From fe06585a4269bef34abedb10190c937b73b5ef84 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 21 Sep 2026 19:36:05 -0400 Subject: [PATCH] Support creating accumulators with already derived dtypes Signed-off-by: Robert Kruszewski --- vortex-array/src/aggregate_fn/accumulator.rs | 13 +++- vortex-array/src/aggregate_fn/combined.rs | 75 +++++++++++++++++--- vortex-array/src/aggregate_fn/vtable.rs | 15 ++-- 3 files changed, 85 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index c6042308dab..6604afb8a7b 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -44,15 +44,24 @@ pub struct Accumulator { impl Accumulator { pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult { let dtypes = AggregateDTypes::try_new(&vtable, &options, dtype)?; + + Ok(Self::from_dtypes(vtable, options, dtypes)) + } + + /// Build an accumulator over dtypes that are already resolved. + /// + /// A nested accumulator takes the dtypes its parent already derived, instead of resolving + /// them again for every batch. + pub fn from_dtypes(vtable: V, options: V::Options, dtypes: AggregateDTypes) -> Self { let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased(); - Ok(Self { + Self { vtable, options, aggregate_fn, dtypes, partial: None, - }) + } } /// The state of a group with no accumulated values. diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 4bc6d10b1e7..f49df80bf7b 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -138,6 +138,21 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone { } } +/// Resolve a child's return dtype, the one dtype the combined parent does not already hold. +fn child_return_dtype( + vtable: &V, + options: &V::Options, + input_dtype: &DType, +) -> VortexResult { + vtable.return_dtype(options, input_dtype).ok_or_else(|| { + vortex_err!( + "Aggregate function {} cannot be applied to dtype {}", + vtable.id(), + input_dtype + ) + }) +} + /// Adapter that exposes any [`BinaryCombined`] as an [`AggregateFnVTable`]. #[derive(Clone, Debug)] pub struct Combined(pub T); @@ -148,15 +163,50 @@ impl Combined { Self(inner) } - /// Construct a pair of empty child accumulators. + /// Derive both children's dtypes from the parent's, rather than resolving them again. + /// + /// Both children read the parent's input, and the parent's partial dtype is the struct the + /// children's partials were resolved into: `{left: , right: }`. + /// Only the children's return dtypes are left to resolve. + fn child_dtypes( + &self, + args: AggregateArgs<'_, CombinedOptions>, + ) -> VortexResult<(AggregateDTypes, AggregateDTypes)> { + let partials = args.partial_dtype.as_struct_fields_opt().ok_or_else(|| { + vortex_err!( + "Combined partial dtype must be a struct, got {}", + args.partial_dtype + ) + })?; + let [Some(l_partial), Some(r_partial)] = + [partials.field_by_index(0), partials.field_by_index(1)] + else { + vortex_bail!( + "Combined partial dtype {} must have two fields", + args.partial_dtype + ); + }; + + let l_return = child_return_dtype(&self.0.left(), &args.options.0, args.dtype)?; + let r_return = child_return_dtype(&self.0.right(), &args.options.1, args.dtype)?; + + Ok(( + AggregateDTypes::new(args.dtype.clone(), l_return, l_partial), + AggregateDTypes::new(args.dtype.clone(), r_return, r_partial), + )) + } + + /// Construct a pair of empty child accumulators over the derived child dtypes. fn new_child_accumulators( &self, args: AggregateArgs<'_, CombinedOptions>, ) -> VortexResult> { - 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)) + let (l_dtypes, r_dtypes) = self.child_dtypes(args)?; + + Ok(( + Accumulator::from_dtypes(self.0.left(), args.options.0.clone(), l_dtypes), + Accumulator::from_dtypes(self.0.right(), args.options.1.clone(), r_dtypes), + )) } } @@ -283,12 +333,15 @@ impl AggregateFnVTable for Combined { ) -> VortexResult { let l_field = states.get_item(FieldName::from(self.0.left_name()))?; let r_field = states.get_item(FieldName::from(self.0.right_name()))?; - let left = self.0.left(); - let right = self.0.right(); - let l_dtypes = 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)?; + let (l_dtypes, r_dtypes) = self.child_dtypes(args)?; + let l_finalized = self + .0 + .left() + .finalize(l_dtypes.args(&args.options.0), l_field)?; + let r_finalized = self + .0 + .right() + .finalize(r_dtypes.args(&args.options.1), r_field)?; BinaryCombined::finalize(&self.0, args, l_finalized, r_finalized) } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index d1b39ee1b0b..76a2fd10415 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -40,6 +40,15 @@ pub struct AggregateDTypes { } impl AggregateDTypes { + /// Pair already-resolved dtypes, e.g. the dtypes a nested aggregate derives for its children. + pub fn new(dtype: DType, return_dtype: DType, partial_dtype: DType) -> Self { + Self { + dtype, + return_dtype, + partial_dtype, + } + } + /// Resolve the return and partial dtypes of `vtable` bound to `options` over `dtype`. /// /// Fails if the aggregate cannot be applied to `dtype`. @@ -62,11 +71,7 @@ impl AggregateDTypes { dtype ) })?; - Ok(Self { - dtype, - return_dtype, - partial_dtype, - }) + Ok(Self::new(dtype, return_dtype, partial_dtype)) } /// Pair these dtypes with `options` for one aggregate execution call.