diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 69bae4e1053..117ac35f89c 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -1,13 +1,17 @@ // 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; 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; @@ -27,46 +31,69 @@ 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. - 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: AggregateDTypes, /// 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 { 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 partial = vtable.empty_partial(&options, &dtype)?; - let aggregate_fn = AggregateFn::new(vtable.clone(), options).erased(); + let dtypes = AggregateDTypes::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, - partial, + dtypes, + partial: None, }) } + + /// The identity partial state: the state of a group with no accumulated values. + pub(crate) fn empty_partial(&self) -> VortexResult { + self.vtable + .empty_partial(&self.options, self.dtypes.borrow()) + } + + /// 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(()) + } + + /// 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() { + // Merging the incoming partial with the empty state is the identity. + None => other, + Some(current) => { + self.vtable + .merge_partials(&self.options, self.dtypes.borrow(), current, other)? + } + }); + Ok(()) + } + + /// 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<()> { + 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 @@ -75,11 +102,19 @@ 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 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<()>; + + /// 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 the - /// options and input dtype used to construct this accumulator. - fn combine_partials(&mut self, other: Scalar) -> VortexResult<()>; + /// 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(&mut self, partial: Scalar) -> 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 { @@ -113,9 +158,9 @@ impl DynAccumulator for Accumulator { } vortex_ensure!( - batch.dtype() == &self.dtype, + batch.dtype() == &self.dtypes.dtype, "Input DType mismatch: expected {}, got {}", - self.dtype, + self.dtypes.dtype, batch.dtype() ); @@ -124,20 +169,22 @@ 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_dtype { partial } else { vortex_ensure!( - partial.dtype().eq_ignore_nullability(&self.partial_dtype), + partial + .dtype() + .eq_ignore_nullability(&self.dtypes.partial_dtype), "Aggregate {} read legacy stat {} with dtype {}, expected {}", self.aggregate_fn, stat, partial.dtype(), - self.partial_dtype, + self.dtypes.partial_dtype, ); - partial.cast(&self.partial_dtype)? + partial.cast(&self.dtypes.partial_dtype)? }; - self.vtable.combine_partials(&mut self.partial, partial)?; + self.fold_partial_scalar(partial)?; return Ok(()); } @@ -156,18 +203,23 @@ 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_dtype, "Aggregate kernel returned {}, expected {}", result.dtype(), - self.partial_dtype, + self.dtypes.partial_dtype, ); - self.vtable.combine_partials(&mut self.partial, result)?; + self.fold_partial_scalar(result)?; 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(&self.options, self.dtypes.borrow(), partial, batch, ctx)? + { return Ok(()); } @@ -188,12 +240,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_dtype, "Aggregate kernel returned {}, expected {}", result.dtype(), - self.partial_dtype, + self.dtypes.partial_dtype, ); - self.vtable.combine_partials(&mut self.partial, result)?; + self.fold_partial_scalar(result)?; return Ok(()); } @@ -203,30 +255,71 @@ 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(&self.options, self.dtypes.borrow(), 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.dtypes.dtype == self.dtypes.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 combine_partial(&mut self, partial: Scalar) -> VortexResult<()> { + vortex_ensure!( + partial.dtype() == &self.dtypes.partial_dtype, + "Partial DType mismatch for {}: expected {}, got {}", + self.aggregate_fn, + self.dtypes.partial_dtype, + partial.dtype(), + ); + self.fold_partial_scalar(partial) + } + + 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(&self.options, self.dtypes.borrow(), 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 dtypes = self.dtypes.borrow(); + let partial = match &self.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_dtype, "Aggregate returned incorrect DType on partial_scalar: expected {}, got {}", - self.partial_dtype, + dtypes.partial_dtype, partial.dtype(), ); } @@ -235,12 +328,20 @@ impl DynAccumulator for Accumulator { } fn final_scalar(&self) -> VortexResult { - let result = self.vtable.finalize_scalar(&self.partial)?; + let dtypes = self.dtypes.borrow(); + let result = match &self.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.return_dtype, "Aggregate returned incorrect DType on final_scalar: expected {}, got {}", - self.return_dtype, + dtypes.return_dtype, result.dtype(), ); @@ -272,6 +373,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 +381,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 +391,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 +426,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; @@ -361,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.partial_dtype, vec![sum, count]) + Scalar::struct_(acc.dtypes.partial_dtype, vec![sum, count]) } /// Kernel registered for `(Dict, Combined)` fires in preference to @@ -448,4 +554,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/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index eda0acdad97..c7643db135d 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; @@ -17,6 +16,7 @@ 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; @@ -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: AggregateDTypes, /// 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 = AggregateDTypes::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.dtype, "Input DType mismatch: expected {}, got {}", - self.dtype, + self.dtypes.dtype, 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_dtype.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.return_dtype, "Return DType mismatch: expected {}, got {}", - self.return_dtype, + self.dtypes.return_dtype, 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.dtype.clone(), )?; let mut states = - builder_with_capacity_in(&self.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)?; @@ -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_dtype, "State DType mismatch: expected {}, got {}", - self.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 70a385add48..589ac5a4ef7 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::AccumulatorRef; +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::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; @@ -48,6 +50,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 { @@ -76,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.return_dtype`. + fn finalize( + &self, + options: &CombinedOptions, + dtypes: AggregateDTypesRef<'_>, + left: ArrayRef, + right: ArrayRef, + ) -> VortexResult; + + /// Combine the finalized child scalars into a result of dtype `dtypes.return_dtype`. + fn finalize_scalar( + &self, + options: &CombinedOptions, + dtypes: AggregateDTypesRef<'_>, + left_scalar: Scalar, + right_scalar: Scalar, + ) -> VortexResult; /// Serialize the options for this combined aggregate. Default: not serializable. fn serialize(&self, options: &CombinedOptions) -> VortexResult>> { @@ -126,14 +149,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() @@ -147,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 { @@ -160,49 +194,69 @@ impl AggregateFnVTable for Combined { fn empty_partial( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypesRef<'_>, ) -> 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, - )) + self.new_child_accumulators(options, dtypes.dtype) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); + fn partial_from_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + 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(); + 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.combine_partial(l_field)?; + right.combine_partial(r_field)?; } - 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(()) - } - - fn to_scalar(&self, 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((left, right)) } - fn reset(&self, partial: &mut Self::Partial) { - partial.0.reset(); - partial.1.reset(); + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + (mut left, mut right): Self::Partial, + (mut other_left, mut other_right): Self::Partial, + ) -> VortexResult { + // 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 is_saturated(&self, partial: &Self::Partial) -> bool { + fn to_scalar( + &self, + _options: &Self::Options, + 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_dtype.clone(), + vec![l_scalar, r_scalar], + )) + } + + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { partial.0.is_saturated() && partial.1.is_saturated() } @@ -213,6 +267,8 @@ impl AggregateFnVTable for Combined { /// `true` so [`Self::accumulate`] is unreachable. fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -224,6 +280,8 @@ impl AggregateFnVTable for Combined { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -231,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: 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 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 = 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) } - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + 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 68a58908018..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,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -63,30 +64,52 @@ impl AggregateFnVTable for AllNan { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + bool::try_from(&scalar) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::bool(*partial, Nullability::Nullable)) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first && second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -102,6 +125,8 @@ impl AggregateFnVTable for AllNan { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -121,12 +146,22 @@ impl AggregateFnVTable for AllNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: 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 a8392f3fa74..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,6 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -150,42 +151,62 @@ impl AggregateFnVTable for AllNonDistinct { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(AllNonDistinctPartial { all_non_distinct: true, }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if !partial.all_non_distinct { - return Ok(()); - } + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + Ok(AllNonDistinctPartial { + all_non_distinct: scalar.as_bool().value().unwrap_or(false), + }) + } - if !other.as_bool().value().unwrap_or(false) { - partial.all_non_distinct = false; - } - Ok(()) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(AllNonDistinctPartial { + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::bool( partial.all_non_distinct, Nullability::NonNullable, )) } - fn reset(&self, partial: &mut Self::Partial) { - partial.all_non_distinct = true; - } - #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !partial.all_non_distinct } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -235,11 +256,21 @@ impl AggregateFnVTable for AllNonDistinct { } } - fn finalize(&self, _partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _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: AggregateDTypesRef<'_>, + 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 fe8527da966..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,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -63,30 +64,52 @@ impl AggregateFnVTable for AllNonNan { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + bool::try_from(&scalar) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::bool(*partial, Nullability::Nullable)) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first && second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -97,6 +120,8 @@ impl AggregateFnVTable for AllNonNan { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -111,12 +136,22 @@ impl AggregateFnVTable for AllNonNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: 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 c07dbb907c9..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,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -54,30 +55,52 @@ impl AggregateFnVTable for AllNonNull { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + bool::try_from(&scalar) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::bool(*partial, Nullability::NonNullable)) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first && second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -88,6 +111,8 @@ impl AggregateFnVTable for AllNonNull { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -101,12 +126,22 @@ impl AggregateFnVTable for AllNonNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: 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 ec64e3d5c43..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,6 +9,7 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; @@ -54,30 +55,52 @@ impl AggregateFnVTable for AllNull { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(true) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - *partial &= bool::try_from(&other)?; - Ok(()) + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + bool::try_from(&scalar) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::bool(*partial, Nullability::NonNullable)) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first && second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = true; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !*partial } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -88,6 +111,8 @@ impl AggregateFnVTable for AllNull { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -104,12 +129,22 @@ impl AggregateFnVTable for AllNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: 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 bb27d9f9ab1..abaaa0a2791 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -71,8 +72,6 @@ enum BoundedMaxState { /// Partial accumulator state for the bounded maximum aggregate. pub struct BoundedMaxPartial { state: BoundedMaxState, - element_dtype: DType, - max_bytes: NonZeroUsize, } impl BoundedMaxPartial { @@ -94,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: 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,45 +190,74 @@ impl AggregateFnVTable for BoundedMax { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(BoundedMaxPartial { state: BoundedMaxState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); - } + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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"); + }; - 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 { + BoundedMaxState::Unknown + } else if bound.is_null() { + BoundedMaxState::Empty + } else { + BoundedMaxState::Value(bound) + } }; + Ok(BoundedMaxPartial { state }) + } - if unknown { - partial.unknown(); - } else { - partial.merge_bound(bound); + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + match second.state { + BoundedMaxState::Empty => {} + BoundedMaxState::Value(max) => first.merge_bound(max), + BoundedMaxState::Unknown => first.unknown(), } - Ok(()) + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + 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_( @@ -249,16 +277,19 @@ impl AggregateFnVTable for BoundedMax { } } - fn reset(&self, partial: &mut Self::Partial) { - partial.state = BoundedMaxState::Empty; - } - - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { matches!(partial.state, BoundedMaxState::Unknown) } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -272,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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + partial.final_scalar(dtypes) } } @@ -332,7 +373,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 +484,8 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(make_bounded_max_partial_dtype(values.dtype())))?; + let empty = acc.empty_partial()?; + acc.fold_partial(empty)?; assert_eq!( acc.finish()?, @@ -463,17 +506,10 @@ 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, + })?; 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..149d4a87f2b 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -56,8 +57,6 @@ enum BoundedMinState { /// Partial accumulator state for the bounded minimum aggregate. pub struct BoundedMinPartial { state: BoundedMinState, - element_dtype: DType, - max_bytes: NonZeroUsize, } impl BoundedMinPartial { @@ -145,39 +144,68 @@ impl AggregateFnVTable for BoundedMin { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(BoundedMinPartial { state: BoundedMinState::Empty, - element_dtype: input_dtype.clone(), - max_bytes: options.max_bytes, }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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 }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + if let BoundedMinState::Value(min) = second.state { + first.merge(min); + } + Ok(first) + } + + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.dtype.as_nullable(); match &partial.state { BoundedMinState::Empty => Ok(Scalar::null(dtype)), BoundedMinState::Value(min) => min.cast(&dtype), } } - fn reset(&self, partial: &mut Self::Partial) { - partial.state = BoundedMinState::Empty; - } - - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -191,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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -322,7 +360,8 @@ mod tests { )?; acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(values.dtype().as_nullable()))?; + let empty = acc.empty_partial()?; + acc.fold_partial(empty)?; 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..31929b3318b 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::AggregateDTypesRef; 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"); @@ -60,54 +54,75 @@ impl AggregateFnVTable for Count { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - Ok(CountPartial { - count: 0, - exclude_nans: options.skip_nans && input_dtype.is_float(), - }) + Ok(0) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let val = other + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + Ok(scalar .as_primitive() .typed_value::() - .vortex_expect("count partial should not be null"); - partial.count += val; - Ok(()) + .vortex_expect("count partial should not be null")) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(partial.count, Nullability::NonNullable)) + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first + second) } - fn reset(&self, partial: &mut Self::Partial) { - partial.count = 0; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, 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.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); } - state.count += count; + *state += count; Ok(true) } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -115,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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -138,6 +163,7 @@ 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; @@ -246,16 +272,12 @@ 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 dtypes = AggregateDTypes::try_new(&Count, &options, dtype)?; - let scalar2 = Scalar::primitive(3u64, Nullability::NonNullable); - Count.combine_partials(&mut state, scalar2)?; + let state = Count.merge_partials(&options, dtypes.borrow(), 5, 3)?; - let result = Count.to_scalar(&state)?; - Count.reset(&mut 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 c41e2575057..5062430f49b 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::AggregateDTypesRef; 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, } @@ -60,40 +59,62 @@ impl AggregateFnVTable for First { fn empty_partial( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { + Ok(FirstPartial { value: None }) + } + + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + // The earlier non-empty partial wins; the later one is ignored. + Ok(FirstPartial { + value: first.value.or(second.value), + }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(partial.return_dtype.clone()), + None => Scalar::null(dtypes.return_dtype.clone()), }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.value = None; - } - #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { partial.value.is_some() } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -110,6 +131,8 @@ impl AggregateFnVTable for First { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -117,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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -134,6 +167,7 @@ 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; @@ -259,18 +293,26 @@ 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))?; - 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)); + let owned = AggregateDTypes::try_new(&First, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); + 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 = 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(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)); + 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 96017c3ed92..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,6 +31,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -219,10 +220,17 @@ 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() -> Self { + Self { + is_constant: true, + first_value: None, + } + } + fn check_value(&mut self, value: Scalar) { if !self.is_constant { return; @@ -286,76 +294,87 @@ impl AggregateFnVTable for IsConstant { fn empty_partial( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - Ok(IsConstantPartial { - is_constant: true, - first_value: None, - element_dtype: input_dtype.clone(), - }) + Ok(IsConstantPartial::empty()) } - 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(()); + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + // A null struct means the producing accumulator was empty. + if scalar.is_null() { + return Ok(IsConstantPartial::empty()); } - 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); - - if let Some(other_val) = other_value { - partial.check_value(other_val); - } - - Ok(()) + Ok(IsConstantPartial { + is_constant, + first_value: scalar.as_struct().field_by_idx(1), + }) } - 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut acc: Self::Partial, + partial: Self::Partial, + ) -> VortexResult { + if !partial.is_constant { + acc.is_constant = false; + } else if let Some(value) = partial.first_value { + acc.check_value(value); + } + Ok(acc) } - fn reset(&self, partial: &mut Self::Partial) { - partial.is_constant = true; - partial.first_value = None; + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + 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 { + 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] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !partial.is_constant } fn accumulate( &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -379,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.dtype.as_nullable())); return Ok(()); } @@ -426,11 +445,21 @@ impl AggregateFnVTable for IsConstant { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { if partial.first_value.is_none() { // Empty accumulator → return false. return Ok(Scalar::bool(false, Nullability::NonNullable)); @@ -448,6 +477,13 @@ 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::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 +789,50 @@ 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 owned = AggregateDTypes::try_new(&IsConstant, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); + let partial = IsConstantPartial { + is_constant: false, + first_value: None, + }; + + let scalar = IsConstant.to_scalar(&EmptyOptions, dtypes, &partial)?; + assert!(!scalar.is_null()); + let parsed = IsConstant.partial_from_scalar(&EmptyOptions, dtypes, scalar)?; + assert_eq!( + 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 99a7b99128a..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,6 +26,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -202,11 +203,20 @@ 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() -> Self { + Self { + is_sorted: true, + first_value: None, + last_value: None, + } + } } static NAMES: std::sync::LazyLock = std::sync::LazyLock::new(|| { @@ -279,130 +289,142 @@ impl AggregateFnVTable for IsSorted { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - Ok(IsSortedPartial { - is_sorted: true, - strict: options.strict, - first_value: None, - last_value: None, - element_dtype: input_dtype.clone(), - }) + Ok(IsSortedPartial::empty()) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if !partial.is_sorted { - return Ok(()); - } - - // Null struct means the other accumulator was empty, skip it. - if other.is_null() { - return Ok(()); + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + // A null struct means the producing accumulator was empty. + if scalar.is_null() { + return Ok(IsSortedPartial::empty()); } - 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, + first_value: scalar.as_struct().field_by_idx(2), + last_value: scalar.as_struct().field_by_idx(3), + }) + } - 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); + fn merge_partials( + &self, + options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut acc: Self::Partial, + partial: Self::Partial, + ) -> VortexResult { + if !acc.is_sorted { + return Ok(acc); + } + + 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); } - return Ok(()); + if acc.first_value.is_none() { + acc.first_value = partial.first_value; + } + return Ok(acc); } - // 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 + // 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 { - *self_last <= *other_first_val + *acc_last <= first }; if !boundary_ok { - partial.is_sorted = false; + acc.is_sorted = false; } - } else if !self_last.is_null() && other_first_val.is_null() { + } else if !acc_last.is_null() && first.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 { + acc.is_sorted = false; + } else if acc_last.is_null() && first.is_null() && options.strict { // both null with strict: violates strict sort - partial.is_sorted = false; + acc.is_sorted = false; } } - // 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); + // 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(()) - } - - 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(), - ], - ) - } - } - }) + Ok(acc) } - fn reset(&self, partial: &mut Self::Partial) { - partial.is_sorted = true; - partial.first_value = None; - partial.last_value = None; + fn to_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + 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() { + return Ok(Scalar::null(dtype)); + } + let first_value = partial + .first_value + .clone() + .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 + .clone() + .unwrap_or_else(|| first_value.clone()); + Ok(Scalar::struct_( + dtype, + vec![ + Scalar::bool(partial.is_sorted, Nullability::NonNullable), + Scalar::bool(options.strict, Nullability::NonNullable), + first_value, + last_value, + ], + )) } #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { !partial.is_sorted } fn accumulate( &self, + options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -415,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 @@ -431,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; } @@ -454,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 @@ -472,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( @@ -489,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!(), }; @@ -516,15 +538,22 @@ impl AggregateFnVTable for IsSorted { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partials: ArrayRef, + ) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } - 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)); - } + fn finalize_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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)) } } @@ -557,17 +586,29 @@ 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::AggregateDTypes; + 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 +759,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 owned = AggregateDTypes::try_new(&IsSorted, &options, dtype)?; + let dtypes = owned.borrow(); + let partial = IsSortedPartial { + is_sorted: false, + first_value: None, + last_value: None, + }; + + let scalar = IsSorted.to_scalar(&options, dtypes, &partial)?; + assert!(!scalar.is_null()); + let parsed = IsSorted.partial_from_scalar(&options, dtypes, scalar)?; + assert_eq!( + 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 63ee54a6efa..52987c6190f 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::AggregateDTypesRef; 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, } @@ -60,41 +59,63 @@ impl AggregateFnVTable for Last { fn empty_partial( &self, _options: &Self::Options, - input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { + Ok(LastPartial { value: None }) + } + + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + // The later non-empty partial wins; an empty later partial changes nothing. + Ok(LastPartial { + value: second.value.or(first.value), + }) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.value { Some(v) => v.clone(), - None => Scalar::null(partial.return_dtype.clone()), + None => Scalar::null(dtypes.return_dtype.clone()), }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.value = None; - } - #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _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: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -108,6 +129,8 @@ impl AggregateFnVTable for Last { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _partial: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -115,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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -132,6 +165,7 @@ 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; @@ -256,18 +290,24 @@ 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()))?; - assert_eq!(Last.to_scalar(&state)?, Scalar::primitive(7i32, Nullable)); + let owned = AggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?; + let dtypes = owned.borrow(); + let partial_of = |value: i32| { + Last.partial_from_scalar(&EmptyOptions, dtypes, Scalar::primitive(value, Nullable)) + }; + + let five = partial_of(5)?; + let seven = partial_of(7)?; + // An empty partial must not clobber a prior value. + 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); + let state = merge(merge(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 ccf9f4d899d..d8bbaa2b308 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -38,12 +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 { - fn merge(&mut self, max: Scalar) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypesRef<'_>, + max: Scalar, + ) { if max.is_null() { return; } @@ -51,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; } @@ -63,12 +67,12 @@ impl MaxPartial { }); } - fn poison(&mut self) { - self.max = Some(nan_scalar(&self.element_dtype)); + fn poison(&mut self, dtypes: AggregateDTypesRef<'_>) { + self.max = Some(nan_scalar(dtypes.dtype)); } 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) } } @@ -122,48 +126,72 @@ impl AggregateFnVTable for Max { } fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + ) -> VortexResult { + Ok(MaxPartial { max: None }) + } + + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, ) -> VortexResult { - Ok(MaxPartial { - max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - }) + let mut partial = MaxPartial { max: None }; + // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. + partial.merge(options, dtypes, scalar); + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn merge_partials( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + if let Some(max) = second.max { + first.merge(options, dtypes, max); + } + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.dtype.as_nullable(); match &partial.max { Some(max) => max.cast(&dtype), None => Ok(Scalar::null(dtype)), } } - fn reset(&self, partial: &mut Self::Partial) { - partial.max = None; - } - - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including maximum is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + 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 partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -171,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), @@ -186,6 +214,8 @@ impl AggregateFnVTable for Max { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -196,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: AggregateDTypesRef<'_>, + 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: 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 a3a0f77cb45..e30e90524e8 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::AggregateDTypesRef; 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: AggregateDTypesRef<'_>, + 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.return_dtype.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: AggregateDTypesRef<'_>, + 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.return_dtype, + ); } - let target = DType::Primitive(PType::F64, Nullability::Nullable); + let target = dtypes.return_dtype.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 488405e8f14..17446b6bfa9 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; @@ -38,12 +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 { - fn merge(&mut self, min: Scalar) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypesRef<'_>, + min: Scalar, + ) { if min.is_null() { return; } @@ -51,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; } @@ -63,12 +67,12 @@ impl MinPartial { }); } - fn poison(&mut self) { - self.min = Some(nan_scalar(&self.element_dtype)); + fn poison(&mut self, dtypes: AggregateDTypesRef<'_>) { + self.min = Some(nan_scalar(dtypes.dtype)); } 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) } } @@ -122,48 +126,72 @@ impl AggregateFnVTable for Min { } fn empty_partial( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + ) -> VortexResult { + Ok(MinPartial { min: None }) + } + + fn partial_from_scalar( &self, options: &Self::Options, - input_dtype: &DType, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, ) -> VortexResult { - Ok(MinPartial { - min: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, - }) + let mut partial = MinPartial { min: None }; + // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop. + partial.merge(options, dtypes, scalar); + Ok(partial) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); - Ok(()) + fn merge_partials( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + if let Some(min) = second.min { + first.merge(options, dtypes, min); + } + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + let dtype = dtypes.dtype.as_nullable(); match &partial.min { Some(min) => min.cast(&dtype), None => Ok(Scalar::null(dtype)), } } - fn reset(&self, partial: &mut Self::Partial) { - partial.min = None; - } - - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including minimum is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + 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 partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -171,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), @@ -186,6 +214,8 @@ impl AggregateFnVTable for Min { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -196,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: AggregateDTypesRef<'_>, + 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: 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 adca6453862..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,12 +9,16 @@ use vortex_mask::AllOr; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypesRef; +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: AggregateDTypesRef<'_>, 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..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,6 +8,8 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypesRef; +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: AggregateDTypesRef<'_>, 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..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,6 +7,7 @@ use super::MinMaxPartial; use super::MinMaxResult; use super::min_max; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypesRef; 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: AggregateDTypesRef<'_>, 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 a5091cf4a9f..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,6 +25,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -197,13 +198,16 @@ pub struct MinMax; pub struct MinMaxPartial { min: Option, max: Option, - element_dtype: DType, - skip_nans: bool, } impl MinMaxPartial { /// Merge a local `MinMaxResult` into this partial state. - fn merge(&mut self, local: Option) { + fn merge( + &mut self, + options: &NumericalAggregateOpts, + dtypes: AggregateDTypesRef<'_>, + local: Option, + ) { let Some(MinMaxResult { min, max }) = local else { return; }; @@ -212,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; } @@ -230,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: AggregateDTypesRef<'_>) { + let nan = nan_scalar(dtypes.dtype); 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) } } @@ -310,51 +314,78 @@ impl AggregateFnVTable for MinMax { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(MinMaxPartial { min: None, max: None, - element_dtype: input_dtype.clone(), - skip_nans: options.skip_nans, }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let local = MinMaxResult::from_scalar(other)?; - partial.merge(local); - Ok(()) + fn partial_from_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + let mut partial = MinMaxPartial { + min: None, + max: None, + }; + // `merge` normalizes the parsed extrema: nulls stay empty and NaNs poison or drop. + partial.merge(options, dtypes, MinMaxResult::from_scalar(scalar)?); + Ok(partial) + } + + fn merge_partials( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + if let (Some(min), Some(max)) = (second.min, second.max) { + first.merge(options, dtypes, Some(MinMaxResult { min, max })); + } + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + 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), }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.min = None; - partial.max = None; - } - #[inline] - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { // A poisoned NaN-including min/max is fully determined. partial.is_poisoned() } fn try_accumulate( &self, + options: &Self::Options, + 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 partial.skip_nans || !partial.element_dtype.is_float() { + if options.skip_nans || !dtypes.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -365,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.dtype.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), @@ -385,6 +420,8 @@ impl AggregateFnVTable for MinMax { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -397,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") @@ -431,12 +472,22 @@ impl AggregateFnVTable for MinMax { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -454,10 +505,12 @@ 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::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,19 +709,18 @@ 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 owned = AggregateDTypes::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)), + }; - let scalar2 = Scalar::struct_(struct_dtype, vec![Scalar::from(2i32), Scalar::from(10i32)]); - MinMax.combine_partials(&mut state, scalar2)?; + let state = + MinMax.merge_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..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,6 +8,8 @@ use vortex_mask::Mask; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypesRef; +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: AggregateDTypesRef<'_>, 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..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,17 +8,25 @@ use vortex_error::vortex_panic; use super::MinMaxPartial; use super::MinMaxResult; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateDTypesRef; +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: AggregateDTypesRef<'_>, 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 723847160e1..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,6 +16,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -118,35 +119,56 @@ impl AggregateFnVTable for NanCount { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { - Ok(0u64) + Ok(0) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let val = other + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first + second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, 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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -189,6 +221,7 @@ 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; @@ -201,7 +234,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 +279,11 @@ 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 dtypes = AggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?; - let scalar2 = Scalar::primitive(3u64, Nullability::NonNullable); - NanCount.combine_partials(&mut state, scalar2)?; + let state = NanCount.merge_partials(&EmptyOptions, dtypes.borrow(), 5, 3)?; - let result = NanCount.to_scalar(&state)?; - NanCount.reset(&mut 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 031e8f6a09b..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,6 +12,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -87,35 +88,56 @@ impl AggregateFnVTable for NullCount { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let count = other + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first + second) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn try_accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -126,6 +148,8 @@ impl AggregateFnVTable for NullCount { fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, 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: AggregateDTypesRef<'_>, + 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: 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 197b8f10b04..7825e5c692b 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| { @@ -113,9 +118,12 @@ 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::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 +137,13 @@ 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) -> SumPartial { + SumPartial { + current: Some(SumState::Decimal(value)), + } + } + #[test] fn sum_decimal_basic() -> VortexResult<()> { let decimal = DecimalArray::new( @@ -355,22 +370,17 @@ 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 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 = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, small)?; + let small = partial_with_decimal(DecimalValue::from(9i64)); + let state = Sum.merge_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(), @@ -385,23 +395,17 @@ 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 merge_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 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 = - 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)); + let state = Sum.merge_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(), @@ -414,45 +418,37 @@ 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 dtypes = AggregateDTypes::try_new(&Sum, &options, input_dtype)?; - 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)); + let one_more = partial_with_decimal(DecimalValue::from(-1i64)); + let state = Sum.merge_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(()) } #[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 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. // - // We use combine_partials to get the state close to 10^37, then accumulate - // a real array that pushes it over. + // 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 return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + 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) + ); - // 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)); // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -461,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 76eaba22c15..2e384636ca1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -10,7 +10,7 @@ pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_err; +use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -27,6 +27,7 @@ use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -144,81 +145,68 @@ impl AggregateFnVTable for Sum { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, ) -> 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, + current: Some(make_zero_state(dtypes.return_dtype)), }) } - 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(()); + fn partial_from_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + vortex_ensure!( + scalar.dtype().eq_ignore_nullability(dtypes.return_dtype), + "Sum partial has dtype {}, expected {}", + scalar.dtype(), + 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.return_dtype)?) }; - 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, - } - } + Ok(SumPartial { current }) + } + + fn merge_partials( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + 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.return_dtype, &state)?, }; - if saturated { - partial.current = None; + if overflow { + first.current = None; } - Ok(()) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match &partial.current { - None => Scalar::null(partial.return_dtype.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 = *partial + Some(SumState::Decimal(value)) => { + let decimal_dtype = *dtypes .return_dtype .as_decimal_opt() .vortex_expect("return dtype must be decimal"); @@ -227,12 +215,13 @@ 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 { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { match partial.current.as_ref() { None => true, Some(SumState::Float(v)) => v.is_nan(), @@ -242,13 +231,15 @@ impl AggregateFnVTable for Sum { fn try_accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, 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) { @@ -256,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.return_dtype { sum } else { - sum.cast(&partial.return_dtype)? + sum.cast(dtypes.return_dtype)? }; - self.combine_partials(partial, sum)?; + merge_sum_result(partial, dtypes.return_dtype, sum)?; return Ok(true); } Ok(false) @@ -279,23 +270,24 @@ impl AggregateFnVTable for Sum { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, 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()) { + 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)? { - self.combine_partials(partial, product)?; + if let Some(product) = multiply_constant(c.scalar(), c.len(), dtypes.return_dtype)? { + merge_sum_result(partial, dtypes.return_dtype, product)?; } return Ok(()); } - let skip_nans = partial.skip_nans; let mut inner = match partial.current.take() { Some(inner) => inner, None => return Ok(()), @@ -303,9 +295,13 @@ 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.return_dtype, d, ctx) + } _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -322,36 +318,42 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + 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, } /// 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 { @@ -361,14 +363,78 @@ 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"), } } +/// 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(..) => 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 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), + (SumState::Float(acc), SumState::Float(other)) => { + *acc += *other; + false + } + (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; + 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, + 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, + return_dtype, + &sum_state_from_scalar(&result, return_dtype)?, + )?, + }; + if overflow { + partial.current = None; + } + Ok(()) +} + /// Checked add for u64, returning true if overflow occurred. #[allow(clippy::inline_always)] #[inline(always)] @@ -406,12 +472,15 @@ 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::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 +605,40 @@ 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 scalar1 = Scalar::primitive(100i64, Nullable); - Sum.combine_partials(&mut state, scalar1)?; + let options = NumericalAggregateOpts::default(); + let owned = AggregateDTypes::try_new(&Sum, &options, dtype)?; + let dtypes = owned.borrow(); - let scalar2 = Scalar::primitive(50i64, Nullable); - Sum.combine_partials(&mut state, scalar2)?; + let partial_of = |value: i64| SumPartial { + current: Some(SumState::Signed(value)), + }; + let state = Sum.merge_partials(&options, dtypes, partial_of(100), partial_of(50))?; - let result = Sum.to_scalar(&state)?; - Sum.reset(&mut state); + let result = Sum.to_scalar(&options, dtypes, &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 owned = AggregateDTypes::try_new(&Sum, &options, dtype)?; + let dtypes = owned.borrow(); + + let overflowed = Sum.partial_from_scalar( + &options, + dtypes, + Scalar::null(DType::Primitive(PType::I64, Nullable)), + )?; + let five = Sum.partial_from_scalar(&options, dtypes, Scalar::primitive(5i64, Nullable))?; + 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()); + Ok(()) + } + // Stats caching test #[test] 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 e1287448a8f..a3c6b243a46 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -26,6 +28,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; @@ -106,71 +109,92 @@ impl AggregateFnVTable for SumV2 { fn empty_partial( &self, - options: &Self::Options, - input_dtype: &DType, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, ) -> 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, - }) + Ok(SumV2Partial::empty(dtypes.return_dtype)) } - 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)?; + fn partial_from_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + let mut partial = SumV2Partial::empty(dtypes.return_dtype); + let (sum, is_overflow, is_empty) = decode_partial_scalar(scalar)?; + 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.return_dtype, &sum)?; + partial.is_overflow = is_overflow || overflowed; + partial.is_empty = is_empty && !partial.is_overflow; + Ok(partial) + } - if partial.is_overflow { - return Ok(()); + fn merge_partials( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + mut acc: Self::Partial, + partial: Self::Partial, + ) -> VortexResult { + if acc.is_overflow { + return Ok(acc); } - if other_is_overflow { - partial.is_overflow = true; - partial.is_empty = false; - return Ok(()); + 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 = true; + acc.is_empty = false; + return Ok(acc); } - if other_is_empty { - return Ok(()); + if partial.is_empty { + return Ok(acc); } - - partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; - partial.is_empty = false; - Ok(()) + acc.is_overflow = checked_add_sum_states(&mut acc.sum, dtypes.return_dtype, &partial.sum)?; + acc.is_empty = false; + Ok(acc) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(Scalar::struct_( - sum_v2_partial_dtype(partial.return_dtype.clone()), + dtypes.partial_dtype.clone(), vec![ - sum_state_scalar(partial, 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), ], )) } - 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 { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, 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); } @@ -190,6 +214,8 @@ impl AggregateFnVTable for SumV2 { fn accumulate( &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -202,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() @@ -211,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.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, &product)?; + partial.is_overflow = + checked_add_sum_state(&mut partial.sum, dtypes.return_dtype, &product)?; partial.is_empty = false; } } @@ -228,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.return_dtype, array, ctx) + } _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -247,7 +276,12 @@ impl AggregateFnVTable for SumV2 { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partials: ArrayRef, + ) -> VortexResult { if let Some(partials) = partials.as_opt::() { return finalize_struct(partials); } @@ -260,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: AggregateDTypesRef<'_>, + 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.return_dtype.as_nullable())); } - Ok(sum_state_scalar(partial, Nullability::Nullable)) + Ok(sum_state_scalar( + partial, + dtypes.return_dtype, + Nullability::Nullable, + )) } } @@ -289,11 +332,20 @@ 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. + fn empty(return_dtype: &DType) -> Self { + Self { + sum: make_zero_state(return_dtype), + is_overflow: false, + is_empty: true, + } + } } fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult { @@ -349,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)?), @@ -357,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) => { @@ -389,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 65e54a5bb6a..683228b177f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -10,9 +10,11 @@ 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; +use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; @@ -216,30 +218,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 +256,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 +290,15 @@ 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 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!( + SumV2 + .finalize_scalar(&options, dtypes.borrow(), &propagated)? + .is_null() + ); Ok(()) } @@ -373,7 +378,10 @@ fn finalize_struct_applies_partial_and_struct_validity( )? .into_array(); - let result = SumV2.finalize(partials)?; + let options = NumericalAggregateOpts::default(); + let dtypes = + 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!( &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 a23bcfdc2c9..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,6 +38,7 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; @@ -134,37 +135,58 @@ impl AggregateFnVTable for UncompressedSizeInBytes { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(0) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let size = other + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + 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 merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + first + .checked_add(second) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) } - fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -181,12 +203,22 @@ impl AggregateFnVTable for UncompressedSizeInBytes { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -341,6 +373,7 @@ 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; @@ -686,19 +719,11 @@ mod tests { #[test] fn state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = UncompressedSizeInBytes.empty_partial(&EmptyOptions, &dtype)?; + let dtypes = AggregateDTypes::try_new(&UncompressedSizeInBytes, &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.merge_partials(&EmptyOptions, dtypes.borrow(), 5, 3)?; - let result = UncompressedSizeInBytes.to_scalar(&state)?; - UncompressedSizeInBytes.reset(&mut 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 feb07e47175..6955e8879be 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; @@ -80,27 +81,52 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn combine_partials(&self, _partial: &mut Self::Partial, _other: Scalar) -> VortexResult<()> { + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _scalar: Scalar, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _first: Self::Partial, + _second: Self::Partial, + ) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn reset(&self, _partial: &mut Self::Partial) {} + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _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: AggregateDTypesRef<'_>, + _state: &Self::Partial, + ) -> bool { false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &Columnar, _ctx: &mut ExecutionCtx, @@ -108,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: AggregateDTypesRef<'_>, + _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: 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 92fac87892a..509c3d8062b 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::AggregateDTypesRef; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; @@ -117,31 +118,52 @@ mod tests { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(()) } - fn combine_partials( + fn partial_from_scalar( &self, - _partial: &mut Self::Partial, - _other: Scalar, - ) -> VortexResult<()> { + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _scalar: Scalar, + ) -> VortexResult { Ok(()) } - fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult { - vortex_panic!("TestAgg is for serde tests only"); + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _first: Self::Partial, + _second: Self::Partial, + ) -> VortexResult { + Ok(()) } - fn reset(&self, _partial: &mut Self::Partial) {} + fn to_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _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: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { true } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, _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: AggregateDTypesRef<'_>, + partials: ArrayRef, + ) -> VortexResult { Ok(partials) } - fn finalize_scalar(&self, _partial: &Self::Partial) -> VortexResult { + fn finalize_scalar( + &self, + _options: &Self::Options, + _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 24449bc7572..3e5c82c14f6 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,74 @@ 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 as an +/// [`AggregateDTypesRef`], so partial states only hold accumulated values. Combined aggregates +/// resolve a separate set for each child. +#[derive(Clone, Debug)] +pub struct AggregateDTypes { + /// The DType of the input. + 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 { + /// Resolve the return and partial dtypes of `vtable` bound to `options` over `dtype`. + /// + /// Fails if the aggregate cannot be applied to `dtype`. + 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 + ) + })?; + Ok(Self { + dtype, + return_dtype, + partial_dtype, + }) + } + + /// Lend the dtypes to an execution method. + 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 @@ -32,6 +101,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 [`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 { /// Options for this aggregate function. type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; @@ -62,6 +135,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,38 +176,86 @@ 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. + /// The partial state of a group with no accumulated values. + /// + /// 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, - input_dtype: &DType, + dtypes: AggregateDTypesRef<'_>, ) -> VortexResult; - /// Combine partial scalar state into the accumulator. - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()>; + /// Parse a partial scalar into the typed partial state. + /// + /// 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`]. + /// + /// Implementations should only parse the scalar here; combining states belongs in + /// [`merge_partials`]. + /// + /// [`to_scalar`]: AggregateFnVTable::to_scalar + /// [`merge_partials`]: AggregateFnVTable::merge_partials + fn partial_from_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult; - /// Convert the partial state into a partial scalar. + /// 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. /// - /// The returned scalar must have the same DType as specified by `partial_dtype` for the - /// options and input dtype used to construct the state. - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult; + /// [`empty_partial`]: AggregateFnVTable::empty_partial + fn merge_partials( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult; - /// Reset the state of the accumulator to an empty group. - fn reset(&self, partial: &mut Self::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 + /// finalization. + /// + /// [`partial_from_scalar`]: AggregateFnVTable::partial_from_scalar + fn to_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, _state: &mut Self::Partial, _batch: &ArrayRef, _ctx: &mut ExecutionCtx, @@ -137,25 +263,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: AggregateDTypesRef<'_>, 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_dtype`; the result must have dtype `dtypes.return_dtype`. + fn finalize( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + 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.return_dtype`. + fn finalize_scalar( + &self, + options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult; } #[derive(Clone, Debug, PartialEq, Eq, Hash)] 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..3ad0625752b 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::AggregateDTypesRef; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::dtype::DType; @@ -298,48 +299,79 @@ impl AggregateFnVTable for BloomFilter { } /// Returns an empty Bloom filter with all blocks zero-initialized. - fn empty_partial(&self, options: &Self::Options, _: &DType) -> VortexResult { + fn empty_partial( + &self, + options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + ) -> 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, + _dtypes: AggregateDTypesRef<'_>, + 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) + /// Merges two filters by OR-ing their blocks together. + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { let bytes: Vec = partial.serialize(); 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. - fn is_saturated(&self, partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> bool { partial.is_saturated() } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -353,12 +385,22 @@ impl AggregateFnVTable for BloomFilter { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -381,9 +423,11 @@ 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::test_harness::check_metadata; + use super::partial::BLOCK_SIZE; use super::*; pub fn setup() -> VortexResult { @@ -417,78 +461,98 @@ 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 { + AggregateDTypes::try_new( + &BloomFilter, + options, + DType::Binary(Nullability::NonNullable), + ) + } + #[test] fn saturation_false_when_empty() -> VortexResult<()> { let options = BloomOptions::default(); - let partial = - BloomFilter.empty_partial(&options, &DType::Binary(Nullability::NonNullable))?; - assert!(!BloomFilter.is_saturated(&partial)); + let dtypes = binary_dtypes(&options)?; + let partial = BloomFilter.empty_partial(&options, dtypes.borrow())?; + assert!(!BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); Ok(()) } #[test] - fn saturation_true_when_every_block_is_full() { - let blocks = vec![[u32::MAX; 8]; 4]; - let partial = BloomPartial::from(blocks); + fn saturation_true_when_every_block_is_full() -> VortexResult<()> { + let options = BloomOptions::default(); + let dtypes = binary_dtypes(&options)?; + + // 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(&partial)); + assert!(BloomFilter.is_saturated(&options, dtypes.borrow(), &partial)); + Ok(()) } #[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 smaller = BloomOptions::new(NonZeroU32::new(4).unwrap(), HashFn::XxHash3_64); + let dtypes = binary_dtypes(&smaller)?; + let bigger = BloomFilter.empty_partial(&BloomOptions::default(), dtypes.borrow())?; - let bigger_scalar = BloomFilter.to_scalar(&bigger)?; - let result = BloomFilter.combine_partials(&mut smaller, bigger_scalar); + let bigger_scalar = + BloomFilter.to_scalar(&BloomOptions::default(), dtypes.borrow(), &bigger)?; + assert!( + BloomFilter + .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!( - result.is_err(), - "combining partials built with different blocks_count must fail loudly, not corrupt state" + BloomFilter + .merge_partials( + &smaller, + dtypes.borrow(), + BloomFilter.empty_partial(&smaller, dtypes.borrow())?, + 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 merge_partials_unions_two_disjoint_partials() -> VortexResult<()> { + let options = BloomOptions::default(); + 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 = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + 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 = BloomFilter.empty_partial( - &BloomOptions::default(), - &DType::Binary(Nullability::NonNullable), - )?; + let mut expected = BloomFilter.empty_partial(&options, dtypes.borrow())?; 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.merge_partials(&options, dtypes.borrow(), 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-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial/mod.rs index 6f2c2312878..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 @@ -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, AggregateDTypes}; /// /// let filter = BloomFilter {}; +/// let options = BloomOptions::default(); +/// let dtypes = AggregateDTypes::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,30 +242,12 @@ 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, - } - } -} - -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. + // 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], + }, } } } 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 7c3bcaf7659..4b1b0c35bba 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::AggregateDTypesRef; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; @@ -156,36 +157,62 @@ impl AggregateFnVTable for GeometryAabb { fn empty_partial( &self, _options: &Self::Options, - _input_dtype: &DType, + _dtypes: AggregateDTypesRef<'_>, ) -> VortexResult { Ok(AabbPartial { rect: None }) } - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if let Some(rect) = rect_from_storage(&other)? { - partial.merge(rect); + fn partial_from_scalar( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + scalar: Scalar, + ) -> VortexResult { + // A null box is an empty group's AABB. + Ok(AabbPartial { + rect: rect_from_storage(&scalar)?, + }) + } + + fn merge_partials( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + mut first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + if let Some(rect) = second.rect { + first.merge(rect); } - Ok(()) + Ok(first) } - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + fn to_scalar( + &self, + _options: &Self::Options, + dtypes: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { Ok(match partial.rect { Some(rect) => rect_to_storage(rect), - None => Scalar::null(aabb_dtype()), + None => Scalar::null(dtypes.partial_dtype.clone()), }) } - fn reset(&self, partial: &mut Self::Partial) { - partial.rect = None; - } - - fn is_saturated(&self, _partial: &Self::Partial) -> bool { + fn is_saturated( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + _partial: &Self::Partial, + ) -> bool { // An AABB can always grow, so it is never saturated. false } fn accumulate( &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -217,13 +244,23 @@ impl AggregateFnVTable for GeometryAabb { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize( + &self, + _options: &Self::Options, + _dtypes: AggregateDTypesRef<'_>, + 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: AggregateDTypesRef<'_>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(options, dtypes, partial) } } @@ -233,6 +270,7 @@ 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; @@ -245,7 +283,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 +425,40 @@ mod tests { Ok(()) } - /// `combine_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 combine_partials_unions_boxes() -> VortexResult<()> { + fn merge_partials_unions_boxes() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let dtypes = AggregateDTypes::try_new(&GeometryAabb, &EmptyOptions, dtype)?; 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.merge_partials( + &EmptyOptions, + 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(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &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 `merge_partials`. #[test] - fn combine_partials_ignores_null() -> VortexResult<()> { - let mut partial = AabbPartial { + 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 = GeometryAabb.empty_partial(&EmptyOptions, dtypes.borrow())?; + 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.merge_partials(&EmptyOptions, dtypes.borrow(), value, empty)?; assert_eq!( - aabb(&GeometryAabb.to_scalar(&partial)?)?, + aabb(&GeometryAabb.to_scalar(&EmptyOptions, dtypes.borrow(), &reduced)?)?, (0.0, 0.0, 1.0, 1.0) ); Ok(())