vortex_array/aggregate_fn/
combined.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Generic adapter for aggregates whose result is computed from two child
5//! aggregate functions, e.g. `Mean = Sum / Count`.
6
7use std::fmt::Debug;
8use std::fmt::Display;
9use std::fmt::Formatter;
10use std::fmt::{self};
11use std::hash::Hash;
12
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_err;
16use vortex_session::VortexSession;
17
18use crate::ArrayRef;
19use crate::Columnar;
20use crate::ExecutionCtx;
21use crate::aggregate_fn::Accumulator;
22use crate::aggregate_fn::AggregateArgs;
23use crate::aggregate_fn::AggregateDTypes;
24use crate::aggregate_fn::AggregateFnId;
25use crate::aggregate_fn::AggregateFnVTable;
26use crate::aggregate_fn::DynAccumulator;
27use crate::builtins::ArrayBuiltins;
28use crate::dtype::DType;
29use crate::dtype::FieldName;
30use crate::dtype::FieldNames;
31use crate::dtype::Nullability;
32use crate::dtype::StructFields;
33use crate::scalar::Scalar;
34
35/// Pair of options for the two children of a [`BinaryCombined`] aggregate.
36///
37/// Wrapper around `(L, R)` because the [`AggregateFnVTable::Options`] bound
38/// requires `Display`, which tuples don't implement.
39#[derive(Clone, Debug, PartialEq, Eq, Hash)]
40pub struct PairOptions<L, R>(pub L, pub R);
41
42impl<L: Display, R: Display> Display for PairOptions<L, R> {
43    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
44        write!(f, "({}, {})", self.0, self.1)
45    }
46}
47
48// Convenience aliases so signatures stay readable.
49type LeftOptions<T> = <<T as BinaryCombined>::Left as AggregateFnVTable>::Options;
50type RightOptions<T> = <<T as BinaryCombined>::Right as AggregateFnVTable>::Options;
51/// Combined options for a [`BinaryCombined`] aggregate.
52pub type CombinedOptions<T> = PairOptions<LeftOptions<T>, RightOptions<T>>;
53/// Pair of typed child accumulators holding the partial state of a [`BinaryCombined`] aggregate.
54type ChildAccumulators<T> = (
55    Accumulator<<T as BinaryCombined>::Left>,
56    Accumulator<<T as BinaryCombined>::Right>,
57);
58
59/// Declare an aggregate function in terms of two child aggregates.
60pub trait BinaryCombined: 'static + Send + Sync + Clone {
61    /// The left child aggregate vtable.
62    type Left: AggregateFnVTable;
63    /// The right child aggregate vtable.
64    type Right: AggregateFnVTable;
65
66    /// Stable identifier for the combined aggregate.
67    fn id(&self) -> AggregateFnId;
68
69    /// Construct the left child vtable.
70    fn left(&self) -> Self::Left;
71
72    /// Construct the right child vtable.
73    fn right(&self) -> Self::Right;
74
75    /// Field name for the left child in the partial struct dtype.
76    fn left_name(&self) -> &'static str {
77        "left"
78    }
79
80    /// Field name for the right child in the partial struct dtype.
81    fn right_name(&self) -> &'static str {
82        "right"
83    }
84
85    /// Return type of the combined aggregate.
86    fn return_dtype(&self, options: &CombinedOptions<Self>, input_dtype: &DType) -> Option<DType>;
87
88    /// Combine the finalized left and right results into the final aggregate.
89    ///
90    /// `args` describes the combined aggregate; `left` and `right` already have their child's
91    /// result dtype. The returned array must have dtype `args.return_dtype`.
92    fn finalize(
93        &self,
94        args: AggregateArgs<'_, CombinedOptions<Self>>,
95        left: ArrayRef,
96        right: ArrayRef,
97    ) -> VortexResult<ArrayRef>;
98
99    /// Combine the finalized child scalars into a result of dtype `args.return_dtype`.
100    fn finalize_scalar(
101        &self,
102        args: AggregateArgs<'_, CombinedOptions<Self>>,
103        left_scalar: Scalar,
104        right_scalar: Scalar,
105    ) -> VortexResult<Scalar>;
106
107    /// Serialize the options for this combined aggregate. Default: not serializable.
108    fn serialize(&self, options: &CombinedOptions<Self>) -> VortexResult<Option<Vec<u8>>> {
109        let _ = options;
110        Ok(None)
111    }
112
113    /// Deserialize the options for this combined aggregate. Default: bails.
114    fn deserialize(
115        &self,
116        metadata: &[u8],
117        session: &VortexSession,
118    ) -> VortexResult<CombinedOptions<Self>> {
119        let _ = (metadata, session);
120        vortex_bail!(
121            "Combined aggregate function {} is not deserializable",
122            BinaryCombined::id(self)
123        );
124    }
125
126    /// Build the partial struct dtype that wraps the two child partials.
127    fn partial_struct_dtype(&self, left: DType, right: DType) -> DType {
128        DType::Struct(
129            StructFields::new(
130                FieldNames::from_iter([
131                    FieldName::from(self.left_name()),
132                    FieldName::from(self.right_name()),
133                ]),
134                vec![left, right],
135            ),
136            Nullability::NonNullable,
137        )
138    }
139}
140
141/// Resolve a child's return dtype, the one dtype the combined parent does not already hold.
142fn child_return_dtype<V: AggregateFnVTable>(
143    vtable: &V,
144    options: &V::Options,
145    input_dtype: &DType,
146) -> VortexResult<DType> {
147    vtable.return_dtype(options, input_dtype).ok_or_else(|| {
148        vortex_err!(
149            "Aggregate function {} cannot be applied to dtype {}",
150            vtable.id(),
151            input_dtype
152        )
153    })
154}
155
156/// Adapter that exposes any [`BinaryCombined`] as an [`AggregateFnVTable`].
157#[derive(Clone, Debug)]
158pub struct Combined<T: BinaryCombined>(pub T);
159
160impl<T: BinaryCombined> Combined<T> {
161    /// Construct a new combined aggregate vtable.
162    pub fn new(inner: T) -> Self {
163        Self(inner)
164    }
165
166    /// Derive both children's dtypes from the parent's, rather than resolving them again.
167    ///
168    /// Both children read the parent's input, and the parent's partial dtype is the struct the
169    /// children's partials were resolved into: `{left: <left partial>, right: <right partial>}`.
170    /// Only the children's return dtypes are left to resolve.
171    fn child_dtypes(
172        &self,
173        args: AggregateArgs<'_, CombinedOptions<T>>,
174    ) -> VortexResult<(AggregateDTypes, AggregateDTypes)> {
175        let partials = args.partial_dtype.as_struct_fields_opt().ok_or_else(|| {
176            vortex_err!(
177                "Combined partial dtype must be a struct, got {}",
178                args.partial_dtype
179            )
180        })?;
181        let [Some(l_partial), Some(r_partial)] =
182            [partials.field_by_index(0), partials.field_by_index(1)]
183        else {
184            vortex_bail!(
185                "Combined partial dtype {} must have two fields",
186                args.partial_dtype
187            );
188        };
189
190        let l_return = child_return_dtype(&self.0.left(), &args.options.0, args.dtype)?;
191        let r_return = child_return_dtype(&self.0.right(), &args.options.1, args.dtype)?;
192
193        Ok((
194            AggregateDTypes::new(args.dtype.clone(), l_return, l_partial),
195            AggregateDTypes::new(args.dtype.clone(), r_return, r_partial),
196        ))
197    }
198
199    /// Construct a pair of empty child accumulators over the derived child dtypes.
200    fn new_child_accumulators(
201        &self,
202        args: AggregateArgs<'_, CombinedOptions<T>>,
203    ) -> VortexResult<ChildAccumulators<T>> {
204        let (l_dtypes, r_dtypes) = self.child_dtypes(args)?;
205
206        Ok((
207            Accumulator::from_dtypes(self.0.left(), args.options.0.clone(), l_dtypes),
208            Accumulator::from_dtypes(self.0.right(), args.options.1.clone(), r_dtypes),
209        ))
210    }
211}
212
213impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
214    type Options = CombinedOptions<T>;
215    // Each child is held as a fully-fledged `Accumulator` so that batches dispatched through
216    // `try_accumulate` consult the kernel registry per-child (e.g. a `(Dict, Sum)` kernel fires
217    // for the inner `Sum` child of `Combined<Mean>`).
218    type Partial = ChildAccumulators<T>;
219
220    fn id(&self) -> AggregateFnId {
221        self.0.id()
222    }
223
224    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
225        BinaryCombined::serialize(&self.0, options)
226    }
227
228    fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult<Self::Options> {
229        BinaryCombined::deserialize(&self.0, metadata, session)
230    }
231
232    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
233        BinaryCombined::return_dtype(&self.0, options, input_dtype)
234    }
235
236    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
237        let l = self.0.left().partial_dtype(&options.0, input_dtype)?;
238        let r = self.0.right().partial_dtype(&options.1, input_dtype)?;
239        Some(self.0.partial_struct_dtype(l, r))
240    }
241
242    fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Self::Partial> {
243        self.new_child_accumulators(args)
244    }
245
246    fn partial_from_scalar(
247        &self,
248        args: AggregateArgs<'_, Self::Options>,
249        scalar: &Scalar,
250    ) -> VortexResult<Self::Partial> {
251        let (mut left, mut right) = self.new_child_accumulators(args)?;
252        // A null partial represents an empty group and parses to empty child accumulators.
253        if !scalar.is_null() {
254            let s = scalar.as_struct();
255            let lname = self.0.left_name();
256            let rname = self.0.right_name();
257            let l_field = s
258                .field(lname)
259                .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", lname))?;
260            let r_field = s
261                .field(rname)
262                .ok_or_else(|| vortex_err!("BinaryCombined partial missing `{}` field", rname))?;
263            left.combine_partials(&l_field)?;
264            right.combine_partials(&r_field)?;
265        }
266        Ok((left, right))
267    }
268
269    fn merge_partials(
270        &self,
271        _args: AggregateArgs<'_, Self::Options>,
272        (mut left, mut right): Self::Partial,
273        (mut other_left, mut other_right): Self::Partial,
274    ) -> VortexResult<Self::Partial> {
275        // Typed child accumulators merge directly, without any scalar interchange.
276        left.merge_from(&mut other_left)?;
277        right.merge_from(&mut other_right)?;
278        Ok((left, right))
279    }
280
281    fn to_scalar(
282        &self,
283        args: AggregateArgs<'_, Self::Options>,
284        partial: &Self::Partial,
285    ) -> VortexResult<Scalar> {
286        let l_scalar = partial.0.partial_scalar()?;
287        let r_scalar = partial.1.partial_scalar()?;
288        Ok(Scalar::struct_(
289            args.partial_dtype.clone(),
290            vec![l_scalar, r_scalar],
291        ))
292    }
293
294    fn is_saturated(
295        &self,
296        _args: AggregateArgs<'_, Self::Options>,
297        partial: &Self::Partial,
298    ) -> bool {
299        partial.0.is_saturated() && partial.1.is_saturated()
300    }
301
302    /// Delegate the batch to each child's `Accumulator::accumulate`, which consults the
303    /// kernel registry against the child's `aggregate_fn` id. This is what makes
304    /// `(encoding, Child)` kernels reachable through `Combined<Parent>` — without it, a
305    /// `(Dict, Sum)` kernel would be dead code for `Combined<Mean>`. We always return
306    /// `true` so [`Self::accumulate`] is unreachable.
307    fn try_accumulate(
308        &self,
309        _args: AggregateArgs<'_, Self::Options>,
310        state: &mut Self::Partial,
311        batch: &ArrayRef,
312        ctx: &mut ExecutionCtx,
313    ) -> VortexResult<bool> {
314        state.0.accumulate(batch, ctx)?;
315        state.1.accumulate(batch, ctx)?;
316        Ok(true)
317    }
318
319    fn accumulate(
320        &self,
321        _args: AggregateArgs<'_, Self::Options>,
322        _state: &mut Self::Partial,
323        _batch: &Columnar,
324        _ctx: &mut ExecutionCtx,
325    ) -> VortexResult<()> {
326        unreachable!("Combined::try_accumulate handles all batches")
327    }
328
329    fn finalize(
330        &self,
331        args: AggregateArgs<'_, Self::Options>,
332        states: ArrayRef,
333    ) -> VortexResult<ArrayRef> {
334        let l_field = states.get_item(FieldName::from(self.0.left_name()))?;
335        let r_field = states.get_item(FieldName::from(self.0.right_name()))?;
336        let (l_dtypes, r_dtypes) = self.child_dtypes(args)?;
337        let l_finalized = self
338            .0
339            .left()
340            .finalize(l_dtypes.args(&args.options.0), l_field)?;
341        let r_finalized = self
342            .0
343            .right()
344            .finalize(r_dtypes.args(&args.options.1), r_field)?;
345        BinaryCombined::finalize(&self.0, args, l_finalized, r_finalized)
346    }
347
348    fn finalize_scalar(
349        &self,
350        args: AggregateArgs<'_, Self::Options>,
351        partial: &Self::Partial,
352    ) -> VortexResult<Scalar> {
353        let l_scalar = partial.0.final_scalar()?;
354        let r_scalar = partial.1.final_scalar()?;
355        BinaryCombined::finalize_scalar(&self.0, args, l_scalar, r_scalar)
356    }
357}
vortex_array/aggregate_fn/combined.rs

vortex_array/aggregate_fn/
combined.rs
