Rework AggregateVTable

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,41 +187,22 @@ pub struct GroupedAccumulator<V: AggregateFnVTable> {
     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<ArrayRef>,
 }
 
 impl<V: AggregateFnVTable> GroupedAccumulator<V> {
     pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult<Self> {
         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<V: AggregateFnVTable> DynGroupedAccumulator for GroupedAccumulator<V> {
             ),
         };
         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<V: AggregateFnVTable> DynGroupedAccumulator for GroupedAccumulator<V> {
         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<ArrayRef> {
         let states = self.flush()?;
-        let results = self.vtable.finalize(states)?;
+        let results = self
+            .vtable
+            .finalize(self.dtypes.args(&self.options), 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<V: AggregateFnVTable> GroupedAccumulator<V> {
         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<V: AggregateFnVTable> GroupedAccumulator<V> {
 
     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);

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::AggregateArgs;
+use crate::aggregate_fn::AggregateDTypes;
 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<T> = <<T as BinaryCombined>::Left as AggregateFnVTable>::Option
 type RightOptions<T> = <<T as BinaryCombined>::Right as AggregateFnVTable>::Options;
 /// Combined options for a [`BinaryCombined`] aggregate.
 pub type CombinedOptions<T> = PairOptions<LeftOptions<T>, RightOptions<T>>;
+/// Pair of typed child accumulators holding the partial state of a [`BinaryCombined`] aggregate.
+type ChildAccumulators<T> = (
+    Accumulator<<T as BinaryCombined>::Left>,
+    Accumulator<<T as BinaryCombined>::Right>,
+);
 
 /// Declare an aggregate function in terms of two child aggregates.
 pub trait BinaryCombined: 'static + Send + Sync + Clone {
@@ -76,12 +83,26 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone {
     }
 
     /// Return type of the combined aggregate.
-    fn return_dtype(&self, input_dtype: &DType) -> Option<DType>;
+    fn return_dtype(&self, options: &CombinedOptions<Self>, input_dtype: &DType) -> Option<DType>;
 
     /// Combine the finalized left and right results into the final aggregate.
-    fn finalize(&self, left: ArrayRef, right: ArrayRef) -> VortexResult<ArrayRef>;
+    ///
+    /// `args` describes the combined aggregate; `left` and `right` already have their child's
+    /// result dtype. The returned array must have dtype `args.return_dtype`.
+    fn finalize(
+        &self,
+        args: AggregateArgs<'_, CombinedOptions<Self>>,
+        left: ArrayRef,
+        right: ArrayRef,
+    ) -> VortexResult<ArrayRef>;
 
-    fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult<Scalar>;
+    /// Combine the finalized child scalars into a result of dtype `args.return_dtype`.
+    fn finalize_scalar(
+        &self,
+        args: AggregateArgs<'_, CombinedOptions<Self>>,
+        left_scalar: Scalar,
+        right_scalar: Scalar,
+    ) -> VortexResult<Scalar>;
 
     /// Serialize the options for this combined aggregate. Default: not serializable.
     fn serialize(&self, options: &CombinedOptions<Self>) -> VortexResult<Option<Vec<u8>>> {
@@ -126,14 +147,25 @@ impl<T: BinaryCombined> Combined<T> {
     pub fn new(inner: T) -> Self {
         Self(inner)
     }
+
+    /// Construct a pair of empty child accumulators.
+    fn new_child_accumulators(
+        &self,
+        args: AggregateArgs<'_, CombinedOptions<T>>,
+    ) -> VortexResult<ChildAccumulators<T>> {
+        let left = Accumulator::try_new(self.0.left(), args.options.0.clone(), args.dtype.clone())?;
+        let right =
+            Accumulator::try_new(self.0.right(), args.options.1.clone(), args.dtype.clone())?;
+        Ok((left, right))
+    }
 }
 
 impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
     type Options = CombinedOptions<T>;
-    // 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<Mean>`).
-    type Partial = (AccumulatorRef, AccumulatorRef);
+    type Partial = ChildAccumulators<T>;
 
     fn id(&self) -> AggregateFnId {
         self.0.id()
@@ -147,8 +179,8 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
         BinaryCombined::deserialize(&self.0, metadata, session)
     }
 
-    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
-        BinaryCombined::return_dtype(&self.0, input_dtype)
+    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
+        BinaryCombined::return_dtype(&self.0, options, input_dtype)
     }
 
     fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -157,52 +189,63 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
         Some(self.0.partial_struct_dtype(l, r))
     }
 
-    fn empty_partial(
+    fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Self::Partial> {
+        self.new_child_accumulators(args)
+    }
+
+    fn partial_from_scalar(
         &self,
-        options: &Self::Options,
-        input_dtype: &DType,
+        args: AggregateArgs<'_, Self::Options>,
+        scalar: Scalar,
     ) -> VortexResult<Self::Partial> {
-        let left = Accumulator::try_new(self.0.left(), options.0.clone(), input_dtype.clone())?;
-        let right = Accumulator::try_new(self.0.right(), options.1.clone(), input_dtype.clone())?;
-        Ok((
-            Box::new(left) as AccumulatorRef,
-            Box::new(right) as AccumulatorRef,
-        ))
+        let (mut left, mut right) = self.new_child_accumulators(args)?;
+        // 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_partials(l_field)?;
+            right.combine_partials(r_field)?;
+        }
+        Ok((left, right))
     }
 
-    fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
-        if other.is_null() {
-            return Ok(());
-        }
-        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<Scalar> {
-        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]))
+    fn merge_partials(
+        &self,
+        _args: AggregateArgs<'_, Self::Options>,
+        (mut left, mut right): Self::Partial,
+        (mut other_left, mut other_right): Self::Partial,
+    ) -> VortexResult<Self::Partial> {
+        // Typed child accumulators merge directly, without any scalar interchange.
+        left.merge_from(&mut other_left)?;
+        right.merge_from(&mut other_right)?;
+        Ok((left, right))
     }
 
-    fn reset(&self, partial: &mut Self::Partial) {
-        partial.0.reset();
-        partial.1.reset();
+    fn to_scalar(
+        &self,
+        args: AggregateArgs<'_, Self::Options>,
+        partial: &Self::Partial,
+    ) -> VortexResult<Scalar> {
+        let l_scalar = partial.0.partial_scalar()?;
+        let r_scalar = partial.1.partial_scalar()?;
+        Ok(Scalar::struct_(
+            args.partial_dtype.clone(),
+            vec![l_scalar, r_scalar],
+        ))
     }
 
-    fn is_saturated(&self, partial: &Self::Partial) -> bool {
+    fn is_saturated(
+        &self,
+        _args: AggregateArgs<'_, Self::Options>,
+        partial: &Self::Partial,
+    ) -> bool {
         partial.0.is_saturated() && partial.1.is_saturated()
     }
 
@@ -213,6 +256,7 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
     /// `true` so [`Self::accumulate`] is unreachable.
     fn try_accumulate(
         &self,
+        _args: AggregateArgs<'_, Self::Options>,
         state: &mut Self::Partial,
         batch: &ArrayRef,
         ctx: &mut ExecutionCtx,
@@ -224,24 +268,37 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
 
     fn accumulate(
         &self,
+        _args: AggregateArgs<'_, Self::Options>,
         _state: &mut Self::Partial,
         _batch: &Columnar,
         _ctx: &mut ExecutionCtx,
     ) -> VortexResult<()> {
         unreachable!("Combined::try_accumulate handles all batches")
     }
 
-    fn finalize(&self, states: ArrayRef) -> VortexResult<ArrayRef> {
+    fn finalize(
+        &self,
+        args: AggregateArgs<'_, Self::Options>,
+        states: ArrayRef,
+    ) -> VortexResult<ArrayRef> {
         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, &args.options.0, args.dtype.clone())?;
+        let r_dtypes = AggregateDTypes::try_new(&right, &args.options.1, args.dtype.clone())?;
+        let l_finalized = left.finalize(l_dtypes.args(&args.options.0), l_field)?;
+        let r_finalized = right.finalize(r_dtypes.args(&args.options.1), r_field)?;
+        BinaryCombined::finalize(&self.0, args, l_finalized, r_finalized)
     }
 
-    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+    fn finalize_scalar(
+        &self,
+        args: AggregateArgs<'_, Self::Options>,
+        partial: &Self::Partial,
+    ) -> VortexResult<Scalar> {
         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, args, l_scalar, r_scalar)
     }
 }