From e22530f662320bee0eb0007edc55a8b957082724 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 8 Sep 2026 15:39:29 +0000 Subject: [PATCH 1/6] Replace AggregateFnVTable empty/merge with partial_from_scalar and reduce_partials The vtable's `empty_partial`, `combine_partials` (typed partial plus untyped scalar) and `reset` are replaced by two primitives so each aggregate either parses scalars or operates on typed state: - `partial_from_scalar(options, input_dtype, scalar) -> Partial` parses a partial scalar (kernel results, cached statistics, other accumulators' `to_scalar`) into the typed state, and is its inverse. - `reduce_partials(options, input_dtype, impl IntoIterator)` reduces owned partials in iteration order; the empty sequence is the identity (the state of a group with no values). Accumulator changes: - `Accumulator.partial` is an `Option`, with `None` as the cheap empty state, so folds and empty accumulators never construct an identity partial. - `DynAccumulator::combine_partials` is removed. `merge_from` downcasts the other accumulator (`DynAccumulator::downcast_mut`), checks that the aggregate, options and input dtype match, and folds its typed partial directly, so merging never round-trips through scalars. - `Combined` holds typed child accumulators and merges children with `merge_from`; its `reduce_partials` seeds from the first pair instead of building fresh children per fold. Partial states: - Min, Max, MinMax, BoundedMin, BoundedMax, IsSorted, IsConstant, Sum and SumV2 build their identity from one `empty` constructor, and reduce bodies move values instead of cloning them. - IsSorted and IsConstant no longer serialize a false verdict without a first value as the null (empty) struct, which previously lost the verdict when merged into a materialized empty state or persisted as a zone stat. IsSorted's reduce honors an unsorted partial before the emptiness check. - BloomPartial gains a typed `union`, so reducing bloom partials ORs blocks without serializing. `can_satisfy` documents that satisfaction is a claim about stored state: a partial for the requested aggregate must be creatable from this aggregate's partial. DuckDB's `finalize_scan` merges thread-local accumulators with `merge_from`. Tests construct partials directly where parsing is not the subject, and cover merging into a materialized empty state, boundary-less false verdicts, type-erased Combined merges, and SumV2 identity/overflow merges. Signed-off-by: "Claude" Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- vortex-array/src/aggregate_fn/accumulator.rs | 215 ++++++++++++-- vortex-array/src/aggregate_fn/combined.rs | 91 ++++-- .../src/aggregate_fn/fns/all_nan/mod.rs | 19 +- .../aggregate_fn/fns/all_non_distinct/mod.rs | 27 +- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 19 +- .../src/aggregate_fn/fns/all_non_null/mod.rs | 19 +- .../src/aggregate_fn/fns/all_null/mod.rs | 19 +- .../src/aggregate_fn/fns/bounded_max/mod.rs | 110 +++++--- .../src/aggregate_fn/fns/bounded_min/mod.rs | 51 +++- .../src/aggregate_fn/fns/count/mod.rs | 43 +-- .../src/aggregate_fn/fns/first/mod.rs | 51 ++-- .../src/aggregate_fn/fns/is_constant/mod.rs | 150 ++++++---- .../src/aggregate_fn/fns/is_sorted/mod.rs | 263 +++++++++++------- vortex-array/src/aggregate_fn/fns/last/mod.rs | 54 ++-- vortex-array/src/aggregate_fn/fns/max/mod.rs | 41 ++- vortex-array/src/aggregate_fn/fns/min/mod.rs | 41 ++- .../src/aggregate_fn/fns/min_max/mod.rs | 65 +++-- .../src/aggregate_fn/fns/nan_count/mod.rs | 35 +-- .../src/aggregate_fn/fns/null_count/mod.rs | 26 +- .../src/aggregate_fn/fns/sum/decimal.rs | 84 +++--- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 209 +++++++++----- .../src/aggregate_fn/fns/sum_v2/mod.rs | 91 +++--- .../src/aggregate_fn/fns/sum_v2/tests.rs | 63 +++-- .../fns/uncompressed_size_in_bytes/mod.rs | 42 ++- vortex-array/src/aggregate_fn/foreign.rs | 12 +- vortex-array/src/aggregate_fn/proto.rs | 14 +- vortex-array/src/aggregate_fn/vtable.rs | 41 ++- .../src/arrays/chunked/compute/aggregate.rs | 2 +- vortex-duckdb/src/table_function.rs | 2 +- .../zoned/aggregates/bloom_filter/constant.rs | 25 +- .../zoned/aggregates/bloom_filter/mod.rs | 104 +++---- .../bloom_filter/partial/aggregate.rs | 31 ++- vortex-spatial/src/aggregate_fn/aabb.rs | 60 ++-- 33 files changed, 1334 insertions(+), 785 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 69bae4e1053..0735cdaaa73 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; + +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -27,6 +31,8 @@ pub type AccumulatorRef = Box; pub struct Accumulator { /// The vtable of the aggregate function. vtable: V, + /// The options of the aggregate function. + options: V::Options, /// Type-erased aggregate function used for kernel dispatch. aggregate_fn: AggregateFnRef, /// The DType of the input. @@ -36,7 +42,10 @@ pub struct Accumulator { /// The DType of the accumulator state. partial_dtype: DType, /// The partial state of the accumulator, updated after each accumulate/merge call. - partial: V::Partial, + /// + /// `None` is the empty-group state; a live partial is only materialized when a batch is + /// accumulated in place, so empty accumulators and folds never construct one. + partial: Option, } impl Accumulator { @@ -55,18 +64,44 @@ impl Accumulator { dtype ) })?; - let partial = vtable.empty_partial(&options, &dtype)?; - let aggregate_fn = AggregateFn::new(vtable.clone(), options).erased(); + let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); Ok(Self { vtable, + options, aggregate_fn, dtype, return_dtype, partial_dtype, - partial, + partial: None, }) } + + /// The identity partial state: the state of a group with no accumulated values. + fn empty_partial(&self) -> VortexResult { + self.vtable.reduce_partials(&self.options, &self.dtype, []) + } + + /// Materialize the partial state in place so a batch can be accumulated into it. + fn ensure_partial(&mut self) -> VortexResult<()> { + if self.partial.is_none() { + self.partial = Some(self.empty_partial()?); + } + Ok(()) + } + + /// Reduce an incoming partial state into the accumulator's current state. + pub(crate) fn fold_partial(&mut self, other: V::Partial) -> VortexResult<()> { + self.partial = Some(match self.partial.take() { + // Reducing the incoming partial with the empty state is the identity. + None => other, + Some(current) => { + self.vtable + .reduce_partials(&self.options, &self.dtype, [current, other])? + } + }); + Ok(()) + } } /// A trait object for type-erased accumulators, used for dynamic dispatch when the aggregate @@ -75,11 +110,11 @@ pub trait DynAccumulator: 'static + Send { /// Accumulate a new array into the accumulator's state. fn accumulate(&mut self, batch: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()>; - /// Fold an external partial-state scalar into this accumulator's state. + /// Drain another accumulator's state into this one, resetting `other`. /// - /// The scalar must have the dtype reported by the vtable's `partial_dtype` for the - /// options and input dtype used to construct this accumulator. - fn combine_partials(&mut self, other: Scalar) -> VortexResult<()>; + /// The other accumulator must have been constructed for the same aggregate function, + /// options, and input dtype as this one. + fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()>; /// Whether the accumulator's result is fully determined. fn is_saturated(&self) -> bool; @@ -104,6 +139,16 @@ pub trait DynAccumulator: 'static + Send { /// /// Resets the accumulator state back to the initial state. fn finish(&mut self) -> VortexResult; + + /// Access the accumulator as [`Any`], so it can be downcast to a typed [`Accumulator`]. + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +impl dyn DynAccumulator { + /// Downcast to the typed [`Accumulator`] of the aggregate vtable `V`. + pub fn downcast_mut(&mut self) -> Option<&mut Accumulator> { + self.as_any_mut().downcast_mut() + } } impl DynAccumulator for Accumulator { @@ -137,7 +182,10 @@ impl DynAccumulator for Accumulator { ); partial.cast(&self.partial_dtype)? }; - self.vtable.combine_partials(&mut self.partial, partial)?; + let parsed = self + .vtable + .partial_from_scalar(&self.options, &self.dtype, partial)?; + self.fold_partial(parsed)?; return Ok(()); } @@ -161,13 +209,18 @@ impl DynAccumulator for Accumulator { result.dtype(), self.partial_dtype, ); - self.vtable.combine_partials(&mut self.partial, result)?; + let parsed = self + .vtable + .partial_from_scalar(&self.options, &self.dtype, result)?; + self.fold_partial(parsed)?; return Ok(()); } } // 2. Allow the vtable to short-circuit on the raw array before decompression. - if self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { + self.ensure_partial()?; + let partial = self.partial.as_mut().vortex_expect("partial materialized"); + if self.vtable.try_accumulate(partial, batch, ctx)? { return Ok(()); } @@ -193,7 +246,10 @@ impl DynAccumulator for Accumulator { result.dtype(), self.partial_dtype, ); - self.vtable.combine_partials(&mut self.partial, result)?; + let parsed = self + .vtable + .partial_from_scalar(&self.options, &self.dtype, result)?; + self.fold_partial(parsed)?; return Ok(()); } @@ -203,23 +259,48 @@ impl DynAccumulator for Accumulator { // 4. Otherwise, execute the batch until it is columnar and accumulate it into the state. let columnar = batch.execute::(ctx)?; - self.vtable.accumulate(&mut self.partial, &columnar, ctx) + self.ensure_partial()?; + let partial = self.partial.as_mut().vortex_expect("partial materialized"); + self.vtable.accumulate(partial, &columnar, ctx) } - fn combine_partials(&mut self, other: Scalar) -> VortexResult<()> { - self.vtable.combine_partials(&mut self.partial, other) + fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()> { + let Some(other) = other.downcast_mut::() else { + vortex_bail!( + "Cannot merge into a {} accumulator from an accumulator of a different aggregate", + self.aggregate_fn, + ); + }; + vortex_ensure!( + other.options == self.options && other.dtype == self.dtype, + "Cannot merge {} accumulators with different options or input dtypes", + self.aggregate_fn, + ); + match other.partial.take() { + Some(partial) => self.fold_partial(partial), + None => Ok(()), + } + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self } fn is_saturated(&self) -> bool { - self.vtable.is_saturated(&self.partial) + self.partial + .as_ref() + .is_some_and(|partial| self.vtable.is_saturated(partial)) } fn reset(&mut self) { - self.vtable.reset(&mut self.partial); + self.partial = None; } fn partial_scalar(&self) -> VortexResult { - let partial = self.vtable.to_scalar(&self.partial)?; + let partial = match &self.partial { + Some(partial) => self.vtable.to_scalar(partial)?, + None => self.vtable.to_scalar(&self.empty_partial()?)?, + }; #[cfg(debug_assertions)] { @@ -235,7 +316,10 @@ impl DynAccumulator for Accumulator { } fn final_scalar(&self) -> VortexResult { - let result = self.vtable.finalize_scalar(&self.partial)?; + let result = match &self.partial { + Some(partial) => self.vtable.finalize_scalar(partial)?, + None => self.vtable.finalize_scalar(&self.empty_partial()?)?, + }; vortex_ensure!( result.dtype() == &self.return_dtype, @@ -272,6 +356,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AccumulatorRef; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -279,6 +364,7 @@ mod tests { use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; + use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; @@ -288,7 +374,10 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; use crate::scalar::Scalar; + use crate::scalar::ScalarValue; /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` @@ -320,7 +409,7 @@ mod tests { } } - /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of + /// Sum partial sentinel `{sum: 42.0, is_overflow: false, is_empty: false}` — distinguishable from the natural Sum of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -448,4 +537,90 @@ mod tests { ); Ok(()) } + + #[test] + fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { + static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let batch = dict_of_seven(); + batch + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(11.0f64))); + + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + acc.accumulate(&batch, &mut ctx)?; + + assert_eq!(acc.finish()?.as_primitive().as_::(), Some(11.0)); + Ok(()) + } + + fn sum_i32_accumulator(options: NumericalAggregateOpts) -> VortexResult { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + Ok(Box::new(Accumulator::try_new(Sum, options, dtype)?)) + } + + #[test] + fn merge_from_drains_other_accumulator() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let mut global = sum_i32_accumulator(NumericalAggregateOpts::default())?; + let mut local = sum_i32_accumulator(NumericalAggregateOpts::default())?; + + global.accumulate(&buffer![10i32, 20].into_array(), &mut ctx)?; + local.accumulate(&buffer![5i32].into_array(), &mut ctx)?; + global.merge_from(local.as_mut())?; + + assert_eq!( + global.finish()?, + Scalar::primitive(35i64, Nullability::Nullable) + ); + // The merged-from accumulator is reset back to the empty state. + assert_eq!( + local.finish()?, + Scalar::primitive(0i64, Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn merge_from_rejects_a_different_aggregate() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut sum = sum_i32_accumulator(NumericalAggregateOpts::default())?; + let mut min: AccumulatorRef = Box::new(Accumulator::try_new( + Min, + NumericalAggregateOpts::default(), + dtype, + )?); + + assert!(sum.merge_from(min.as_mut()).is_err()); + Ok(()) + } + + #[test] + fn merge_from_rejects_mismatched_options() -> VortexResult<()> { + let mut skipping = sum_i32_accumulator(NumericalAggregateOpts::skip_nans())?; + let mut including = sum_i32_accumulator(NumericalAggregateOpts::include_nans())?; + + assert!(skipping.merge_from(including.as_mut()).is_err()); + Ok(()) + } + + #[test] + fn merge_from_combines_child_accumulators() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let mut global: AccumulatorRef = Box::new(mean_f64_accumulator()?); + let mut local: AccumulatorRef = Box::new(mean_f64_accumulator()?); + + global.accumulate(&buffer![1.0f64, 2.0].into_array(), &mut ctx)?; + local.accumulate(&buffer![6.0f64].into_array(), &mut ctx)?; + global.merge_from(local.as_mut())?; + + assert_eq!(global.finish()?.as_primitive().as_::(), Some(3.0)); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 70a385add48..f600d6a802a 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -19,9 +19,9 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AccumulatorRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynAccumulator; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; @@ -48,6 +48,11 @@ type LeftOptions = <::Left as AggregateFnVTable>::Option type RightOptions = <::Right as AggregateFnVTable>::Options; /// Combined options for a [`BinaryCombined`] aggregate. pub type CombinedOptions = PairOptions, RightOptions>; +/// Pair of typed child accumulators holding the partial state of a [`BinaryCombined`] aggregate. +type ChildAccumulators = ( + Accumulator<::Left>, + Accumulator<::Right>, +); /// Declare an aggregate function in terms of two child aggregates. pub trait BinaryCombined: 'static + Send + Sync + Clone { @@ -126,14 +131,25 @@ impl Combined { pub fn new(inner: T) -> Self { Self(inner) } + + /// Construct a pair of empty child accumulators. + fn new_child_accumulators( + &self, + options: &CombinedOptions, + input_dtype: &DType, + ) -> VortexResult> { + let left = Accumulator::try_new(self.0.left(), options.0.clone(), input_dtype.clone())?; + let right = Accumulator::try_new(self.0.right(), options.1.clone(), input_dtype.clone())?; + Ok((left, right)) + } } impl AggregateFnVTable for Combined { type Options = CombinedOptions; - // Each child is held as a fully-fledged `AccumulatorRef` so that batches dispatched through + // Each child is held as a fully-fledged `Accumulator` so that batches dispatched through // `try_accumulate` consult the kernel registry per-child (e.g. a `(Dict, Sum)` kernel fires // for the inner `Sum` child of `Combined`). - type Partial = (AccumulatorRef, AccumulatorRef); + type Partial = ChildAccumulators; fn id(&self) -> AggregateFnId { self.0.id() @@ -157,35 +173,55 @@ impl AggregateFnVTable for Combined { Some(self.0.partial_struct_dtype(l, r)) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - let left = Accumulator::try_new(self.0.left(), options.0.clone(), input_dtype.clone())?; - let right = Accumulator::try_new(self.0.right(), options.1.clone(), input_dtype.clone())?; - Ok(( - Box::new(left) as AccumulatorRef, - Box::new(right) as AccumulatorRef, - )) + let (mut left, mut right) = self.new_child_accumulators(options, input_dtype)?; + // A null partial represents an empty group and parses to empty child accumulators. + if !scalar.is_null() { + let s = scalar.as_struct(); + let lname = self.0.left_name(); + let rname = self.0.right_name(); + let l_field = s + .field(lname) + .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", lname))?; + let r_field = s + .field(rname) + .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", rname))?; + left.fold_partial(self.0.left().partial_from_scalar( + &options.0, + input_dtype, + l_field, + )?)?; + right.fold_partial(self.0.right().partial_from_scalar( + &options.1, + input_dtype, + r_field, + )?)?; + } + Ok((left, right)) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut partials = partials.into_iter(); + let Some((mut left, mut right)) = partials.next() else { + return self.new_child_accumulators(options, input_dtype); + }; + // The children are typed accumulators of the same child aggregates, so the remaining + // partials merge state directly without any scalar interchange. + for (mut l, mut r) in partials { + left.merge_from(&mut l)?; + right.merge_from(&mut r)?; } - let s = other.as_struct(); - let lname = self.0.left_name(); - let rname = self.0.right_name(); - let l_field = s - .field(lname) - .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", lname))?; - let r_field = s - .field(rname) - .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", rname))?; - partial.0.combine_partials(l_field)?; - partial.1.combine_partials(r_field)?; - Ok(()) + Ok((left, right)) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -197,11 +233,6 @@ impl AggregateFnVTable for Combined { Ok(Scalar::struct_(dtype, vec![l_scalar, r_scalar])) } - fn reset(&self, partial: &mut Self::Partial) { - partial.0.reset(); - partial.1.reset(); - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { partial.0.is_saturated() && partial.1.is_saturated() } diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index 68a58908018..a3d3198c9b6 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -60,27 +60,28 @@ impl AggregateFnVTable for AllNan { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(true) + bool::try_from(&scalar) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().all(|partial| partial)) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { !*partial } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index a8392f3fa74..d9b399a9e4a 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -147,25 +147,26 @@ impl AggregateFnVTable for AllNonDistinct { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { Ok(AllNonDistinctPartial { - all_non_distinct: true, + all_non_distinct: scalar.as_bool().value().unwrap_or(false), }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if !partial.all_non_distinct { - return Ok(()); - } - - if !other.as_bool().value().unwrap_or(false) { - partial.all_non_distinct = false; - } - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(AllNonDistinctPartial { + all_non_distinct: partials.into_iter().all(|partial| partial.all_non_distinct), + }) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -175,10 +176,6 @@ impl AggregateFnVTable for AllNonDistinct { )) } - fn reset(&self, partial: &mut Self::Partial) { - partial.all_non_distinct = true; - } - #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { !partial.all_non_distinct diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index fe8527da966..5ee577b58f9 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -60,27 +60,28 @@ impl AggregateFnVTable for AllNonNan { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(true) + bool::try_from(&scalar) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().all(|partial| partial)) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { !*partial } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index c07dbb907c9..fd95c5016d2 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -51,27 +51,28 @@ impl AggregateFnVTable for AllNonNull { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(true) + bool::try_from(&scalar) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().all(|partial| partial)) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { !*partial } diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index ec64e3d5c43..1bff71d5ca8 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -51,27 +51,28 @@ impl AggregateFnVTable for AllNull { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(true) + bool::try_from(&scalar) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().all(|partial| partial)) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { !*partial } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index bb27d9f9ab1..c0d0c838770 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -76,6 +76,15 @@ pub struct BoundedMaxPartial { } impl BoundedMaxPartial { + /// The state of a group with no accumulated values. + fn empty(options: &BoundedMaxOptions, input_dtype: &DType) -> Self { + Self { + state: BoundedMaxState::Empty, + element_dtype: input_dtype.clone(), + max_bytes: options.max_bytes, + } + } + fn merge_bound(&mut self, max: Scalar) { if max.is_null() { return; @@ -189,42 +198,61 @@ impl AggregateFnVTable for BoundedMax { supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { + // A null partial means the producing accumulator saw nothing valid. + let state = if scalar.is_null() { + BoundedMaxState::Empty + } else { + let Some(fields) = scalar.as_struct_opt() else { + vortex_bail!( + "BoundedMax partial must be a struct, got {}", + scalar.dtype() + ); + }; + let Some(bound) = fields.field_by_idx(0) else { + vortex_bail!("BoundedMax partial is missing its bound field"); + }; + let Some(unknown) = fields + .field_by_idx(1) + .and_then(|unknown| unknown.as_bool().value()) + else { + vortex_bail!("BoundedMax partial is missing its non-null unknown field"); + }; + + if unknown { + BoundedMaxState::Unknown + } else if bound.is_null() { + BoundedMaxState::Empty + } else { + BoundedMaxState::Value(bound) + } + }; Ok(BoundedMaxPartial { - state: BoundedMaxState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, + state, + ..BoundedMaxPartial::empty(options, input_dtype) }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); - } - - let Some(other) = other.as_struct_opt() else { - vortex_bail!("BoundedMax partial must be a struct, got {}", other.dtype()); - }; - let Some(bound) = other.field_by_idx(0) else { - vortex_bail!("BoundedMax partial is missing its bound field"); - }; - let Some(unknown) = other - .field_by_idx(1) - .and_then(|unknown| unknown.as_bool().value()) - else { - vortex_bail!("BoundedMax partial is missing its non-null unknown field"); - }; - - if unknown { - partial.unknown(); - } else { - partial.merge_bound(bound); + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = BoundedMaxPartial::empty(options, input_dtype); + for partial in partials { + match partial.state { + BoundedMaxState::Empty => {} + BoundedMaxState::Value(max) => acc.merge_bound(max), + BoundedMaxState::Unknown => acc.unknown(), + } } - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -249,10 +277,6 @@ impl AggregateFnVTable for BoundedMax { } } - fn reset(&self, partial: &mut Self::Partial) { - partial.state = BoundedMaxState::Empty; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { matches!(partial.state, BoundedMaxState::Unknown) } @@ -332,7 +356,8 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_max::BoundedMax; use crate::aggregate_fn::fns::bounded_max::BoundedMaxOptions; - use crate::aggregate_fn::fns::bounded_max::make_bounded_max_partial_dtype; + use crate::aggregate_fn::fns::bounded_max::BoundedMaxPartial; + use crate::aggregate_fn::fns::bounded_max::BoundedMaxState; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::array_session; @@ -442,7 +467,11 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(make_bounded_max_partial_dtype(values.dtype())))?; + acc.fold_partial(BoundedMaxPartial { + state: BoundedMaxState::Empty, + element_dtype: values.dtype().clone(), + max_bytes: max_bytes(2), + })?; assert_eq!( acc.finish()?, @@ -463,17 +492,12 @@ mod tests { values.dtype().clone(), )?; - let partial_dtype = make_bounded_max_partial_dtype(values.dtype()); - let unknown = Scalar::struct_( - partial_dtype, - vec![ - Scalar::null(values.dtype().as_nullable()), - Scalar::bool(true, Nullability::NonNullable), - ], - ); - acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(unknown)?; + acc.fold_partial(BoundedMaxPartial { + state: BoundedMaxState::Unknown, + element_dtype: values.dtype().clone(), + max_bytes: max_bytes(2), + })?; assert_eq!(acc.finish()?, Scalar::null(values.dtype().as_nullable())); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 7a72442550b..8a70d924ce6 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -61,6 +61,15 @@ pub struct BoundedMinPartial { } impl BoundedMinPartial { + /// The state of a group with no accumulated values. + fn empty(options: &BoundedMinOptions, input_dtype: &DType) -> Self { + Self { + state: BoundedMinState::Empty, + element_dtype: input_dtype.clone(), + max_bytes: options.max_bytes, + } + } + fn merge(&mut self, min: Scalar) { if min.is_null() { return; @@ -143,21 +152,37 @@ impl AggregateFnVTable for BoundedMin { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { + // A null partial means the producing accumulator saw nothing valid. + let state = if scalar.is_null() { + BoundedMinState::Empty + } else { + BoundedMinState::Value(scalar) + }; Ok(BoundedMinPartial { - state: BoundedMinState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, + state, + ..BoundedMinPartial::empty(options, input_dtype) }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = BoundedMinPartial::empty(options, input_dtype); + for partial in partials { + if let BoundedMinState::Value(min) = partial.state { + acc.merge(min); + } + } + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -168,10 +193,6 @@ impl AggregateFnVTable for BoundedMin { } } - fn reset(&self, partial: &mut Self::Partial) { - partial.state = BoundedMinState::Empty; - } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { false } @@ -249,6 +270,8 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_min::BoundedMin; use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions; + use crate::aggregate_fn::fns::bounded_min::BoundedMinPartial; + use crate::aggregate_fn::fns::bounded_min::BoundedMinState; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::array_session; @@ -322,7 +345,11 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(values.dtype().as_nullable()))?; + acc.fold_partial(BoundedMinPartial { + state: BoundedMinState::Empty, + element_dtype: values.dtype().clone(), + max_bytes: max_bytes(2), + })?; assert_eq!( acc.finish()?, diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 8f7d68027bc..64a27d9fd79 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -58,34 +58,37 @@ impl AggregateFnVTable for Count { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { Ok(CountPartial { - count: 0, + count: scalar + .as_primitive() + .typed_value::() + .vortex_expect("count partial should not be null"), exclude_nans: options.skip_nans && input_dtype.is_float(), }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let val = other - .as_primitive() - .typed_value::() - .vortex_expect("count partial should not be null"); - partial.count += val; - Ok(()) + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(CountPartial { + count: partials.into_iter().map(|partial| partial.count).sum(), + exclude_nans: options.skip_nans && input_dtype.is_float(), + }) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { Ok(Scalar::primitive(partial.count, Nullability::NonNullable)) } - fn reset(&self, partial: &mut Self::Partial) { - partial.count = 0; - } - #[inline] fn is_saturated(&self, _partial: &Self::Partial) -> bool { false @@ -142,6 +145,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; + use crate::aggregate_fn::fns::count::CountPartial; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -246,16 +250,15 @@ mod tests { #[test] fn count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Count.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - - let scalar1 = Scalar::primitive(5u64, Nullability::NonNullable); - Count.combine_partials(&mut state, scalar1)?; + let options = NumericalAggregateOpts::default(); - let scalar2 = Scalar::primitive(3u64, Nullability::NonNullable); - Count.combine_partials(&mut state, scalar2)?; + let partial_of = |count: u64| CountPartial { + count, + exclude_nans: false, + }; + let state = Count.reduce_partials(&options, &dtype, [partial_of(5), partial_of(3)])?; let result = Count.to_scalar(&state)?; - Count.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index c41e2575057..028545ea031 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -57,23 +57,30 @@ impl AggregateFnVTable for First { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { + // A null partial means the producing accumulator saw nothing valid. Ok(FirstPartial { return_dtype: input_dtype.as_nullable(), - value: None, + value: (!scalar.is_null()).then_some(scalar), }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - // Only the first non-null partial wins; later ones are ignored. - if partial.value.is_none() && !other.is_null() { - partial.value = Some(other); - } - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + // The first non-empty partial in iteration order wins; later ones are ignored. + Ok(FirstPartial { + return_dtype: input_dtype.as_nullable(), + value: partials.into_iter().find_map(|partial| partial.value), + }) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -83,10 +90,6 @@ impl AggregateFnVTable for First { }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.value = None; - } - #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { partial.value.is_some() @@ -138,6 +141,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::first::First; + use crate::aggregate_fn::fns::first::FirstPartial; use crate::aggregate_fn::fns::first::first; use crate::array_session; use crate::arrays::ChunkedArray; @@ -259,17 +263,20 @@ mod tests { #[test] fn first_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = First.empty_partial(&EmptyOptions, &dtype)?; - - // A null partial means the sub-accumulator saw nothing valid - should be ignored. - First.combine_partials(&mut state, Scalar::null(dtype.as_nullable()))?; - assert!(!First.is_saturated(&state)); - - First.combine_partials(&mut state, Scalar::primitive(5i32, Nullable))?; + let partial_of = |value: Option| FirstPartial { + return_dtype: dtype.as_nullable(), + value, + }; + + // An empty partial means the sub-accumulator saw nothing valid - it is ignored. + let empty = partial_of(None); + assert!(!First.is_saturated(&empty)); + + // The first non-empty partial wins; subsequent valid partials are dropped. + let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); + let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); + let state = First.reduce_partials(&EmptyOptions, &dtype, [empty, five, seven])?; assert!(First.is_saturated(&state)); - - // Subsequent valid partials are dropped. - First.combine_partials(&mut state, Scalar::primitive(7i32, Nullable))?; assert_eq!(First.to_scalar(&state)?, Scalar::primitive(5i32, Nullable)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 96017c3ed92..a0d478bd40b 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -223,6 +223,15 @@ pub struct IsConstantPartial { } impl IsConstantPartial { + /// The state of a group with no accumulated values. + fn empty(input_dtype: &DType) -> Self { + Self { + is_constant: true, + first_value: None, + element_dtype: input_dtype.clone(), + } + } + fn check_value(&mut self, value: Scalar) { if !self.is_constant { return; @@ -283,70 +292,66 @@ impl AggregateFnVTable for IsConstant { } } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(IsConstantPartial { - is_constant: true, - first_value: None, - element_dtype: input_dtype.clone(), - }) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if !partial.is_constant { - return Ok(()); - } - - // Null struct means the other accumulator was empty, skip it. - if other.is_null() { - return Ok(()); + // A null struct means the producing accumulator was empty. + if scalar.is_null() { + return Ok(IsConstantPartial::empty(input_dtype)); } - let other_is_constant = other + let is_constant = scalar .as_struct() .field_by_idx(0) .map(|s| s.as_bool().value().unwrap_or(false)) .unwrap_or(false); - if !other_is_constant { - partial.is_constant = false; - return Ok(()); - } - - let other_value = other.as_struct().field_by_idx(1); + Ok(IsConstantPartial { + is_constant, + first_value: scalar.as_struct().field_by_idx(1), + element_dtype: input_dtype.clone(), + }) + } - if let Some(other_val) = other_value { - partial.check_value(other_val); + fn reduce_partials( + &self, + _options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = IsConstantPartial::empty(input_dtype); + for partial in partials { + if !partial.is_constant { + acc.is_constant = false; + break; + } + if let Some(value) = partial.first_value { + acc.check_value(value); + } } - - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { let dtype = make_is_constant_partial_dtype(&partial.element_dtype); - Ok(match &partial.first_value { - None => { - // Empty accumulator — return null struct. - Scalar::null(dtype) - } - Some(first_value) => Scalar::struct_( - dtype, - vec![ - Scalar::bool(partial.is_constant, Nullability::NonNullable), - first_value - .clone() - .cast(&partial.element_dtype.as_nullable())?, - ], - ), - }) - } - - fn reset(&self, partial: &mut Self::Partial) { - partial.is_constant = true; - partial.first_value = None; + let element_dtype = partial.element_dtype.as_nullable(); + // Only a constant partial that saw no values is the empty (null) state: a non-constant + // verdict stands regardless of whether a value was observed. + let first_value = match &partial.first_value { + Some(first_value) => first_value.clone().cast(&element_dtype)?, + None if partial.is_constant => return Ok(Scalar::null(dtype)), + None => Scalar::null(element_dtype), + }; + Ok(Scalar::struct_( + dtype, + vec![ + Scalar::bool(partial.is_constant, Nullability::NonNullable), + first_value, + ], + )) } #[inline] @@ -448,6 +453,12 @@ mod tests { use crate::IntoArray as _; use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::fns::is_constant::IsConstant; + use crate::aggregate_fn::fns::is_constant::IsConstantPartial; use crate::aggregate_fn::fns::is_constant::is_constant; use crate::array_session; use crate::arrays::BoolArray; @@ -753,4 +764,49 @@ mod tests { Ok(()) } + + /// Merging a non-constant partial into a materialized empty one must keep the false verdict + /// in the partial scalar rather than collapsing it to the empty state. + #[test] + fn non_constant_merged_into_empty_keeps_verdict() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut empty = Accumulator::try_new(IsConstant, EmptyOptions, dtype.clone())?; + let mut varying = Accumulator::try_new(IsConstant, EmptyOptions, dtype)?; + + // An empty batch materializes the empty partial in place. + empty.accumulate( + &PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(), + &mut ctx, + )?; + varying.accumulate(&buffer![1i32, 2].into_array(), &mut ctx)?; + empty.merge_from(&mut varying)?; + + assert!(!empty.partial_scalar()?.is_null()); + assert_eq!( + empty.finish()?, + Scalar::bool(false, Nullability::NonNullable) + ); + Ok(()) + } + + /// A non-constant verdict without an observed value is not the empty state. + #[test] + fn non_constant_partial_without_value_is_not_empty() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let partial = IsConstantPartial { + is_constant: false, + first_value: None, + element_dtype: dtype.clone(), + }; + + let scalar = IsConstant.to_scalar(&partial)?; + assert!(!scalar.is_null()); + let parsed = IsConstant.partial_from_scalar(&EmptyOptions, &dtype, scalar)?; + assert_eq!( + IsConstant.finalize_scalar(&parsed)?, + Scalar::bool(false, Nullability::NonNullable) + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 99a7b99128a..9ab0501a4ce 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -209,6 +209,19 @@ pub struct IsSortedPartial { element_dtype: DType, } +impl IsSortedPartial { + /// The state of a group with no accumulated values. + fn empty(options: &IsSortedOptions, input_dtype: &DType) -> Self { + Self { + is_sorted: true, + strict: options.strict, + first_value: None, + last_value: None, + element_dtype: input_dtype.clone(), + } + } +} + static NAMES: std::sync::LazyLock = std::sync::LazyLock::new(|| { FieldNames::from(["is_sorted", "strict", "first_value", "last_value"]) }); @@ -277,123 +290,120 @@ impl AggregateFnVTable for IsSorted { } } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(IsSortedPartial { - is_sorted: true, - strict: options.strict, - first_value: None, - last_value: None, - element_dtype: input_dtype.clone(), - }) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if !partial.is_sorted { - return Ok(()); + // A null struct means the producing accumulator was empty. + if scalar.is_null() { + return Ok(IsSortedPartial::empty(options, input_dtype)); } - // Null struct means the other accumulator was empty, skip it. - if other.is_null() { - return Ok(()); - } - - let other_is_sorted = other + let is_sorted = scalar .as_struct() .field_by_idx(0) .map(|s| s.as_bool().value().unwrap_or(false)) .unwrap_or(false); - let other_first = other.as_struct().field_by_idx(2); - let other_last = other.as_struct().field_by_idx(3); + // The scalar's own strict flag is ignored: strictness comes from the options. + Ok(IsSortedPartial { + is_sorted, + strict: options.strict, + first_value: scalar.as_struct().field_by_idx(2), + last_value: scalar.as_struct().field_by_idx(3), + element_dtype: input_dtype.clone(), + }) + } + + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = IsSortedPartial::empty(options, input_dtype); - if !other_is_sorted { - partial.is_sorted = false; - // Still update last_value for correctness if needed, but we're done. - if let Some(last) = other_last { - partial.last_value = Some(last); + for partial in partials { + if !acc.is_sorted { + break; } - return Ok(()); - } - // Check boundary: self.last_value vs other.first_value - if let Some(self_last) = &partial.last_value - && let Some(other_first_val) = &other_first - { - if !self_last.is_null() && !other_first_val.is_null() { - let boundary_ok = if partial.strict { - *self_last < *other_first_val - } else { - *self_last <= *other_first_val - }; - if !boundary_ok { - partial.is_sorted = false; + if !partial.is_sorted { + // An unsorted partial settles the verdict whether or not it observed boundaries. + acc.is_sorted = false; + if let Some(last) = partial.last_value.or_else(|| partial.first_value.clone()) { + acc.last_value = Some(last); } - } else if !self_last.is_null() && other_first_val.is_null() { - // non-null before null violates sort order - partial.is_sorted = false; - } else if self_last.is_null() && other_first_val.is_null() && partial.strict { - // both null with strict: violates strict sort - partial.is_sorted = false; + if acc.first_value.is_none() { + acc.first_value = partial.first_value; + } + break; } - } - // Update first_value if this is the first non-empty chunk. - if partial.first_value.is_none() { - partial.first_value = other_first; - } - if let Some(last) = other_last { - partial.last_value = Some(last); + // A sorted partial without a first value is empty and contributes nothing. + let Some(first) = partial.first_value else { + continue; + }; + // A partial that saw a single value carries it as both boundaries. + let last = partial.last_value.unwrap_or_else(|| first.clone()); + + // Check boundary: acc.last_value vs partial.first_value + if let Some(acc_last) = &acc.last_value { + if !acc_last.is_null() && !first.is_null() { + let boundary_ok = if acc.strict { + *acc_last < first + } else { + *acc_last <= first + }; + if !boundary_ok { + acc.is_sorted = false; + } + } else if !acc_last.is_null() && first.is_null() { + // non-null before null violates sort order + acc.is_sorted = false; + } else if acc_last.is_null() && first.is_null() && acc.strict { + // both null with strict: violates strict sort + acc.is_sorted = false; + } + } + + // Update first_value if this is the first non-empty partial. + if acc.first_value.is_none() { + acc.first_value = Some(first); + } + acc.last_value = Some(last); } - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { let dtype = make_is_sorted_partial_dtype(&partial.element_dtype); - Ok(match (&partial.first_value, &partial.last_value) { - (None, _) => { - // Empty accumulator — return null struct. - Scalar::null(dtype) - } - (Some(first_value), Some(last_value)) => { - // SAFETY: We constructed partial_dtype and the children match its field dtypes. - unsafe { - Scalar::struct_unchecked( - dtype, - [ - Scalar::bool(partial.is_sorted, Nullability::NonNullable), - Scalar::bool(partial.strict, Nullability::NonNullable), - first_value.clone(), - last_value.clone(), - ], - ) - } - } - (Some(first_value), None) => { - // SAFETY: We constructed partial_dtype and the children match its field dtypes. - unsafe { - Scalar::struct_unchecked( - dtype, - [ - Scalar::bool(partial.is_sorted, Nullability::NonNullable), - Scalar::bool(partial.strict, Nullability::NonNullable), - first_value.clone(), - first_value.clone(), - ], - ) - } - } - }) - } - - fn reset(&self, partial: &mut Self::Partial) { - partial.is_sorted = true; - partial.first_value = None; - partial.last_value = None; + // Only a sorted partial that saw no values is the empty (null) state: an unsorted verdict + // stands regardless of which boundaries were observed. + if partial.is_sorted && partial.first_value.is_none() { + return Ok(Scalar::null(dtype)); + } + let first_value = partial + .first_value + .clone() + .unwrap_or_else(|| Scalar::null(partial.element_dtype.as_nullable())); + // A partial that saw a single value carries it as both boundaries. + let last_value = partial + .last_value + .clone() + .unwrap_or_else(|| first_value.clone()); + Ok(Scalar::struct_( + dtype, + vec![ + Scalar::bool(partial.is_sorted, Nullability::NonNullable), + Scalar::bool(partial.strict, Nullability::NonNullable), + first_value, + last_value, + ], + )) } #[inline] @@ -521,10 +531,7 @@ impl AggregateFnVTable for IsSorted { } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - if partial.first_value.is_none() { - // Empty accumulator → vacuously sorted. - return Ok(Scalar::bool(true, Nullability::NonNullable)); - } + // The empty state is vacuously sorted, so the verdict stands on its own. Ok(Scalar::bool(partial.is_sorted, Nullability::NonNullable)) } } @@ -557,17 +564,28 @@ where #[cfg(test)] mod tests { use rstest::rstest; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::fns::is_sorted::IsSorted; + use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; + use crate::aggregate_fn::fns::is_sorted::IsSortedPartial; use crate::aggregate_fn::fns::is_sorted::is_sorted; use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::scalar::Scalar; use crate::validity::Validity; // Tests migrated from compute/is_sorted.rs @@ -718,4 +736,53 @@ mod tests { Ok(()) } + + /// Merging an unsorted partial into a materialized empty one must keep the false verdict in + /// both the finalized result and the partial scalar. + #[test] + fn unsorted_merged_into_empty_keeps_verdict() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let options = IsSortedOptions { strict: false }; + let mut empty = Accumulator::try_new(IsSorted, options.clone(), dtype.clone())?; + let mut unsorted = Accumulator::try_new(IsSorted, options, dtype)?; + + // An empty batch materializes the empty partial in place. + empty.accumulate( + &PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(), + &mut ctx, + )?; + unsorted.accumulate(&buffer![3i32, 1].into_array(), &mut ctx)?; + empty.merge_from(&mut unsorted)?; + + assert!(!empty.partial_scalar()?.is_null()); + assert_eq!( + empty.finish()?, + Scalar::bool(false, Nullability::NonNullable) + ); + Ok(()) + } + + /// An unsorted verdict without observed boundaries is not the empty state. + #[test] + fn unsorted_partial_without_boundaries_is_not_empty() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let options = IsSortedOptions { strict: false }; + let partial = IsSortedPartial { + is_sorted: false, + strict: false, + first_value: None, + last_value: None, + element_dtype: dtype.clone(), + }; + + let scalar = IsSorted.to_scalar(&partial)?; + assert!(!scalar.is_null()); + let parsed = IsSorted.partial_from_scalar(&options, &dtype, scalar)?; + assert_eq!( + IsSorted.finalize_scalar(&parsed)?, + Scalar::bool(false, Nullability::NonNullable) + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 63ee54a6efa..b119f4aa18a 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -57,23 +57,33 @@ impl AggregateFnVTable for Last { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { + // A null partial means the producing accumulator saw nothing valid. Ok(LastPartial { return_dtype: input_dtype.as_nullable(), - value: None, + value: (!scalar.is_null()).then_some(scalar), }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - // Each new non-null partial replaces the previous one; nulls are ignored. - if !other.is_null() { - partial.value = Some(other); - } - Ok(()) + fn reduce_partials( + &self, + _options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + // The last non-empty partial in iteration order wins; empty ones are ignored. + Ok(LastPartial { + return_dtype: input_dtype.as_nullable(), + value: partials + .into_iter() + .filter_map(|partial| partial.value) + .last(), + }) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -83,10 +93,6 @@ impl AggregateFnVTable for Last { }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.value = None; - } - #[inline] fn is_saturated(&self, _partial: &Self::Partial) -> bool { // Last can never short-circuit: a later batch can always supersede the current value. @@ -136,6 +142,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::last::Last; + use crate::aggregate_fn::fns::last::LastPartial; use crate::aggregate_fn::fns::last::last; use crate::array_session; use crate::arrays::ChunkedArray; @@ -256,17 +263,18 @@ mod tests { #[test] fn last_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Last.empty_partial(&EmptyOptions, &dtype)?; - - Last.combine_partials(&mut state, Scalar::primitive(5i32, Nullable))?; - assert_eq!(Last.to_scalar(&state)?, Scalar::primitive(5i32, Nullable)); - - // A later non-null partial replaces the prior value. - Last.combine_partials(&mut state, Scalar::primitive(7i32, Nullable))?; - assert_eq!(Last.to_scalar(&state)?, Scalar::primitive(7i32, Nullable)); - - // A null partial must not clobber the stored value. - Last.combine_partials(&mut state, Scalar::null(dtype.as_nullable()))?; + let partial_of = |value: Option| LastPartial { + return_dtype: dtype.as_nullable(), + value, + }; + + let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); + let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); + // An empty partial must not clobber a prior value. + let empty = partial_of(None); + + // The last non-empty partial in order replaces the prior values. + let state = Last.reduce_partials(&EmptyOptions, &dtype, [five, seven, empty])?; assert_eq!(Last.to_scalar(&state)?, Scalar::primitive(7i32, Nullable)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index ccf9f4d899d..76713db16c3 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -43,6 +43,15 @@ pub struct MaxPartial { } impl MaxPartial { + /// The state of a group with no accumulated values. + fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { + Self { + max: None, + element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, + } + } + fn merge(&mut self, max: Scalar) { if max.is_null() { return; @@ -121,21 +130,31 @@ impl AggregateFnVTable for Max { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(MaxPartial { - max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - }) + let mut partial = MaxPartial::empty(options, input_dtype); + // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. + partial.merge(scalar); + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = MaxPartial::empty(options, input_dtype); + for partial in partials { + if let Some(max) = partial.max { + acc.merge(max); + } + } + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -146,10 +165,6 @@ impl AggregateFnVTable for Max { } } - fn reset(&self, partial: &mut Self::Partial) { - partial.max = None; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { // A poisoned NaN-including maximum is fully determined. partial.is_poisoned() diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index 488405e8f14..a9443020551 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -43,6 +43,15 @@ pub struct MinPartial { } impl MinPartial { + /// The state of a group with no accumulated values. + fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { + Self { + min: None, + element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, + } + } + fn merge(&mut self, min: Scalar) { if min.is_null() { return; @@ -121,21 +130,31 @@ impl AggregateFnVTable for Min { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(MinPartial { - min: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - }) + let mut partial = MinPartial::empty(options, input_dtype); + // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. + partial.merge(scalar); + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = MinPartial::empty(options, input_dtype); + for partial in partials { + if let Some(min) = partial.min { + acc.merge(min); + } + } + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -146,10 +165,6 @@ impl AggregateFnVTable for Min { } } - fn reset(&self, partial: &mut Self::Partial) { - partial.min = None; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { // A poisoned NaN-including minimum is fully determined. partial.is_poisoned() diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index a5091cf4a9f..737cb6eb716 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -202,6 +202,16 @@ pub struct MinMaxPartial { } impl MinMaxPartial { + /// The state of a group with no accumulated values. + fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { + Self { + min: None, + max: None, + element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, + } + } + /// Merge a local `MinMaxResult` into this partial state. fn merge(&mut self, local: Option) { let Some(MinMaxResult { min, max }) = local else { @@ -308,23 +318,31 @@ impl AggregateFnVTable for MinMax { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(MinMaxPartial { - min: None, - max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - }) + let mut partial = MinMaxPartial::empty(options, input_dtype); + // `merge` normalizes the parsed extrema: nulls stay empty and NaNs poison or drop. + partial.merge(MinMaxResult::from_scalar(scalar)?); + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let local = MinMaxResult::from_scalar(other)?; - partial.merge(local); - Ok(()) + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = MinMaxPartial::empty(options, input_dtype); + for partial in partials { + if let (Some(min), Some(max)) = (partial.min, partial.max) { + acc.merge(Some(MinMaxResult { min, max })); + } + } + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -335,11 +353,6 @@ impl AggregateFnVTable for MinMax { }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.min = None; - partial.max = None; - } - #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { // A poisoned NaN-including min/max is fully determined. @@ -458,6 +471,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMax; + use crate::aggregate_fn::fns::min_max::MinMaxPartial; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::make_minmax_dtype; use crate::aggregate_fn::fns::min_max::min_max; @@ -656,17 +670,16 @@ mod tests { #[test] fn test_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = MinMax.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - - let struct_dtype = make_minmax_dtype(&dtype); - let scalar1 = Scalar::struct_( - struct_dtype.clone(), - vec![Scalar::from(5i32), Scalar::from(15i32)], - ); - MinMax.combine_partials(&mut state, scalar1)?; + let options = NumericalAggregateOpts::default(); + let partial_of = |min: i32, max: i32| MinMaxPartial { + min: Some(Scalar::from(min)), + max: Some(Scalar::from(max)), + element_dtype: dtype.clone(), + skip_nans: true, + }; - let scalar2 = Scalar::struct_(struct_dtype, vec![Scalar::from(2i32), Scalar::from(10i32)]); - MinMax.combine_partials(&mut state, scalar2)?; + let state = + MinMax.reduce_partials(&options, &dtype, [partial_of(5, 15), partial_of(2, 10)])?; let result = MinMaxResult::from_scalar(MinMax.to_scalar(&state)?)? .vortex_expect("should have result"); diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index 723847160e1..79c7a7e51b2 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -115,29 +115,29 @@ impl AggregateFnVTable for NanCount { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(0u64) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let val = other + Ok(scalar .as_primitive() .typed_value::() - .vortex_expect("nan_count partial should not be null"); - *partial += val; - Ok(()) + .vortex_expect("nan_count partial should not be null")) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(*partial, NonNullable)) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().sum()) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] @@ -201,7 +201,6 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::scalar::Scalar; use crate::validity::Validity; #[test] @@ -247,16 +246,10 @@ mod tests { #[test] fn nan_count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut state = NanCount.empty_partial(&EmptyOptions, &dtype)?; - - let scalar1 = Scalar::primitive(5u64, Nullability::NonNullable); - NanCount.combine_partials(&mut state, scalar1)?; - let scalar2 = Scalar::primitive(3u64, Nullability::NonNullable); - NanCount.combine_partials(&mut state, scalar2)?; + let state = NanCount.reduce_partials(&EmptyOptions, &dtype, [5, 3])?; let result = NanCount.to_scalar(&state)?; - NanCount.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index 031e8f6a09b..608e1f7704f 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -84,29 +84,29 @@ impl AggregateFnVTable for NullCount { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(0) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let count = other + Ok(scalar .as_primitive() .typed_value::() - .vortex_expect("null_count partial should not be null"); - *partial += count; - Ok(()) + .vortex_expect("null_count partial should not be null")) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(*partial, NonNullable)) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + Ok(partials.into_iter().sum()) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 197b8f10b04..3d919525ce0 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -116,6 +116,8 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumPartial; + use crate::aggregate_fn::fns::sum::SumState; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::DecimalArray; @@ -129,6 +131,18 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; + /// A partial whose running decimal sum is `value` (test-only helper bypassing scalar parsing). + fn partial_with_decimal(value: DecimalValue, decimal_dtype: DecimalDType) -> SumPartial { + SumPartial { + return_dtype: DType::Decimal(decimal_dtype, Nullable), + current: Some(SumState::Decimal { + value, + dtype: decimal_dtype, + }), + skip_nans: true, + } + } + #[test] fn sum_decimal_basic() -> VortexResult<()> { let decimal = DecimalArray::new( @@ -355,20 +369,16 @@ mod tests { fn sum_decimal_near_precision_boundary() -> VortexResult<()> { // Input precision 4 → return precision min(76, 4+10) = 14. // Native type for precision 14 is I64 (max precision 18), so 14 < 18. - // Use combine_partials to push state near (but under) 10^14. + // Reduce partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; - - let near_limit = Scalar::decimal( - DecimalValue::from(99_999_999_999_990i64), - DecimalDType::new(14, 0), - Nullable, - ); - Sum.combine_partials(&mut state, near_limit)?; + let options = NumericalAggregateOpts::default(); + let sum_decimal = DecimalDType::new(14, 0); + let near_limit = + partial_with_decimal(DecimalValue::from(99_999_999_999_990i64), sum_decimal); // Add a small value that keeps us just under 10^14. - let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, small)?; + let small = partial_with_decimal(DecimalValue::from(9i64), sum_decimal); + let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, small])?; let result = Sum.to_scalar(&state)?; assert!(!result.is_null()); @@ -385,21 +395,16 @@ mod tests { // The max representable value for precision 14 is 10^14 - 1. // When the sum reaches exactly 10^14, fits_in_precision fails even though // i256 arithmetic does not overflow. This tests the precision-based - // saturation path in combine_partials. + // saturation path in reduce_partials. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; - - let near_limit = Scalar::decimal( - DecimalValue::from(99_999_999_999_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - Sum.combine_partials(&mut state, near_limit)?; + let options = NumericalAggregateOpts::default(); + let sum_decimal = DecimalDType::new(14, 0); + let near_limit = + partial_with_decimal(DecimalValue::from(99_999_999_999_999i64), sum_decimal); // Push the sum to exactly 10^14, exceeding precision 14. - let one_more = - Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, one_more)?; + let one_more = partial_with_decimal(DecimalValue::from(1i64), sum_decimal); + let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, one_more])?; let result = Sum.to_scalar(&state)?; assert!(result.is_null()); @@ -414,21 +419,13 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; - - let near_limit = Scalar::decimal( - DecimalValue::from(-99_999_999_999_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - Sum.combine_partials(&mut state, near_limit)?; + let options = NumericalAggregateOpts::default(); + let sum_decimal = DecimalDType::new(14, 0); - let one_more = Scalar::decimal( - DecimalValue::from(-1i64), - DecimalDType::new(14, 0), - Nullable, - ); - Sum.combine_partials(&mut state, one_more)?; + let near_limit = + partial_with_decimal(DecimalValue::from(-99_999_999_999_999i64), sum_decimal); + let one_more = partial_with_decimal(DecimalValue::from(-1i64), sum_decimal); + let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, one_more])?; let result = Sum.to_scalar(&state)?; assert!(result.is_null()); @@ -437,22 +434,17 @@ mod tests { #[test] fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { - // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Test precision overflow via the accumulate_decimal path (not reduce_partials). // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. // Native for precision 38 is I128 (max 38), so 38 = 38. // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. // - // We use combine_partials to get the state close to 10^37, then accumulate - // a real array that pushes it over. - let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); + // We seed the state close to 10^37, then accumulate a real array that pushes it over. let return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; - // Set state to 10^37 - 1 via combine_partials. + // Set state to 10^37 - 1. let near_limit_val: i128 = 10i128.pow(37) - 1; - let near_limit = - Scalar::decimal(DecimalValue::from(near_limit_val), return_dtype, Nullable); - Sum.combine_partials(&mut state, near_limit)?; + let mut state = partial_with_decimal(DecimalValue::from(near_limit_val), return_dtype); // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 76eaba22c15..2f2886f9ad0 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -10,6 +10,7 @@ pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; @@ -142,73 +143,48 @@ impl AggregateFnVTable for Sum { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - let return_dtype = self - .return_dtype(options, input_dtype) - .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; - let initial = make_zero_state(&return_dtype); - - Ok(SumPartial { - return_dtype, - current: Some(initial), - skip_nans: options.skip_nans, - }) + let mut partial = SumPartial::empty(options, input_dtype)?; + vortex_ensure!( + scalar.dtype().eq_ignore_nullability(&partial.return_dtype), + "Sum partial has dtype {}, expected {}", + scalar.dtype(), + partial.return_dtype + ); + // A null partial means the producing accumulator saturated (overflow). + partial.current = if scalar.is_null() { + None + } else { + Some(sum_state_from_scalar(&scalar, &partial.return_dtype)?) + }; + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - // A null partial means the sub-accumulator saturated (overflow). - partial.current = None; - return Ok(()); - } - let Some(ref mut inner) = partial.current else { - return Ok(()); - }; - let saturated = match inner { - SumState::Unsigned(acc) => { - let val = other - .as_primitive() - .typed_value::() - .vortex_expect("checked non-null"); - checked_add_u64(acc, val) - } - SumState::Signed(acc) => { - let val = other - .as_primitive() - .typed_value::() - .vortex_expect("checked non-null"); - checked_add_i64(acc, val) - } - SumState::Float(acc) => { - let val = other - .as_primitive() - .typed_value::() - .vortex_expect("checked non-null"); - *acc += val; - false - } - SumState::Decimal { value, dtype } => { - let val = other - .as_decimal() - .decimal_value() - .vortex_expect("checked non-null"); - match value.checked_add(&val) { - Some(r) => { - *value = r; - !value.fits_in_precision(*dtype) - } - None => true, - } + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = SumPartial::empty(options, input_dtype)?; + for partial in partials { + let overflow = match (acc.current.as_mut(), partial.current) { + (None, _) => break, + // A saturated (overflowed) partial poisons the reduction. + (Some(_), None) => true, + (Some(acc_state), Some(state)) => checked_add_sum_states(acc_state, &state)?, + }; + if overflow { + acc.current = None; + break; } - }; - if saturated { - partial.current = None; } - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -227,10 +203,6 @@ impl AggregateFnVTable for Sum { }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.current = Some(make_zero_state(&partial.return_dtype)); - } - #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { match partial.current.as_ref() { @@ -261,7 +233,7 @@ impl AggregateFnVTable for Sum { } else { sum.cast(&partial.return_dtype)? }; - self.combine_partials(partial, sum)?; + merge_sum_result(partial, sum)?; return Ok(true); } Ok(false) @@ -283,14 +255,14 @@ impl AggregateFnVTable for Sum { batch: &Columnar, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - // Constants compute scalar * len and combine via combine_partials. + // Constants compute scalar * len and merge the product into the running state. if let Columnar::Constant(c) = batch { // NaN constants are treated as missing when skipping NaNs. if partial.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { return Ok(()); } if let Some(product) = multiply_constant(c.scalar(), c.len(), &partial.return_dtype)? { - self.combine_partials(partial, product)?; + merge_sum_result(partial, product)?; } return Ok(()); } @@ -341,6 +313,20 @@ pub struct SumPartial { skip_nans: bool, } +impl SumPartial { + /// The state of a group with no accumulated values, or an error for unsupported input dtypes. + fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> VortexResult { + let return_dtype = Sum + .return_dtype(options, input_dtype) + .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; + Ok(Self { + current: Some(make_zero_state(&return_dtype)), + return_dtype, + skip_nans: options.skip_nans, + }) + } +} + /// The accumulated sum value. // TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the // input type every time? Perhaps? @@ -369,6 +355,63 @@ pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { } } +/// Parse a non-null sum value of `return_dtype` into its accumulated state. +fn sum_state_from_scalar(scalar: &Scalar, return_dtype: &DType) -> VortexResult { + Ok(match return_dtype { + DType::Primitive(ptype, _) if ptype.is_unsigned_int() => { + SumState::Unsigned(u64::try_from(scalar)?) + } + DType::Primitive(ptype, _) if ptype.is_signed_int() => { + SumState::Signed(i64::try_from(scalar)?) + } + DType::Primitive(..) => SumState::Float(f64::try_from(scalar)?), + DType::Decimal(dtype, _) => SumState::Decimal { + value: DecimalValue::try_from(scalar)?, + dtype: *dtype, + }, + _ => vortex_bail!("Unsupported sum type {}", return_dtype), + }) +} + +/// Checked add of one sum state into another, returning true if overflow occurred. +/// +/// A decimal sum that no longer fits its precision counts as an overflow. +pub(crate) fn checked_add_sum_states(state: &mut SumState, other: &SumState) -> VortexResult { + Ok(match (state, other) { + (SumState::Unsigned(acc), SumState::Unsigned(other)) => checked_add_u64(acc, *other), + (SumState::Signed(acc), SumState::Signed(other)) => checked_add_i64(acc, *other), + (SumState::Float(acc), SumState::Float(other)) => { + *acc += *other; + false + } + (SumState::Decimal { value, dtype }, SumState::Decimal { value: other, .. }) => { + match value.checked_add(other) { + Some(result) if result.fits_in_precision(*dtype) => { + *value = result; + false + } + Some(_) | None => true, + } + } + _ => vortex_bail!("Mismatched sum partial states"), + }) +} + +/// Merge a finalized sum result (nullable; null means overflow) into the partial. +fn merge_sum_result(partial: &mut SumPartial, result: Scalar) -> VortexResult<()> { + let overflow = match partial.current.as_mut() { + None => return Ok(()), + Some(_) if result.is_null() => true, + Some(acc) => { + checked_add_sum_states(acc, &sum_state_from_scalar(&result, &partial.return_dtype)?)? + } + }; + if overflow { + partial.current = None; + } + Ok(()) +} + /// Checked add for u64, returning true if overflow occurred. #[allow(clippy::inline_always)] #[inline(always)] @@ -412,6 +455,8 @@ mod tests { use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumPartial; + use crate::aggregate_fn::fns::sum::SumState; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::BoolArray; @@ -536,20 +581,38 @@ mod tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let options = NumericalAggregateOpts::default(); - let scalar1 = Scalar::primitive(100i64, Nullable); - Sum.combine_partials(&mut state, scalar1)?; - - let scalar2 = Scalar::primitive(50i64, Nullable); - Sum.combine_partials(&mut state, scalar2)?; + let partial_of = |value: i64| SumPartial { + return_dtype: DType::Primitive(PType::I64, Nullable), + current: Some(SumState::Signed(value)), + skip_nans: true, + }; + let state = Sum.reduce_partials(&options, &dtype, [partial_of(100), partial_of(50)])?; let result = Sum.to_scalar(&state)?; - Sum.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) } + #[test] + fn sum_overflowed_partial_poisons_reduction() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let options = NumericalAggregateOpts::default(); + + let overflowed = Sum.partial_from_scalar( + &options, + &dtype, + Scalar::null(DType::Primitive(PType::I64, Nullable)), + )?; + let five = Sum.partial_from_scalar(&options, &dtype, Scalar::primitive(5i64, Nullable))?; + let state = Sum.reduce_partials(&options, &dtype, [five, overflowed])?; + + assert!(Sum.is_saturated(&state)); + assert!(Sum.to_scalar(&state)?.is_null()); + Ok(()) + } + // Stats caching test #[test] diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index e1287448a8f..678779dc127 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -26,6 +26,7 @@ use crate::aggregate_fn::fns::sum::SumState; use crate::aggregate_fn::fns::sum::accumulate_bool; use crate::aggregate_fn::fns::sum::accumulate_decimal; use crate::aggregate_fn::fns::sum::accumulate_primitive; +use crate::aggregate_fn::fns::sum::checked_add_sum_states; use crate::aggregate_fn::fns::sum::make_zero_state; use crate::aggregate_fn::fns::sum::multiply_constant; use crate::arrays::Struct; @@ -104,43 +105,51 @@ impl AggregateFnVTable for SumV2 { .map(sum_v2_partial_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - let return_dtype = self - .return_dtype(options, input_dtype) - .ok_or_else(|| vortex_err!("Unsupported sum_v2 dtype: {}", input_dtype))?; - let sum = make_zero_state(&return_dtype); - Ok(SumV2Partial { - return_dtype, - sum, - is_overflow: false, - is_empty: true, - skip_nans: options.skip_nans, - }) + let mut partial = SumV2Partial::empty(options, input_dtype)?; + let (sum, is_overflow, is_empty) = decode_partial_scalar(scalar)?; + validate_sum_field_dtype(&sum, &partial.return_dtype)?; + + // Adding the parsed value to the zero state cannot overflow; treat a decimal value that + // no longer fits its precision as an already-overflowed partial. + let overflowed = checked_add_sum_state(&mut partial.sum, &sum)?; + partial.is_overflow = is_overflow || overflowed; + partial.is_empty = is_empty && !partial.is_overflow; + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let (other_sum, other_is_overflow, other_is_empty) = decode_partial_scalar(other)?; - validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; - - if partial.is_overflow { - return Ok(()); - } - if other_is_overflow { - partial.is_overflow = true; - partial.is_empty = false; - return Ok(()); - } - if other_is_empty { - return Ok(()); + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + // Seed from the first partial so an overflowed state keeps its last valid sum. + let mut partials = partials.into_iter(); + let Some(mut acc) = partials.next() else { + return SumV2Partial::empty(options, input_dtype); + }; + for partial in partials { + if acc.is_overflow { + break; + } + if partial.is_overflow { + acc.is_overflow = true; + acc.is_empty = false; + continue; + } + if partial.is_empty { + continue; + } + acc.is_overflow = checked_add_sum_states(&mut acc.sum, &partial.sum)?; + acc.is_empty = false; } - - partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; - partial.is_empty = false; - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -154,12 +163,6 @@ impl AggregateFnVTable for SumV2 { )) } - fn reset(&self, partial: &mut Self::Partial) { - partial.sum = make_zero_state(&partial.return_dtype); - partial.is_overflow = false; - partial.is_empty = true; - } - fn is_saturated(&self, partial: &Self::Partial) -> bool { partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) } @@ -296,6 +299,22 @@ pub struct SumV2Partial { skip_nans: bool, } +impl SumV2Partial { + /// The state of a group with no accumulated values, or an error for unsupported input dtypes. + fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> VortexResult { + let return_dtype = SumV2 + .return_dtype(options, input_dtype) + .ok_or_else(|| vortex_err!("Unsupported sum_v2 dtype: {}", input_dtype))?; + Ok(Self { + sum: make_zero_state(&return_dtype), + return_dtype, + is_overflow: false, + is_empty: true, + skip_nans: options.skip_nans, + }) + } +} + fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult { let (validity, len) = match batch { Columnar::Canonical(Canonical::Primitive(array)) => { diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 65e54a5bb6a..160a5640616 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -10,6 +10,7 @@ use vortex_proto::expr as pb; use super::SumV2; use super::sum_v2; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; @@ -216,30 +217,33 @@ fn empty_chunk_is_a_merge_identity() -> VortexResult<()> { Ok(()) } +/// An accumulator over the single `i64` value `value`. +fn accumulated(value: i64, ctx: &mut ExecutionCtx) -> VortexResult> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + acc.accumulate(&PrimitiveArray::from_iter([value]).into_array(), ctx)?; + Ok(acc) +} + #[test] -fn combine_partials_empty_is_identity() -> VortexResult<()> { +fn merge_from_empty_is_identity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut empty = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; - let empty_partial = empty.partial_scalar()?; + let new_accumulator = + || Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone()); - empty.combine_partials(empty_partial.clone())?; + let mut empty = new_accumulator()?; + empty.merge_from(&mut new_accumulator()?)?; assert!(empty.final_scalar()?.is_null()); - let mut value = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; - let batch = PrimitiveArray::from_iter([7i64]).into_array(); - value.accumulate(&batch, &mut array_session().create_execution_ctx())?; - let value_partial = value.partial_scalar()?; - - empty.combine_partials(value_partial.clone())?; + empty.merge_from(&mut accumulated(7, &mut ctx)?)?; assert_eq!( empty.final_scalar()?.as_primitive().typed_value::(), Some(7) ); - let mut value_then_empty = - Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; - value_then_empty.combine_partials(value_partial)?; - value_then_empty.combine_partials(empty_partial)?; + let mut value_then_empty = accumulated(7, &mut ctx)?; + value_then_empty.merge_from(&mut new_accumulator()?)?; assert_eq!( value_then_empty .final_scalar()? @@ -251,24 +255,19 @@ fn combine_partials_empty_is_identity() -> VortexResult<()> { } #[test] -fn combine_partials_overflow_is_absorbing() -> VortexResult<()> { +fn merge_from_overflow_is_absorbing() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut max = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; - let max_batch = PrimitiveArray::from_iter([i64::MAX]).into_array(); - max.accumulate(&max_batch, &mut array_session().create_execution_ctx())?; - - let mut one = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; - let one_batch = PrimitiveArray::from_iter([1i64]).into_array(); - one.accumulate(&one_batch, &mut array_session().create_execution_ctx())?; - - let mut combined = - Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; - combined.combine_partials(max.partial_scalar()?)?; - combined.combine_partials(one.partial_scalar()?)?; + let options = NumericalAggregateOpts::default(); + + let mut combined = Accumulator::try_new(SumV2, options, dtype.clone())?; + combined.merge_from(&mut accumulated(i64::MAX, &mut ctx)?)?; + combined.merge_from(&mut accumulated(1, &mut ctx)?)?; assert!(combined.is_saturated()); assert!(combined.final_scalar()?.is_null()); - combined.combine_partials(max.partial_scalar()?)?; + // Further input cannot revive an overflowed sum, and the state keeps the last valid sum. + combined.merge_from(&mut accumulated(i64::MAX, &mut ctx)?)?; let overflow_partial = combined.partial_scalar()?; let fields = overflow_partial.as_struct(); assert_eq!( @@ -290,10 +289,10 @@ fn combine_partials_overflow_is_absorbing() -> VortexResult<()> { Some(false) ); - let mut propagated = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; - propagated.combine_partials(overflow_partial)?; - assert!(propagated.is_saturated()); - assert!(propagated.final_scalar()?.is_null()); + // The overflow flag survives the scalar round trip. + let propagated = SumV2.partial_from_scalar(&options, &dtype, overflow_partial)?; + assert!(SumV2.is_saturated(&propagated)); + assert!(SumV2.finalize_scalar(&propagated)?.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index a23bcfdc2c9..b775b58cdc5 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -131,31 +131,32 @@ impl AggregateFnVTable for UncompressedSizeInBytes { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(0) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let size = other + Ok(scalar .as_primitive() .typed_value::() - .vortex_expect("uncompressed_size_in_bytes partial should not be null"); - *partial = partial - .checked_add(size) - .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?; - Ok(()) + .vortex_expect("uncompressed_size_in_bytes partial should not be null")) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(*partial, NonNullable)) + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + partials.into_iter().try_fold(0u64, |acc, partial| { + acc.checked_add(partial) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) + }) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] @@ -686,19 +687,10 @@ mod tests { #[test] fn state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = UncompressedSizeInBytes.empty_partial(&EmptyOptions, &dtype)?; - UncompressedSizeInBytes.combine_partials( - &mut state, - Scalar::primitive(5u64, Nullability::NonNullable), - )?; - UncompressedSizeInBytes.combine_partials( - &mut state, - Scalar::primitive(3u64, Nullability::NonNullable), - )?; + let state = UncompressedSizeInBytes.reduce_partials(&EmptyOptions, &dtype, [5, 3])?; let result = UncompressedSizeInBytes.to_scalar(&state)?; - UncompressedSizeInBytes.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs index feb07e47175..e56ed84a754 100644 --- a/vortex-array/src/aggregate_fn/foreign.rs +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -77,15 +77,21 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { None } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + _scalar: Scalar, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn combine_partials(&self, _partial: &mut Self::Partial, _other: Scalar) -> VortexResult<()> { + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + _partials: impl IntoIterator, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } @@ -93,8 +99,6 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn reset(&self, _partial: &mut Self::Partial) {} - fn is_saturated(&self, _state: &Self::Partial) -> bool { false } diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 92fac87892a..d74aabefde9 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -114,19 +114,21 @@ mod tests { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + _scalar: Scalar, ) -> VortexResult { Ok(()) } - fn combine_partials( + fn reduce_partials( &self, - _partial: &mut Self::Partial, - _other: Scalar, - ) -> VortexResult<()> { + _options: &Self::Options, + _input_dtype: &DType, + _partials: impl IntoIterator, + ) -> VortexResult { Ok(()) } @@ -134,8 +136,6 @@ mod tests { vortex_panic!("TestAgg is for serde tests only"); } - fn reset(&self, _partial: &mut Self::Partial) {} - fn is_saturated(&self, _partial: &Self::Partial) -> bool { true } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 24449bc7572..4816a78c73e 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -62,6 +62,11 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Return whether this stored aggregate can satisfy `requested`. /// + /// Satisfaction is a claim about stored state, not just result semantics: consumers read this + /// aggregate's persisted partial state in place of ever accumulating `requested`, so anything + /// other than [`AggregateFnSatisfaction::No`] requires that a partial state for `requested` + /// can be created from this aggregate's partial state. + /// /// The default implementation only treats exactly equal aggregate functions as satisfying the /// request. Approximate pruning aggregates can override this to expose looser-but-sound bounds. fn can_satisfy( @@ -98,25 +103,45 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns `None` if the aggregate function cannot be applied to the input dtype. fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option; - /// Return the partial accumulator state for an empty group. - fn empty_partial( + /// Parse a partial scalar into the typed partial accumulator state. + /// + /// The scalar must have the DType specified by `partial_dtype` for the given options and + /// input dtype; this is the inverse of [`to_scalar`]. Partial scalars are produced by + /// aggregate kernels, cached statistics, and other accumulators' [`to_scalar`]. + /// + /// Implementations should only parse the scalar here; combining states belongs in + /// [`reduce_partials`]. + /// + /// [`to_scalar`]: AggregateFnVTable::to_scalar + /// [`reduce_partials`]: AggregateFnVTable::reduce_partials + fn partial_from_scalar( &self, options: &Self::Options, input_dtype: &DType, + scalar: Scalar, ) -> VortexResult; - /// Combine partial scalar state into the accumulator. - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()>; + /// Reduce a sequence of partial states into a single partial state. + /// + /// Partials must be reduced in iteration order, since some aggregates (e.g. first/last or + /// is_sorted) are order-dependent. Reducing an empty sequence returns the identity: the + /// partial state of a group with no accumulated values. + fn reduce_partials( + &self, + options: &Self::Options, + input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult; /// Convert the partial state into a partial scalar. /// /// The returned scalar must have the same DType as specified by `partial_dtype` for the - /// options and input dtype used to construct the state. + /// options and input dtype used to construct the state. This is the inverse of + /// [`partial_from_scalar`]. + /// + /// [`partial_from_scalar`]: AggregateFnVTable::partial_from_scalar fn to_scalar(&self, partial: &Self::Partial) -> VortexResult; - /// Reset the state of the accumulator to an empty group. - fn reset(&self, partial: &mut Self::Partial); - /// Is the partial accumulator state is "saturated", i.e. has it reached a state where the /// final result is fully determined. fn is_saturated(&self, state: &Self::Partial) -> bool; diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 149c72f499e..db3eda9f87f 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -31,7 +31,7 @@ impl DynAggregateKernel for ChunkedArrayAggregate { acc.accumulate(chunk, ctx)?; } // Return the partial (not finalized) result, since the outer accumulator - // will call combine_partials() on this value. + // will merge this partial scalar into its own state. Ok(Some(acc.flush()?)) } } diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 88aa9f9bd1c..292fce69f1f 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -209,7 +209,7 @@ pub fn finalize_scan(global: &GlobalState, chunk: &mut DataChunkRef) -> VortexRe .vortex_expect("no local state"); for other in rest.iter_mut() { for ((_, acc), (_, part)) in base.iter_mut().zip(other.iter_mut()) { - acc.combine_partials(part.flush()?)?; + acc.merge_from(part.as_mut())?; } } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs index 61377dca882..8457569bb45 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs @@ -30,7 +30,6 @@ pub(super) fn accumulate_constant( #[cfg(test)] mod tests { - use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -40,19 +39,13 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_error::VortexResult; - use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; + use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; use crate::layouts::zoned::aggregates::bloom_filter::constant::accumulate_constant; #[test] fn nulls_are_omitted() { - let bloom = BloomFilter; - let mut zone_partial = bloom - .empty_partial( - &BloomOptions::default(), - &DType::Primitive(PType::I32, Nullability::Nullable), - ) - .unwrap(); + let mut zone_partial = BloomPartial::from(&BloomOptions::default()); assert!( accumulate_constant( @@ -69,13 +62,7 @@ mod tests { #[test] fn null_always_returns_false() { - let bloom = BloomFilter; - let zone_partial = bloom - .empty_partial( - &BloomOptions::default(), - &DType::Primitive(PType::I32, Nullability::Nullable), - ) - .unwrap(); + let zone_partial = BloomPartial::from(&BloomOptions::default()); assert!( !zone_partial @@ -92,12 +79,10 @@ mod tests { fn valid_extension_is_a_member() -> VortexResult<()> { let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); let scalar = Scalar::extension_ref( - ext_dtype.clone(), + ext_dtype, Scalar::primitive(1_000i64, Nullability::NonNullable), ); - let bloom = BloomFilter; - let mut zone_partial = - bloom.empty_partial(&BloomOptions::default(), &DType::Extension(ext_dtype))?; + let mut zone_partial = BloomPartial::from(&BloomOptions::default()); accumulate_constant(&ConstantArray::new(scalar.clone(), 1), &mut zone_partial)?; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index a5abad8007a..3599b3444bb 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -297,26 +297,43 @@ impl AggregateFnVTable for BloomFilter { self.return_dtype(options, input_dtype) } - /// Returns an empty Bloom filter with all blocks zero-initialized. - fn empty_partial(&self, options: &Self::Options, _: &DType) -> VortexResult { - Ok(BloomPartial::from(options)) - } - - // Combination happens by doing an OR between both filters bits - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); + /// Parses a serialized filter into a partial with the configured block count. + /// + /// A null scalar is an empty filter. This assumes that `scalar` was created using the same + /// hash function as `options`. Ideally, an assertion here about which `hash_fn` was used to + /// create `scalar` would catch this invariant. + fn partial_from_scalar( + &self, + options: &Self::Options, + _input_dtype: &DType, + scalar: Scalar, + ) -> VortexResult { + let mut partial = BloomPartial::from(options); + if scalar.is_null() { + return Ok(partial); } - let other_as_bytes = other + let bytes = scalar .as_binary() .value() .ok_or_else(|| vortex_err!("non-null bloom partial has no bytes"))?; + partial.merge(bytes)?; + Ok(partial) + } - // This assumes that `other` was created using the same hash function as - // `partial`. Ideally, an assertion here about which `hash_fn` was used to create `other` - // would catch this invariant. - partial.merge(other_as_bytes) + /// Reduces partials by OR-ing their blocks together; an empty sequence is an empty filter + /// with all blocks zero-initialized. + fn reduce_partials( + &self, + options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = BloomPartial::from(options); + for partial in partials { + acc.union(&partial)?; + } + Ok(acc) } /// Returns the non-nullable binary representation of a bloom filter @@ -327,10 +344,6 @@ impl AggregateFnVTable for BloomFilter { Ok(Scalar::binary(bytes, Nullability::NonNullable)) } - fn reset(&self, partial: &mut Self::Partial) { - partial.reset(); - } - /// Returns true if all the blocks are full. /// /// When a bloom filter is saturated, it cannot rule out any values. @@ -421,7 +434,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { fn saturation_false_when_empty() -> VortexResult<()> { let options = BloomOptions::default(); let partial = - BloomFilter.empty_partial(&options, &DType::Binary(Nullability::NonNullable))?; + BloomFilter.reduce_partials(&options, &DType::Binary(Nullability::NonNullable), [])?; assert!(!BloomFilter.is_saturated(&partial)); Ok(()) } @@ -435,60 +448,57 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { } #[test] - fn combine_partials_rejects_mismatched_block_counts() -> VortexResult<()> { - let mut smaller = BloomFilter.empty_partial( - &BloomOptions::new(NonZeroU32::new(4).unwrap(), HashFn::XxHash3_64), - &DType::Binary(Nullability::NonNullable), - )?; - let bigger = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + fn mismatched_block_counts_are_rejected() -> VortexResult<()> { + let dtype = DType::Binary(Nullability::NonNullable); + let smaller = BloomOptions::new(NonZeroU32::new(4).unwrap(), HashFn::XxHash3_64); + let bigger = BloomPartial::from(&BloomOptions::default()); let bigger_scalar = BloomFilter.to_scalar(&bigger)?; - let result = BloomFilter.combine_partials(&mut smaller, bigger_scalar); + assert!( + BloomFilter + .partial_from_scalar(&smaller, &dtype, bigger_scalar) + .is_err(), + "parsing a partial built with a different blocks_count must fail loudly, not corrupt state" + ); assert!( - result.is_err(), - "combining partials built with different blocks_count must fail loudly, not corrupt state" + BloomFilter + .reduce_partials(&smaller, &dtype, [BloomPartial::from(&smaller), bigger]) + .is_err(), + "reducing partials built with different blocks_count must fail loudly, not corrupt state" ); Ok(()) } #[test] - fn combine_partials_unions_two_disjoint_partials() -> VortexResult<()> { - let mut partial = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + fn reduce_partials_unions_two_disjoint_partials() -> VortexResult<()> { + let options = BloomOptions::default(); + let mut partial = BloomPartial::from(&options); for i in 0..50i64 { partial.insert(i.to_le_bytes()); } - let mut secondary_partial = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + let mut secondary_partial = BloomPartial::from(&options); for i in 50..100i64 { secondary_partial.insert(i.to_le_bytes()); } // The following expected works because seed is equal for all. // If the seed is different for both partials, then this will fail. - let mut expected = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + let mut expected = BloomPartial::from(&options); for i in 0..100i64 { expected.insert(i.to_le_bytes()); } - let secondary_partial_as_scalar = BloomFilter.to_scalar(&secondary_partial)?; - BloomFilter.combine_partials(&mut partial, secondary_partial_as_scalar)?; + let partial = BloomFilter.reduce_partials( + &options, + &DType::Binary(Nullability::NonNullable), + [partial, secondary_partial], + )?; assert!( partial == expected, - "merging via combine_partials should equal a single filter built from the union of inputs" + "reducing partials should equal a single filter built from the union of inputs" ); for i in 0..100i64 { diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/aggregate.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/aggregate.rs index e3da98c11a9..34b75eecfbe 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/aggregate.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/aggregate.rs @@ -15,12 +15,6 @@ use super::BloomPartial; /// Practical implementation to avoid having to share blocks impl BloomPartial { - /// Resets and empties all blocks. - #[inline] - pub(in crate::layouts::zoned) fn reset(&mut self) { - self.blocks.fill([0; 8]); - } - /// Returns true if all the blocks are saturated, in other words, /// all bits are `1`. #[inline] @@ -28,6 +22,31 @@ impl BloomPartial { self.blocks.iter().all(|byte| *byte == [u32::MAX; 8]) } + /// Merges a compatible partial into this one. + /// + /// The merge is a bitwise OR, which represents the union of two split-block + /// Bloom filters when they use the same block count. + /// + /// _Notice_ This method only validates the block count. + /// Merging a filter created with a different hash function + /// will produce an invalid filter and introduce false negatives. + #[inline] + pub(in crate::layouts::zoned) fn union(&mut self, other: &BloomPartial) -> VortexResult<()> { + vortex_ensure_eq!( + self.len(), + other.len(), + "bloom partial block count mismatch" + ); + + for (dst_block, src_block) in self.blocks.iter_mut().zip(&other.blocks) { + for (dst_split, src_split) in dst_block.iter_mut().zip(src_block) { + *dst_split |= *src_split; + } + } + + Ok(()) + } + /// Merges a compatible serialized Bloom filter into this partial. /// /// The merge is a bitwise OR, which represents the union of two split-block diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index 7c3bcaf7659..f5d98e7b4e2 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/vortex-spatial/src/aggregate_fn/aabb.rs @@ -153,19 +153,31 @@ impl AggregateFnVTable for GeometryAabb { self.return_dtype(options, input_dtype) } - fn empty_partial( + fn partial_from_scalar( &self, _options: &Self::Options, _input_dtype: &DType, + scalar: Scalar, ) -> VortexResult { - Ok(AabbPartial { rect: None }) + // A null box is an empty group's AABB. + Ok(AabbPartial { + rect: rect_from_storage(&scalar)?, + }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if let Some(rect) = rect_from_storage(&other)? { - partial.merge(rect); + fn reduce_partials( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partials: impl IntoIterator, + ) -> VortexResult { + let mut acc = AabbPartial { rect: None }; + for partial in partials { + if let Some(rect) = partial.rect { + acc.merge(rect); + } } - Ok(()) + Ok(acc) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -175,10 +187,6 @@ impl AggregateFnVTable for GeometryAabb { }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.rect = None; - } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { // An AABB can always grow, so it is never saturated. false @@ -245,7 +253,6 @@ mod tests { use super::AabbPartial; use super::GeometryAabb; - use super::aabb_dtype; use super::rect_from_storage; use crate::test_harness::linestring_column; use crate::test_harness::multilinestring_column; @@ -388,38 +395,37 @@ mod tests { Ok(()) } - /// `combine_partials` unions partial boxes - the path the zoned writer takes when a zone's + /// `reduce_partials` unions partial boxes - the path the zoned writer takes when a zone's /// array is chunked. #[test] - fn combine_partials_unions_boxes() -> VortexResult<()> { + fn reduce_partials_unions_boxes() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); let bbox = |xmin, ymin, xmax, ymax| AabbPartial { rect: Some(SpatialRect::new((xmin, ymin), (xmax, ymax))), }; - let mut partial = AabbPartial { rect: None }; - GeometryAabb.combine_partials( - &mut partial, - GeometryAabb.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, - )?; - GeometryAabb.combine_partials( - &mut partial, - GeometryAabb.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, + let reduced = GeometryAabb.reduce_partials( + &EmptyOptions, + &dtype, + [bbox(0.0, 0.0, 1.0, 1.0), bbox(5.0, -2.0, 7.0, 3.0)], )?; assert_eq!( - aabb(&GeometryAabb.to_scalar(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&reduced)?)?, (0.0, -2.0, 7.0, 3.0) ); Ok(()) } - /// A null partial (an empty group's AABB) is a no-op in `combine_partials`. + /// An empty partial (an empty group's AABB) is a no-op in `reduce_partials`. #[test] - fn combine_partials_ignores_null() -> VortexResult<()> { - let mut partial = AabbPartial { + fn reduce_partials_ignores_empty() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let empty = AabbPartial { rect: None }; + let value = AabbPartial { rect: Some(SpatialRect::new((0.0, 0.0), (1.0, 1.0))), }; - GeometryAabb.combine_partials(&mut partial, Scalar::null(aabb_dtype()))?; + let reduced = GeometryAabb.reduce_partials(&EmptyOptions, &dtype, [value, empty])?; assert_eq!( - aabb(&GeometryAabb.to_scalar(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&reduced)?)?, (0.0, 0.0, 1.0, 1.0) ); Ok(()) From 62d35d1f743a2863e75cd2bb9afe4772e3348d9f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 16:49:03 +0000 Subject: [PATCH 2/6] Pass options and AggregateDTypes into aggregate execution methods Every AggregateFnVTable execution method now receives the bound options and the resolved input/partial/result dtypes, so partial states only hold accumulated values instead of copies of the options and dtypes. - Add `AggregateDTypes<'a>` (borrowed) and `OwnedAggregateDTypes`, resolved once by `Accumulator`/`GroupedAccumulator` and lent to every vtable call. - Split reduction into typed primitives: `empty_partial` (the identity) and binary `merge_partials`; `reduce_partials` becomes a provided fold. - Add `DynAccumulator::combine_partial_scalar`, which parses a partial scalar and merges it through the typed vtable in one monomorphized call; `Accumulator` uses the same path for kernel results and cached statistics. - Strip options/dtype fields from every partial (Sum, SumV2, Count, Min/Max/MinMax, BoundedMin/Max, IsSorted, IsConstant, First/Last, BloomPartial's hash_fn, ...). `Count`'s partial is now a plain `u64`. - `BinaryCombined::return_dtype`/`finalize`/`finalize_scalar` take the combined options and dtypes; `Mean` derives its target dtype from them. Signed-off-by: "Claude" Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- vortex-array/src/aggregate_fn/accumulator.rs | 145 ++++++------ .../src/aggregate_fn/accumulator_grouped.rs | 51 ++--- vortex-array/src/aggregate_fn/combined.rs | 129 +++++++---- .../src/aggregate_fn/fns/all_nan/mod.rs | 54 ++++- .../aggregate_fn/fns/all_non_distinct/mod.rs | 52 ++++- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 54 ++++- .../src/aggregate_fn/fns/all_non_null/mod.rs | 54 ++++- .../src/aggregate_fn/fns/all_null/mod.rs | 54 ++++- .../src/aggregate_fn/fns/bounded_max/mod.rs | 99 +++++---- .../src/aggregate_fn/fns/bounded_min/mod.rs | 87 +++++--- .../src/aggregate_fn/fns/count/mod.rs | 99 +++++---- .../src/aggregate_fn/fns/first/mod.rs | 81 +++++-- .../src/aggregate_fn/fns/is_constant/mod.rs | 84 ++++--- .../src/aggregate_fn/fns/is_sorted/mod.rs | 181 ++++++++------- vortex-array/src/aggregate_fn/fns/last/mod.rs | 80 ++++--- vortex-array/src/aggregate_fn/fns/max/mod.rs | 108 +++++---- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 45 +++- vortex-array/src/aggregate_fn/fns/min/mod.rs | 108 +++++---- .../src/aggregate_fn/fns/min_max/bool.rs | 16 +- .../src/aggregate_fn/fns/min_max/decimal.rs | 6 +- .../src/aggregate_fn/fns/min_max/extension.rs | 5 +- .../src/aggregate_fn/fns/min_max/mod.rs | 155 ++++++++----- .../src/aggregate_fn/fns/min_max/primitive.rs | 8 +- .../src/aggregate_fn/fns/min_max/varbin.rs | 10 +- .../src/aggregate_fn/fns/nan_count/mod.rs | 58 ++++- .../src/aggregate_fn/fns/null_count/mod.rs | 54 ++++- .../src/aggregate_fn/fns/sum/decimal.rs | 64 +++--- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 208 ++++++++++-------- .../src/aggregate_fn/fns/sum/primitive.rs | 4 +- .../src/aggregate_fn/fns/sum_v2/mod.rs | 162 +++++++++----- .../src/aggregate_fn/fns/sum_v2/tests.rs | 17 +- .../fns/uncompressed_size_in_bytes/mod.rs | 64 ++++-- vortex-array/src/aggregate_fn/foreign.rs | 48 +++- vortex-array/src/aggregate_fn/proto.rs | 48 +++- vortex-array/src/aggregate_fn/vtable.rs | 200 ++++++++++++++--- .../zoned/aggregates/bloom_filter/mod.rs | 107 ++++++--- .../aggregates/bloom_filter/partial/mod.rs | 41 ++-- .../aggregates/bloom_filter/partial/serde.rs | 6 +- vortex-spatial/src/aggregate_fn/aabb.rs | 73 ++++-- 39 files changed, 1936 insertions(+), 983 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 0735cdaaa73..c2b88ad6e6a 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -7,7 +7,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::Columnar; @@ -15,6 +14,7 @@ use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::columnar::AnyColumnar; use crate::dtype::DType; @@ -35,12 +35,8 @@ pub struct Accumulator { options: V::Options, /// Type-erased aggregate function used for kernel dispatch. aggregate_fn: AggregateFnRef, - /// The DType of the input. - dtype: DType, - /// The DType of the aggregate. - return_dtype: DType, - /// The DType of the accumulator state. - partial_dtype: DType, + /// The input, partial, and result dtypes lent to every vtable call. + dtypes: OwnedAggregateDTypes, /// The partial state of the accumulator, updated after each accumulate/merge call. /// /// `None` is the empty-group state; a live partial is only materialized when a batch is @@ -50,36 +46,22 @@ pub struct Accumulator { impl Accumulator { pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult { - let return_dtype = vtable.return_dtype(&options, &dtype).ok_or_else(|| { - vortex_err!( - "Aggregate function {} cannot be applied to dtype {}", - vtable.id(), - dtype - ) - })?; - let partial_dtype = vtable.partial_dtype(&options, &dtype).ok_or_else(|| { - vortex_err!( - "Aggregate function {} cannot be applied to dtype {}", - vtable.id(), - dtype - ) - })?; + let dtypes = OwnedAggregateDTypes::try_new(&vtable, &options, dtype)?; let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); Ok(Self { vtable, options, aggregate_fn, - dtype, - return_dtype, - partial_dtype, + dtypes, partial: None, }) } /// The identity partial state: the state of a group with no accumulated values. fn empty_partial(&self) -> VortexResult { - self.vtable.reduce_partials(&self.options, &self.dtype, []) + self.vtable + .empty_partial(&self.options, self.dtypes.borrow()) } /// Materialize the partial state in place so a batch can be accumulated into it. @@ -90,18 +72,28 @@ impl Accumulator { Ok(()) } - /// Reduce an incoming partial state into the accumulator's current state. + /// Merge an incoming partial state into the accumulator's current state. pub(crate) fn fold_partial(&mut self, other: V::Partial) -> VortexResult<()> { self.partial = Some(match self.partial.take() { - // Reducing the incoming partial with the empty state is the identity. + // Merging the incoming partial with the empty state is the identity. None => other, Some(current) => { self.vtable - .reduce_partials(&self.options, &self.dtype, [current, other])? + .merge_partials(&self.options, self.dtypes.borrow(), current, other)? } }); Ok(()) } + + /// Parse a partial scalar of dtype `dtypes.partial` and merge it into the current state. + /// + /// Both steps go through the typed vtable of `V`, so they inline into one monomorphized call. + fn fold_partial_scalar(&mut self, scalar: Scalar) -> VortexResult<()> { + let other = self + .vtable + .partial_from_scalar(&self.options, self.dtypes.borrow(), scalar)?; + self.fold_partial(other) + } } /// A trait object for type-erased accumulators, used for dynamic dispatch when the aggregate @@ -116,6 +108,14 @@ pub trait DynAccumulator: 'static + Send { /// options, and input dtype as this one. fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()>; + /// Parse a partial scalar and merge it into this accumulator's state. + /// + /// The scalar must have the dtype reported by the vtable's `partial_dtype` for this + /// accumulator's options and input dtype, and represents input following the input already + /// accumulated. Parsing and merging both run through the typed vtable, so they inline into + /// a single monomorphized call per aggregate. + fn combine_partial_scalar(&mut self, partial: Scalar) -> VortexResult<()>; + /// Whether the accumulator's result is fully determined. fn is_saturated(&self) -> bool; @@ -158,9 +158,9 @@ impl DynAccumulator for Accumulator { } vortex_ensure!( - batch.dtype() == &self.dtype, + batch.dtype() == self.dtypes.input(), "Input DType mismatch: expected {}, got {}", - self.dtype, + self.dtypes.input(), batch.dtype() ); @@ -169,23 +169,20 @@ impl DynAccumulator for Accumulator { if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - let partial = if partial.dtype() == &self.partial_dtype { + let partial = if partial.dtype() == self.dtypes.partial() { partial } else { vortex_ensure!( - partial.dtype().eq_ignore_nullability(&self.partial_dtype), + partial.dtype().eq_ignore_nullability(self.dtypes.partial()), "Aggregate {} read legacy stat {} with dtype {}, expected {}", self.aggregate_fn, stat, partial.dtype(), - self.partial_dtype, + self.dtypes.partial(), ); - partial.cast(&self.partial_dtype)? + partial.cast(self.dtypes.partial())? }; - let parsed = self - .vtable - .partial_from_scalar(&self.options, &self.dtype, partial)?; - self.fold_partial(parsed)?; + self.fold_partial_scalar(partial)?; return Ok(()); } @@ -204,15 +201,12 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)? { vortex_ensure!( - result.dtype() == &self.partial_dtype, + result.dtype() == self.dtypes.partial(), "Aggregate kernel returned {}, expected {}", result.dtype(), - self.partial_dtype, + self.dtypes.partial(), ); - let parsed = self - .vtable - .partial_from_scalar(&self.options, &self.dtype, result)?; - self.fold_partial(parsed)?; + self.fold_partial_scalar(result)?; return Ok(()); } } @@ -220,7 +214,10 @@ impl DynAccumulator for Accumulator { // 2. Allow the vtable to short-circuit on the raw array before decompression. self.ensure_partial()?; let partial = self.partial.as_mut().vortex_expect("partial materialized"); - if self.vtable.try_accumulate(partial, batch, ctx)? { + if self + .vtable + .try_accumulate(&self.options, self.dtypes.borrow(), partial, batch, ctx)? + { return Ok(()); } @@ -241,15 +238,12 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)? { vortex_ensure!( - result.dtype() == &self.partial_dtype, + result.dtype() == self.dtypes.partial(), "Aggregate kernel returned {}, expected {}", result.dtype(), - self.partial_dtype, + self.dtypes.partial(), ); - let parsed = self - .vtable - .partial_from_scalar(&self.options, &self.dtype, result)?; - self.fold_partial(parsed)?; + self.fold_partial_scalar(result)?; return Ok(()); } @@ -261,7 +255,8 @@ impl DynAccumulator for Accumulator { self.ensure_partial()?; let partial = self.partial.as_mut().vortex_expect("partial materialized"); - self.vtable.accumulate(partial, &columnar, ctx) + self.vtable + .accumulate(&self.options, self.dtypes.borrow(), partial, &columnar, ctx) } fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()> { @@ -272,7 +267,7 @@ impl DynAccumulator for Accumulator { ); }; vortex_ensure!( - other.options == self.options && other.dtype == self.dtype, + other.options == self.options && other.dtypes.input() == self.dtypes.input(), "Cannot merge {} accumulators with different options or input dtypes", self.aggregate_fn, ); @@ -282,14 +277,26 @@ impl DynAccumulator for Accumulator { } } + fn combine_partial_scalar(&mut self, partial: Scalar) -> VortexResult<()> { + vortex_ensure!( + partial.dtype() == self.dtypes.partial(), + "Partial DType mismatch for {}: expected {}, got {}", + self.aggregate_fn, + self.dtypes.partial(), + partial.dtype(), + ); + self.fold_partial_scalar(partial) + } + fn as_any_mut(&mut self) -> &mut dyn Any { self } fn is_saturated(&self) -> bool { - self.partial - .as_ref() - .is_some_and(|partial| self.vtable.is_saturated(partial)) + self.partial.as_ref().is_some_and(|partial| { + self.vtable + .is_saturated(&self.options, self.dtypes.borrow(), partial) + }) } fn reset(&mut self) { @@ -297,17 +304,20 @@ impl DynAccumulator for Accumulator { } fn partial_scalar(&self) -> VortexResult { + let dtypes = self.dtypes.borrow(); let partial = match &self.partial { - Some(partial) => self.vtable.to_scalar(partial)?, - None => self.vtable.to_scalar(&self.empty_partial()?)?, + Some(partial) => self.vtable.to_scalar(&self.options, dtypes, partial)?, + None => self + .vtable + .to_scalar(&self.options, dtypes, &self.empty_partial()?)?, }; #[cfg(debug_assertions)] { vortex_ensure!( - partial.dtype() == &self.partial_dtype, + partial.dtype() == dtypes.partial, "Aggregate returned incorrect DType on partial_scalar: expected {}, got {}", - self.partial_dtype, + dtypes.partial, partial.dtype(), ); } @@ -316,15 +326,20 @@ impl DynAccumulator for Accumulator { } fn final_scalar(&self) -> VortexResult { + let dtypes = self.dtypes.borrow(); let result = match &self.partial { - Some(partial) => self.vtable.finalize_scalar(partial)?, - None => self.vtable.finalize_scalar(&self.empty_partial()?)?, + Some(partial) => self + .vtable + .finalize_scalar(&self.options, dtypes, partial)?, + None => self + .vtable + .finalize_scalar(&self.options, dtypes, &self.empty_partial()?)?, }; vortex_ensure!( - result.dtype() == &self.return_dtype, + result.dtype() == dtypes.result, "Aggregate returned incorrect DType on final_scalar: expected {}, got {}", - self.return_dtype, + dtypes.result, result.dtype(), ); @@ -450,7 +465,7 @@ mod tests { let acc = mean_f64_accumulator().expect("build accumulator"); let sum = Scalar::primitive(42.0f64, Nullability::Nullable); let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.partial_dtype, vec![sum, count]) + Scalar::struct_(acc.dtypes.partial().clone(), vec![sum, count]) } /// Kernel registered for `(Dict, Combined)` fires in preference to diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index eda0acdad97..8483daf028e 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -7,7 +7,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -21,6 +20,7 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; @@ -187,12 +187,8 @@ pub struct GroupedAccumulator { options: V::Options, /// Type-erased aggregate function used for kernel dispatch. aggregate_fn: AggregateFnRef, - /// The DType of the input. - dtype: DType, - /// The DType of the aggregate. - return_dtype: DType, - /// The DType of the partial accumulator state. - partial_dtype: DType, + /// The input, partial, and result dtypes lent to every vtable call. + dtypes: OwnedAggregateDTypes, /// The accumulated state for prior batches of groups. partials: Vec, } @@ -200,28 +196,13 @@ pub struct GroupedAccumulator { impl GroupedAccumulator { pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult { let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); - let return_dtype = vtable.return_dtype(&options, &dtype).ok_or_else(|| { - vortex_err!( - "Aggregate function {} cannot be applied to dtype {}", - vtable.id(), - dtype - ) - })?; - let partial_dtype = vtable.partial_dtype(&options, &dtype).ok_or_else(|| { - vortex_err!( - "Aggregate function {} cannot be applied to dtype {}", - vtable.id(), - dtype - ) - })?; + let dtypes = OwnedAggregateDTypes::try_new(&vtable, &options, dtype)?; Ok(Self { vtable, options, aggregate_fn, - dtype, - return_dtype, - partial_dtype, + dtypes, partials: vec![], }) } @@ -253,9 +234,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { ), }; vortex_ensure!( - elements_dtype.as_ref() == &self.dtype, + elements_dtype.as_ref() == self.dtypes.input(), "Input DType mismatch: expected {}, got {}", - self.dtype, + self.dtypes.input(), elements_dtype ); @@ -277,17 +258,19 @@ impl DynGroupedAccumulator for GroupedAccumulator { if states.len() == 1 { return Ok(states.pop().vortex_expect("checked one partial")); } - Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) + Ok(ChunkedArray::try_new(states, self.dtypes.partial().clone())?.into_array()) } fn finish(&mut self) -> VortexResult { let states = self.flush()?; - let results = self.vtable.finalize(states)?; + let results = self + .vtable + .finalize(&self.options, self.dtypes.borrow(), states)?; vortex_ensure!( - results.dtype() == &self.return_dtype, + results.dtype() == self.dtypes.result(), "Return DType mismatch: expected {}, got {}", - self.return_dtype, + self.dtypes.result(), results.dtype() ); @@ -359,10 +342,10 @@ impl GroupedAccumulator { let mut accumulator = Accumulator::try_new( self.vtable.clone(), self.options.clone(), - self.dtype.clone(), + self.dtypes.input().clone(), )?; let mut states = - builder_with_capacity_in(&self.partial_dtype, grouped.len(), ctx.allocator()); + builder_with_capacity_in(self.dtypes.partial(), grouped.len(), ctx.allocator()); let group_ranges = grouped.group_ranges(ctx)?; let group_validity = grouped.group_validity(ctx)?; @@ -381,9 +364,9 @@ impl GroupedAccumulator { fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> { vortex_ensure!( - state.dtype() == &self.partial_dtype, + state.dtype() == self.dtypes.partial(), "State DType mismatch: expected {}, got {}", - self.partial_dtype, + self.dtypes.partial(), state.dtype() ); self.partials.push(state); diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index f600d6a802a..38f432055e7 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -19,9 +19,11 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::OwnedAggregateDTypes; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; @@ -81,12 +83,28 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone { } /// Return type of the combined aggregate. - fn return_dtype(&self, input_dtype: &DType) -> Option; + fn return_dtype(&self, options: &CombinedOptions, input_dtype: &DType) -> Option; /// Combine the finalized left and right results into the final aggregate. - fn finalize(&self, left: ArrayRef, right: ArrayRef) -> VortexResult; - - fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult; + /// + /// `dtypes` describes the combined aggregate; `left` and `right` already have their child's + /// result dtype. The returned array must have dtype `dtypes.result`. + fn finalize( + &self, + options: &CombinedOptions, + dtypes: AggregateDTypes<'_>, + left: ArrayRef, + right: ArrayRef, + ) -> VortexResult; + + /// Combine the finalized child scalars into a result of dtype `dtypes.result`. + fn finalize_scalar( + &self, + options: &CombinedOptions, + dtypes: AggregateDTypes<'_>, + left_scalar: Scalar, + right_scalar: Scalar, + ) -> VortexResult; /// Serialize the options for this combined aggregate. Default: not serializable. fn serialize(&self, options: &CombinedOptions) -> VortexResult>> { @@ -163,8 +181,8 @@ impl AggregateFnVTable for Combined { BinaryCombined::deserialize(&self.0, metadata, session) } - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - BinaryCombined::return_dtype(&self.0, input_dtype) + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + BinaryCombined::return_dtype(&self.0, options, input_dtype) } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { @@ -173,13 +191,21 @@ impl AggregateFnVTable for Combined { Some(self.0.partial_struct_dtype(l, r)) } + fn empty_partial( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + self.new_child_accumulators(options, dtypes.input) + } + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let (mut left, mut right) = self.new_child_accumulators(options, input_dtype)?; + let (mut left, mut right) = self.new_child_accumulators(options, dtypes.input)?; // A null partial represents an empty group and parses to empty child accumulators. if !scalar.is_null() { let s = scalar.as_struct(); @@ -191,49 +217,46 @@ impl AggregateFnVTable for Combined { let r_field = s .field(rname) .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", rname))?; - left.fold_partial(self.0.left().partial_from_scalar( - &options.0, - input_dtype, - l_field, - )?)?; - right.fold_partial(self.0.right().partial_from_scalar( - &options.1, - input_dtype, - r_field, - )?)?; + left.combine_partial_scalar(l_field)?; + right.combine_partial_scalar(r_field)?; } Ok((left, right)) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + (mut left, mut right): Self::Partial, + (mut other_left, mut other_right): Self::Partial, ) -> VortexResult { - let mut partials = partials.into_iter(); - let Some((mut left, mut right)) = partials.next() else { - return self.new_child_accumulators(options, input_dtype); - }; - // The children are typed accumulators of the same child aggregates, so the remaining - // partials merge state directly without any scalar interchange. - for (mut l, mut r) in partials { - left.merge_from(&mut l)?; - right.merge_from(&mut r)?; - } + // The children are typed accumulators of the same child aggregates, so they merge state + // directly without any scalar interchange. + left.merge_from(&mut other_left)?; + right.merge_from(&mut other_right)?; Ok((left, right)) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { let l_scalar = partial.0.partial_scalar()?; let r_scalar = partial.1.partial_scalar()?; - let dtype = self - .0 - .partial_struct_dtype(l_scalar.dtype().clone(), r_scalar.dtype().clone()); - Ok(Scalar::struct_(dtype, vec![l_scalar, r_scalar])) + Ok(Scalar::struct_( + dtypes.partial.clone(), + vec![l_scalar, r_scalar], + )) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { partial.0.is_saturated() && partial.1.is_saturated() } @@ -244,6 +267,8 @@ impl AggregateFnVTable for Combined { /// `true` so [`Self::accumulate`] is unreachable. fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -255,6 +280,8 @@ impl AggregateFnVTable for Combined { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -262,17 +289,31 @@ impl AggregateFnVTable for Combined { unreachable!("Combined::try_accumulate handles all batches") } - fn finalize(&self, states: ArrayRef) -> VortexResult { + fn finalize( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + states: ArrayRef, + ) -> VortexResult { let l_field = states.get_item(FieldName::from(self.0.left_name()))?; let r_field = states.get_item(FieldName::from(self.0.right_name()))?; - let l_finalized = self.0.left().finalize(l_field)?; - let r_finalized = self.0.right().finalize(r_field)?; - BinaryCombined::finalize(&self.0, l_finalized, r_finalized) + let left = self.0.left(); + let right = self.0.right(); + let l_dtypes = OwnedAggregateDTypes::try_new(&left, &options.0, dtypes.input.clone())?; + let r_dtypes = OwnedAggregateDTypes::try_new(&right, &options.1, dtypes.input.clone())?; + let l_finalized = left.finalize(&options.0, l_dtypes.borrow(), l_field)?; + let r_finalized = right.finalize(&options.1, r_dtypes.borrow(), r_field)?; + BinaryCombined::finalize(&self.0, options, dtypes, l_finalized, r_finalized) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { let l_scalar = partial.0.final_scalar()?; let r_scalar = partial.1.final_scalar()?; - BinaryCombined::finalize_scalar(&self.0, l_scalar, r_scalar) + BinaryCombined::finalize_scalar(&self.0, options, dtypes, l_scalar, r_scalar) } } diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index a3d3198c9b6..743b9fe8064 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -9,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -60,34 +61,55 @@ impl AggregateFnVTable for AllNan { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(true) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().all(|partial| partial)) + Ok(first && second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -103,6 +125,8 @@ impl AggregateFnVTable for AllNan { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -122,12 +146,22 @@ impl AggregateFnVTable for AllNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index d9b399a9e4a..7be5fa20b67 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -38,6 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -147,10 +148,20 @@ impl AggregateFnVTable for AllNonDistinct { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(AllNonDistinctPartial { + all_non_distinct: true, + }) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { Ok(AllNonDistinctPartial { @@ -158,18 +169,24 @@ impl AggregateFnVTable for AllNonDistinct { }) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { Ok(AllNonDistinctPartial { - all_non_distinct: partials.into_iter().all(|partial| partial.all_non_distinct), + all_non_distinct: first.all_non_distinct && second.all_non_distinct, }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool( partial.all_non_distinct, Nullability::NonNullable, @@ -177,12 +194,19 @@ impl AggregateFnVTable for AllNonDistinct { } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !partial.all_non_distinct } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -232,11 +256,21 @@ impl AggregateFnVTable for AllNonDistinct { } } - fn finalize(&self, _partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partials: ArrayRef, + ) -> VortexResult { vortex_bail!("AllNonDistinct does not support array finalization"); } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool( partial.all_non_distinct, Nullability::NonNullable, diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index 5ee577b58f9..0bb45773d69 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -9,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -60,34 +61,55 @@ impl AggregateFnVTable for AllNonNan { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(true) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().all(|partial| partial)) + Ok(first && second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -98,6 +120,8 @@ impl AggregateFnVTable for AllNonNan { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -112,12 +136,22 @@ impl AggregateFnVTable for AllNonNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index fd95c5016d2..bd42570abf1 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -9,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -51,34 +52,55 @@ impl AggregateFnVTable for AllNonNull { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(true) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().all(|partial| partial)) + Ok(first && second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -89,6 +111,8 @@ impl AggregateFnVTable for AllNonNull { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -102,12 +126,22 @@ impl AggregateFnVTable for AllNonNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index 1bff71d5ca8..ad61771138f 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -9,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -51,34 +52,55 @@ impl AggregateFnVTable for AllNull { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(true) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().all(|partial| partial)) + Ok(first && second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -89,6 +111,8 @@ impl AggregateFnVTable for AllNull { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -105,12 +129,22 @@ impl AggregateFnVTable for AllNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index c0d0c838770..e7630411b39 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -19,6 +19,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -71,20 +72,9 @@ enum BoundedMaxState { /// Partial accumulator state for the bounded maximum aggregate. pub struct BoundedMaxPartial { state: BoundedMaxState, - element_dtype: DType, - max_bytes: NonZeroUsize, } impl BoundedMaxPartial { - /// The state of a group with no accumulated values. - fn empty(options: &BoundedMaxOptions, input_dtype: &DType) -> Self { - Self { - state: BoundedMaxState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, - } - } - fn merge_bound(&mut self, max: Scalar) { if max.is_null() { return; @@ -103,8 +93,8 @@ impl BoundedMaxPartial { self.state = BoundedMaxState::Unknown; } - fn final_scalar(&self) -> VortexResult { - let dtype = self.element_dtype.as_nullable(); + fn final_scalar(&self, dtypes: AggregateDTypes<'_>) -> VortexResult { + let dtype = dtypes.result.clone(); match &self.state { BoundedMaxState::Value(max) => max.cast(&dtype), BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), @@ -198,10 +188,20 @@ impl AggregateFnVTable for BoundedMax { supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(BoundedMaxPartial { + state: BoundedMaxState::Empty, + }) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -232,32 +232,32 @@ impl AggregateFnVTable for BoundedMax { BoundedMaxState::Value(bound) } }; - Ok(BoundedMaxPartial { - state, - ..BoundedMaxPartial::empty(options, input_dtype) - }) + Ok(BoundedMaxPartial { state }) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = BoundedMaxPartial::empty(options, input_dtype); - for partial in partials { - match partial.state { - BoundedMaxState::Empty => {} - BoundedMaxState::Value(max) => acc.merge_bound(max), - BoundedMaxState::Unknown => acc.unknown(), - } + match second.state { + BoundedMaxState::Empty => {} + BoundedMaxState::Value(max) => first.merge_bound(max), + BoundedMaxState::Unknown => first.unknown(), } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = make_bounded_max_partial_dtype(&partial.element_dtype); - let bound_dtype = partial.element_dtype.as_nullable(); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.partial.clone(); + let bound_dtype = dtypes.input.as_nullable(); match &partial.state { BoundedMaxState::Empty => Ok(Scalar::null(dtype)), BoundedMaxState::Value(max) => Ok(Scalar::struct_( @@ -277,12 +277,19 @@ impl AggregateFnVTable for BoundedMax { } } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { matches!(partial.state, BoundedMaxState::Unknown) } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -296,19 +303,29 @@ impl AggregateFnVTable for BoundedMax { let Some(result) = min_max(&array, ctx, NumericalAggregateOpts::default())? else { return Ok(()); }; - match truncate_max(result.max, partial.max_bytes.get())? { + match truncate_max(result.max, options.max_bytes.get())? { Some(bound) => partial.merge_bound(bound), None => partial.unknown(), } Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { partials.get_item(BOUNDED_MAX_BOUND) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - partial.final_scalar() + fn finalize_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + partial.final_scalar(dtypes) } } @@ -469,8 +486,6 @@ mod tests { acc.accumulate(&values, &mut ctx)?; acc.fold_partial(BoundedMaxPartial { state: BoundedMaxState::Empty, - element_dtype: values.dtype().clone(), - max_bytes: max_bytes(2), })?; assert_eq!( @@ -495,8 +510,6 @@ mod tests { acc.accumulate(&values, &mut ctx)?; acc.fold_partial(BoundedMaxPartial { state: BoundedMaxState::Unknown, - element_dtype: values.dtype().clone(), - max_bytes: max_bytes(2), })?; assert_eq!(acc.finish()?, Scalar::null(values.dtype().as_nullable())); diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 8a70d924ce6..4cd3acbd086 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -17,6 +17,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -56,20 +57,9 @@ enum BoundedMinState { /// Partial accumulator state for the bounded minimum aggregate. pub struct BoundedMinPartial { state: BoundedMinState, - element_dtype: DType, - max_bytes: NonZeroUsize, } impl BoundedMinPartial { - /// The state of a group with no accumulated values. - fn empty(options: &BoundedMinOptions, input_dtype: &DType) -> Self { - Self { - state: BoundedMinState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, - } - } - fn merge(&mut self, min: Scalar) { if min.is_null() { return; @@ -152,10 +142,20 @@ impl AggregateFnVTable for BoundedMin { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(BoundedMinPartial { + state: BoundedMinState::Empty, + }) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -164,41 +164,48 @@ impl AggregateFnVTable for BoundedMin { } else { BoundedMinState::Value(scalar) }; - Ok(BoundedMinPartial { - state, - ..BoundedMinPartial::empty(options, input_dtype) - }) + Ok(BoundedMinPartial { state }) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = BoundedMinPartial::empty(options, input_dtype); - for partial in partials { - if let BoundedMinState::Value(min) = partial.state { - acc.merge(min); - } + if let BoundedMinState::Value(min) = second.state { + first.merge(min); } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.input.as_nullable(); match &partial.state { BoundedMinState::Empty => Ok(Scalar::null(dtype)), BoundedMinState::Value(min) => min.cast(&dtype), } } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -212,18 +219,28 @@ impl AggregateFnVTable for BoundedMin { let Some(result) = min_max(&array, ctx, NumericalAggregateOpts::default())? else { return Ok(()); }; - if let Some(bound) = truncate_min(result.min, partial.max_bytes.get())? { + if let Some(bound) = truncate_min(result.min, options.max_bytes.get())? { partial.merge(bound); } Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -347,8 +364,6 @@ mod tests { acc.accumulate(&values, &mut ctx)?; acc.fold_partial(BoundedMinPartial { state: BoundedMinState::Empty, - element_dtype: values.dtype().clone(), - max_bytes: max_bytes(2), })?; assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 64a27d9fd79..3c3917c74cb 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -10,6 +10,7 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; @@ -30,16 +31,9 @@ use crate::scalar::Scalar; #[derive(Clone, Debug)] pub struct Count; -/// Partial accumulator state for the count aggregate. -pub struct CountPartial { - count: u64, - /// Whether NaN values must be excluded from the count (float input with `skip_nans`). - exclude_nans: bool, -} - impl AggregateFnVTable for Count { type Options = NumericalAggregateOpts; - type Partial = CountPartial; + type Partial = u64; fn id(&self) -> AggregateFnId { static ID: CachedId = CachedId::new("vortex.count"); @@ -58,59 +52,77 @@ impl AggregateFnVTable for Count { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(0) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - Ok(CountPartial { - count: scalar - .as_primitive() - .typed_value::() - .vortex_expect("count partial should not be null"), - exclude_nans: options.skip_nans && input_dtype.is_float(), - }) + Ok(scalar + .as_primitive() + .typed_value::() + .vortex_expect("count partial should not be null")) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(CountPartial { - count: partials.into_iter().map(|partial| partial.count).sum(), - exclude_nans: options.skip_nans && input_dtype.is_float(), - }) + Ok(first + second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(partial.count, Nullability::NonNullable)) + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + Ok(Scalar::primitive(*partial, Nullability::NonNullable)) } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { let mut count = batch.valid_count(ctx)? as u64; - if state.exclude_nans { + // NaN values are excluded from the count of a float input when they are skipped. + if options.skip_nans && dtypes.input.is_float() { // `nan_count` shortcircuits on an exact `Stat::NaNCount` before scanning the batch. count = count.saturating_sub(nan_count(batch, ctx)? as u64); } - state.count += count; + *state += count; Ok(true) } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -118,12 +130,22 @@ impl AggregateFnVTable for Count { unreachable!("Count::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -144,8 +166,8 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::count::Count; - use crate::aggregate_fn::fns::count::CountPartial; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -251,14 +273,11 @@ mod tests { fn count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); + let dtypes = OwnedAggregateDTypes::try_new(&Count, &options, dtype)?; - let partial_of = |count: u64| CountPartial { - count, - exclude_nans: false, - }; - let state = Count.reduce_partials(&options, &dtype, [partial_of(5), partial_of(3)])?; + let state = Count.reduce_partials(&options, dtypes.borrow(), [5, 3])?; - let result = Count.to_scalar(&state)?; + let result = Count.to_scalar(&options, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 028545ea031..5e9642f78ab 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -8,6 +8,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -30,8 +31,6 @@ pub struct First; /// Partial accumulator state for the [`First`] aggregate. pub struct FirstPartial { - /// The nullable version of the input dtype, used for the result and for empty/all-null inputs. - return_dtype: DType, /// The first non-null value seen so far, or `None` if no non-null value has been observed. value: Option, } @@ -57,46 +56,65 @@ impl AggregateFnVTable for First { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(FirstPartial { value: None }) + } + fn partial_from_scalar( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. Ok(FirstPartial { - return_dtype: input_dtype.as_nullable(), value: (!scalar.is_null()).then_some(scalar), }) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - // The first non-empty partial in iteration order wins; later ones are ignored. + // The earlier non-empty partial wins; the later one is ignored. Ok(FirstPartial { - return_dtype: input_dtype.as_nullable(), - value: partials.into_iter().find_map(|partial| partial.value), + value: first.value.or(second.value), }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(partial.return_dtype.clone()), + None => Scalar::null(dtypes.result.clone()), }) } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { partial.value.is_some() } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -113,6 +131,8 @@ impl AggregateFnVTable for First { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -120,12 +140,22 @@ impl AggregateFnVTable for First { unreachable!("First::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -140,6 +170,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::first::FirstPartial; use crate::aggregate_fn::fns::first::first; @@ -263,21 +294,23 @@ mod tests { #[test] fn first_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let partial_of = |value: Option| FirstPartial { - return_dtype: dtype.as_nullable(), - value, - }; + let owned = OwnedAggregateDTypes::try_new(&First, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); + let partial_of = |value: Option| FirstPartial { value }; // An empty partial means the sub-accumulator saw nothing valid - it is ignored. let empty = partial_of(None); - assert!(!First.is_saturated(&empty)); + assert!(!First.is_saturated(&EmptyOptions, dtypes, &empty)); // The first non-empty partial wins; subsequent valid partials are dropped. let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); - let state = First.reduce_partials(&EmptyOptions, &dtype, [empty, five, seven])?; - assert!(First.is_saturated(&state)); - assert_eq!(First.to_scalar(&state)?, Scalar::primitive(5i32, Nullable)); + let state = First.reduce_partials(&EmptyOptions, dtypes, [empty, five, seven])?; + assert!(First.is_saturated(&EmptyOptions, dtypes, &state)); + assert_eq!( + First.to_scalar(&EmptyOptions, dtypes, &state)?, + Scalar::primitive(5i32, Nullable) + ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index a0d478bd40b..dc17f7540e9 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -31,6 +31,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -219,16 +220,14 @@ pub struct IsConstantPartial { is_constant: bool, /// None = empty (no values seen), Some(null) = all nulls, Some(v) = first value seen. first_value: Option, - element_dtype: DType, } impl IsConstantPartial { /// The state of a group with no accumulated values. - fn empty(input_dtype: &DType) -> Self { + fn empty() -> Self { Self { is_constant: true, first_value: None, - element_dtype: input_dtype.clone(), } } @@ -292,15 +291,23 @@ impl AggregateFnVTable for IsConstant { } } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(IsConstantPartial::empty()) + } + fn partial_from_scalar( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null struct means the producing accumulator was empty. if scalar.is_null() { - return Ok(IsConstantPartial::empty(input_dtype)); + return Ok(IsConstantPartial::empty()); } let is_constant = scalar @@ -312,32 +319,32 @@ impl AggregateFnVTable for IsConstant { Ok(IsConstantPartial { is_constant, first_value: scalar.as_struct().field_by_idx(1), - element_dtype: input_dtype.clone(), }) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + mut acc: Self::Partial, + partial: Self::Partial, ) -> VortexResult { - let mut acc = IsConstantPartial::empty(input_dtype); - for partial in partials { - if !partial.is_constant { - acc.is_constant = false; - break; - } - if let Some(value) = partial.first_value { - acc.check_value(value); - } + if !partial.is_constant { + acc.is_constant = false; + } else if let Some(value) = partial.first_value { + acc.check_value(value); } Ok(acc) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = make_is_constant_partial_dtype(&partial.element_dtype); - let element_dtype = partial.element_dtype.as_nullable(); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.partial.clone(); + let element_dtype = dtypes.input.as_nullable(); // Only a constant partial that saw no values is the empty (null) state: a non-constant // verdict stands regardless of whether a value was observed. let first_value = match &partial.first_value { @@ -355,12 +362,19 @@ impl AggregateFnVTable for IsConstant { } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !partial.is_constant } fn accumulate( &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -384,7 +398,7 @@ impl AggregateFnVTable for IsConstant { let all_invalid = array_ref.all_invalid(ctx)?; if all_invalid { - partial.check_value(Scalar::null(partial.element_dtype.as_nullable())); + partial.check_value(Scalar::null(dtypes.input.as_nullable())); return Ok(()); } @@ -431,11 +445,21 @@ impl AggregateFnVTable for IsConstant { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { if partial.first_value.is_none() { // Empty accumulator → return false. return Ok(Scalar::bool(false, Nullability::NonNullable)); @@ -457,6 +481,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::is_constant::IsConstant; use crate::aggregate_fn::fns::is_constant::IsConstantPartial; use crate::aggregate_fn::fns::is_constant::is_constant; @@ -794,17 +819,18 @@ mod tests { #[test] fn non_constant_partial_without_value_is_not_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let owned = OwnedAggregateDTypes::try_new(&IsConstant, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); let partial = IsConstantPartial { is_constant: false, first_value: None, - element_dtype: dtype.clone(), }; - let scalar = IsConstant.to_scalar(&partial)?; + let scalar = IsConstant.to_scalar(&EmptyOptions, dtypes, &partial)?; assert!(!scalar.is_null()); - let parsed = IsConstant.partial_from_scalar(&EmptyOptions, &dtype, scalar)?; + let parsed = IsConstant.partial_from_scalar(&EmptyOptions, dtypes, scalar)?; assert_eq!( - IsConstant.finalize_scalar(&parsed)?, + IsConstant.finalize_scalar(&EmptyOptions, dtypes, &parsed)?, Scalar::bool(false, Nullability::NonNullable) ); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 9ab0501a4ce..3c9ebdd3389 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -26,6 +26,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -202,22 +203,18 @@ impl IsSorted { /// Partial accumulator state for is_sorted. pub struct IsSortedPartial { is_sorted: bool, - strict: bool, /// None = empty (no values seen). first_value: Option, last_value: Option, - element_dtype: DType, } impl IsSortedPartial { /// The state of a group with no accumulated values. - fn empty(options: &IsSortedOptions, input_dtype: &DType) -> Self { + fn empty() -> Self { Self { is_sorted: true, - strict: options.strict, first_value: None, last_value: None, - element_dtype: input_dtype.clone(), } } } @@ -290,15 +287,23 @@ impl AggregateFnVTable for IsSorted { } } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(IsSortedPartial::empty()) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null struct means the producing accumulator was empty. if scalar.is_null() { - return Ok(IsSortedPartial::empty(options, input_dtype)); + return Ok(IsSortedPartial::empty()); } let is_sorted = scalar @@ -310,77 +315,77 @@ impl AggregateFnVTable for IsSorted { // The scalar's own strict flag is ignored: strictness comes from the options. Ok(IsSortedPartial { is_sorted, - strict: options.strict, first_value: scalar.as_struct().field_by_idx(2), last_value: scalar.as_struct().field_by_idx(3), - element_dtype: input_dtype.clone(), }) } - fn reduce_partials( + fn merge_partials( &self, options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + mut acc: Self::Partial, + partial: Self::Partial, ) -> VortexResult { - let mut acc = IsSortedPartial::empty(options, input_dtype); + if !acc.is_sorted { + return Ok(acc); + } - for partial in partials { - if !acc.is_sorted { - break; + if !partial.is_sorted { + // An unsorted partial settles the verdict whether or not it observed boundaries. + acc.is_sorted = false; + if let Some(last) = partial.last_value.or_else(|| partial.first_value.clone()) { + acc.last_value = Some(last); } - - if !partial.is_sorted { - // An unsorted partial settles the verdict whether or not it observed boundaries. - acc.is_sorted = false; - if let Some(last) = partial.last_value.or_else(|| partial.first_value.clone()) { - acc.last_value = Some(last); - } - if acc.first_value.is_none() { - acc.first_value = partial.first_value; - } - break; + if acc.first_value.is_none() { + acc.first_value = partial.first_value; } + return Ok(acc); + } - // A sorted partial without a first value is empty and contributes nothing. - let Some(first) = partial.first_value else { - continue; - }; - // A partial that saw a single value carries it as both boundaries. - let last = partial.last_value.unwrap_or_else(|| first.clone()); - - // Check boundary: acc.last_value vs partial.first_value - if let Some(acc_last) = &acc.last_value { - if !acc_last.is_null() && !first.is_null() { - let boundary_ok = if acc.strict { - *acc_last < first - } else { - *acc_last <= first - }; - if !boundary_ok { - acc.is_sorted = false; - } - } else if !acc_last.is_null() && first.is_null() { - // non-null before null violates sort order - acc.is_sorted = false; - } else if acc_last.is_null() && first.is_null() && acc.strict { - // both null with strict: violates strict sort + // A sorted partial without a first value is empty and contributes nothing. + let Some(first) = partial.first_value else { + return Ok(acc); + }; + // A partial that saw a single value carries it as both boundaries. + let last = partial.last_value.unwrap_or_else(|| first.clone()); + + // Check boundary: acc.last_value vs partial.first_value + if let Some(acc_last) = &acc.last_value { + if !acc_last.is_null() && !first.is_null() { + let boundary_ok = if options.strict { + *acc_last < first + } else { + *acc_last <= first + }; + if !boundary_ok { acc.is_sorted = false; } + } else if !acc_last.is_null() && first.is_null() { + // non-null before null violates sort order + acc.is_sorted = false; + } else if acc_last.is_null() && first.is_null() && options.strict { + // both null with strict: violates strict sort + acc.is_sorted = false; } + } - // Update first_value if this is the first non-empty partial. - if acc.first_value.is_none() { - acc.first_value = Some(first); - } - acc.last_value = Some(last); + // Update first_value if this is the first non-empty partial. + if acc.first_value.is_none() { + acc.first_value = Some(first); } + acc.last_value = Some(last); Ok(acc) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = make_is_sorted_partial_dtype(&partial.element_dtype); + fn to_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.partial.clone(); // Only a sorted partial that saw no values is the empty (null) state: an unsorted verdict // stands regardless of which boundaries were observed. if partial.is_sorted && partial.first_value.is_none() { @@ -389,7 +394,7 @@ impl AggregateFnVTable for IsSorted { let first_value = partial .first_value .clone() - .unwrap_or_else(|| Scalar::null(partial.element_dtype.as_nullable())); + .unwrap_or_else(|| Scalar::null(dtypes.input.as_nullable())); // A partial that saw a single value carries it as both boundaries. let last_value = partial .last_value @@ -399,7 +404,7 @@ impl AggregateFnVTable for IsSorted { dtype, vec![ Scalar::bool(partial.is_sorted, Nullability::NonNullable), - Scalar::bool(partial.strict, Nullability::NonNullable), + Scalar::bool(options.strict, Nullability::NonNullable), first_value, last_value, ], @@ -407,12 +412,19 @@ impl AggregateFnVTable for IsSorted { } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { !partial.is_sorted } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -425,14 +437,14 @@ impl AggregateFnVTable for IsSorted { Columnar::Constant(c) => { // Constant arrays are sorted but not strict sorted (if len > 1). let value = c.scalar().clone().into_nullable(); - if partial.strict && c.len() > 1 { + if options.strict && c.len() > 1 { partial.is_sorted = false; } // Check boundary with previous chunk. if let Some(self_last) = &partial.last_value { if !self_last.is_null() && !value.is_null() { - let boundary_ok = if partial.strict { + let boundary_ok = if options.strict { *self_last < value } else { *self_last <= value @@ -441,7 +453,7 @@ impl AggregateFnVTable for IsSorted { partial.is_sorted = false; } } else if (!self_last.is_null() && value.is_null()) - || (self_last.is_null() && value.is_null() && partial.strict) + || (self_last.is_null() && value.is_null() && options.strict) { partial.is_sorted = false; } @@ -464,7 +476,7 @@ impl AggregateFnVTable for IsSorted { let first_value = array_ref.execute_scalar(0, ctx)?.into_nullable(); if let Some(self_last) = &partial.last_value { if !self_last.is_null() && !first_value.is_null() { - let boundary_ok = if partial.strict { + let boundary_ok = if options.strict { *self_last < first_value } else { *self_last <= first_value @@ -482,7 +494,7 @@ impl AggregateFnVTable for IsSorted { return Ok(()); } } else if (!self_last.is_null() && first_value.is_null()) - || (self_last.is_null() && first_value.is_null() && partial.strict) + || (self_last.is_null() && first_value.is_null() && options.strict) { partial.is_sorted = false; partial.last_value = Some( @@ -499,12 +511,12 @@ impl AggregateFnVTable for IsSorted { // Check within-batch sortedness. let batch_is_sorted = match c { - Canonical::Primitive(p) => check_primitive_sorted(p, partial.strict, ctx)?, - Canonical::Bool(b) => check_bool_sorted(b, partial.strict, ctx)?, - Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict, ctx)?, - Canonical::Decimal(d) => check_decimal_sorted(d, partial.strict, ctx)?, - Canonical::Extension(e) => check_extension_sorted(e, partial.strict, ctx)?, - Canonical::Null(_) => !partial.strict, + Canonical::Primitive(p) => check_primitive_sorted(p, options.strict, ctx)?, + Canonical::Bool(b) => check_bool_sorted(b, options.strict, ctx)?, + Canonical::VarBinView(v) => check_varbinview_sorted(v, options.strict, ctx)?, + Canonical::Decimal(d) => check_decimal_sorted(d, options.strict, ctx)?, + Canonical::Extension(e) => check_extension_sorted(e, options.strict, ctx)?, + Canonical::Null(_) => !options.strict, // Struct, List, FixedSizeList should have been filtered out by return_dtype _ => unreachable!(), }; @@ -526,11 +538,21 @@ impl AggregateFnVTable for IsSorted { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { // The empty state is vacuously sorted, so the verdict stands on its own. Ok(Scalar::bool(partial.is_sorted, Nullability::NonNullable)) } @@ -574,6 +596,7 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::is_sorted::IsSorted; use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; use crate::aggregate_fn::fns::is_sorted::IsSortedPartial; @@ -768,19 +791,19 @@ mod tests { fn unsorted_partial_without_boundaries_is_not_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = IsSortedOptions { strict: false }; + let owned = OwnedAggregateDTypes::try_new(&IsSorted, &options, dtype)?; + let dtypes = owned.borrow(); let partial = IsSortedPartial { is_sorted: false, - strict: false, first_value: None, last_value: None, - element_dtype: dtype.clone(), }; - let scalar = IsSorted.to_scalar(&partial)?; + let scalar = IsSorted.to_scalar(&options, dtypes, &partial)?; assert!(!scalar.is_null()); - let parsed = IsSorted.partial_from_scalar(&options, &dtype, scalar)?; + let parsed = IsSorted.partial_from_scalar(&options, dtypes, scalar)?; assert_eq!( - IsSorted.finalize_scalar(&parsed)?, + IsSorted.finalize_scalar(&options, dtypes, &parsed)?, Scalar::bool(false, Nullability::NonNullable) ); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index b119f4aa18a..c81b66165cb 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -8,6 +8,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -30,8 +31,6 @@ pub struct Last; /// Partial accumulator state for the [`Last`] aggregate. pub struct LastPartial { - /// The nullable version of the input dtype, used for the result and for empty/all-null inputs. - return_dtype: DType, /// The last non-null value seen so far, or `None` if no non-null value has been observed. value: Option, } @@ -57,50 +56,66 @@ impl AggregateFnVTable for Last { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(LastPartial { value: None }) + } + fn partial_from_scalar( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. Ok(LastPartial { - return_dtype: input_dtype.as_nullable(), value: (!scalar.is_null()).then_some(scalar), }) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - // The last non-empty partial in iteration order wins; empty ones are ignored. + // The later non-empty partial wins; an empty later partial changes nothing. Ok(LastPartial { - return_dtype: input_dtype.as_nullable(), - value: partials - .into_iter() - .filter_map(|partial| partial.value) - .last(), + value: second.value.or(first.value), }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(partial.return_dtype.clone()), + None => Scalar::null(dtypes.result.clone()), }) } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { // Last can never short-circuit: a later batch can always supersede the current value. false } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -114,6 +129,8 @@ impl AggregateFnVTable for Last { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -121,12 +138,22 @@ impl AggregateFnVTable for Last { unreachable!("Last::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -141,6 +168,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::last::LastPartial; use crate::aggregate_fn::fns::last::last; @@ -263,10 +291,9 @@ mod tests { #[test] fn last_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let partial_of = |value: Option| LastPartial { - return_dtype: dtype.as_nullable(), - value, - }; + let owned = OwnedAggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); + let partial_of = |value: Option| LastPartial { value }; let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); @@ -274,8 +301,11 @@ mod tests { let empty = partial_of(None); // The last non-empty partial in order replaces the prior values. - let state = Last.reduce_partials(&EmptyOptions, &dtype, [five, seven, empty])?; - assert_eq!(Last.to_scalar(&state)?, Scalar::primitive(7i32, Nullable)); + let state = Last.reduce_partials(&EmptyOptions, dtypes, [five, seven, empty])?; + assert_eq!( + Last.to_scalar(&EmptyOptions, dtypes, &state)?, + Scalar::primitive(7i32, Nullable) + ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index 76713db16c3..c292fa77253 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -10,6 +10,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -38,21 +39,15 @@ pub struct Max; /// Partial accumulator state for the maximum aggregate. pub struct MaxPartial { max: Option, - element_dtype: DType, - skip_nans: bool, } impl MaxPartial { - /// The state of a group with no accumulated values. - fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { - Self { - max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - } - } - - fn merge(&mut self, max: Scalar) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, + max: Scalar, + ) { if max.is_null() { return; } @@ -60,8 +55,8 @@ impl MaxPartial { // NaN scalars are incomparable under `partial_max`; they poison the maximum when NaNs // participate, and are dropped when they are skipped. if scalar_is_nan(&max) || self.is_poisoned() { - if !self.skip_nans { - self.poison(); + if !options.skip_nans { + self.poison(dtypes); } return; } @@ -72,12 +67,12 @@ impl MaxPartial { }); } - fn poison(&mut self) { - self.max = Some(nan_scalar(&self.element_dtype)); + fn poison(&mut self, dtypes: AggregateDTypes<'_>) { + self.max = Some(nan_scalar(dtypes.input)); } fn is_poisoned(&self) -> bool { - self.element_dtype.is_float() && self.max.as_ref().is_some_and(scalar_is_nan) + self.max.as_ref().is_some_and(scalar_is_nan) } } @@ -130,55 +125,73 @@ impl AggregateFnVTable for Max { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(MaxPartial { max: None }) + } + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = MaxPartial::empty(options, input_dtype); + let mut partial = MaxPartial { max: None }; // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. - partial.merge(scalar); + partial.merge(options, dtypes, scalar); Ok(partial) } - fn reduce_partials( + fn merge_partials( &self, options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = MaxPartial::empty(options, input_dtype); - for partial in partials { - if let Some(max) = partial.max { - acc.merge(max); - } + if let Some(max) = second.max { + first.merge(options, dtypes, max); } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.input.as_nullable(); match &partial.max { Some(max) => max.cast(&dtype), None => Ok(Scalar::null(dtype)), } } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including maximum is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to the NaN-including float maximum; everything else // takes the default dispatch path. - if partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.input.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -186,13 +199,13 @@ impl AggregateFnVTable for Max { // NaN-free batch: the cached NaN-skipping maximum (if any) is valid. `to_scalar` // re-casts to the result dtype, so the cached scalar can merge as-is. if let Some(max) = batch.statistics().get(Stat::Max).as_exact() { - partial.merge(max); + partial.merge(options, dtypes, max); return Ok(true); } Ok(false) } Precision::Exact(_) => { - partial.poison(); + partial.poison(dtypes); Ok(true) } _ => Ok(false), @@ -201,6 +214,8 @@ impl AggregateFnVTable for Max { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -211,21 +226,28 @@ impl AggregateFnVTable for Max { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - let options = NumericalAggregateOpts { - skip_nans: partial.skip_nans, - }; - if let Some(result) = min_max(&array, ctx, options)? { - partial.merge(result.max); + if let Some(result) = min_max(&array, ctx, *options)? { + partial.merge(options, dtypes, result.max); } Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index a3a0f77cb45..bb6bd51260d 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; @@ -9,6 +10,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; @@ -88,15 +90,21 @@ impl BinaryCombined for Mean { "count" } - fn return_dtype(&self, input_dtype: &DType) -> Option { + fn return_dtype(&self, _options: &CombinedOptions, input_dtype: &DType) -> Option { Some(mean_output_dtype(input_dtype)?.with_nullability(Nullability::Nullable)) } - fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &CombinedOptions, + dtypes: AggregateDTypes<'_>, + sum: ArrayRef, + count: ArrayRef, + ) -> VortexResult { if let DType::Decimal(..) = sum.dtype() { vortex_bail!("grouped mean over decimals is not yet supported"); } - let target = DType::Primitive(PType::F64, Nullability::Nullable); + let target = dtypes.result.clone(); let sum = sum.cast(target.clone())?; let count = count.cast(target.clone())?; @@ -114,12 +122,23 @@ impl BinaryCombined for Mean { sum.binary(count, Operator::Div) } - fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult { + fn finalize_scalar( + &self, + _options: &CombinedOptions, + dtypes: AggregateDTypes<'_>, + left_scalar: Scalar, + right_scalar: Scalar, + ) -> VortexResult { if let DType::Decimal(decimal_dtype, _) = *left_scalar.dtype() { - return finalize_decimal_scalar(&left_scalar, &right_scalar, decimal_dtype); + return finalize_decimal_scalar( + &left_scalar, + &right_scalar, + decimal_dtype, + dtypes.result, + ); } - let target = DType::Primitive(PType::F64, Nullability::Nullable); + let target = dtypes.result.clone(); let sum_cast = left_scalar.cast(&target)?; let count_cast = right_scalar.cast(&target)?; @@ -163,18 +182,20 @@ fn finalize_decimal_scalar( sum: &Scalar, count: &Scalar, sum_decimal: DecimalDType, + target_dtype: &DType, ) -> VortexResult { - let target_decimal_dtype = mean_decimal_dtype(&sum_decimal); - let target_dtype = DType::Decimal(target_decimal_dtype, Nullability::Nullable); + let target_decimal_dtype = *target_dtype + .as_decimal_opt() + .vortex_expect("decimal mean result dtype"); // overflow let Some(sum_value) = sum.as_decimal().decimal_value() else { - return Ok(Scalar::null(target_dtype)); + return Ok(Scalar::null(target_dtype.clone())); }; // empty input let count = count.as_primitive().typed_value::().unwrap_or(0); if count == 0 { - return Ok(Scalar::null(target_dtype)); + return Ok(Scalar::null(target_dtype.clone())); } let Ok(sum) = DecimalValue::rescale_i256( @@ -182,12 +203,12 @@ fn finalize_decimal_scalar( sum_decimal.scale(), target_decimal_dtype.scale(), ) else { - return Ok(Scalar::null(target_dtype)); + return Ok(Scalar::null(target_dtype.clone())); }; let mean = sum / i256::from_i128(i128::from(count)); let Ok(mean) = DecimalValue::try_from_i256(mean, target_decimal_dtype) else { - return Ok(Scalar::null(target_dtype)); + return Ok(Scalar::null(target_dtype.clone())); }; Ok(Scalar::decimal( mean, diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index a9443020551..c7a606237fa 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -10,6 +10,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -38,21 +39,15 @@ pub struct Min; /// Partial accumulator state for the minimum aggregate. pub struct MinPartial { min: Option, - element_dtype: DType, - skip_nans: bool, } impl MinPartial { - /// The state of a group with no accumulated values. - fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { - Self { - min: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - } - } - - fn merge(&mut self, min: Scalar) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, + min: Scalar, + ) { if min.is_null() { return; } @@ -60,8 +55,8 @@ impl MinPartial { // NaN scalars are incomparable under `partial_min`; they poison the minimum when NaNs // participate, and are dropped when they are skipped. if scalar_is_nan(&min) || self.is_poisoned() { - if !self.skip_nans { - self.poison(); + if !options.skip_nans { + self.poison(dtypes); } return; } @@ -72,12 +67,12 @@ impl MinPartial { }); } - fn poison(&mut self) { - self.min = Some(nan_scalar(&self.element_dtype)); + fn poison(&mut self, dtypes: AggregateDTypes<'_>) { + self.min = Some(nan_scalar(dtypes.input)); } fn is_poisoned(&self) -> bool { - self.element_dtype.is_float() && self.min.as_ref().is_some_and(scalar_is_nan) + self.min.as_ref().is_some_and(scalar_is_nan) } } @@ -130,55 +125,73 @@ impl AggregateFnVTable for Min { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(MinPartial { min: None }) + } + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = MinPartial::empty(options, input_dtype); + let mut partial = MinPartial { min: None }; // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. - partial.merge(scalar); + partial.merge(options, dtypes, scalar); Ok(partial) } - fn reduce_partials( + fn merge_partials( &self, options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = MinPartial::empty(options, input_dtype); - for partial in partials { - if let Some(min) = partial.min { - acc.merge(min); - } + if let Some(min) = second.min { + first.merge(options, dtypes, min); } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.input.as_nullable(); match &partial.min { Some(min) => min.cast(&dtype), None => Ok(Scalar::null(dtype)), } } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including minimum is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to the NaN-including float minimum; everything else // takes the default dispatch path. - if partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.input.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -186,13 +199,13 @@ impl AggregateFnVTable for Min { // NaN-free batch: the cached NaN-skipping minimum (if any) is valid. `to_scalar` // re-casts to the result dtype, so the cached scalar can merge as-is. if let Some(min) = batch.statistics().get(Stat::Min).as_exact() { - partial.merge(min); + partial.merge(options, dtypes, min); return Ok(true); } Ok(false) } Precision::Exact(_) => { - partial.poison(); + partial.poison(dtypes); Ok(true) } _ => Ok(false), @@ -201,6 +214,8 @@ impl AggregateFnVTable for Min { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -211,21 +226,28 @@ impl AggregateFnVTable for Min { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - let options = NumericalAggregateOpts { - skip_nans: partial.skip_nans, - }; - if let Some(result) = min_max(&array, ctx, options)? { - partial.merge(result.min); + if let Some(result) = min_max(&array, ctx, *options)? { + partial.merge(options, dtypes, result.min); } Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs b/vortex-array/src/aggregate_fn/fns/min_max/bool.rs index adca6453862..2aebb7f1941 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/bool.rs @@ -9,12 +9,16 @@ use vortex_mask::AllOr; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::dtype::Nullability::NonNullable; use crate::scalar::Scalar; pub(super) fn accumulate_bool( + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, partial: &mut MinMaxPartial, array: &BoolArray, ctx: &mut ExecutionCtx, @@ -44,9 +48,13 @@ pub(super) fn accumulate_bool( (false, true) }; - partial.merge(Some(MinMaxResult { - min: Scalar::bool(min, NonNullable), - max: Scalar::bool(max, NonNullable), - })); + partial.merge( + options, + dtypes, + Some(MinMaxResult { + min: Scalar::bool(min, NonNullable), + max: Scalar::bool(max, NonNullable), + }), + ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs index 019136e2584..a5c2e0071dd 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs @@ -8,6 +8,8 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; use crate::dtype::NativeDecimalType; @@ -17,13 +19,15 @@ use crate::scalar::DecimalValue; use crate::scalar::Scalar; pub(super) fn accumulate_decimal( + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, partial: &mut MinMaxPartial, array: &DecimalArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { match_each_decimal_value_type!(array.values_type(), |T| { let local = compute_min_max_with_validity::(array, ctx)?; - partial.merge(local); + partial.merge(options, dtypes, local); Ok(()) }) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index e981c7c6af1..1014bea111c 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -7,6 +7,7 @@ use super::MinMaxPartial; use super::MinMaxResult; use super::min_max; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; @@ -14,6 +15,8 @@ use crate::dtype::Nullability; use crate::scalar::Scalar; pub(super) fn accumulate_extension( + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, partial: &mut MinMaxPartial, array: &ExtensionArray, ctx: &mut ExecutionCtx, @@ -28,6 +31,6 @@ pub(super) fn accumulate_extension( min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), max: Scalar::extension_ref(non_nullable_ext_dtype, max), }); - partial.merge(local); + partial.merge(options, dtypes, local); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 737cb6eb716..674e1e42a3a 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -25,6 +25,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -197,23 +198,16 @@ pub struct MinMax; pub struct MinMaxPartial { min: Option, max: Option, - element_dtype: DType, - skip_nans: bool, } impl MinMaxPartial { - /// The state of a group with no accumulated values. - fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> Self { - Self { - min: None, - max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - } - } - /// Merge a local `MinMaxResult` into this partial state. - fn merge(&mut self, local: Option) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, + local: Option, + ) { let Some(MinMaxResult { min, max }) = local else { return; }; @@ -222,8 +216,8 @@ impl MinMaxPartial { // explicitly: a NaN extremum poisons the partial state when NaNs participate, and is // dropped when they are skipped. if scalar_is_nan(&min) || scalar_is_nan(&max) || self.is_poisoned() { - if !self.skip_nans { - self.poison(); + if !options.skip_nans { + self.poison(dtypes); } return; } @@ -240,15 +234,15 @@ impl MinMaxPartial { } /// Poison the partial state to `{min: NaN, max: NaN}`. - fn poison(&mut self) { - let nan = nan_scalar(&self.element_dtype); + fn poison(&mut self, dtypes: AggregateDTypes<'_>) { + let nan = nan_scalar(dtypes.input); self.min = Some(nan.clone()); self.max = Some(nan); } /// Whether the partial state is poisoned to NaN. fn is_poisoned(&self) -> bool { - self.element_dtype.is_float() && self.min.as_ref().is_some_and(scalar_is_nan) + self.min.as_ref().is_some_and(scalar_is_nan) } } @@ -318,35 +312,52 @@ impl AggregateFnVTable for MinMax { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(MinMaxPartial { + min: None, + max: None, + }) + } + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = MinMaxPartial::empty(options, input_dtype); + let mut partial = MinMaxPartial { + min: None, + max: None, + }; // `merge` normalizes the parsed extrema: nulls stay empty and NaNs poison or drop. - partial.merge(MinMaxResult::from_scalar(scalar)?); + partial.merge(options, dtypes, MinMaxResult::from_scalar(scalar)?); Ok(partial) } - fn reduce_partials( + fn merge_partials( &self, options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = MinMaxPartial::empty(options, input_dtype); - for partial in partials { - if let (Some(min), Some(max)) = (partial.min, partial.max) { - acc.merge(Some(MinMaxResult { min, max })); - } + if let (Some(min), Some(max)) = (second.min, second.max) { + first.merge(options, dtypes, Some(MinMaxResult { min, max })); } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = make_minmax_dtype(&partial.element_dtype); + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.partial.clone(); Ok(match (&partial.min, &partial.max) { (Some(min), Some(max)) => Scalar::struct_(dtype, vec![min.clone(), max.clone()]), _ => Scalar::null(dtype), @@ -354,20 +365,27 @@ impl AggregateFnVTable for MinMax { } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including min/max is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to NaN-including float min/max; everything else // takes the default dispatch path. - if partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.input.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -378,18 +396,22 @@ impl AggregateFnVTable for MinMax { if let Some((min, max)) = cached_min.zip(cached_max) { // Cached float stats carry the (possibly nullable) array dtype; `to_scalar` // builds a struct with non-nullable fields, so normalise here. - let non_nullable_dtype = partial.element_dtype.as_nonnullable(); - partial.merge(Some(MinMaxResult { - min: min.cast(&non_nullable_dtype)?, - max: max.cast(&non_nullable_dtype)?, - })); + let non_nullable_dtype = dtypes.input.as_nonnullable(); + partial.merge( + options, + dtypes, + Some(MinMaxResult { + min: min.cast(&non_nullable_dtype)?, + max: max.cast(&non_nullable_dtype)?, + }), + ); return Ok(true); } Ok(false) } Precision::Exact(_) => { // At least one NaN value poisons both extrema without scanning the batch. - partial.poison(); + partial.poison(dtypes); Ok(true) } _ => Ok(false), @@ -398,6 +420,8 @@ impl AggregateFnVTable for MinMax { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -410,25 +434,29 @@ impl AggregateFnVTable for MinMax { } // NaN float constants are skipped or poison the extrema, per the options. if scalar_is_nan(scalar) { - if !partial.skip_nans { - partial.poison(); + if !options.skip_nans { + partial.poison(dtypes); } return Ok(()); } let non_nullable_dtype = scalar.dtype().as_nonnullable(); let cast = scalar.cast(&non_nullable_dtype)?; - partial.merge(Some(MinMaxResult { - min: cast.clone(), - max: cast, - })); + partial.merge( + options, + dtypes, + Some(MinMaxResult { + min: cast.clone(), + max: cast, + }), + ); Ok(()) } Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx), - Canonical::Bool(b) => accumulate_bool(partial, b, ctx), - Canonical::VarBinView(v) => accumulate_varbinview(partial, v, ctx), - Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx), - Canonical::Extension(e) => accumulate_extension(partial, e, ctx), + Canonical::Primitive(p) => accumulate_primitive(options, dtypes, partial, p, ctx), + Canonical::Bool(b) => accumulate_bool(options, dtypes, partial, b, ctx), + Canonical::VarBinView(v) => accumulate_varbinview(options, dtypes, partial, v, ctx), + Canonical::Decimal(d) => accumulate_decimal(options, dtypes, partial, d, ctx), + Canonical::Extension(e) => accumulate_extension(options, dtypes, partial, e, ctx), Canonical::Null(_) => Ok(()), Canonical::Union(_) => { todo!("TODO(connor)[Union]: implement min_max for Union arrays") @@ -444,12 +472,22 @@ impl AggregateFnVTable for MinMax { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -470,6 +508,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::MinMaxPartial; use crate::aggregate_fn::fns::min_max::MinMaxResult; @@ -671,17 +710,17 @@ mod tests { fn test_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); + let owned = OwnedAggregateDTypes::try_new(&MinMax, &options, dtype)?; + let dtypes = owned.borrow(); let partial_of = |min: i32, max: i32| MinMaxPartial { min: Some(Scalar::from(min)), max: Some(Scalar::from(max)), - element_dtype: dtype.clone(), - skip_nans: true, }; let state = - MinMax.reduce_partials(&options, &dtype, [partial_of(5, 15), partial_of(2, 10)])?; + MinMax.reduce_partials(&options, dtypes, [partial_of(5, 15), partial_of(2, 10)])?; - let result = MinMaxResult::from_scalar(MinMax.to_scalar(&state)?)? + let result = MinMaxResult::from_scalar(MinMax.to_scalar(&options, dtypes, &state)?)? .vortex_expect("should have result"); assert_eq!(result.min, Scalar::from(2i32)); assert_eq!(result.max, Scalar::from(15i32)); diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index 7ca97712477..a35fedf7000 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -8,6 +8,8 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::PrimitiveArray; use crate::dtype::NativePType; use crate::dtype::Nullability::NonNullable; @@ -16,14 +18,16 @@ use crate::scalar::PValue; use crate::scalar::Scalar; pub(super) fn accumulate_primitive( + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, partial: &mut MinMaxPartial, p: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let skip_nans = partial.skip_nans; + let skip_nans = options.skip_nans; match_each_native_ptype!(p.ptype(), |T| { let local = compute_min_max_with_validity::(p, ctx, skip_nans)?; - partial.merge(local); + partial.merge(options, dtypes, local); Ok(()) }) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs index 5d4de12cd11..9e0846eb553 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs @@ -8,17 +8,25 @@ use vortex_error::vortex_panic; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::VarBinViewArray; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::scalar::Scalar; pub(super) fn accumulate_varbinview( + options: &NumericalAggregateOpts, + dtypes: AggregateDTypes<'_>, partial: &mut MinMaxPartial, array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - partial.merge(varbin_compute_min_max(array, array.dtype(), ctx)?); + partial.merge( + options, + dtypes, + varbin_compute_min_max(array, array.dtype(), ctx)?, + ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index 79c7a7e51b2..dcd49f56652 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -16,6 +16,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -115,10 +116,18 @@ impl AggregateFnVTable for NanCount { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(0) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -127,26 +136,39 @@ impl AggregateFnVTable for NanCount { .vortex_expect("nan_count partial should not be null")) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().sum()) + Ok(first + second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -172,12 +194,22 @@ impl AggregateFnVTable for NanCount { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -192,6 +224,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::nan_count::nan_count; use crate::array_session; @@ -246,10 +279,11 @@ mod tests { #[test] fn nan_count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let dtypes = OwnedAggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?; - let state = NanCount.reduce_partials(&EmptyOptions, &dtype, [5, 3])?; + let state = NanCount.reduce_partials(&EmptyOptions, dtypes.borrow(), [5, 3])?; - let result = NanCount.to_scalar(&state)?; + let result = NanCount.to_scalar(&EmptyOptions, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index 608e1f7704f..db139ae8d9a 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -12,6 +12,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -84,10 +85,18 @@ impl AggregateFnVTable for NullCount { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(0) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -96,26 +105,39 @@ impl AggregateFnVTable for NullCount { .vortex_expect("null_count partial should not be null")) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - Ok(partials.into_iter().sum()) + Ok(first + second) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -126,6 +148,8 @@ impl AggregateFnVTable for NullCount { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -143,12 +167,22 @@ impl AggregateFnVTable for NullCount { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 3d919525ce0..b970954549f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -15,6 +15,7 @@ use vortex_mask::Mask; use super::SumState; use crate::ExecutionCtx; use crate::arrays::DecimalArray; +use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; @@ -25,6 +26,7 @@ use crate::scalar::DecimalValue; /// Returns Ok(true) if saturated (overflow), Ok(false) if not. pub(crate) fn accumulate_decimal( inner: &mut SumState, + return_dtype: &DType, d: &DecimalArray, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -37,9 +39,12 @@ pub(crate) fn accumulate_decimal( } }; - let SumState::Decimal { value, dtype } = inner else { + let SumState::Decimal(value) = inner else { vortex_panic!("expected decimal sum state for decimal input"); }; + let dtype = return_dtype + .as_decimal_opt() + .vortex_expect("decimal sum result dtype"); let values_type = DecimalType::smallest_decimal_value_type(dtype); match_each_decimal_value_type!(d.values_type(), |T| { @@ -115,6 +120,7 @@ mod tests { use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumPartial; use crate::aggregate_fn::fns::sum::SumState; @@ -132,14 +138,9 @@ mod tests { use crate::validity::Validity; /// A partial whose running decimal sum is `value` (test-only helper bypassing scalar parsing). - fn partial_with_decimal(value: DecimalValue, decimal_dtype: DecimalDType) -> SumPartial { + fn partial_with_decimal(value: DecimalValue) -> SumPartial { SumPartial { - return_dtype: DType::Decimal(decimal_dtype, Nullable), - current: Some(SumState::Decimal { - value, - dtype: decimal_dtype, - }), - skip_nans: true, + current: Some(SumState::Decimal(value)), } } @@ -372,15 +373,14 @@ mod tests { // Reduce partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let sum_decimal = DecimalDType::new(14, 0); + let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; - let near_limit = - partial_with_decimal(DecimalValue::from(99_999_999_999_990i64), sum_decimal); + let near_limit = partial_with_decimal(DecimalValue::from(99_999_999_999_990i64)); // Add a small value that keeps us just under 10^14. - let small = partial_with_decimal(DecimalValue::from(9i64), sum_decimal); - let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, small])?; + let small = partial_with_decimal(DecimalValue::from(9i64)); + let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, small])?; - let result = Sum.to_scalar(&state)?; + let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -398,15 +398,14 @@ mod tests { // saturation path in reduce_partials. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let sum_decimal = DecimalDType::new(14, 0); + let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; - let near_limit = - partial_with_decimal(DecimalValue::from(99_999_999_999_999i64), sum_decimal); + let near_limit = partial_with_decimal(DecimalValue::from(99_999_999_999_999i64)); // Push the sum to exactly 10^14, exceeding precision 14. - let one_more = partial_with_decimal(DecimalValue::from(1i64), sum_decimal); - let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, one_more])?; + let one_more = partial_with_decimal(DecimalValue::from(1i64)); + let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, one_more])?; - let result = Sum.to_scalar(&state)?; + let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -420,14 +419,13 @@ mod tests { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let sum_decimal = DecimalDType::new(14, 0); + let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; - let near_limit = - partial_with_decimal(DecimalValue::from(-99_999_999_999_999i64), sum_decimal); - let one_more = partial_with_decimal(DecimalValue::from(-1i64), sum_decimal); - let state = Sum.reduce_partials(&options, &input_dtype, [near_limit, one_more])?; + let near_limit = partial_with_decimal(DecimalValue::from(-99_999_999_999_999i64)); + let one_more = partial_with_decimal(DecimalValue::from(-1i64)); + let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, one_more])?; - let result = Sum.to_scalar(&state)?; + let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(result.is_null()); Ok(()) } @@ -440,11 +438,17 @@ mod tests { // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. // // We seed the state close to 10^37, then accumulate a real array that pushes it over. - let return_dtype = DecimalDType::new(37, 0); + let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); + let options = NumericalAggregateOpts::default(); + let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; + assert_eq!( + dtypes.result(), + &DType::Decimal(DecimalDType::new(37, 0), Nullable) + ); // Set state to 10^37 - 1. let near_limit_val: i128 = 10i128.pow(37) - 1; - let mut state = partial_with_decimal(DecimalValue::from(near_limit_val), return_dtype); + let mut state = partial_with_decimal(DecimalValue::from(near_limit_val)); // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -453,9 +457,9 @@ mod tests { // Drive accumulate through the vtable directly. let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); let mut ctx = array_session().create_execution_ctx(); - Sum.accumulate(&mut state, &columnar, &mut ctx)?; + Sum.accumulate(&options, dtypes.borrow(), &mut state, &columnar, &mut ctx)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 2f2886f9ad0..86eaf221f0b 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -11,7 +11,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -28,6 +27,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -143,59 +143,71 @@ impl AggregateFnVTable for Sum { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(SumPartial { + current: Some(make_zero_state(dtypes.result)), + }) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = SumPartial::empty(options, input_dtype)?; vortex_ensure!( - scalar.dtype().eq_ignore_nullability(&partial.return_dtype), + scalar.dtype().eq_ignore_nullability(dtypes.result), "Sum partial has dtype {}, expected {}", scalar.dtype(), - partial.return_dtype + dtypes.result ); // A null partial means the producing accumulator saturated (overflow). - partial.current = if scalar.is_null() { + let current = if scalar.is_null() { None } else { - Some(sum_state_from_scalar(&scalar, &partial.return_dtype)?) + Some(sum_state_from_scalar(&scalar, dtypes.result)?) }; - Ok(partial) + Ok(SumPartial { current }) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = SumPartial::empty(options, input_dtype)?; - for partial in partials { - let overflow = match (acc.current.as_mut(), partial.current) { - (None, _) => break, - // A saturated (overflowed) partial poisons the reduction. - (Some(_), None) => true, - (Some(acc_state), Some(state)) => checked_add_sum_states(acc_state, &state)?, - }; - if overflow { - acc.current = None; - break; - } + let overflow = match (first.current.as_mut(), second.current) { + // A saturated sum stays saturated. + (None, _) => false, + // A saturated (overflowed) partial poisons the merge. + (Some(_), None) => true, + (Some(acc), Some(state)) => checked_add_sum_states(acc, dtypes.result, &state)?, + }; + if overflow { + first.current = None; } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.current { - None => Scalar::null(partial.return_dtype.as_nullable()), + None => Scalar::null(dtypes.result.as_nullable()), Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable), Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable), Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable), - Some(SumState::Decimal { value, .. }) => { - let decimal_dtype = *partial - .return_dtype + Some(SumState::Decimal(value)) => { + let decimal_dtype = *dtypes + .result .as_decimal_opt() .vortex_expect("return dtype must be decimal"); Scalar::decimal(*value, decimal_dtype, Nullability::Nullable) @@ -204,7 +216,12 @@ impl AggregateFnVTable for Sum { } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { match partial.current.as_ref() { None => true, Some(SumState::Float(v)) => v.is_nan(), @@ -214,13 +231,15 @@ impl AggregateFnVTable for Sum { fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes // the default dispatch path. - if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + if options.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -228,12 +247,12 @@ impl AggregateFnVTable for Sum { // NaN-free batch: the cached NaN-skipping sum (if any) equals the // NaN-including sum. if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { - let sum = if sum.dtype() == &partial.return_dtype { + let sum = if sum.dtype() == dtypes.result { sum } else { - sum.cast(&partial.return_dtype)? + sum.cast(dtypes.result)? }; - merge_sum_result(partial, sum)?; + merge_sum_result(partial, dtypes.result, sum)?; return Ok(true); } Ok(false) @@ -251,6 +270,8 @@ impl AggregateFnVTable for Sum { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -258,16 +279,15 @@ impl AggregateFnVTable for Sum { // Constants compute scalar * len and merge the product into the running state. if let Columnar::Constant(c) = batch { // NaN constants are treated as missing when skipping NaNs. - if partial.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { + if options.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { return Ok(()); } - if let Some(product) = multiply_constant(c.scalar(), c.len(), &partial.return_dtype)? { - merge_sum_result(partial, product)?; + if let Some(product) = multiply_constant(c.scalar(), c.len(), dtypes.result)? { + merge_sum_result(partial, dtypes.result, product)?; } return Ok(()); } - let skip_nans = partial.skip_nans; let mut inner = match partial.current.take() { Some(inner) => inner, None => return Ok(()), @@ -275,9 +295,11 @@ impl AggregateFnVTable for Sum { let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans), + Canonical::Primitive(p) => { + accumulate_primitive(&mut inner, p, ctx, options.skip_nans) + } Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, dtypes.result, d, ctx), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -294,50 +316,42 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } -/// The group state for a sum aggregate, containing the accumulated value and configuration -/// needed for reset/result without external context. +/// The accumulated sum, or the overflow state. pub struct SumPartial { - return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, - /// Whether NaN values in float inputs are skipped. - skip_nans: bool, -} - -impl SumPartial { - /// The state of a group with no accumulated values, or an error for unsupported input dtypes. - fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> VortexResult { - let return_dtype = Sum - .return_dtype(options, input_dtype) - .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; - Ok(Self { - current: Some(make_zero_state(&return_dtype)), - return_dtype, - skip_nans: options.skip_nans, - }) - } } /// The accumulated sum value. +/// +/// Decimal sums do not carry their dtype: the result dtype of the aggregate provides the +/// precision and scale whenever they are needed. // TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the // input type every time? Perhaps? pub enum SumState { Unsigned(u64), Signed(i64), Float(f64), - Decimal { - value: DecimalValue, - dtype: DecimalDType, - }, + Decimal(DecimalValue), } pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { @@ -347,10 +361,7 @@ pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0), PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0), }, - DType::Decimal(decimal, _) => SumState::Decimal { - value: DecimalValue::zero(decimal), - dtype: *decimal, - }, + DType::Decimal(decimal, _) => SumState::Decimal(DecimalValue::zero(decimal)), _ => vortex_panic!("Unsupported sum type"), } } @@ -365,18 +376,19 @@ fn sum_state_from_scalar(scalar: &Scalar, return_dtype: &DType) -> VortexResult< SumState::Signed(i64::try_from(scalar)?) } DType::Primitive(..) => SumState::Float(f64::try_from(scalar)?), - DType::Decimal(dtype, _) => SumState::Decimal { - value: DecimalValue::try_from(scalar)?, - dtype: *dtype, - }, + DType::Decimal(..) => SumState::Decimal(DecimalValue::try_from(scalar)?), _ => vortex_bail!("Unsupported sum type {}", return_dtype), }) } /// Checked add of one sum state into another, returning true if overflow occurred. /// -/// A decimal sum that no longer fits its precision counts as an overflow. -pub(crate) fn checked_add_sum_states(state: &mut SumState, other: &SumState) -> VortexResult { +/// A decimal sum that no longer fits the precision of `return_dtype` counts as an overflow. +pub(crate) fn checked_add_sum_states( + state: &mut SumState, + return_dtype: &DType, + other: &SumState, +) -> VortexResult { Ok(match (state, other) { (SumState::Unsigned(acc), SumState::Unsigned(other)) => checked_add_u64(acc, *other), (SumState::Signed(acc), SumState::Signed(other)) => checked_add_i64(acc, *other), @@ -384,7 +396,10 @@ pub(crate) fn checked_add_sum_states(state: &mut SumState, other: &SumState) -> *acc += *other; false } - (SumState::Decimal { value, dtype }, SumState::Decimal { value: other, .. }) => { + (SumState::Decimal(value), SumState::Decimal(other)) => { + let dtype = return_dtype + .as_decimal_opt() + .vortex_expect("decimal sum result dtype"); match value.checked_add(other) { Some(result) if result.fits_in_precision(*dtype) => { *value = result; @@ -398,13 +413,19 @@ pub(crate) fn checked_add_sum_states(state: &mut SumState, other: &SumState) -> } /// Merge a finalized sum result (nullable; null means overflow) into the partial. -fn merge_sum_result(partial: &mut SumPartial, result: Scalar) -> VortexResult<()> { +fn merge_sum_result( + partial: &mut SumPartial, + return_dtype: &DType, + result: Scalar, +) -> VortexResult<()> { let overflow = match partial.current.as_mut() { None => return Ok(()), Some(_) if result.is_null() => true, - Some(acc) => { - checked_add_sum_states(acc, &sum_state_from_scalar(&result, &partial.return_dtype)?)? - } + Some(acc) => checked_add_sum_states( + acc, + return_dtype, + &sum_state_from_scalar(&result, return_dtype)?, + )?, }; if overflow { partial.current = None; @@ -454,6 +475,7 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumPartial; use crate::aggregate_fn::fns::sum::SumState; @@ -582,15 +604,15 @@ mod tests { fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); + let owned = OwnedAggregateDTypes::try_new(&Sum, &options, dtype)?; + let dtypes = owned.borrow(); let partial_of = |value: i64| SumPartial { - return_dtype: DType::Primitive(PType::I64, Nullable), current: Some(SumState::Signed(value)), - skip_nans: true, }; - let state = Sum.reduce_partials(&options, &dtype, [partial_of(100), partial_of(50)])?; + let state = Sum.reduce_partials(&options, dtypes, [partial_of(100), partial_of(50)])?; - let result = Sum.to_scalar(&state)?; + let result = Sum.to_scalar(&options, dtypes, &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) } @@ -599,17 +621,19 @@ mod tests { fn sum_overflowed_partial_poisons_reduction() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); + let owned = OwnedAggregateDTypes::try_new(&Sum, &options, dtype)?; + let dtypes = owned.borrow(); let overflowed = Sum.partial_from_scalar( &options, - &dtype, + dtypes, Scalar::null(DType::Primitive(PType::I64, Nullable)), )?; - let five = Sum.partial_from_scalar(&options, &dtype, Scalar::primitive(5i64, Nullable))?; - let state = Sum.reduce_partials(&options, &dtype, [five, overflowed])?; + let five = Sum.partial_from_scalar(&options, dtypes, Scalar::primitive(5i64, Nullable))?; + let state = Sum.reduce_partials(&options, dtypes, [five, overflowed])?; - assert!(Sum.is_saturated(&state)); - assert!(Sum.to_scalar(&state)?.is_null()); + assert!(Sum.is_saturated(&options, dtypes, &state)); + assert!(Sum.to_scalar(&options, dtypes, &state)?.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 930a353f6e1..1638be27bb9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -59,7 +59,7 @@ fn accumulate_primitive_all( Ok(false) } ), - SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), + SumState::Decimal(_) => vortex_panic!("decimal sum state with primitive input"), } } @@ -174,7 +174,7 @@ fn accumulate_primitive_valid( Ok(false) } ), - SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), + SumState::Decimal(_) => vortex_panic!("decimal sum state with primitive input"), } } diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index 678779dc127..96e71b7fc7c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -4,6 +4,7 @@ mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -17,6 +18,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -105,75 +107,94 @@ impl AggregateFnVTable for SumV2 { .map(sum_v2_partial_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(SumV2Partial::empty(dtypes.result)) + } + fn partial_from_scalar( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = SumV2Partial::empty(options, input_dtype)?; + let mut partial = SumV2Partial::empty(dtypes.result); let (sum, is_overflow, is_empty) = decode_partial_scalar(scalar)?; - validate_sum_field_dtype(&sum, &partial.return_dtype)?; + validate_sum_field_dtype(&sum, dtypes.result)?; // Adding the parsed value to the zero state cannot overflow; treat a decimal value that // no longer fits its precision as an already-overflowed partial. - let overflowed = checked_add_sum_state(&mut partial.sum, &sum)?; + let overflowed = checked_add_sum_state(&mut partial.sum, dtypes.result, &sum)?; partial.is_overflow = is_overflow || overflowed; partial.is_empty = is_empty && !partial.is_overflow; Ok(partial) } - fn reduce_partials( + fn merge_partials( &self, - options: &Self::Options, - input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + mut acc: Self::Partial, + partial: Self::Partial, ) -> VortexResult { - // Seed from the first partial so an overflowed state keeps its last valid sum. - let mut partials = partials.into_iter(); - let Some(mut acc) = partials.next() else { - return SumV2Partial::empty(options, input_dtype); - }; - for partial in partials { - if acc.is_overflow { - break; - } - if partial.is_overflow { - acc.is_overflow = true; - acc.is_empty = false; - continue; - } - if partial.is_empty { - continue; + if acc.is_overflow { + return Ok(acc); + } + if partial.is_overflow { + // An overflowed state keeps its last valid sum, so an empty accumulator adopts the + // incoming sum rather than reporting zero. + if acc.is_empty { + acc.sum = partial.sum; } - acc.is_overflow = checked_add_sum_states(&mut acc.sum, &partial.sum)?; + acc.is_overflow = true; acc.is_empty = false; + return Ok(acc); } + if partial.is_empty { + return Ok(acc); + } + acc.is_overflow = checked_add_sum_states(&mut acc.sum, dtypes.result, &partial.sum)?; + acc.is_empty = false; Ok(acc) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::struct_( - sum_v2_partial_dtype(partial.return_dtype.clone()), + dtypes.partial.clone(), vec![ - sum_state_scalar(partial, Nullability::NonNullable), + sum_state_scalar(partial, dtypes.result, Nullability::NonNullable), Scalar::bool(partial.is_overflow, Nullability::NonNullable), Scalar::bool(partial.is_empty, Nullability::NonNullable), ], )) } - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) } fn try_accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - if partial.skip_nans || !matches!(&partial.sum, SumState::Float(_)) { + if options.skip_nans || !matches!(&partial.sum, SumState::Float(_)) { return Ok(false); } @@ -193,6 +214,8 @@ impl AggregateFnVTable for SumV2 { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -205,7 +228,7 @@ impl AggregateFnVTable for SumV2 { if !constant.scalar().is_null() && !constant.is_empty() { partial.is_empty = false; } - if partial.skip_nans + if options.skip_nans && constant .scalar() .as_primitive_opt() @@ -214,13 +237,14 @@ impl AggregateFnVTable for SumV2 { return Ok(()); } if let Some(product) = - multiply_constant(constant.scalar(), constant.len(), &partial.return_dtype)? + multiply_constant(constant.scalar(), constant.len(), dtypes.result)? { if product.is_null() { partial.is_overflow = true; partial.is_empty = false; } else { - partial.is_overflow = checked_add_sum_state(&mut partial.sum, &product)?; + partial.is_overflow = + checked_add_sum_state(&mut partial.sum, dtypes.result, &product)?; partial.is_empty = false; } } @@ -231,10 +255,12 @@ impl AggregateFnVTable for SumV2 { let result = match batch { Columnar::Canonical(canonical) => match canonical { Canonical::Primitive(array) => { - accumulate_primitive(&mut partial.sum, array, ctx, partial.skip_nans) + accumulate_primitive(&mut partial.sum, array, ctx, options.skip_nans) } Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx), - Canonical::Decimal(array) => accumulate_decimal(&mut partial.sum, array, ctx), + Canonical::Decimal(array) => { + accumulate_decimal(&mut partial.sum, dtypes.result, array, ctx) + } _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -250,7 +276,12 @@ impl AggregateFnVTable for SumV2 { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { if let Some(partials) = partials.as_opt::() { return finalize_struct(partials); } @@ -263,11 +294,20 @@ impl AggregateFnVTable for SumV2 { sum.mask(is_invalid.not()?) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { if partial.is_overflow || partial.is_empty { - return Ok(Scalar::null(partial.return_dtype.as_nullable())); + return Ok(Scalar::null(dtypes.result.as_nullable())); } - Ok(sum_state_scalar(partial, Nullability::Nullable)) + Ok(sum_state_scalar( + partial, + dtypes.result, + Nullability::Nullable, + )) } } @@ -292,26 +332,19 @@ fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult { /// In-memory state for SumV2 accumulation. pub struct SumV2Partial { - return_dtype: DType, sum: SumState, is_overflow: bool, is_empty: bool, - skip_nans: bool, } impl SumV2Partial { - /// The state of a group with no accumulated values, or an error for unsupported input dtypes. - fn empty(options: &NumericalAggregateOpts, input_dtype: &DType) -> VortexResult { - let return_dtype = SumV2 - .return_dtype(options, input_dtype) - .ok_or_else(|| vortex_err!("Unsupported sum_v2 dtype: {}", input_dtype))?; - Ok(Self { - sum: make_zero_state(&return_dtype), - return_dtype, + /// The state of a group with no accumulated values. + fn empty(return_dtype: &DType) -> Self { + Self { + sum: make_zero_state(return_dtype), is_overflow: false, is_empty: true, - skip_nans: options.skip_nans, - }) + } } } @@ -368,7 +401,11 @@ fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult< Ok(()) } -fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { +fn checked_add_sum_state( + state: &mut SumState, + return_dtype: &DType, + other: &Scalar, +) -> VortexResult { Ok(match state { SumState::Unsigned(sum) => checked_add_u64(sum, u64::try_from(other)?), SumState::Signed(sum) => checked_add_i64(sum, i64::try_from(other)?), @@ -376,7 +413,10 @@ fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { + SumState::Decimal(value) => { + let dtype = return_dtype + .as_decimal_opt() + .vortex_expect("decimal sum result dtype"); let other = DecimalValue::try_from(other)?; match value.checked_add(&other) { Some(result) if result.fits_in_precision(*dtype) => { @@ -408,12 +448,22 @@ fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields { ) } -fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar { +fn sum_state_scalar( + partial: &SumV2Partial, + return_dtype: &DType, + nullability: Nullability, +) -> Scalar { match &partial.sum { SumState::Unsigned(value) => Scalar::primitive(*value, nullability), SumState::Signed(value) => Scalar::primitive(*value, nullability), SumState::Float(value) => Scalar::primitive(*value, nullability), - SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), + SumState::Decimal(value) => Scalar::decimal( + *value, + *return_dtype + .as_decimal_opt() + .vortex_expect("decimal sum result dtype"), + nullability, + ), } } diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 160a5640616..6d7701e5b49 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -21,6 +21,7 @@ use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::array_session; use crate::arrays::BoolArray; @@ -290,9 +291,14 @@ fn merge_from_overflow_is_absorbing() -> VortexResult<()> { ); // The overflow flag survives the scalar round trip. - let propagated = SumV2.partial_from_scalar(&options, &dtype, overflow_partial)?; - assert!(SumV2.is_saturated(&propagated)); - assert!(SumV2.finalize_scalar(&propagated)?.is_null()); + let dtypes = OwnedAggregateDTypes::try_new(&SumV2, &options, dtype)?; + let propagated = SumV2.partial_from_scalar(&options, dtypes.borrow(), overflow_partial)?; + assert!(SumV2.is_saturated(&options, dtypes.borrow(), &propagated)); + assert!( + SumV2 + .finalize_scalar(&options, dtypes.borrow(), &propagated)? + .is_null() + ); Ok(()) } @@ -372,7 +378,10 @@ fn finalize_struct_applies_partial_and_struct_validity( )? .into_array(); - let result = SumV2.finalize(partials)?; + let options = NumericalAggregateOpts::default(); + let dtypes = + OwnedAggregateDTypes::try_new(&SumV2, &options, DType::Primitive(PType::I64, Nullable))?; + let result = SumV2.finalize(&options, dtypes.borrow(), partials)?; let expected = PrimitiveArray::from_option_iter(expected).into_array(); assert_arrays_eq!( &result, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index b775b58cdc5..31457051f7a 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -38,6 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -131,10 +132,18 @@ impl AggregateFnVTable for UncompressedSizeInBytes { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(0) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -143,29 +152,41 @@ impl AggregateFnVTable for UncompressedSizeInBytes { .vortex_expect("uncompressed_size_in_bytes partial should not be null")) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - partials.into_iter().try_fold(0u64, |acc, partial| { - acc.checked_add(partial) - .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) - }) + first + .checked_add(second) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -182,12 +203,22 @@ impl AggregateFnVTable for UncompressedSizeInBytes { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -345,6 +376,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::array_session; @@ -687,10 +719,12 @@ mod tests { #[test] fn state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let dtypes = OwnedAggregateDTypes::try_new(&UncompressedSizeInBytes, &EmptyOptions, dtype)?; - let state = UncompressedSizeInBytes.reduce_partials(&EmptyOptions, &dtype, [5, 3])?; + let state = + UncompressedSizeInBytes.reduce_partials(&EmptyOptions, dtypes.borrow(), [5, 3])?; - let result = UncompressedSizeInBytes.to_scalar(&state)?; + let result = UncompressedSizeInBytes.to_scalar(&EmptyOptions, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs index e56ed84a754..5e5d7ac549c 100644 --- a/vortex-array/src/aggregate_fn/foreign.rs +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -12,6 +12,7 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; @@ -77,34 +78,55 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { None } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, _scalar: Scalar, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - _partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + _first: Self::Partial, + _second: Self::Partial, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn is_saturated(&self, _state: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _state: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -112,11 +134,21 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn finalize(&self, _states: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _states: ArrayRef, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn finalize_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } } diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index d74aabefde9..93b8e7a5605 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -69,6 +69,7 @@ mod tests { use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; @@ -114,34 +115,55 @@ mod tests { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(()) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, _scalar: Scalar, ) -> VortexResult { Ok(()) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - _partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + _first: Self::Partial, + _second: Self::Partial, ) -> VortexResult { Ok(()) } - fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> VortexResult { vortex_panic!("TestAgg is for serde tests only"); } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { true } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -149,11 +171,21 @@ mod tests { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> VortexResult { vortex_panic!("TestAgg is for serde tests only"); } } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 4816a78c73e..23c89471a91 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -10,6 +10,7 @@ use std::hash::Hash; use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_proto::expr as pb; use vortex_session::VortexSession; @@ -23,6 +24,83 @@ use crate::aggregate_fn::AggregateFnSatisfaction; use crate::dtype::DType; use crate::scalar::Scalar; +/// Resolved dtypes of one aggregate function bound to its options and input. +/// +/// Accumulators resolve these once and lend them to every execution method, so partial states +/// only hold accumulated values. Combined aggregates resolve a separate set for each child. +#[derive(Clone, Copy, Debug)] +pub struct AggregateDTypes<'a> { + /// The dtype of the values being aggregated. + pub input: &'a DType, + /// The dtype of the partial state scalar, as reported by [`AggregateFnVTable::partial_dtype`]. + pub partial: &'a DType, + /// The dtype of the final aggregate result, as reported by [`AggregateFnVTable::return_dtype`]. + pub result: &'a DType, +} + +/// Owned [`AggregateDTypes`], resolved once from an aggregate's options and input dtype. +#[derive(Clone, Debug)] +pub struct OwnedAggregateDTypes { + input: DType, + partial: DType, + result: DType, +} + +impl OwnedAggregateDTypes { + /// Resolve the partial and result dtypes of `vtable` bound to `options` over `input`. + /// + /// Fails if the aggregate cannot be applied to `input`. + pub fn try_new( + vtable: &V, + options: &V::Options, + input: DType, + ) -> VortexResult { + let result = vtable.return_dtype(options, &input).ok_or_else(|| { + vortex_err!( + "Aggregate function {} cannot be applied to dtype {}", + vtable.id(), + input + ) + })?; + let partial = vtable.partial_dtype(options, &input).ok_or_else(|| { + vortex_err!( + "Aggregate function {} cannot be applied to dtype {}", + vtable.id(), + input + ) + })?; + Ok(Self { + input, + partial, + result, + }) + } + + /// The dtype of the values being aggregated. + pub fn input(&self) -> &DType { + &self.input + } + + /// The dtype of the partial state scalar. + pub fn partial(&self) -> &DType { + &self.partial + } + + /// The dtype of the final aggregate result. + pub fn result(&self) -> &DType { + &self.result + } + + /// Lend the dtypes to an execution method. + pub fn borrow(&self) -> AggregateDTypes<'_> { + AggregateDTypes { + input: &self.input, + partial: &self.partial, + result: &self.result, + } + } +} + /// Defines the interface for aggregate function vtables. /// /// This trait is non-object-safe and allows the implementer to make use of associated types @@ -32,6 +110,10 @@ use crate::scalar::Scalar; /// The [`AggregateFnVTable`] trait should be implemented for a struct that holds global data across /// all instances of the aggregate. In almost all cases, this struct will be an empty unit /// struct, since most aggregates do not require any global state. +/// +/// Execution methods receive the options and the resolved [`AggregateDTypes`] of the aggregate +/// they operate on, so partial states only hold accumulated values. Callers must pass the same +/// options and dtypes for the whole lifetime of a partial state. pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Options for this aggregate function. type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; @@ -103,58 +185,102 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns `None` if the aggregate function cannot be applied to the input dtype. fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option; - /// Parse a partial scalar into the typed partial accumulator state. + /// The partial state of a group with no accumulated values. /// - /// The scalar must have the DType specified by `partial_dtype` for the given options and - /// input dtype; this is the inverse of [`to_scalar`]. Partial scalars are produced by - /// aggregate kernels, cached statistics, and other accumulators' [`to_scalar`]. + /// This is the identity of [`merge_partials`]: merging it with any partial state, on either + /// side, yields that state. + /// + /// [`merge_partials`]: AggregateFnVTable::merge_partials + fn empty_partial( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + ) -> VortexResult; + + /// Parse a partial scalar into the typed partial state. + /// + /// The scalar must have dtype `dtypes.partial`; this is the inverse of [`to_scalar`]. Partial + /// scalars are produced by aggregate kernels, cached statistics, and other accumulators' + /// [`to_scalar`]. /// /// Implementations should only parse the scalar here; combining states belongs in - /// [`reduce_partials`]. + /// [`merge_partials`]. /// /// [`to_scalar`]: AggregateFnVTable::to_scalar - /// [`reduce_partials`]: AggregateFnVTable::reduce_partials + /// [`merge_partials`]: AggregateFnVTable::merge_partials fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult; - /// Reduce a sequence of partial states into a single partial state. + /// Merge two partial states into one. + /// + /// `first` accumulated the input preceding `second`'s. Order-dependent aggregates (e.g. + /// first/last or is_sorted) rely on this, so the merge need not be commutative. Merging with + /// [`empty_partial`] on either side is the identity. /// - /// Partials must be reduced in iteration order, since some aggregates (e.g. first/last or - /// is_sorted) are order-dependent. Reducing an empty sequence returns the identity: the - /// partial state of a group with no accumulated values. + /// [`empty_partial`]: AggregateFnVTable::empty_partial + fn merge_partials( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult; + + /// Reduce a sequence of partial states in iteration order, starting from [`empty_partial`]. + /// + /// [`empty_partial`]: AggregateFnVTable::empty_partial fn reduce_partials( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypes<'_>, partials: impl IntoIterator, - ) -> VortexResult; + ) -> VortexResult { + partials + .into_iter() + .try_fold(self.empty_partial(options, dtypes)?, |acc, partial| { + self.merge_partials(options, dtypes, acc, partial) + }) + } - /// Convert the partial state into a partial scalar. + /// Convert the partial state into a partial scalar of dtype `dtypes.partial`. /// - /// The returned scalar must have the same DType as specified by `partial_dtype` for the - /// options and input dtype used to construct the state. This is the inverse of - /// [`partial_from_scalar`]. + /// This is the inverse of [`partial_from_scalar`]: parsing the returned scalar must + /// reconstruct a state that behaves identically under accumulation, merging, saturation, and + /// finalization. /// /// [`partial_from_scalar`]: AggregateFnVTable::partial_from_scalar - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult; + fn to_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult; - /// Is the partial accumulator state is "saturated", i.e. has it reached a state where the - /// final result is fully determined. - fn is_saturated(&self, state: &Self::Partial) -> bool; + /// Is the partial state "saturated", i.e. has it reached a state where the final result is + /// fully determined. + fn is_saturated( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool; /// Try to accumulate the raw array before decompression. /// - /// Returns `true` if the array was handled, `false` to fall through to - /// the default kernel dispatch and canonicalization path. + /// Returns `true` if the array was handled, `false` to fall through to the default kernel + /// dispatch and canonicalization path. When returning `false`, the partial state must be + /// unchanged, since the same batch is then accumulated by the fallback path. /// /// This is useful for aggregates that only depend on array metadata (e.g., validity) /// rather than the encoded data, avoiding unnecessary decompression. fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, _state: &mut Self::Partial, _batch: &ArrayRef, _ctx: &mut ExecutionCtx, @@ -162,25 +288,33 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { Ok(false) } - /// Accumulate a new canonical array into the accumulator state. + /// Accumulate a new canonical array into the partial state. fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, state: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, ) -> VortexResult<()>; - /// Finalize an array of accumulator states into an array of aggregate results. + /// Finalize an array of partial states into an array of aggregate results. /// - /// The provides `states` array has dtype as specified by `state_dtype`, the result array - /// must have dtype as specified by `return_dtype`. - fn finalize(&self, states: ArrayRef) -> VortexResult; + /// The `states` array has dtype `dtypes.partial`; the result must have dtype `dtypes.result`. + fn finalize( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + states: ArrayRef, + ) -> VortexResult; - /// Finalize a scalar accumulator state into an aggregate result. - /// - /// The provided `state` has dtype as specified by `state_dtype`, the result scalar must have - /// dtype as specified by `return_dtype`. - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult; + /// Finalize a partial state into an aggregate result of dtype `dtypes.result`. + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult; } #[derive(Clone, Debug, PartialEq, Eq, Hash)] diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 3599b3444bb..68bd20a6ed9 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -16,6 +16,7 @@ use std::num::NonZeroU32; use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateDTypes; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::dtype::DType; @@ -297,6 +298,15 @@ impl AggregateFnVTable for BloomFilter { self.return_dtype(options, input_dtype) } + /// Returns an empty Bloom filter with all blocks zero-initialized. + fn empty_partial( + &self, + options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(BloomPartial::from(options)) + } + /// Parses a serialized filter into a partial with the configured block count. /// /// A null scalar is an empty filter. This assumes that `scalar` was created using the same @@ -305,7 +315,7 @@ impl AggregateFnVTable for BloomFilter { fn partial_from_scalar( &self, options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { let mut partial = BloomPartial::from(options); @@ -321,25 +331,27 @@ impl AggregateFnVTable for BloomFilter { Ok(partial) } - /// Reduces partials by OR-ing their blocks together; an empty sequence is an empty filter - /// with all blocks zero-initialized. - fn reduce_partials( + /// Merges two filters by OR-ing their blocks together. + fn merge_partials( &self, - options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = BloomPartial::from(options); - for partial in partials { - acc.union(&partial)?; - } - Ok(acc) + first.union(&second)?; + Ok(first) } /// Returns the non-nullable binary representation of a bloom filter /// /// Basically turns each block into a single byte sequence. - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { let bytes: Vec = partial.serialize(); Ok(Scalar::binary(bytes, Nullability::NonNullable)) } @@ -347,12 +359,19 @@ impl AggregateFnVTable for BloomFilter { /// Returns true if all the blocks are full. /// /// When a bloom filter is saturated, it cannot rule out any values. - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> bool { partial.is_saturated() } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -366,12 +385,22 @@ impl AggregateFnVTable for BloomFilter { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -395,6 +424,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::OwnedAggregateDTypes; use vortex_array::test_harness::check_metadata; use super::*; @@ -430,40 +460,57 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { Ok(bloom_filter) } + /// Aggregate dtypes for a binary input; the Bloom filter's dtypes do not depend on options. + fn binary_dtypes(options: &BloomOptions) -> VortexResult { + OwnedAggregateDTypes::try_new( + &BloomFilter, + options, + DType::Binary(Nullability::NonNullable), + ) + } + #[test] fn saturation_false_when_empty() -> VortexResult<()> { let options = BloomOptions::default(); - let partial = - BloomFilter.reduce_partials(&options, &DType::Binary(Nullability::NonNullable), [])?; - assert!(!BloomFilter.is_saturated(&partial)); + let dtypes = binary_dtypes(&options)?; + let partial = BloomFilter.reduce_partials(&options, dtypes.borrow(), [])?; + assert!(!BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); Ok(()) } #[test] - fn saturation_true_when_every_block_is_full() { + fn saturation_true_when_every_block_is_full() -> VortexResult<()> { + let options = BloomOptions::default(); + let dtypes = binary_dtypes(&options)?; let blocks = vec![[u32::MAX; 8]; 4]; let partial = BloomPartial::from(blocks); - assert!(BloomFilter.is_saturated(&partial)); + assert!(BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); + Ok(()) } #[test] fn mismatched_block_counts_are_rejected() -> VortexResult<()> { - let dtype = DType::Binary(Nullability::NonNullable); let smaller = BloomOptions::new(NonZeroU32::new(4).unwrap(), HashFn::XxHash3_64); + let dtypes = binary_dtypes(&smaller)?; let bigger = BloomPartial::from(&BloomOptions::default()); - let bigger_scalar = BloomFilter.to_scalar(&bigger)?; + let bigger_scalar = + BloomFilter.to_scalar(&BloomOptions::default(), dtypes.borrow(), &bigger)?; assert!( BloomFilter - .partial_from_scalar(&smaller, &dtype, bigger_scalar) + .partial_from_scalar(&smaller, dtypes.borrow(), bigger_scalar) .is_err(), "parsing a partial built with a different blocks_count must fail loudly, not corrupt state" ); assert!( BloomFilter - .reduce_partials(&smaller, &dtype, [BloomPartial::from(&smaller), bigger]) + .reduce_partials( + &smaller, + dtypes.borrow(), + [BloomPartial::from(&smaller), bigger] + ) .is_err(), "reducing partials built with different blocks_count must fail loudly, not corrupt state" ); @@ -490,11 +537,9 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { expected.insert(i.to_le_bytes()); } - let partial = BloomFilter.reduce_partials( - &options, - &DType::Binary(Nullability::NonNullable), - [partial, secondary_partial], - )?; + let dtypes = binary_dtypes(&options)?; + let partial = + BloomFilter.reduce_partials(&options, dtypes.borrow(), [partial, secondary_partial])?; assert!( partial == expected, diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs index 6f2c2312878..b72ae7a6f84 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs @@ -93,14 +93,18 @@ impl TryFrom for HashFn { /// ```rust /// use vortex_array::dtype::{DType, Nullability}; /// use vortex_layout::layouts::zoned::aggregates::bloom_filter::{BloomFilter, BloomOptions}; -/// use vortex_array::aggregate_fn::AggregateFnVTable; +/// use vortex_array::aggregate_fn::{AggregateFnVTable, OwnedAggregateDTypes}; /// /// let filter = BloomFilter {}; +/// let options = BloomOptions::default(); +/// let dtypes = OwnedAggregateDTypes::try_new( +/// &filter, +/// &options, +/// DType::Binary(Nullability::NonNullable), +/// ) +/// .expect("valid input dtype"); /// let mut zone = filter -/// .empty_partial( -/// &BloomOptions::default(), -/// &DType::Binary(Nullability::NonNullable), -/// ) +/// .empty_partial(&options, dtypes.borrow()) /// .expect("valid partial"); /// /// zone.insert(b"Denmark"); @@ -109,9 +113,9 @@ impl TryFrom for HashFn { /// assert_eq!(zone.contains(b"Japan"), false); /// assert_eq!(zone.contains(b"Brazil"), false); /// ``` +#[derive(PartialEq, Eq)] pub struct BloomPartial { blocks: Vec<[u32; 8]>, - hash_fn: HashFn, } impl BloomPartial { @@ -238,31 +242,20 @@ impl BloomPartial { /// start an empty partial from [`super::BloomFilter`]. impl From<&BloomOptions> for BloomPartial { fn from(options: &BloomOptions) -> Self { - Self { - blocks: vec![[0u32; 8]; options.blocks_count.get() as usize], - hash_fn: options.hash_fn, + // The hash function lives in the options, not the partial: adding a variant must + // revisit how partials are hashed, so match exhaustively here. + match options.hash_fn { + HashFn::XxHash3_64 => Self { + blocks: vec![[0u32; 8]; options.blocks_count.get() as usize], + }, } } } -impl PartialEq for BloomPartial { - fn eq(&self, other: &Self) -> bool { - // Currently, the Bloom filter only supports one hash function, - // so two partials with the same blocks are equal. - // If the filter supports more hash functions in the future, - // this would no longer be true, because the same blocks could represent - // different values. - self.blocks == other.blocks && self.hash_fn == other.hash_fn - } -} - #[cfg(test)] impl From> for BloomPartial { fn from(value: Vec<[u32; 8]>) -> Self { - BloomPartial { - blocks: value, - hash_fn: HashFn::XxHash3_64, // Default. Only used for tests. - } + BloomPartial { blocks: value } } } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/serde.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/serde.rs index f233cb96d7a..d1a8841b8be 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/serde.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/serde.rs @@ -14,7 +14,6 @@ use vortex_error::vortex_ensure; use super::BLOCK_SIZE; use super::BYTES_PER_SPLIT; use super::BloomPartial; -use crate::layouts::zoned::aggregates::bloom_filter::HashFn; impl BloomPartial { /// Deserialize a partial from its byte representation. @@ -50,10 +49,7 @@ impl BloomPartial { "bloom blocks length must be non-zero and lower than u32::MAX", ); - Ok(BloomPartial { - blocks, - hash_fn: HashFn::XxHash3_64, // Default option - }) + Ok(BloomPartial { blocks }) } /// Serialize partial filter into its bytes format (little endian) diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index f5d98e7b4e2..e6132ee14aa 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/vortex-spatial/src/aggregate_fn/aabb.rs @@ -8,6 +8,7 @@ use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateDTypes; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; @@ -153,10 +154,18 @@ impl AggregateFnVTable for GeometryAabb { self.return_dtype(options, input_dtype) } + fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + ) -> VortexResult { + Ok(AabbPartial { rect: None }) + } + fn partial_from_scalar( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypes<'_>, scalar: Scalar, ) -> VortexResult { // A null box is an empty group's AABB. @@ -165,35 +174,45 @@ impl AggregateFnVTable for GeometryAabb { }) } - fn reduce_partials( + fn merge_partials( &self, _options: &Self::Options, - _input_dtype: &DType, - partials: impl IntoIterator, + _dtypes: AggregateDTypes<'_>, + mut first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - let mut acc = AabbPartial { rect: None }; - for partial in partials { - if let Some(rect) = partial.rect { - acc.merge(rect); - } + if let Some(rect) = second.rect { + first.merge(rect); } - Ok(acc) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match partial.rect { Some(rect) => rect_to_storage(rect), - None => Scalar::null(aabb_dtype()), + None => Scalar::null(dtypes.partial.clone()), }) } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + _partial: &Self::Partial, + ) -> bool { // An AABB can always grow, so it is never saturated. false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -225,13 +244,23 @@ impl AggregateFnVTable for GeometryAabb { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypes<'_>, + partials: ArrayRef, + ) -> VortexResult { // The stored partial is already the AABB struct, so finalizing is the identity. Ok(partials) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypes<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -244,6 +273,7 @@ mod tests { use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::OwnedAggregateDTypes; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -400,16 +430,17 @@ mod tests { #[test] fn reduce_partials_unions_boxes() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let dtypes = OwnedAggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; let bbox = |xmin, ymin, xmax, ymax| AabbPartial { rect: Some(SpatialRect::new((xmin, ymin), (xmax, ymax))), }; let reduced = GeometryAabb.reduce_partials( &EmptyOptions, - &dtype, + dtypes.borrow(), [bbox(0.0, 0.0, 1.0, 1.0), bbox(5.0, -2.0, 7.0, 3.0)], )?; assert_eq!( - aabb(&GeometryAabb.to_scalar(&reduced)?)?, + aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &reduced)?)?, (0.0, -2.0, 7.0, 3.0) ); Ok(()) @@ -419,13 +450,15 @@ mod tests { #[test] fn reduce_partials_ignores_empty() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let dtypes = OwnedAggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; let empty = AabbPartial { rect: None }; let value = AabbPartial { rect: Some(SpatialRect::new((0.0, 0.0), (1.0, 1.0))), }; - let reduced = GeometryAabb.reduce_partials(&EmptyOptions, &dtype, [value, empty])?; + let reduced = + GeometryAabb.reduce_partials(&EmptyOptions, dtypes.borrow(), [value, empty])?; assert_eq!( - aabb(&GeometryAabb.to_scalar(&reduced)?)?, + aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &reduced)?)?, (0.0, 0.0, 1.0, 1.0) ); Ok(()) From 2c3a95d8cf67ea41fa750614dae340116dcde822 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 10 Sep 2026 17:45:33 +0000 Subject: [PATCH 3/6] Name the owned aggregate dtypes AggregateDTypes and drop reduce_partials Renames the resolved-dtype pair so the owned type carries the plain name and the borrowed view is suffixed: - `OwnedAggregateDTypes` -> `AggregateDTypes` - `AggregateDTypes<'a>` -> `AggregateDTypesRef<'a>` Its fields and accessors keep the names they had on the accumulators they were lifted from: `dtype`, `return_dtype`, and `partial_dtype`. Also removes `AggregateFnVTable::reduce_partials`. It was a provided fold over `empty_partial` and `merge_partials` with no non-test callers, so the tests now merge directly, which exercises the binary operation rather than the fold wrapper. Merging with the empty partial is the identity, so a two-element reduction is exactly one `merge_partials`. Signed-off-by: "Claude" Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- vortex-array/src/aggregate_fn/accumulator.rs | 46 +++--- .../src/aggregate_fn/accumulator_grouped.rs | 24 ++-- vortex-array/src/aggregate_fn/combined.rs | 38 ++--- .../src/aggregate_fn/fns/all_nan/mod.rs | 20 +-- .../aggregate_fn/fns/all_non_distinct/mod.rs | 18 +-- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 20 +-- .../src/aggregate_fn/fns/all_non_null/mod.rs | 20 +-- .../src/aggregate_fn/fns/all_null/mod.rs | 20 +-- .../src/aggregate_fn/fns/bounded_max/mod.rs | 26 ++-- .../src/aggregate_fn/fns/bounded_min/mod.rs | 20 +-- .../src/aggregate_fn/fns/count/mod.rs | 28 ++-- .../src/aggregate_fn/fns/first/mod.rs | 29 ++-- .../src/aggregate_fn/fns/is_constant/mod.rs | 28 ++-- .../src/aggregate_fn/fns/is_sorted/mod.rs | 26 ++-- vortex-array/src/aggregate_fn/fns/last/mod.rs | 29 ++-- vortex-array/src/aggregate_fn/fns/max/mod.rs | 30 ++-- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 12 +- vortex-array/src/aggregate_fn/fns/min/mod.rs | 30 ++-- .../src/aggregate_fn/fns/min_max/bool.rs | 4 +- .../src/aggregate_fn/fns/min_max/decimal.rs | 4 +- .../src/aggregate_fn/fns/min_max/extension.rs | 4 +- .../src/aggregate_fn/fns/min_max/mod.rs | 38 ++--- .../src/aggregate_fn/fns/min_max/primitive.rs | 4 +- .../src/aggregate_fn/fns/min_max/varbin.rs | 4 +- .../src/aggregate_fn/fns/nan_count/mod.rs | 24 ++-- .../src/aggregate_fn/fns/null_count/mod.rs | 20 +-- .../src/aggregate_fn/fns/sum/decimal.rs | 22 +-- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 58 ++++---- .../src/aggregate_fn/fns/sum_v2/mod.rs | 44 +++--- .../src/aggregate_fn/fns/sum_v2/tests.rs | 6 +- .../fns/uncompressed_size_in_bytes/mod.rs | 25 ++-- vortex-array/src/aggregate_fn/foreign.rs | 18 +-- vortex-array/src/aggregate_fn/proto.rs | 18 +-- vortex-array/src/aggregate_fn/vtable.rs | 133 ++++++++---------- .../zoned/aggregates/bloom_filter/mod.rs | 35 ++--- .../aggregates/bloom_filter/partial/mod.rs | 4 +- vortex-spatial/src/aggregate_fn/aabb.rs | 42 +++--- 37 files changed, 483 insertions(+), 488 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index c2b88ad6e6a..7a15d7a01f8 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -11,10 +11,10 @@ use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::columnar::AnyColumnar; use crate::dtype::DType; @@ -36,7 +36,7 @@ pub struct Accumulator { /// Type-erased aggregate function used for kernel dispatch. aggregate_fn: AggregateFnRef, /// The input, partial, and result dtypes lent to every vtable call. - dtypes: OwnedAggregateDTypes, + dtypes: AggregateDTypes, /// The partial state of the accumulator, updated after each accumulate/merge call. /// /// `None` is the empty-group state; a live partial is only materialized when a batch is @@ -46,7 +46,7 @@ pub struct Accumulator { impl Accumulator { pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult { - let dtypes = OwnedAggregateDTypes::try_new(&vtable, &options, dtype)?; + let dtypes = AggregateDTypes::try_new(&vtable, &options, dtype)?; let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); Ok(Self { @@ -85,7 +85,7 @@ impl Accumulator { Ok(()) } - /// Parse a partial scalar of dtype `dtypes.partial` and merge it into the current state. + /// Parse a partial scalar of dtype `dtypes.partial_dtype` and merge it into the current state. /// /// Both steps go through the typed vtable of `V`, so they inline into one monomorphized call. fn fold_partial_scalar(&mut self, scalar: Scalar) -> VortexResult<()> { @@ -158,9 +158,9 @@ impl DynAccumulator for Accumulator { } vortex_ensure!( - batch.dtype() == self.dtypes.input(), + batch.dtype() == self.dtypes.dtype(), "Input DType mismatch: expected {}, got {}", - self.dtypes.input(), + self.dtypes.dtype(), batch.dtype() ); @@ -169,18 +169,20 @@ impl DynAccumulator for Accumulator { if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - let partial = if partial.dtype() == self.dtypes.partial() { + let partial = if partial.dtype() == self.dtypes.partial_dtype() { partial } else { vortex_ensure!( - partial.dtype().eq_ignore_nullability(self.dtypes.partial()), + partial + .dtype() + .eq_ignore_nullability(self.dtypes.partial_dtype()), "Aggregate {} read legacy stat {} with dtype {}, expected {}", self.aggregate_fn, stat, partial.dtype(), - self.dtypes.partial(), + self.dtypes.partial_dtype(), ); - partial.cast(self.dtypes.partial())? + partial.cast(self.dtypes.partial_dtype())? }; self.fold_partial_scalar(partial)?; return Ok(()); @@ -201,10 +203,10 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)? { vortex_ensure!( - result.dtype() == self.dtypes.partial(), + result.dtype() == self.dtypes.partial_dtype(), "Aggregate kernel returned {}, expected {}", result.dtype(), - self.dtypes.partial(), + self.dtypes.partial_dtype(), ); self.fold_partial_scalar(result)?; return Ok(()); @@ -238,10 +240,10 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)? { vortex_ensure!( - result.dtype() == self.dtypes.partial(), + result.dtype() == self.dtypes.partial_dtype(), "Aggregate kernel returned {}, expected {}", result.dtype(), - self.dtypes.partial(), + self.dtypes.partial_dtype(), ); self.fold_partial_scalar(result)?; return Ok(()); @@ -267,7 +269,7 @@ impl DynAccumulator for Accumulator { ); }; vortex_ensure!( - other.options == self.options && other.dtypes.input() == self.dtypes.input(), + other.options == self.options && other.dtypes.dtype() == self.dtypes.dtype(), "Cannot merge {} accumulators with different options or input dtypes", self.aggregate_fn, ); @@ -279,10 +281,10 @@ impl DynAccumulator for Accumulator { fn combine_partial_scalar(&mut self, partial: Scalar) -> VortexResult<()> { vortex_ensure!( - partial.dtype() == self.dtypes.partial(), + partial.dtype() == self.dtypes.partial_dtype(), "Partial DType mismatch for {}: expected {}, got {}", self.aggregate_fn, - self.dtypes.partial(), + self.dtypes.partial_dtype(), partial.dtype(), ); self.fold_partial_scalar(partial) @@ -315,9 +317,9 @@ impl DynAccumulator for Accumulator { #[cfg(debug_assertions)] { vortex_ensure!( - partial.dtype() == dtypes.partial, + partial.dtype() == dtypes.partial_dtype, "Aggregate returned incorrect DType on partial_scalar: expected {}, got {}", - dtypes.partial, + dtypes.partial_dtype, partial.dtype(), ); } @@ -337,9 +339,9 @@ impl DynAccumulator for Accumulator { }; vortex_ensure!( - result.dtype() == dtypes.result, + result.dtype() == dtypes.return_dtype, "Aggregate returned incorrect DType on final_scalar: expected {}, got {}", - dtypes.result, + dtypes.return_dtype, result.dtype(), ); @@ -465,7 +467,7 @@ mod tests { let acc = mean_f64_accumulator().expect("build accumulator"); let sum = Scalar::primitive(42.0f64, Nullability::Nullable); let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.dtypes.partial().clone(), vec![sum, count]) + Scalar::struct_(acc.dtypes.partial_dtype().clone(), vec![sum, count]) } /// Kernel registered for `(Dict, Combined)` fires in preference to diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 8483daf028e..1eedfdd6ce1 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -16,11 +16,11 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; @@ -188,7 +188,7 @@ pub struct GroupedAccumulator { /// Type-erased aggregate function used for kernel dispatch. aggregate_fn: AggregateFnRef, /// The input, partial, and result dtypes lent to every vtable call. - dtypes: OwnedAggregateDTypes, + dtypes: AggregateDTypes, /// The accumulated state for prior batches of groups. partials: Vec, } @@ -196,7 +196,7 @@ pub struct GroupedAccumulator { impl GroupedAccumulator { pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult { let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); - let dtypes = OwnedAggregateDTypes::try_new(&vtable, &options, dtype)?; + let dtypes = AggregateDTypes::try_new(&vtable, &options, dtype)?; Ok(Self { vtable, @@ -234,9 +234,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { ), }; vortex_ensure!( - elements_dtype.as_ref() == self.dtypes.input(), + elements_dtype.as_ref() == self.dtypes.dtype(), "Input DType mismatch: expected {}, got {}", - self.dtypes.input(), + self.dtypes.dtype(), elements_dtype ); @@ -258,7 +258,7 @@ impl DynGroupedAccumulator for GroupedAccumulator { if states.len() == 1 { return Ok(states.pop().vortex_expect("checked one partial")); } - Ok(ChunkedArray::try_new(states, self.dtypes.partial().clone())?.into_array()) + Ok(ChunkedArray::try_new(states, self.dtypes.partial_dtype().clone())?.into_array()) } fn finish(&mut self) -> VortexResult { @@ -268,9 +268,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { .finalize(&self.options, self.dtypes.borrow(), states)?; vortex_ensure!( - results.dtype() == self.dtypes.result(), + results.dtype() == self.dtypes.return_dtype(), "Return DType mismatch: expected {}, got {}", - self.dtypes.result(), + self.dtypes.return_dtype(), results.dtype() ); @@ -342,10 +342,10 @@ impl GroupedAccumulator { let mut accumulator = Accumulator::try_new( self.vtable.clone(), self.options.clone(), - self.dtypes.input().clone(), + self.dtypes.dtype().clone(), )?; let mut states = - builder_with_capacity_in(self.dtypes.partial(), grouped.len(), ctx.allocator()); + builder_with_capacity_in(self.dtypes.partial_dtype(), grouped.len(), ctx.allocator()); let group_ranges = grouped.group_ranges(ctx)?; let group_validity = grouped.group_validity(ctx)?; @@ -364,9 +364,9 @@ impl GroupedAccumulator { fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> { vortex_ensure!( - state.dtype() == self.dtypes.partial(), + state.dtype() == self.dtypes.partial_dtype(), "State DType mismatch: expected {}, got {}", - self.dtypes.partial(), + self.dtypes.partial_dtype(), state.dtype() ); self.partials.push(state); diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 38f432055e7..8e8c7d3cb62 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -20,10 +20,10 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::OwnedAggregateDTypes; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; @@ -88,20 +88,20 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone { /// Combine the finalized left and right results into the final aggregate. /// /// `dtypes` describes the combined aggregate; `left` and `right` already have their child's - /// result dtype. The returned array must have dtype `dtypes.result`. + /// result dtype. The returned array must have dtype `dtypes.return_dtype`. fn finalize( &self, options: &CombinedOptions, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, left: ArrayRef, right: ArrayRef, ) -> VortexResult; - /// Combine the finalized child scalars into a result of dtype `dtypes.result`. + /// Combine the finalized child scalars into a result of dtype `dtypes.return_dtype`. fn finalize_scalar( &self, options: &CombinedOptions, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, left_scalar: Scalar, right_scalar: Scalar, ) -> VortexResult; @@ -194,18 +194,18 @@ impl AggregateFnVTable for Combined { fn empty_partial( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - self.new_child_accumulators(options, dtypes.input) + self.new_child_accumulators(options, dtypes.dtype) } fn partial_from_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { - let (mut left, mut right) = self.new_child_accumulators(options, dtypes.input)?; + let (mut left, mut right) = self.new_child_accumulators(options, dtypes.dtype)?; // A null partial represents an empty group and parses to empty child accumulators. if !scalar.is_null() { let s = scalar.as_struct(); @@ -226,7 +226,7 @@ impl AggregateFnVTable for Combined { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, (mut left, mut right): Self::Partial, (mut other_left, mut other_right): Self::Partial, ) -> VortexResult { @@ -240,13 +240,13 @@ impl AggregateFnVTable for Combined { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { let l_scalar = partial.0.partial_scalar()?; let r_scalar = partial.1.partial_scalar()?; Ok(Scalar::struct_( - dtypes.partial.clone(), + dtypes.partial_dtype.clone(), vec![l_scalar, r_scalar], )) } @@ -254,7 +254,7 @@ impl AggregateFnVTable for Combined { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { partial.0.is_saturated() && partial.1.is_saturated() @@ -268,7 +268,7 @@ impl AggregateFnVTable for Combined { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -281,7 +281,7 @@ impl AggregateFnVTable for Combined { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -292,15 +292,15 @@ impl AggregateFnVTable for Combined { fn finalize( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, states: ArrayRef, ) -> VortexResult { let l_field = states.get_item(FieldName::from(self.0.left_name()))?; let r_field = states.get_item(FieldName::from(self.0.right_name()))?; let left = self.0.left(); let right = self.0.right(); - let l_dtypes = OwnedAggregateDTypes::try_new(&left, &options.0, dtypes.input.clone())?; - let r_dtypes = OwnedAggregateDTypes::try_new(&right, &options.1, dtypes.input.clone())?; + let l_dtypes = AggregateDTypes::try_new(&left, &options.0, dtypes.dtype.clone())?; + let r_dtypes = AggregateDTypes::try_new(&right, &options.1, dtypes.dtype.clone())?; let l_finalized = left.finalize(&options.0, l_dtypes.borrow(), l_field)?; let r_finalized = right.finalize(&options.1, r_dtypes.borrow(), r_field)?; BinaryCombined::finalize(&self.0, options, dtypes, l_finalized, r_finalized) @@ -309,7 +309,7 @@ impl AggregateFnVTable for Combined { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { let l_scalar = partial.0.final_scalar()?; diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index 743b9fe8064..1dccd337e22 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -9,7 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -64,7 +64,7 @@ impl AggregateFnVTable for AllNan { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } @@ -72,7 +72,7 @@ impl AggregateFnVTable for AllNan { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) @@ -81,7 +81,7 @@ impl AggregateFnVTable for AllNan { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -91,7 +91,7 @@ impl AggregateFnVTable for AllNan { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) @@ -100,7 +100,7 @@ impl AggregateFnVTable for AllNan { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !*partial @@ -109,7 +109,7 @@ impl AggregateFnVTable for AllNan { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -126,7 +126,7 @@ impl AggregateFnVTable for AllNan { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -149,7 +149,7 @@ impl AggregateFnVTable for AllNan { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -158,7 +158,7 @@ impl AggregateFnVTable for AllNan { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index 7be5fa20b67..92682e5a545 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -38,7 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -151,7 +151,7 @@ impl AggregateFnVTable for AllNonDistinct { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(AllNonDistinctPartial { all_non_distinct: true, @@ -161,7 +161,7 @@ impl AggregateFnVTable for AllNonDistinct { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { Ok(AllNonDistinctPartial { @@ -172,7 +172,7 @@ impl AggregateFnVTable for AllNonDistinct { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -184,7 +184,7 @@ impl AggregateFnVTable for AllNonDistinct { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool( @@ -197,7 +197,7 @@ impl AggregateFnVTable for AllNonDistinct { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !partial.all_non_distinct @@ -206,7 +206,7 @@ impl AggregateFnVTable for AllNonDistinct { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -259,7 +259,7 @@ impl AggregateFnVTable for AllNonDistinct { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partials: ArrayRef, ) -> VortexResult { vortex_bail!("AllNonDistinct does not support array finalization"); @@ -268,7 +268,7 @@ impl AggregateFnVTable for AllNonDistinct { fn finalize_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool( diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index 0bb45773d69..a5c2750d261 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -9,7 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -64,7 +64,7 @@ impl AggregateFnVTable for AllNonNan { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } @@ -72,7 +72,7 @@ impl AggregateFnVTable for AllNonNan { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) @@ -81,7 +81,7 @@ impl AggregateFnVTable for AllNonNan { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -91,7 +91,7 @@ impl AggregateFnVTable for AllNonNan { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::Nullable)) @@ -100,7 +100,7 @@ impl AggregateFnVTable for AllNonNan { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !*partial @@ -109,7 +109,7 @@ impl AggregateFnVTable for AllNonNan { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -121,7 +121,7 @@ impl AggregateFnVTable for AllNonNan { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -139,7 +139,7 @@ impl AggregateFnVTable for AllNonNan { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -148,7 +148,7 @@ impl AggregateFnVTable for AllNonNan { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index bd42570abf1..30910d9653c 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -9,7 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -55,7 +55,7 @@ impl AggregateFnVTable for AllNonNull { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } @@ -63,7 +63,7 @@ impl AggregateFnVTable for AllNonNull { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) @@ -72,7 +72,7 @@ impl AggregateFnVTable for AllNonNull { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -82,7 +82,7 @@ impl AggregateFnVTable for AllNonNull { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) @@ -91,7 +91,7 @@ impl AggregateFnVTable for AllNonNull { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !*partial @@ -100,7 +100,7 @@ impl AggregateFnVTable for AllNonNull { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -112,7 +112,7 @@ impl AggregateFnVTable for AllNonNull { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -129,7 +129,7 @@ impl AggregateFnVTable for AllNonNull { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -138,7 +138,7 @@ impl AggregateFnVTable for AllNonNull { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index ad61771138f..8a4e090ed8a 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -9,7 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -55,7 +55,7 @@ impl AggregateFnVTable for AllNull { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } @@ -63,7 +63,7 @@ impl AggregateFnVTable for AllNull { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { bool::try_from(&scalar) @@ -72,7 +72,7 @@ impl AggregateFnVTable for AllNull { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -82,7 +82,7 @@ impl AggregateFnVTable for AllNull { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::bool(*partial, Nullability::NonNullable)) @@ -91,7 +91,7 @@ impl AggregateFnVTable for AllNull { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !*partial @@ -100,7 +100,7 @@ impl AggregateFnVTable for AllNull { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -112,7 +112,7 @@ impl AggregateFnVTable for AllNull { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -132,7 +132,7 @@ impl AggregateFnVTable for AllNull { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -141,7 +141,7 @@ impl AggregateFnVTable for AllNull { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index e7630411b39..8ef939b3475 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -19,7 +19,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -93,8 +93,8 @@ impl BoundedMaxPartial { self.state = BoundedMaxState::Unknown; } - fn final_scalar(&self, dtypes: AggregateDTypes<'_>) -> VortexResult { - let dtype = dtypes.result.clone(); + fn final_scalar(&self, dtypes: AggregateDTypesRef<'_>) -> VortexResult { + let dtype = dtypes.return_dtype.clone(); match &self.state { BoundedMaxState::Value(max) => max.cast(&dtype), BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), @@ -191,7 +191,7 @@ impl AggregateFnVTable for BoundedMax { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(BoundedMaxPartial { state: BoundedMaxState::Empty, @@ -201,7 +201,7 @@ impl AggregateFnVTable for BoundedMax { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -238,7 +238,7 @@ impl AggregateFnVTable for BoundedMax { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -253,11 +253,11 @@ impl AggregateFnVTable for BoundedMax { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.partial.clone(); - let bound_dtype = dtypes.input.as_nullable(); + let dtype = dtypes.partial_dtype.clone(); + let bound_dtype = dtypes.dtype.as_nullable(); match &partial.state { BoundedMaxState::Empty => Ok(Scalar::null(dtype)), BoundedMaxState::Value(max) => Ok(Scalar::struct_( @@ -280,7 +280,7 @@ impl AggregateFnVTable for BoundedMax { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { matches!(partial.state, BoundedMaxState::Unknown) @@ -289,7 +289,7 @@ impl AggregateFnVTable for BoundedMax { fn accumulate( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -313,7 +313,7 @@ impl AggregateFnVTable for BoundedMax { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { partials.get_item(BOUNDED_MAX_BOUND) @@ -322,7 +322,7 @@ impl AggregateFnVTable for BoundedMax { fn finalize_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { partial.final_scalar(dtypes) diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 4cd3acbd086..1ec0a17a72c 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -17,7 +17,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -145,7 +145,7 @@ impl AggregateFnVTable for BoundedMin { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(BoundedMinPartial { state: BoundedMinState::Empty, @@ -155,7 +155,7 @@ impl AggregateFnVTable for BoundedMin { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -170,7 +170,7 @@ impl AggregateFnVTable for BoundedMin { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -183,10 +183,10 @@ impl AggregateFnVTable for BoundedMin { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.input.as_nullable(); + let dtype = dtypes.dtype.as_nullable(); match &partial.state { BoundedMinState::Empty => Ok(Scalar::null(dtype)), BoundedMinState::Value(min) => min.cast(&dtype), @@ -196,7 +196,7 @@ impl AggregateFnVTable for BoundedMin { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { false @@ -205,7 +205,7 @@ impl AggregateFnVTable for BoundedMin { fn accumulate( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -228,7 +228,7 @@ impl AggregateFnVTable for BoundedMin { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -237,7 +237,7 @@ impl AggregateFnVTable for BoundedMin { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 3c3917c74cb..31929b3318b 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -10,7 +10,7 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; @@ -55,7 +55,7 @@ impl AggregateFnVTable for Count { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } @@ -63,7 +63,7 @@ impl AggregateFnVTable for Count { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -75,7 +75,7 @@ impl AggregateFnVTable for Count { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -85,7 +85,7 @@ impl AggregateFnVTable for Count { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::primitive(*partial, Nullability::NonNullable)) @@ -95,7 +95,7 @@ impl AggregateFnVTable for Count { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { false @@ -104,14 +104,14 @@ impl AggregateFnVTable for Count { fn try_accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { let mut count = batch.valid_count(ctx)? as u64; // NaN values are excluded from the count of a float input when they are skipped. - if options.skip_nans && dtypes.input.is_float() { + if options.skip_nans && dtypes.dtype.is_float() { // `nan_count` shortcircuits on an exact `Stat::NaNCount` before scanning the batch. count = count.saturating_sub(nan_count(batch, ctx)? as u64); } @@ -122,7 +122,7 @@ impl AggregateFnVTable for Count { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -133,7 +133,7 @@ impl AggregateFnVTable for Count { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -142,7 +142,7 @@ impl AggregateFnVTable for Count { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -163,10 +163,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::count::Count; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -273,9 +273,9 @@ mod tests { fn count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let dtypes = OwnedAggregateDTypes::try_new(&Count, &options, dtype)?; + let dtypes = AggregateDTypes::try_new(&Count, &options, dtype)?; - let state = Count.reduce_partials(&options, dtypes.borrow(), [5, 3])?; + let state = Count.merge_partials(&options, dtypes.borrow(), 5, 3)?; let result = Count.to_scalar(&options, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 5e9642f78ab..88214ca7e86 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -8,7 +8,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -59,7 +59,7 @@ impl AggregateFnVTable for First { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(FirstPartial { value: None }) } @@ -67,7 +67,7 @@ impl AggregateFnVTable for First { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -79,7 +79,7 @@ impl AggregateFnVTable for First { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -92,12 +92,12 @@ impl AggregateFnVTable for First { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(dtypes.result.clone()), + None => Scalar::null(dtypes.return_dtype.clone()), }) } @@ -105,7 +105,7 @@ impl AggregateFnVTable for First { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { partial.value.is_some() @@ -114,7 +114,7 @@ impl AggregateFnVTable for First { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -132,7 +132,7 @@ impl AggregateFnVTable for First { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -143,7 +143,7 @@ impl AggregateFnVTable for First { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -152,7 +152,7 @@ impl AggregateFnVTable for First { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -167,10 +167,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::first::FirstPartial; use crate::aggregate_fn::fns::first::first; @@ -294,7 +294,7 @@ mod tests { #[test] fn first_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let owned = OwnedAggregateDTypes::try_new(&First, &EmptyOptions, dtype)?; + let owned = AggregateDTypes::try_new(&First, &EmptyOptions, dtype)?; let dtypes = owned.borrow(); let partial_of = |value: Option| FirstPartial { value }; @@ -305,7 +305,8 @@ mod tests { // The first non-empty partial wins; subsequent valid partials are dropped. let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); - let state = First.reduce_partials(&EmptyOptions, dtypes, [empty, five, seven])?; + let merge = |first, second| First.merge_partials(&EmptyOptions, dtypes, first, second); + let state = merge(merge(empty, five)?, seven)?; assert!(First.is_saturated(&EmptyOptions, dtypes, &state)); assert_eq!( First.to_scalar(&EmptyOptions, dtypes, &state)?, diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index dc17f7540e9..b3907a7ff74 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -31,7 +31,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -294,7 +294,7 @@ impl AggregateFnVTable for IsConstant { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(IsConstantPartial::empty()) } @@ -302,7 +302,7 @@ impl AggregateFnVTable for IsConstant { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null struct means the producing accumulator was empty. @@ -325,7 +325,7 @@ impl AggregateFnVTable for IsConstant { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut acc: Self::Partial, partial: Self::Partial, ) -> VortexResult { @@ -340,11 +340,11 @@ impl AggregateFnVTable for IsConstant { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.partial.clone(); - let element_dtype = dtypes.input.as_nullable(); + let dtype = dtypes.partial_dtype.clone(); + let element_dtype = dtypes.dtype.as_nullable(); // Only a constant partial that saw no values is the empty (null) state: a non-constant // verdict stands regardless of whether a value was observed. let first_value = match &partial.first_value { @@ -365,7 +365,7 @@ impl AggregateFnVTable for IsConstant { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !partial.is_constant @@ -374,7 +374,7 @@ impl AggregateFnVTable for IsConstant { fn accumulate( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -398,7 +398,7 @@ impl AggregateFnVTable for IsConstant { let all_invalid = array_ref.all_invalid(ctx)?; if all_invalid { - partial.check_value(Scalar::null(dtypes.input.as_nullable())); + partial.check_value(Scalar::null(dtypes.dtype.as_nullable())); return Ok(()); } @@ -448,7 +448,7 @@ impl AggregateFnVTable for IsConstant { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) @@ -457,7 +457,7 @@ impl AggregateFnVTable for IsConstant { fn finalize_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { if partial.first_value.is_none() { @@ -478,10 +478,10 @@ mod tests { use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::is_constant::IsConstant; use crate::aggregate_fn::fns::is_constant::IsConstantPartial; use crate::aggregate_fn::fns::is_constant::is_constant; @@ -819,7 +819,7 @@ mod tests { #[test] fn non_constant_partial_without_value_is_not_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let owned = OwnedAggregateDTypes::try_new(&IsConstant, &EmptyOptions, dtype)?; + let owned = AggregateDTypes::try_new(&IsConstant, &EmptyOptions, dtype)?; let dtypes = owned.borrow(); let partial = IsConstantPartial { is_constant: false, diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 3c9ebdd3389..fe711678d72 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -26,7 +26,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -290,7 +290,7 @@ impl AggregateFnVTable for IsSorted { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(IsSortedPartial::empty()) } @@ -298,7 +298,7 @@ impl AggregateFnVTable for IsSorted { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null struct means the producing accumulator was empty. @@ -323,7 +323,7 @@ impl AggregateFnVTable for IsSorted { fn merge_partials( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut acc: Self::Partial, partial: Self::Partial, ) -> VortexResult { @@ -382,10 +382,10 @@ impl AggregateFnVTable for IsSorted { fn to_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.partial.clone(); + let dtype = dtypes.partial_dtype.clone(); // Only a sorted partial that saw no values is the empty (null) state: an unsorted verdict // stands regardless of which boundaries were observed. if partial.is_sorted && partial.first_value.is_none() { @@ -394,7 +394,7 @@ impl AggregateFnVTable for IsSorted { let first_value = partial .first_value .clone() - .unwrap_or_else(|| Scalar::null(dtypes.input.as_nullable())); + .unwrap_or_else(|| Scalar::null(dtypes.dtype.as_nullable())); // A partial that saw a single value carries it as both boundaries. let last_value = partial .last_value @@ -415,7 +415,7 @@ impl AggregateFnVTable for IsSorted { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { !partial.is_sorted @@ -424,7 +424,7 @@ impl AggregateFnVTable for IsSorted { fn accumulate( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -541,7 +541,7 @@ impl AggregateFnVTable for IsSorted { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) @@ -550,7 +550,7 @@ impl AggregateFnVTable for IsSorted { fn finalize_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { // The empty state is vacuously sorted, so the verdict stands on its own. @@ -594,9 +594,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::is_sorted::IsSorted; use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; use crate::aggregate_fn::fns::is_sorted::IsSortedPartial; @@ -791,7 +791,7 @@ mod tests { fn unsorted_partial_without_boundaries_is_not_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = IsSortedOptions { strict: false }; - let owned = OwnedAggregateDTypes::try_new(&IsSorted, &options, dtype)?; + let owned = AggregateDTypes::try_new(&IsSorted, &options, dtype)?; let dtypes = owned.borrow(); let partial = IsSortedPartial { is_sorted: false, diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index c81b66165cb..96d4e03e725 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -8,7 +8,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -59,7 +59,7 @@ impl AggregateFnVTable for Last { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(LastPartial { value: None }) } @@ -67,7 +67,7 @@ impl AggregateFnVTable for Last { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null partial means the producing accumulator saw nothing valid. @@ -79,7 +79,7 @@ impl AggregateFnVTable for Last { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -92,12 +92,12 @@ impl AggregateFnVTable for Last { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(dtypes.result.clone()), + None => Scalar::null(dtypes.return_dtype.clone()), }) } @@ -105,7 +105,7 @@ impl AggregateFnVTable for Last { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { // Last can never short-circuit: a later batch can always supersede the current value. @@ -115,7 +115,7 @@ impl AggregateFnVTable for Last { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -130,7 +130,7 @@ impl AggregateFnVTable for Last { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -141,7 +141,7 @@ impl AggregateFnVTable for Last { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -150,7 +150,7 @@ impl AggregateFnVTable for Last { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -165,10 +165,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::last::LastPartial; use crate::aggregate_fn::fns::last::last; @@ -291,7 +291,7 @@ mod tests { #[test] fn last_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let owned = OwnedAggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?; + let owned = AggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?; let dtypes = owned.borrow(); let partial_of = |value: Option| LastPartial { value }; @@ -301,7 +301,8 @@ mod tests { let empty = partial_of(None); // The last non-empty partial in order replaces the prior values. - let state = Last.reduce_partials(&EmptyOptions, dtypes, [five, seven, empty])?; + let merge = |first, second| Last.merge_partials(&EmptyOptions, dtypes, first, second); + let state = merge(merge(five, seven)?, empty)?; assert_eq!( Last.to_scalar(&EmptyOptions, dtypes, &state)?, Scalar::primitive(7i32, Nullable) diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index c292fa77253..d8bbaa2b308 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -10,7 +10,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -45,7 +45,7 @@ impl MaxPartial { fn merge( &mut self, options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, max: Scalar, ) { if max.is_null() { @@ -67,8 +67,8 @@ impl MaxPartial { }); } - fn poison(&mut self, dtypes: AggregateDTypes<'_>) { - self.max = Some(nan_scalar(dtypes.input)); + fn poison(&mut self, dtypes: AggregateDTypesRef<'_>) { + self.max = Some(nan_scalar(dtypes.dtype)); } fn is_poisoned(&self) -> bool { @@ -128,7 +128,7 @@ impl AggregateFnVTable for Max { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(MaxPartial { max: None }) } @@ -136,7 +136,7 @@ impl AggregateFnVTable for Max { fn partial_from_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { let mut partial = MaxPartial { max: None }; @@ -148,7 +148,7 @@ impl AggregateFnVTable for Max { fn merge_partials( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -161,10 +161,10 @@ impl AggregateFnVTable for Max { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.input.as_nullable(); + let dtype = dtypes.dtype.as_nullable(); match &partial.max { Some(max) => max.cast(&dtype), None => Ok(Scalar::null(dtype)), @@ -174,7 +174,7 @@ impl AggregateFnVTable for Max { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { // A poisoned NaN-including maximum is fully determined. @@ -184,14 +184,14 @@ impl AggregateFnVTable for Max { fn try_accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to the NaN-including float maximum; everything else // takes the default dispatch path. - if options.skip_nans || !dtypes.input.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -215,7 +215,7 @@ impl AggregateFnVTable for Max { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -235,7 +235,7 @@ impl AggregateFnVTable for Max { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -244,7 +244,7 @@ impl AggregateFnVTable for Max { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index bb6bd51260d..e30e90524e8 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -10,7 +10,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; @@ -97,14 +97,14 @@ impl BinaryCombined for Mean { fn finalize( &self, _options: &CombinedOptions, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, sum: ArrayRef, count: ArrayRef, ) -> VortexResult { if let DType::Decimal(..) = sum.dtype() { vortex_bail!("grouped mean over decimals is not yet supported"); } - let target = dtypes.result.clone(); + let target = dtypes.return_dtype.clone(); let sum = sum.cast(target.clone())?; let count = count.cast(target.clone())?; @@ -125,7 +125,7 @@ impl BinaryCombined for Mean { fn finalize_scalar( &self, _options: &CombinedOptions, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, left_scalar: Scalar, right_scalar: Scalar, ) -> VortexResult { @@ -134,11 +134,11 @@ impl BinaryCombined for Mean { &left_scalar, &right_scalar, decimal_dtype, - dtypes.result, + dtypes.return_dtype, ); } - let target = dtypes.result.clone(); + let target = dtypes.return_dtype.clone(); let sum_cast = left_scalar.cast(&target)?; let count_cast = right_scalar.cast(&target)?; diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index c7a606237fa..17446b6bfa9 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -10,7 +10,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -45,7 +45,7 @@ impl MinPartial { fn merge( &mut self, options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, min: Scalar, ) { if min.is_null() { @@ -67,8 +67,8 @@ impl MinPartial { }); } - fn poison(&mut self, dtypes: AggregateDTypes<'_>) { - self.min = Some(nan_scalar(dtypes.input)); + fn poison(&mut self, dtypes: AggregateDTypesRef<'_>) { + self.min = Some(nan_scalar(dtypes.dtype)); } fn is_poisoned(&self) -> bool { @@ -128,7 +128,7 @@ impl AggregateFnVTable for Min { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(MinPartial { min: None }) } @@ -136,7 +136,7 @@ impl AggregateFnVTable for Min { fn partial_from_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { let mut partial = MinPartial { min: None }; @@ -148,7 +148,7 @@ impl AggregateFnVTable for Min { fn merge_partials( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -161,10 +161,10 @@ impl AggregateFnVTable for Min { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.input.as_nullable(); + let dtype = dtypes.dtype.as_nullable(); match &partial.min { Some(min) => min.cast(&dtype), None => Ok(Scalar::null(dtype)), @@ -174,7 +174,7 @@ impl AggregateFnVTable for Min { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { // A poisoned NaN-including minimum is fully determined. @@ -184,14 +184,14 @@ impl AggregateFnVTable for Min { fn try_accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to the NaN-including float minimum; everything else // takes the default dispatch path. - if options.skip_nans || !dtypes.input.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -215,7 +215,7 @@ impl AggregateFnVTable for Min { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -235,7 +235,7 @@ impl AggregateFnVTable for Min { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -244,7 +244,7 @@ impl AggregateFnVTable for Min { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs b/vortex-array/src/aggregate_fn/fns/min_max/bool.rs index 2aebb7f1941..1c25baa0d2f 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/bool.rs @@ -9,7 +9,7 @@ use vortex_mask::AllOr; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; @@ -18,7 +18,7 @@ use crate::scalar::Scalar; pub(super) fn accumulate_bool( options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut MinMaxPartial, array: &BoolArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs index a5c2e0071dd..f00ef587dc5 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs @@ -8,7 +8,7 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; @@ -20,7 +20,7 @@ use crate::scalar::Scalar; pub(super) fn accumulate_decimal( options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut MinMaxPartial, array: &DecimalArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index 1014bea111c..513b9a98d29 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -7,7 +7,7 @@ use super::MinMaxPartial; use super::MinMaxResult; use super::min_max; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; @@ -16,7 +16,7 @@ use crate::scalar::Scalar; pub(super) fn accumulate_extension( options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut MinMaxPartial, array: &ExtensionArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 674e1e42a3a..13cee4a8a06 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -25,7 +25,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -205,7 +205,7 @@ impl MinMaxPartial { fn merge( &mut self, options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, local: Option, ) { let Some(MinMaxResult { min, max }) = local else { @@ -234,8 +234,8 @@ impl MinMaxPartial { } /// Poison the partial state to `{min: NaN, max: NaN}`. - fn poison(&mut self, dtypes: AggregateDTypes<'_>) { - let nan = nan_scalar(dtypes.input); + fn poison(&mut self, dtypes: AggregateDTypesRef<'_>) { + let nan = nan_scalar(dtypes.dtype); self.min = Some(nan.clone()); self.max = Some(nan); } @@ -315,7 +315,7 @@ impl AggregateFnVTable for MinMax { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(MinMaxPartial { min: None, @@ -326,7 +326,7 @@ impl AggregateFnVTable for MinMax { fn partial_from_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { let mut partial = MinMaxPartial { @@ -341,7 +341,7 @@ impl AggregateFnVTable for MinMax { fn merge_partials( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -354,10 +354,10 @@ impl AggregateFnVTable for MinMax { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { - let dtype = dtypes.partial.clone(); + let dtype = dtypes.partial_dtype.clone(); Ok(match (&partial.min, &partial.max) { (Some(min), Some(max)) => Scalar::struct_(dtype, vec![min.clone(), max.clone()]), _ => Scalar::null(dtype), @@ -368,7 +368,7 @@ impl AggregateFnVTable for MinMax { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { // A poisoned NaN-including min/max is fully determined. @@ -378,14 +378,14 @@ impl AggregateFnVTable for MinMax { fn try_accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { // NaN-aware shortcircuits only apply to NaN-including float min/max; everything else // takes the default dispatch path. - if options.skip_nans || !dtypes.input.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -396,7 +396,7 @@ impl AggregateFnVTable for MinMax { if let Some((min, max)) = cached_min.zip(cached_max) { // Cached float stats carry the (possibly nullable) array dtype; `to_scalar` // builds a struct with non-nullable fields, so normalise here. - let non_nullable_dtype = dtypes.input.as_nonnullable(); + let non_nullable_dtype = dtypes.dtype.as_nonnullable(); partial.merge( options, dtypes, @@ -421,7 +421,7 @@ impl AggregateFnVTable for MinMax { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -475,7 +475,7 @@ impl AggregateFnVTable for MinMax { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -484,7 +484,7 @@ impl AggregateFnVTable for MinMax { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -505,10 +505,10 @@ mod tests { use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::MinMaxPartial; use crate::aggregate_fn::fns::min_max::MinMaxResult; @@ -710,7 +710,7 @@ mod tests { fn test_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let owned = OwnedAggregateDTypes::try_new(&MinMax, &options, dtype)?; + let owned = AggregateDTypes::try_new(&MinMax, &options, dtype)?; let dtypes = owned.borrow(); let partial_of = |min: i32, max: i32| MinMaxPartial { min: Some(Scalar::from(min)), @@ -718,7 +718,7 @@ mod tests { }; let state = - MinMax.reduce_partials(&options, dtypes, [partial_of(5, 15), partial_of(2, 10)])?; + MinMax.merge_partials(&options, dtypes, partial_of(5, 15), partial_of(2, 10))?; let result = MinMaxResult::from_scalar(MinMax.to_scalar(&options, dtypes, &state)?)? .vortex_expect("should have result"); diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index a35fedf7000..7f3f183edd9 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -8,7 +8,7 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::PrimitiveArray; use crate::dtype::NativePType; @@ -19,7 +19,7 @@ use crate::scalar::Scalar; pub(super) fn accumulate_primitive( options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut MinMaxPartial, p: &PrimitiveArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs index 9e0846eb553..c793e527e0b 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs @@ -8,7 +8,7 @@ use vortex_error::vortex_panic; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::VarBinViewArray; use crate::dtype::DType; @@ -17,7 +17,7 @@ use crate::scalar::Scalar; pub(super) fn accumulate_varbinview( options: &NumericalAggregateOpts, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut MinMaxPartial, array: &VarBinViewArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index dcd49f56652..a9e881806b0 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -16,7 +16,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -119,7 +119,7 @@ impl AggregateFnVTable for NanCount { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } @@ -127,7 +127,7 @@ impl AggregateFnVTable for NanCount { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -139,7 +139,7 @@ impl AggregateFnVTable for NanCount { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -149,7 +149,7 @@ impl AggregateFnVTable for NanCount { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) @@ -159,7 +159,7 @@ impl AggregateFnVTable for NanCount { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { false @@ -168,7 +168,7 @@ impl AggregateFnVTable for NanCount { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -197,7 +197,7 @@ impl AggregateFnVTable for NanCount { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -206,7 +206,7 @@ impl AggregateFnVTable for NanCount { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -221,10 +221,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::nan_count::nan_count; use crate::array_session; @@ -279,9 +279,9 @@ mod tests { #[test] fn nan_count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let dtypes = OwnedAggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?; + let dtypes = AggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?; - let state = NanCount.reduce_partials(&EmptyOptions, dtypes.borrow(), [5, 3])?; + let state = NanCount.merge_partials(&EmptyOptions, dtypes.borrow(), 5, 3)?; let result = NanCount.to_scalar(&EmptyOptions, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index db139ae8d9a..4999e8bfebb 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -12,7 +12,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -88,7 +88,7 @@ impl AggregateFnVTable for NullCount { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } @@ -96,7 +96,7 @@ impl AggregateFnVTable for NullCount { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -108,7 +108,7 @@ impl AggregateFnVTable for NullCount { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -118,7 +118,7 @@ impl AggregateFnVTable for NullCount { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) @@ -128,7 +128,7 @@ impl AggregateFnVTable for NullCount { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { false @@ -137,7 +137,7 @@ impl AggregateFnVTable for NullCount { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -149,7 +149,7 @@ impl AggregateFnVTable for NullCount { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -170,7 +170,7 @@ impl AggregateFnVTable for NullCount { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -179,7 +179,7 @@ impl AggregateFnVTable for NullCount { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index b970954549f..1a92e48d1c5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -118,9 +118,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumPartial; use crate::aggregate_fn::fns::sum::SumState; @@ -373,12 +373,12 @@ mod tests { // Reduce partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; + let dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; let near_limit = partial_with_decimal(DecimalValue::from(99_999_999_999_990i64)); // Add a small value that keeps us just under 10^14. let small = partial_with_decimal(DecimalValue::from(9i64)); - let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, small])?; + let state = Sum.merge_partials(&options, dtypes.borrow(), near_limit, small)?; let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(!result.is_null()); @@ -395,15 +395,15 @@ mod tests { // The max representable value for precision 14 is 10^14 - 1. // When the sum reaches exactly 10^14, fits_in_precision fails even though // i256 arithmetic does not overflow. This tests the precision-based - // saturation path in reduce_partials. + // saturation path in merge_partials. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; + let dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; let near_limit = partial_with_decimal(DecimalValue::from(99_999_999_999_999i64)); // Push the sum to exactly 10^14, exceeding precision 14. let one_more = partial_with_decimal(DecimalValue::from(1i64)); - let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, one_more])?; + let state = Sum.merge_partials(&options, dtypes.borrow(), near_limit, one_more)?; let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(result.is_null()); @@ -419,11 +419,11 @@ mod tests { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; + let dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; let near_limit = partial_with_decimal(DecimalValue::from(-99_999_999_999_999i64)); let one_more = partial_with_decimal(DecimalValue::from(-1i64)); - let state = Sum.reduce_partials(&options, dtypes.borrow(), [near_limit, one_more])?; + let state = Sum.merge_partials(&options, dtypes.borrow(), near_limit, one_more)?; let result = Sum.to_scalar(&options, dtypes.borrow(), &state)?; assert!(result.is_null()); @@ -432,7 +432,7 @@ mod tests { #[test] fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { - // Test precision overflow via the accumulate_decimal path (not reduce_partials). + // Test precision overflow via the accumulate_decimal path (not merge_partials). // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. // Native for precision 38 is I128 (max 38), so 38 = 38. // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. @@ -440,9 +440,9 @@ mod tests { // We seed the state close to 10^37, then accumulate a real array that pushes it over. let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let dtypes = OwnedAggregateDTypes::try_new(&Sum, &options, input_dtype)?; + let dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; assert_eq!( - dtypes.result(), + dtypes.return_dtype(), &DType::Decimal(DecimalDType::new(37, 0), Nullable) ); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 86eaf221f0b..2e384636ca1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -27,7 +27,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -146,30 +146,30 @@ impl AggregateFnVTable for Sum { fn empty_partial( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(SumPartial { - current: Some(make_zero_state(dtypes.result)), + current: Some(make_zero_state(dtypes.return_dtype)), }) } fn partial_from_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { vortex_ensure!( - scalar.dtype().eq_ignore_nullability(dtypes.result), + scalar.dtype().eq_ignore_nullability(dtypes.return_dtype), "Sum partial has dtype {}, expected {}", scalar.dtype(), - dtypes.result + dtypes.return_dtype ); // A null partial means the producing accumulator saturated (overflow). let current = if scalar.is_null() { None } else { - Some(sum_state_from_scalar(&scalar, dtypes.result)?) + Some(sum_state_from_scalar(&scalar, dtypes.return_dtype)?) }; Ok(SumPartial { current }) } @@ -177,7 +177,7 @@ impl AggregateFnVTable for Sum { fn merge_partials( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -186,7 +186,7 @@ impl AggregateFnVTable for Sum { (None, _) => false, // A saturated (overflowed) partial poisons the merge. (Some(_), None) => true, - (Some(acc), Some(state)) => checked_add_sum_states(acc, dtypes.result, &state)?, + (Some(acc), Some(state)) => checked_add_sum_states(acc, dtypes.return_dtype, &state)?, }; if overflow { first.current = None; @@ -197,17 +197,17 @@ impl AggregateFnVTable for Sum { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(match &partial.current { - None => Scalar::null(dtypes.result.as_nullable()), + None => Scalar::null(dtypes.return_dtype.as_nullable()), Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable), Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable), Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable), Some(SumState::Decimal(value)) => { let decimal_dtype = *dtypes - .result + .return_dtype .as_decimal_opt() .vortex_expect("return dtype must be decimal"); Scalar::decimal(*value, decimal_dtype, Nullability::Nullable) @@ -219,7 +219,7 @@ impl AggregateFnVTable for Sum { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { match partial.current.as_ref() { @@ -232,7 +232,7 @@ impl AggregateFnVTable for Sum { fn try_accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, @@ -247,12 +247,12 @@ impl AggregateFnVTable for Sum { // NaN-free batch: the cached NaN-skipping sum (if any) equals the // NaN-including sum. if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { - let sum = if sum.dtype() == dtypes.result { + let sum = if sum.dtype() == dtypes.return_dtype { sum } else { - sum.cast(dtypes.result)? + sum.cast(dtypes.return_dtype)? }; - merge_sum_result(partial, dtypes.result, sum)?; + merge_sum_result(partial, dtypes.return_dtype, sum)?; return Ok(true); } Ok(false) @@ -271,7 +271,7 @@ impl AggregateFnVTable for Sum { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -282,8 +282,8 @@ impl AggregateFnVTable for Sum { if options.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { return Ok(()); } - if let Some(product) = multiply_constant(c.scalar(), c.len(), dtypes.result)? { - merge_sum_result(partial, dtypes.result, product)?; + if let Some(product) = multiply_constant(c.scalar(), c.len(), dtypes.return_dtype)? { + merge_sum_result(partial, dtypes.return_dtype, product)?; } return Ok(()); } @@ -299,7 +299,9 @@ impl AggregateFnVTable for Sum { accumulate_primitive(&mut inner, p, ctx, options.skip_nans) } Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, dtypes.result, d, ctx), + Canonical::Decimal(d) => { + accumulate_decimal(&mut inner, dtypes.return_dtype, d, ctx) + } _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -319,7 +321,7 @@ impl AggregateFnVTable for Sum { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -328,7 +330,7 @@ impl AggregateFnVTable for Sum { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -470,12 +472,12 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumPartial; use crate::aggregate_fn::fns::sum::SumState; @@ -604,13 +606,13 @@ mod tests { fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let owned = OwnedAggregateDTypes::try_new(&Sum, &options, dtype)?; + let owned = AggregateDTypes::try_new(&Sum, &options, dtype)?; let dtypes = owned.borrow(); let partial_of = |value: i64| SumPartial { current: Some(SumState::Signed(value)), }; - let state = Sum.reduce_partials(&options, dtypes, [partial_of(100), partial_of(50)])?; + let state = Sum.merge_partials(&options, dtypes, partial_of(100), partial_of(50))?; let result = Sum.to_scalar(&options, dtypes, &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(150)); @@ -621,7 +623,7 @@ mod tests { fn sum_overflowed_partial_poisons_reduction() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let options = NumericalAggregateOpts::default(); - let owned = OwnedAggregateDTypes::try_new(&Sum, &options, dtype)?; + let owned = AggregateDTypes::try_new(&Sum, &options, dtype)?; let dtypes = owned.borrow(); let overflowed = Sum.partial_from_scalar( @@ -630,7 +632,7 @@ mod tests { Scalar::null(DType::Primitive(PType::I64, Nullable)), )?; let five = Sum.partial_from_scalar(&options, dtypes, Scalar::primitive(5i64, Nullable))?; - let state = Sum.reduce_partials(&options, dtypes, [five, overflowed])?; + let state = Sum.merge_partials(&options, dtypes, five, overflowed)?; assert!(Sum.is_saturated(&options, dtypes, &state)); assert!(Sum.to_scalar(&options, dtypes, &state)?.is_null()); diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index 96e71b7fc7c..a3c6b243a46 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -18,7 +18,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -110,24 +110,24 @@ impl AggregateFnVTable for SumV2 { fn empty_partial( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - Ok(SumV2Partial::empty(dtypes.result)) + Ok(SumV2Partial::empty(dtypes.return_dtype)) } fn partial_from_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { - let mut partial = SumV2Partial::empty(dtypes.result); + let mut partial = SumV2Partial::empty(dtypes.return_dtype); let (sum, is_overflow, is_empty) = decode_partial_scalar(scalar)?; - validate_sum_field_dtype(&sum, dtypes.result)?; + validate_sum_field_dtype(&sum, dtypes.return_dtype)?; // Adding the parsed value to the zero state cannot overflow; treat a decimal value that // no longer fits its precision as an already-overflowed partial. - let overflowed = checked_add_sum_state(&mut partial.sum, dtypes.result, &sum)?; + let overflowed = checked_add_sum_state(&mut partial.sum, dtypes.return_dtype, &sum)?; partial.is_overflow = is_overflow || overflowed; partial.is_empty = is_empty && !partial.is_overflow; Ok(partial) @@ -136,7 +136,7 @@ impl AggregateFnVTable for SumV2 { fn merge_partials( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, mut acc: Self::Partial, partial: Self::Partial, ) -> VortexResult { @@ -156,7 +156,7 @@ impl AggregateFnVTable for SumV2 { if partial.is_empty { return Ok(acc); } - acc.is_overflow = checked_add_sum_states(&mut acc.sum, dtypes.result, &partial.sum)?; + acc.is_overflow = checked_add_sum_states(&mut acc.sum, dtypes.return_dtype, &partial.sum)?; acc.is_empty = false; Ok(acc) } @@ -164,13 +164,13 @@ impl AggregateFnVTable for SumV2 { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::struct_( - dtypes.partial.clone(), + dtypes.partial_dtype.clone(), vec![ - sum_state_scalar(partial, dtypes.result, Nullability::NonNullable), + sum_state_scalar(partial, dtypes.return_dtype, Nullability::NonNullable), Scalar::bool(partial.is_overflow, Nullability::NonNullable), Scalar::bool(partial.is_empty, Nullability::NonNullable), ], @@ -180,7 +180,7 @@ impl AggregateFnVTable for SumV2 { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) @@ -189,7 +189,7 @@ impl AggregateFnVTable for SumV2 { fn try_accumulate( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, _ctx: &mut ExecutionCtx, @@ -215,7 +215,7 @@ impl AggregateFnVTable for SumV2 { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -237,14 +237,14 @@ impl AggregateFnVTable for SumV2 { return Ok(()); } if let Some(product) = - multiply_constant(constant.scalar(), constant.len(), dtypes.result)? + multiply_constant(constant.scalar(), constant.len(), dtypes.return_dtype)? { if product.is_null() { partial.is_overflow = true; partial.is_empty = false; } else { partial.is_overflow = - checked_add_sum_state(&mut partial.sum, dtypes.result, &product)?; + checked_add_sum_state(&mut partial.sum, dtypes.return_dtype, &product)?; partial.is_empty = false; } } @@ -259,7 +259,7 @@ impl AggregateFnVTable for SumV2 { } Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx), Canonical::Decimal(array) => { - accumulate_decimal(&mut partial.sum, dtypes.result, array, ctx) + accumulate_decimal(&mut partial.sum, dtypes.return_dtype, array, ctx) } _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()), }, @@ -279,7 +279,7 @@ impl AggregateFnVTable for SumV2 { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { if let Some(partials) = partials.as_opt::() { @@ -297,15 +297,15 @@ impl AggregateFnVTable for SumV2 { fn finalize_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { if partial.is_overflow || partial.is_empty { - return Ok(Scalar::null(dtypes.result.as_nullable())); + return Ok(Scalar::null(dtypes.return_dtype.as_nullable())); } Ok(sum_state_scalar( partial, - dtypes.result, + dtypes.return_dtype, Nullability::Nullable, )) } diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 6d7701e5b49..683228b177f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -14,6 +14,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; @@ -21,7 +22,6 @@ use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::sum::Sum; use crate::array_session; use crate::arrays::BoolArray; @@ -291,7 +291,7 @@ fn merge_from_overflow_is_absorbing() -> VortexResult<()> { ); // The overflow flag survives the scalar round trip. - let dtypes = OwnedAggregateDTypes::try_new(&SumV2, &options, dtype)?; + let dtypes = AggregateDTypes::try_new(&SumV2, &options, dtype)?; let propagated = SumV2.partial_from_scalar(&options, dtypes.borrow(), overflow_partial)?; assert!(SumV2.is_saturated(&options, dtypes.borrow(), &propagated)); assert!( @@ -380,7 +380,7 @@ fn finalize_struct_applies_partial_and_struct_validity( let options = NumericalAggregateOpts::default(); let dtypes = - OwnedAggregateDTypes::try_new(&SumV2, &options, DType::Primitive(PType::I64, Nullable))?; + AggregateDTypes::try_new(&SumV2, &options, DType::Primitive(PType::I64, Nullable))?; let result = SumV2.finalize(&options, dtypes.borrow(), partials)?; let expected = PrimitiveArray::from_option_iter(expected).into_array(); assert_arrays_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 31457051f7a..fdf20a8238d 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -38,7 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -135,7 +135,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } @@ -143,7 +143,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { Ok(scalar @@ -155,7 +155,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -167,7 +167,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(Scalar::primitive(*partial, NonNullable)) @@ -177,7 +177,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { false @@ -186,7 +186,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -206,7 +206,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -215,7 +215,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -373,10 +373,10 @@ mod tests { use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::OwnedAggregateDTypes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::array_session; @@ -719,10 +719,9 @@ mod tests { #[test] fn state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let dtypes = OwnedAggregateDTypes::try_new(&UncompressedSizeInBytes, &EmptyOptions, dtype)?; + let dtypes = AggregateDTypes::try_new(&UncompressedSizeInBytes, &EmptyOptions, dtype)?; - let state = - UncompressedSizeInBytes.reduce_partials(&EmptyOptions, dtypes.borrow(), [5, 3])?; + let state = UncompressedSizeInBytes.merge_partials(&EmptyOptions, dtypes.borrow(), 5, 3)?; let result = UncompressedSizeInBytes.to_scalar(&EmptyOptions, dtypes.borrow(), &state)?; assert_eq!(result.as_primitive().typed_value::(), Some(8)); diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs index 5e5d7ac549c..6955e8879be 100644 --- a/vortex-array/src/aggregate_fn/foreign.rs +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -12,7 +12,7 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; @@ -81,7 +81,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } @@ -89,7 +89,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _scalar: Scalar, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) @@ -98,7 +98,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _first: Self::Partial, _second: Self::Partial, ) -> VortexResult { @@ -108,7 +108,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) @@ -117,7 +117,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _state: &Self::Partial, ) -> bool { false @@ -126,7 +126,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -137,7 +137,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _states: ArrayRef, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) @@ -146,7 +146,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn finalize_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 93b8e7a5605..509c3d8062b 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -69,7 +69,7 @@ mod tests { use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; - use crate::aggregate_fn::AggregateDTypes; + use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; @@ -118,7 +118,7 @@ mod tests { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(()) } @@ -126,7 +126,7 @@ mod tests { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _scalar: Scalar, ) -> VortexResult { Ok(()) @@ -135,7 +135,7 @@ mod tests { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _first: Self::Partial, _second: Self::Partial, ) -> VortexResult { @@ -145,7 +145,7 @@ mod tests { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> VortexResult { vortex_panic!("TestAgg is for serde tests only"); @@ -154,7 +154,7 @@ mod tests { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { true @@ -163,7 +163,7 @@ mod tests { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -174,7 +174,7 @@ mod tests { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -183,7 +183,7 @@ mod tests { fn finalize_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> VortexResult { vortex_panic!("TestAgg is for serde tests only"); diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 23c89471a91..199ce4efc76 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -26,81 +26,86 @@ use crate::scalar::Scalar; /// Resolved dtypes of one aggregate function bound to its options and input. /// -/// Accumulators resolve these once and lend them to every execution method, so partial states -/// only hold accumulated values. Combined aggregates resolve a separate set for each child. -#[derive(Clone, Copy, Debug)] -pub struct AggregateDTypes<'a> { - /// The dtype of the values being aggregated. - pub input: &'a DType, - /// The dtype of the partial state scalar, as reported by [`AggregateFnVTable::partial_dtype`]. - pub partial: &'a DType, - /// The dtype of the final aggregate result, as reported by [`AggregateFnVTable::return_dtype`]. - pub result: &'a DType, -} - -/// Owned [`AggregateDTypes`], resolved once from an aggregate's options and input dtype. +/// Accumulators resolve these once and lend them to every execution method as an +/// [`AggregateDTypesRef`], so partial states only hold accumulated values. Combined aggregates +/// resolve a separate set for each child. #[derive(Clone, Debug)] -pub struct OwnedAggregateDTypes { - input: DType, - partial: DType, - result: DType, +pub struct AggregateDTypes { + /// The DType of the input. + dtype: DType, + /// The DType of the aggregate. + return_dtype: DType, + /// The DType of the partial accumulator state. + partial_dtype: DType, } -impl OwnedAggregateDTypes { - /// Resolve the partial and result dtypes of `vtable` bound to `options` over `input`. +impl AggregateDTypes { + /// Resolve the return and partial dtypes of `vtable` bound to `options` over `dtype`. /// - /// Fails if the aggregate cannot be applied to `input`. + /// Fails if the aggregate cannot be applied to `dtype`. pub fn try_new( vtable: &V, options: &V::Options, - input: DType, + dtype: DType, ) -> VortexResult { - let result = vtable.return_dtype(options, &input).ok_or_else(|| { + let return_dtype = vtable.return_dtype(options, &dtype).ok_or_else(|| { vortex_err!( "Aggregate function {} cannot be applied to dtype {}", vtable.id(), - input + dtype ) })?; - let partial = vtable.partial_dtype(options, &input).ok_or_else(|| { + let partial_dtype = vtable.partial_dtype(options, &dtype).ok_or_else(|| { vortex_err!( "Aggregate function {} cannot be applied to dtype {}", vtable.id(), - input + dtype ) })?; Ok(Self { - input, - partial, - result, + dtype, + return_dtype, + partial_dtype, }) } - /// The dtype of the values being aggregated. - pub fn input(&self) -> &DType { - &self.input + /// The DType of the input. + pub fn dtype(&self) -> &DType { + &self.dtype } - /// The dtype of the partial state scalar. - pub fn partial(&self) -> &DType { - &self.partial + /// The DType of the aggregate. + pub fn return_dtype(&self) -> &DType { + &self.return_dtype } - /// The dtype of the final aggregate result. - pub fn result(&self) -> &DType { - &self.result + /// The DType of the partial accumulator state. + pub fn partial_dtype(&self) -> &DType { + &self.partial_dtype } /// Lend the dtypes to an execution method. - pub fn borrow(&self) -> AggregateDTypes<'_> { - AggregateDTypes { - input: &self.input, - partial: &self.partial, - result: &self.result, + pub fn borrow(&self) -> AggregateDTypesRef<'_> { + AggregateDTypesRef { + dtype: &self.dtype, + return_dtype: &self.return_dtype, + partial_dtype: &self.partial_dtype, } } } +/// A borrowed [`AggregateDTypes`], lent to every aggregate execution method. +#[derive(Clone, Copy, Debug)] +pub struct AggregateDTypesRef<'a> { + /// The DType of the input. + pub dtype: &'a DType, + /// The DType of the aggregate, as reported by [`AggregateFnVTable::return_dtype`]. + pub return_dtype: &'a DType, + /// The DType of the partial accumulator state, as reported by + /// [`AggregateFnVTable::partial_dtype`]. + pub partial_dtype: &'a DType, +} + /// Defines the interface for aggregate function vtables. /// /// This trait is non-object-safe and allows the implementer to make use of associated types @@ -111,7 +116,7 @@ impl OwnedAggregateDTypes { /// all instances of the aggregate. In almost all cases, this struct will be an empty unit /// struct, since most aggregates do not require any global state. /// -/// Execution methods receive the options and the resolved [`AggregateDTypes`] of the aggregate +/// Execution methods receive the options and the resolved [`AggregateDTypesRef`] of the aggregate /// they operate on, so partial states only hold accumulated values. Callers must pass the same /// options and dtypes for the whole lifetime of a partial state. pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { @@ -194,12 +199,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn empty_partial( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, ) -> VortexResult; /// Parse a partial scalar into the typed partial state. /// - /// The scalar must have dtype `dtypes.partial`; this is the inverse of [`to_scalar`]. Partial + /// The scalar must have dtype `dtypes.partial_dtype`; this is the inverse of [`to_scalar`]. Partial /// scalars are produced by aggregate kernels, cached statistics, and other accumulators' /// [`to_scalar`]. /// @@ -211,7 +216,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn partial_from_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult; @@ -225,28 +230,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn merge_partials( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, first: Self::Partial, second: Self::Partial, ) -> VortexResult; - /// Reduce a sequence of partial states in iteration order, starting from [`empty_partial`]. - /// - /// [`empty_partial`]: AggregateFnVTable::empty_partial - fn reduce_partials( - &self, - options: &Self::Options, - dtypes: AggregateDTypes<'_>, - partials: impl IntoIterator, - ) -> VortexResult { - partials - .into_iter() - .try_fold(self.empty_partial(options, dtypes)?, |acc, partial| { - self.merge_partials(options, dtypes, acc, partial) - }) - } - - /// Convert the partial state into a partial scalar of dtype `dtypes.partial`. + /// Convert the partial state into a partial scalar of dtype `dtypes.partial_dtype`. /// /// This is the inverse of [`partial_from_scalar`]: parsing the returned scalar must /// reconstruct a state that behaves identically under accumulation, merging, saturation, and @@ -256,7 +245,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn to_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult; @@ -265,7 +254,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn is_saturated( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool; @@ -280,7 +269,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn try_accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &ArrayRef, _ctx: &mut ExecutionCtx, @@ -292,7 +281,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { fn accumulate( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -300,19 +289,19 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Finalize an array of partial states into an array of aggregate results. /// - /// The `states` array has dtype `dtypes.partial`; the result must have dtype `dtypes.result`. + /// The `states` array has dtype `dtypes.partial_dtype`; the result must have dtype `dtypes.return_dtype`. fn finalize( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, states: ArrayRef, ) -> VortexResult; - /// Finalize a partial state into an aggregate result of dtype `dtypes.result`. + /// Finalize a partial state into an aggregate result of dtype `dtypes.return_dtype`. fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult; } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 68bd20a6ed9..f924d4bfe6d 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -16,7 +16,7 @@ use std::num::NonZeroU32; use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateDTypesRef; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::dtype::DType; @@ -302,7 +302,7 @@ impl AggregateFnVTable for BloomFilter { fn empty_partial( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(BloomPartial::from(options)) } @@ -315,7 +315,7 @@ impl AggregateFnVTable for BloomFilter { fn partial_from_scalar( &self, options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { let mut partial = BloomPartial::from(options); @@ -335,7 +335,7 @@ impl AggregateFnVTable for BloomFilter { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -349,7 +349,7 @@ impl AggregateFnVTable for BloomFilter { fn to_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { let bytes: Vec = partial.serialize(); @@ -362,7 +362,7 @@ impl AggregateFnVTable for BloomFilter { fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> bool { partial.is_saturated() @@ -371,7 +371,7 @@ impl AggregateFnVTable for BloomFilter { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -388,7 +388,7 @@ impl AggregateFnVTable for BloomFilter { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { Ok(partials) @@ -397,7 +397,7 @@ impl AggregateFnVTable for BloomFilter { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -423,8 +423,8 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; use vortex_array::aggregate_fn::DynAccumulator; - use vortex_array::aggregate_fn::OwnedAggregateDTypes; use vortex_array::test_harness::check_metadata; use super::*; @@ -461,8 +461,8 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { } /// Aggregate dtypes for a binary input; the Bloom filter's dtypes do not depend on options. - fn binary_dtypes(options: &BloomOptions) -> VortexResult { - OwnedAggregateDTypes::try_new( + fn binary_dtypes(options: &BloomOptions) -> VortexResult { + AggregateDTypes::try_new( &BloomFilter, options, DType::Binary(Nullability::NonNullable), @@ -473,7 +473,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { fn saturation_false_when_empty() -> VortexResult<()> { let options = BloomOptions::default(); let dtypes = binary_dtypes(&options)?; - let partial = BloomFilter.reduce_partials(&options, dtypes.borrow(), [])?; + let partial = BloomFilter.empty_partial(&options, dtypes.borrow())?; assert!(!BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); Ok(()) } @@ -506,10 +506,11 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { assert!( BloomFilter - .reduce_partials( + .merge_partials( &smaller, dtypes.borrow(), - [BloomPartial::from(&smaller), bigger] + BloomPartial::from(&smaller), + bigger ) .is_err(), "reducing partials built with different blocks_count must fail loudly, not corrupt state" @@ -518,7 +519,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { } #[test] - fn reduce_partials_unions_two_disjoint_partials() -> VortexResult<()> { + fn merge_partials_unions_two_disjoint_partials() -> VortexResult<()> { let options = BloomOptions::default(); let mut partial = BloomPartial::from(&options); for i in 0..50i64 { @@ -539,7 +540,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { let dtypes = binary_dtypes(&options)?; let partial = - BloomFilter.reduce_partials(&options, dtypes.borrow(), [partial, secondary_partial])?; + BloomFilter.merge_partials(&options, dtypes.borrow(), partial, secondary_partial)?; assert!( partial == expected, diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs index b72ae7a6f84..bf839216ac0 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs @@ -93,11 +93,11 @@ impl TryFrom for HashFn { /// ```rust /// use vortex_array::dtype::{DType, Nullability}; /// use vortex_layout::layouts::zoned::aggregates::bloom_filter::{BloomFilter, BloomOptions}; -/// use vortex_array::aggregate_fn::{AggregateFnVTable, OwnedAggregateDTypes}; +/// use vortex_array::aggregate_fn::{AggregateFnVTable, AggregateDTypes}; /// /// let filter = BloomFilter {}; /// let options = BloomOptions::default(); -/// let dtypes = OwnedAggregateDTypes::try_new( +/// let dtypes = AggregateDTypes::try_new( /// &filter, /// &options, /// DType::Binary(Nullability::NonNullable), diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index e6132ee14aa..4529da8c73e 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/vortex-spatial/src/aggregate_fn/aabb.rs @@ -8,7 +8,7 @@ use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateDTypesRef; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; @@ -157,7 +157,7 @@ impl AggregateFnVTable for GeometryAabb { fn empty_partial( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(AabbPartial { rect: None }) } @@ -165,7 +165,7 @@ impl AggregateFnVTable for GeometryAabb { fn partial_from_scalar( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, scalar: Scalar, ) -> VortexResult { // A null box is an empty group's AABB. @@ -177,7 +177,7 @@ impl AggregateFnVTable for GeometryAabb { fn merge_partials( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { @@ -190,19 +190,19 @@ impl AggregateFnVTable for GeometryAabb { fn to_scalar( &self, _options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { Ok(match partial.rect { Some(rect) => rect_to_storage(rect), - None => Scalar::null(dtypes.partial.clone()), + None => Scalar::null(dtypes.partial_dtype.clone()), }) } fn is_saturated( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, _partial: &Self::Partial, ) -> bool { // An AABB can always grow, so it is never saturated. @@ -212,7 +212,7 @@ impl AggregateFnVTable for GeometryAabb { fn accumulate( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -247,7 +247,7 @@ impl AggregateFnVTable for GeometryAabb { fn finalize( &self, _options: &Self::Options, - _dtypes: AggregateDTypes<'_>, + _dtypes: AggregateDTypesRef<'_>, partials: ArrayRef, ) -> VortexResult { // The stored partial is already the AABB struct, so finalizing is the identity. @@ -257,7 +257,7 @@ impl AggregateFnVTable for GeometryAabb { fn finalize_scalar( &self, options: &Self::Options, - dtypes: AggregateDTypes<'_>, + dtypes: AggregateDTypesRef<'_>, partial: &Self::Partial, ) -> VortexResult { self.to_scalar(options, dtypes, partial) @@ -270,10 +270,10 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::EmptyOptions; - use vortex_array::aggregate_fn::OwnedAggregateDTypes; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -425,19 +425,20 @@ mod tests { Ok(()) } - /// `reduce_partials` unions partial boxes - the path the zoned writer takes when a zone's + /// `merge_partials` unions partial boxes - the path the zoned writer takes when a zone's /// array is chunked. #[test] - fn reduce_partials_unions_boxes() -> VortexResult<()> { + fn merge_partials_unions_boxes() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let dtypes = OwnedAggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; + let dtypes = AggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; let bbox = |xmin, ymin, xmax, ymax| AabbPartial { rect: Some(SpatialRect::new((xmin, ymin), (xmax, ymax))), }; - let reduced = GeometryAabb.reduce_partials( + let reduced = GeometryAabb.merge_partials( &EmptyOptions, dtypes.borrow(), - [bbox(0.0, 0.0, 1.0, 1.0), bbox(5.0, -2.0, 7.0, 3.0)], + bbox(0.0, 0.0, 1.0, 1.0), + bbox(5.0, -2.0, 7.0, 3.0), )?; assert_eq!( aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &reduced)?)?, @@ -446,17 +447,16 @@ mod tests { Ok(()) } - /// An empty partial (an empty group's AABB) is a no-op in `reduce_partials`. + /// An empty partial (an empty group's AABB) is a no-op in `merge_partials`. #[test] - fn reduce_partials_ignores_empty() -> VortexResult<()> { + fn merge_partials_ignores_empty() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let dtypes = OwnedAggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; + let dtypes = AggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; let empty = AabbPartial { rect: None }; let value = AabbPartial { rect: Some(SpatialRect::new((0.0, 0.0), (1.0, 1.0))), }; - let reduced = - GeometryAabb.reduce_partials(&EmptyOptions, dtypes.borrow(), [value, empty])?; + let reduced = GeometryAabb.merge_partials(&EmptyOptions, dtypes.borrow(), value, empty)?; assert_eq!( aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &reduced)?)?, (0.0, 0.0, 1.0, 1.0) From f4abbfa632187785bb8e7cd9d551a2f5f477ee90 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 10 Sep 2026 21:09:34 +0000 Subject: [PATCH 4/6] Rename combine_partials to combine_partial and expose AggregateDTypes fields `DynAccumulator::combine_partials` took a single `Scalar`, so the plural was misleading; the sibling `merge_from` already distinguishes combining another accumulator from combining one partial scalar. `AggregateDTypes` now exposes `dtype`, `return_dtype` and `partial_dtype` as public fields, matching `AggregateDTypesRef`, and drops the three accessors that only returned them. `try_new` and `borrow` stay, since they resolve the dtypes through the vtable and lend them out rather than just reading a field. Signed-off-by: "Robert" Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- vortex-array/src/aggregate_fn/accumulator.rs | 32 +++++++++---------- .../src/aggregate_fn/accumulator_grouped.rs | 18 +++++------ vortex-array/src/aggregate_fn/combined.rs | 4 +-- .../src/aggregate_fn/fns/sum/decimal.rs | 4 +-- vortex-array/src/aggregate_fn/vtable.rs | 26 ++++----------- 5 files changed, 35 insertions(+), 49 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 7a15d7a01f8..79e6252ef50 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -114,7 +114,7 @@ pub trait DynAccumulator: 'static + Send { /// accumulator's options and input dtype, and represents input following the input already /// accumulated. Parsing and merging both run through the typed vtable, so they inline into /// a single monomorphized call per aggregate. - fn combine_partial_scalar(&mut self, partial: Scalar) -> VortexResult<()>; + fn combine_partial(&mut self, partial: Scalar) -> VortexResult<()>; /// Whether the accumulator's result is fully determined. fn is_saturated(&self) -> bool; @@ -158,9 +158,9 @@ impl DynAccumulator for Accumulator { } vortex_ensure!( - batch.dtype() == self.dtypes.dtype(), + batch.dtype() == &self.dtypes.dtype, "Input DType mismatch: expected {}, got {}", - self.dtypes.dtype(), + self.dtypes.dtype, batch.dtype() ); @@ -169,20 +169,20 @@ impl DynAccumulator for Accumulator { if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - let partial = if partial.dtype() == self.dtypes.partial_dtype() { + let partial = if partial.dtype() == &self.dtypes.partial_dtype { partial } else { vortex_ensure!( partial .dtype() - .eq_ignore_nullability(self.dtypes.partial_dtype()), + .eq_ignore_nullability(&self.dtypes.partial_dtype), "Aggregate {} read legacy stat {} with dtype {}, expected {}", self.aggregate_fn, stat, partial.dtype(), - self.dtypes.partial_dtype(), + self.dtypes.partial_dtype, ); - partial.cast(self.dtypes.partial_dtype())? + partial.cast(&self.dtypes.partial_dtype)? }; self.fold_partial_scalar(partial)?; return Ok(()); @@ -203,10 +203,10 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)? { vortex_ensure!( - result.dtype() == self.dtypes.partial_dtype(), + result.dtype() == &self.dtypes.partial_dtype, "Aggregate kernel returned {}, expected {}", result.dtype(), - self.dtypes.partial_dtype(), + self.dtypes.partial_dtype, ); self.fold_partial_scalar(result)?; return Ok(()); @@ -240,10 +240,10 @@ impl DynAccumulator for Accumulator { && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)? { vortex_ensure!( - result.dtype() == self.dtypes.partial_dtype(), + result.dtype() == &self.dtypes.partial_dtype, "Aggregate kernel returned {}, expected {}", result.dtype(), - self.dtypes.partial_dtype(), + self.dtypes.partial_dtype, ); self.fold_partial_scalar(result)?; return Ok(()); @@ -269,7 +269,7 @@ impl DynAccumulator for Accumulator { ); }; vortex_ensure!( - other.options == self.options && other.dtypes.dtype() == self.dtypes.dtype(), + other.options == self.options && other.dtypes.dtype == self.dtypes.dtype, "Cannot merge {} accumulators with different options or input dtypes", self.aggregate_fn, ); @@ -279,12 +279,12 @@ impl DynAccumulator for Accumulator { } } - fn combine_partial_scalar(&mut self, partial: Scalar) -> VortexResult<()> { + fn combine_partial(&mut self, partial: Scalar) -> VortexResult<()> { vortex_ensure!( - partial.dtype() == self.dtypes.partial_dtype(), + partial.dtype() == &self.dtypes.partial_dtype, "Partial DType mismatch for {}: expected {}, got {}", self.aggregate_fn, - self.dtypes.partial_dtype(), + self.dtypes.partial_dtype, partial.dtype(), ); self.fold_partial_scalar(partial) @@ -467,7 +467,7 @@ mod tests { let acc = mean_f64_accumulator().expect("build accumulator"); let sum = Scalar::primitive(42.0f64, Nullability::Nullable); let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.dtypes.partial_dtype().clone(), vec![sum, count]) + Scalar::struct_(acc.dtypes.partial_dtype, vec![sum, count]) } /// Kernel registered for `(Dict, Combined)` fires in preference to diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 1eedfdd6ce1..c7643db135d 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -234,9 +234,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { ), }; vortex_ensure!( - elements_dtype.as_ref() == self.dtypes.dtype(), + elements_dtype.as_ref() == &self.dtypes.dtype, "Input DType mismatch: expected {}, got {}", - self.dtypes.dtype(), + self.dtypes.dtype, elements_dtype ); @@ -258,7 +258,7 @@ impl DynGroupedAccumulator for GroupedAccumulator { if states.len() == 1 { return Ok(states.pop().vortex_expect("checked one partial")); } - Ok(ChunkedArray::try_new(states, self.dtypes.partial_dtype().clone())?.into_array()) + Ok(ChunkedArray::try_new(states, self.dtypes.partial_dtype.clone())?.into_array()) } fn finish(&mut self) -> VortexResult { @@ -268,9 +268,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { .finalize(&self.options, self.dtypes.borrow(), states)?; vortex_ensure!( - results.dtype() == self.dtypes.return_dtype(), + results.dtype() == &self.dtypes.return_dtype, "Return DType mismatch: expected {}, got {}", - self.dtypes.return_dtype(), + self.dtypes.return_dtype, results.dtype() ); @@ -342,10 +342,10 @@ impl GroupedAccumulator { let mut accumulator = Accumulator::try_new( self.vtable.clone(), self.options.clone(), - self.dtypes.dtype().clone(), + self.dtypes.dtype.clone(), )?; let mut states = - builder_with_capacity_in(self.dtypes.partial_dtype(), grouped.len(), ctx.allocator()); + builder_with_capacity_in(&self.dtypes.partial_dtype, grouped.len(), ctx.allocator()); let group_ranges = grouped.group_ranges(ctx)?; let group_validity = grouped.group_validity(ctx)?; @@ -364,9 +364,9 @@ impl GroupedAccumulator { fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> { vortex_ensure!( - state.dtype() == self.dtypes.partial_dtype(), + state.dtype() == &self.dtypes.partial_dtype, "State DType mismatch: expected {}, got {}", - self.dtypes.partial_dtype(), + self.dtypes.partial_dtype, state.dtype() ); self.partials.push(state); diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 8e8c7d3cb62..589ac5a4ef7 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -217,8 +217,8 @@ impl AggregateFnVTable for Combined { let r_field = s .field(rname) .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", rname))?; - left.combine_partial_scalar(l_field)?; - right.combine_partial_scalar(r_field)?; + left.combine_partial(l_field)?; + right.combine_partial(r_field)?; } Ok((left, right)) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 1a92e48d1c5..7825e5c692b 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -442,8 +442,8 @@ mod tests { let options = NumericalAggregateOpts::default(); let dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; assert_eq!( - dtypes.return_dtype(), - &DType::Decimal(DecimalDType::new(37, 0), Nullable) + dtypes.return_dtype, + DType::Decimal(DecimalDType::new(37, 0), Nullable) ); // Set state to 10^37 - 1. diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 199ce4efc76..3e5c82c14f6 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -32,11 +32,12 @@ use crate::scalar::Scalar; #[derive(Clone, Debug)] pub struct AggregateDTypes { /// The DType of the input. - dtype: DType, - /// The DType of the aggregate. - return_dtype: DType, - /// The DType of the partial accumulator state. - partial_dtype: DType, + pub dtype: DType, + /// The DType of the aggregate, as reported by [`AggregateFnVTable::return_dtype`]. + pub return_dtype: DType, + /// The DType of the partial accumulator state, as reported by + /// [`AggregateFnVTable::partial_dtype`]. + pub partial_dtype: DType, } impl AggregateDTypes { @@ -69,21 +70,6 @@ impl AggregateDTypes { }) } - /// The DType of the input. - pub fn dtype(&self) -> &DType { - &self.dtype - } - - /// The DType of the aggregate. - pub fn return_dtype(&self) -> &DType { - &self.return_dtype - } - - /// The DType of the partial accumulator state. - pub fn partial_dtype(&self) -> &DType { - &self.partial_dtype - } - /// Lend the dtypes to an execution method. pub fn borrow(&self) -> AggregateDTypesRef<'_> { AggregateDTypesRef { From b6a35017d455fef4ef92f03742a33d2c583096f8 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 10 Sep 2026 22:54:43 +0000 Subject: [PATCH 5/6] Build empty partials in tests through empty_partial Tests that needed an empty partial were spelling out the identity state by hand, restating in the test what each vtable already defines. They now call `empty_partial`, so a change to an aggregate's identity state cannot leave a test asserting against a stale hand-written one. Covers First, Last, BoundedMin, BoundedMax, GeometryAabb and BloomFilter. First and Last no longer touch their partial structs at all: the non-empty cases go through `partial_from_scalar`, which for both is exactly a wrapped non-null scalar. `Accumulator::empty_partial` becomes `pub(crate)`, matching the existing visibility of `fold_partial`, so the bounded min/max tests can fold an empty partial without rebuilding the dtypes. Left alone: the is_sorted and is_constant partials, whose tests assert a non-empty verdict is distinguishable from the empty state and so must build it directly; the non-identity partials in sum, min_max and the decimal sum tests; and the BloomPartial unit tests, which exercise the partial type itself rather than the aggregate. Signed-off-by: "Robert" Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- vortex-array/src/aggregate_fn/accumulator.rs | 2 +- .../src/aggregate_fn/fns/bounded_max/mod.rs | 5 ++--- .../src/aggregate_fn/fns/bounded_min/mod.rs | 7 ++----- vortex-array/src/aggregate_fn/fns/first/mod.rs | 11 ++++++----- vortex-array/src/aggregate_fn/fns/last/mod.rs | 11 ++++++----- .../layouts/zoned/aggregates/bloom_filter/mod.rs | 13 +++++++------ vortex-spatial/src/aggregate_fn/aabb.rs | 2 +- 7 files changed, 25 insertions(+), 26 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 79e6252ef50..117ac35f89c 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -59,7 +59,7 @@ impl Accumulator { } /// The identity partial state: the state of a group with no accumulated values. - fn empty_partial(&self) -> VortexResult { + pub(crate) fn empty_partial(&self) -> VortexResult { self.vtable .empty_partial(&self.options, self.dtypes.borrow()) } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 8ef939b3475..abaaa0a2791 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -484,9 +484,8 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.fold_partial(BoundedMaxPartial { - state: BoundedMaxState::Empty, - })?; + let empty = acc.empty_partial()?; + acc.fold_partial(empty)?; assert_eq!( acc.finish()?, diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 1ec0a17a72c..149d4a87f2b 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -287,8 +287,6 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_min::BoundedMin; use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions; - use crate::aggregate_fn::fns::bounded_min::BoundedMinPartial; - use crate::aggregate_fn::fns::bounded_min::BoundedMinState; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::array_session; @@ -362,9 +360,8 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.fold_partial(BoundedMinPartial { - state: BoundedMinState::Empty, - })?; + let empty = acc.empty_partial()?; + acc.fold_partial(empty)?; assert_eq!( acc.finish()?, diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 88214ca7e86..5062430f49b 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -172,7 +172,6 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::first::First; - use crate::aggregate_fn::fns::first::FirstPartial; use crate::aggregate_fn::fns::first::first; use crate::array_session; use crate::arrays::ChunkedArray; @@ -296,15 +295,17 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let owned = AggregateDTypes::try_new(&First, &EmptyOptions, dtype)?; let dtypes = owned.borrow(); - let partial_of = |value: Option| FirstPartial { value }; + let partial_of = |value: i32| { + First.partial_from_scalar(&EmptyOptions, dtypes, Scalar::primitive(value, Nullable)) + }; // An empty partial means the sub-accumulator saw nothing valid - it is ignored. - let empty = partial_of(None); + let empty = First.empty_partial(&EmptyOptions, dtypes)?; assert!(!First.is_saturated(&EmptyOptions, dtypes, &empty)); // The first non-empty partial wins; subsequent valid partials are dropped. - let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); - let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); + let five = partial_of(5)?; + let seven = partial_of(7)?; let merge = |first, second| First.merge_partials(&EmptyOptions, dtypes, first, second); let state = merge(merge(empty, five)?, seven)?; assert!(First.is_saturated(&EmptyOptions, dtypes, &state)); diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 96d4e03e725..52987c6190f 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -170,7 +170,6 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::last::Last; - use crate::aggregate_fn::fns::last::LastPartial; use crate::aggregate_fn::fns::last::last; use crate::array_session; use crate::arrays::ChunkedArray; @@ -293,12 +292,14 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let owned = AggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?; let dtypes = owned.borrow(); - let partial_of = |value: Option| LastPartial { value }; + let partial_of = |value: i32| { + Last.partial_from_scalar(&EmptyOptions, dtypes, Scalar::primitive(value, Nullable)) + }; - let five = partial_of(Some(Scalar::primitive(5i32, Nullable))); - let seven = partial_of(Some(Scalar::primitive(7i32, Nullable))); + let five = partial_of(5)?; + let seven = partial_of(7)?; // An empty partial must not clobber a prior value. - let empty = partial_of(None); + let empty = Last.empty_partial(&EmptyOptions, dtypes)?; // The last non-empty partial in order replaces the prior values. let merge = |first, second| Last.merge_partials(&EmptyOptions, dtypes, first, second); diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index f924d4bfe6d..90c7f74d17a 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -493,7 +493,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { fn mismatched_block_counts_are_rejected() -> VortexResult<()> { let smaller = BloomOptions::new(NonZeroU32::new(4).unwrap(), HashFn::XxHash3_64); let dtypes = binary_dtypes(&smaller)?; - let bigger = BloomPartial::from(&BloomOptions::default()); + let bigger = BloomFilter.empty_partial(&BloomOptions::default(), dtypes.borrow())?; let bigger_scalar = BloomFilter.to_scalar(&BloomOptions::default(), dtypes.borrow(), &bigger)?; @@ -509,7 +509,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { .merge_partials( &smaller, dtypes.borrow(), - BloomPartial::from(&smaller), + BloomFilter.empty_partial(&smaller, dtypes.borrow())?, bigger ) .is_err(), @@ -521,24 +521,25 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { #[test] fn merge_partials_unions_two_disjoint_partials() -> VortexResult<()> { let options = BloomOptions::default(); - let mut partial = BloomPartial::from(&options); + let dtypes = binary_dtypes(&options)?; + + let mut partial = BloomFilter.empty_partial(&options, dtypes.borrow())?; for i in 0..50i64 { partial.insert(i.to_le_bytes()); } - let mut secondary_partial = BloomPartial::from(&options); + let mut secondary_partial = BloomFilter.empty_partial(&options, dtypes.borrow())?; for i in 50..100i64 { secondary_partial.insert(i.to_le_bytes()); } // The following expected works because seed is equal for all. // If the seed is different for both partials, then this will fail. - let mut expected = BloomPartial::from(&options); + let mut expected = BloomFilter.empty_partial(&options, dtypes.borrow())?; for i in 0..100i64 { expected.insert(i.to_le_bytes()); } - let dtypes = binary_dtypes(&options)?; let partial = BloomFilter.merge_partials(&options, dtypes.borrow(), partial, secondary_partial)?; diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index 4529da8c73e..4b1b0c35bba 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/vortex-spatial/src/aggregate_fn/aabb.rs @@ -452,7 +452,7 @@ mod tests { fn merge_partials_ignores_empty() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); let dtypes = AggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; - let empty = AabbPartial { rect: None }; + let empty = GeometryAabb.empty_partial(&EmptyOptions, dtypes.borrow())?; let value = AabbPartial { rect: Some(SpatialRect::new((0.0, 0.0), (1.0, 1.0))), }; From 798246ed5cdd2e930cffcce24e472848ac2b9a87 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 11 Sep 2026 11:14:38 +0000 Subject: [PATCH 6/6] Drop the test-only BloomPartial constructor from raw blocks `From>` existed only so one test could hand-build a saturated filter, which also let that partial disagree with its own options: it carried four blocks while the options said 256. The test now builds the filter through `partial_from_scalar` from an all-ones buffer sized to the options' block count, so the partial is consistent with the options it is checked against and nothing outside the module can bypass `BloomPartial`'s invariants. Signed-off-by: "Robert" Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd --- .../src/layouts/zoned/aggregates/bloom_filter/mod.rs | 11 +++++++++-- .../zoned/aggregates/bloom_filter/partial/mod.rs | 7 ------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 90c7f74d17a..3ad0625752b 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -427,6 +427,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::test_harness::check_metadata; + use super::partial::BLOCK_SIZE; use super::*; pub fn setup() -> VortexResult { @@ -482,8 +483,14 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { fn saturation_true_when_every_block_is_full() -> VortexResult<()> { let options = BloomOptions::default(); let dtypes = binary_dtypes(&options)?; - let blocks = vec![[u32::MAX; 8]; 4]; - let partial = BloomPartial::from(blocks); + + // Every bit of every block set - the only state the filter reports as saturated. + let full = vec![u8::MAX; options.blocks_count().get() as usize * BLOCK_SIZE]; + let partial = BloomFilter.partial_from_scalar( + &options, + dtypes.borrow(), + Scalar::binary(full, Nullability::NonNullable), + )?; assert!(BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); Ok(()) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs index bf839216ac0..2c106dde75e 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs @@ -252,13 +252,6 @@ impl From<&BloomOptions> for BloomPartial { } } -#[cfg(test)] -impl From> for BloomPartial { - fn from(value: Vec<[u32; 8]>) -> Self { - BloomPartial { blocks: value } - } -} - #[cfg(test)] mod tests { use std::num::NonZeroU32;