vortex_array/aggregate_fn/
accumulator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_ensure_eq;
11
12use crate::ArrayRef;
13use crate::Columnar;
14use crate::ExecutionCtx;
15use crate::aggregate_fn::AggregateDTypes;
16use crate::aggregate_fn::AggregateFn;
17use crate::aggregate_fn::AggregateFnRef;
18use crate::aggregate_fn::AggregateFnVTable;
19use crate::columnar::AnyColumnar;
20use crate::dtype::DType;
21use crate::executor::max_iterations;
22use crate::expr::stats::Precision;
23use crate::expr::stats::Stat;
24use crate::expr::stats::StatsProvider;
25use crate::scalar::Scalar;
26
27/// Reference-counted type-erased accumulator.
28pub type AccumulatorRef = Box<dyn DynAccumulator>;
29
30/// An accumulator used for computing aggregates over an entire stream of arrays.
31pub struct Accumulator<V: AggregateFnVTable> {
32    /// The vtable of the aggregate function.
33    vtable: V,
34    /// The options of the aggregate function.
35    options: V::Options,
36    /// Type-erased aggregate function used for kernel dispatch.
37    aggregate_fn: AggregateFnRef,
38    /// The input, partial, and result dtypes lent to every vtable call.
39    dtypes: AggregateDTypes,
40    /// The partial state of the accumulator, with `None` as the empty group.
41    partial: Option<V::Partial>,
42}
43
44impl<V: AggregateFnVTable> Accumulator<V> {
45    pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult<Self> {
46        let dtypes = AggregateDTypes::try_new(&vtable, &options, dtype)?;
47
48        Ok(Self::from_dtypes(vtable, options, dtypes))
49    }
50
51    /// Build an accumulator over dtypes that are already resolved.
52    ///
53    /// A nested accumulator takes the dtypes its parent already derived, instead of resolving
54    /// them again for every batch.
55    pub fn from_dtypes(vtable: V, options: V::Options, dtypes: AggregateDTypes) -> Self {
56        let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased();
57
58        Self {
59            vtable,
60            options,
61            aggregate_fn,
62            dtypes,
63            partial: None,
64        }
65    }
66
67    /// The state of a group with no accumulated values.
68    pub(crate) fn empty_partial(&self) -> VortexResult<V::Partial> {
69        self.vtable.empty_partial(self.dtypes.args(&self.options))
70    }
71
72    /// Take the accumulated partial state, leaving the accumulator empty.
73    ///
74    /// Returns `None` if nothing was accumulated.
75    pub(crate) fn take_partial(&mut self) -> Option<V::Partial> {
76        self.partial.take()
77    }
78
79    fn ensure_partial(&mut self) -> VortexResult<()> {
80        if self.partial.is_none() {
81            self.partial = Some(self.empty_partial()?);
82        }
83        Ok(())
84    }
85
86    pub(crate) fn fold_partial(&mut self, other: V::Partial) -> VortexResult<()> {
87        self.partial = Some(match self.partial.take() {
88            // The empty state is the identity of the merge.
89            None => other,
90            Some(current) => {
91                self.vtable
92                    .merge_partials(self.dtypes.args(&self.options), current, other)?
93            }
94        });
95        Ok(())
96    }
97
98    fn fold_partial_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
99        let other = self
100            .vtable
101            .partial_from_scalar(self.dtypes.args(&self.options), scalar)?;
102        self.fold_partial(other)
103    }
104
105    /// Accumulate `batch` into the partial state through kernel dispatch, the vtable's
106    /// short-circuit, or execution to a columnar array, without consulting cached statistics.
107    fn accumulate_batch(&mut self, batch: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
108        // 1. Kernel registry first: a registered `(encoding, aggregate_fn)` kernel is strictly
109        //    more specific than the vtable's `try_accumulate` short-circuit. Checking the
110        //    registry first gives kernels for `Combined<V>` aggregates a chance to fire —
111        //    `Combined::try_accumulate` always returns true, so a later kernel check would be
112        //    unreachable.
113        {
114            let kernel = ctx
115                .aggregate_kernels()
116                .find_aggregate_kernel(batch.encoding_id(), self.aggregate_fn.id());
117            if let Some(kernel) = kernel
118                && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)?
119            {
120                vortex_ensure_eq!(
121                    result.dtype(),
122                    &self.dtypes.partial_dtype,
123                    "Aggregate kernel returned the wrong partial dtype",
124                );
125                self.fold_partial_scalar(&result)?;
126                return Ok(());
127            }
128        }
129
130        // 2. Allow the vtable to short-circuit on the raw array before decompression.
131        self.ensure_partial()?;
132        let partial = self.partial.as_mut().vortex_expect("partial materialized");
133        if self
134            .vtable
135            .try_accumulate(self.dtypes.args(&self.options), partial, batch, ctx)?
136        {
137            return Ok(());
138        }
139
140        // 3. Iteratively check the registry against each intermediate encoding, executing one
141        //    step between checks. Mirrors the loop in `GroupedAccumulator::accumulate_list_view`.
142        //    Iteration 0 re-checks the initial encoding — a redundant HashMap miss, the price of
143        //    keeping the loop body uniform. Terminates on `AnyColumnar` (Canonical or Constant)
144        //    since the vtable's `accumulate(&Columnar)` handles both cases directly.
145        let mut batch = batch.clone();
146        for _ in 0..max_iterations() {
147            if batch.is::<AnyColumnar>() {
148                break;
149            }
150
151            if let Some(kernel) = ctx
152                .aggregate_kernels()
153                .find_aggregate_kernel(batch.encoding_id(), self.aggregate_fn.id())
154                && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)?
155            {
156                vortex_ensure_eq!(
157                    result.dtype(),
158                    &self.dtypes.partial_dtype,
159                    "Aggregate kernel returned the wrong partial dtype",
160                );
161                self.fold_partial_scalar(&result)?;
162                return Ok(());
163            }
164
165            batch = batch.execute(ctx)?;
166        }
167
168        // 4. Otherwise, execute the batch until it is columnar and accumulate it into the state.
169        let columnar = batch.execute::<Columnar>(ctx)?;
170
171        self.ensure_partial()?;
172        let partial = self.partial.as_mut().vortex_expect("partial materialized");
173        self.vtable
174            .accumulate(self.dtypes.args(&self.options), partial, &columnar, ctx)
175    }
176}
177
178/// A trait object for type-erased accumulators, used for dynamic dispatch when the aggregate
179/// function is not known at compile time.
180pub trait DynAccumulator: 'static + Send {
181    /// Accumulate a new array into the accumulator's state.
182    fn accumulate(&mut self, batch: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()>;
183
184    /// Drain another accumulator's state into this one, resetting `other`.
185    ///
186    /// The other accumulator must have been constructed for the same aggregate function,
187    /// options, and input dtype as this one.
188    fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()>;
189
190    /// Parse a partial scalar and merge it into this accumulator's state.
191    ///
192    /// The scalar must have the dtype reported by the vtable's `partial_dtype`, and represents
193    /// input following the input already accumulated.
194    fn combine_partials(&mut self, partial: &Scalar) -> VortexResult<()>;
195
196    /// Whether the accumulator's result is fully determined.
197    fn is_saturated(&self) -> bool;
198
199    /// Reset the accumulator's state to the empty group.
200    fn reset(&mut self);
201
202    /// Read the current partial state as a scalar without resetting it.
203    ///
204    /// The returned scalar has the dtype reported by the vtable's `partial_dtype`.
205    fn partial_scalar(&self) -> VortexResult<Scalar>;
206
207    /// Compute the final aggregate result as a scalar without resetting state.
208    fn final_scalar(&self) -> VortexResult<Scalar>;
209
210    /// Flush the accumulation state and return the partial aggregate result as a scalar.
211    ///
212    /// Resets the accumulator state back to the initial state.
213    fn flush(&mut self) -> VortexResult<Scalar>;
214
215    /// Finish the accumulation and return the final aggregate result as a scalar.
216    ///
217    /// Resets the accumulator state back to the initial state.
218    fn finish(&mut self) -> VortexResult<Scalar>;
219
220    /// Access the accumulator as [`Any`], so it can be downcast to a typed [`Accumulator`].
221    fn as_any_mut(&mut self) -> &mut dyn Any;
222}
223
224impl dyn DynAccumulator {
225    /// Downcast to the typed [`Accumulator`] of the aggregate vtable `V`.
226    pub fn downcast_mut<V: AggregateFnVTable>(&mut self) -> Option<&mut Accumulator<V>> {
227        self.as_any_mut().downcast_mut()
228    }
229}
230
231impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
232    fn accumulate(&mut self, batch: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
233        if self.is_saturated() {
234            return Ok(());
235        }
236
237        vortex_ensure_eq!(batch.dtype(), &self.dtypes.dtype, "Input DType mismatch");
238
239        // Legacy stats bridge: an aggregate still cached under a legacy `Stat` slot consumes that
240        // exact stat before kernel dispatch or decode, and otherwise caches its result there.
241        let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) else {
242            return self.accumulate_batch(batch, ctx);
243        };
244
245        if let Precision::Exact(partial) = batch.statistics().get(stat) {
246            let partial = if partial.dtype() == &self.dtypes.partial_dtype {
247                partial
248            } else {
249                vortex_ensure!(
250                    partial
251                        .dtype()
252                        .eq_ignore_nullability(&self.dtypes.partial_dtype),
253                    "Aggregate {} read legacy stat {} with dtype {}, expected {}",
254                    self.aggregate_fn,
255                    stat,
256                    partial.dtype(),
257                    self.dtypes.partial_dtype,
258                );
259                partial.cast(&self.dtypes.partial_dtype)?
260            };
261            self.fold_partial_scalar(&partial)?;
262            return Ok(());
263        }
264
265        // Accumulate the batch into an empty state of its own, so that its result is cached on
266        // the requested batch rather than on whatever array the batch executes into.
267        let running = self.partial.take();
268        let accumulated = self.accumulate_batch(batch, ctx);
269        let batch_partial = std::mem::replace(&mut self.partial, running);
270        accumulated?;
271
272        let Some(batch_partial) = batch_partial else {
273            return Ok(());
274        };
275        let batch_result = self
276            .vtable
277            .to_scalar(self.dtypes.args(&self.options), &batch_partial)?;
278        // A null partial, e.g. an overflowed sum or the minimum of an all-null batch, has no exact
279        // stat value.
280        if let Some(value) = batch_result.into_value() {
281            batch.statistics().set(stat, Precision::Exact(value));
282        }
283        self.fold_partial(batch_partial)
284    }
285
286    fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()> {
287        let Some(other) = other.downcast_mut::<V>() else {
288            vortex_bail!(
289                "Cannot merge into a {} accumulator from an accumulator of a different aggregate",
290                self.aggregate_fn,
291            );
292        };
293        vortex_ensure!(
294            other.options == self.options && other.dtypes.dtype == self.dtypes.dtype,
295            "Cannot merge {} accumulators with different options or input dtypes",
296            self.aggregate_fn,
297        );
298        match other.partial.take() {
299            Some(partial) => self.fold_partial(partial),
300            None => Ok(()),
301        }
302    }
303
304    fn combine_partials(&mut self, partial: &Scalar) -> VortexResult<()> {
305        vortex_ensure_eq!(
306            partial.dtype(),
307            &self.dtypes.partial_dtype,
308            "Partial DType mismatch for {}",
309            self.aggregate_fn,
310        );
311        self.fold_partial_scalar(partial)
312    }
313
314    fn as_any_mut(&mut self) -> &mut dyn Any {
315        self
316    }
317
318    fn is_saturated(&self) -> bool {
319        self.partial.as_ref().is_some_and(|partial| {
320            self.vtable
321                .is_saturated(self.dtypes.args(&self.options), partial)
322        })
323    }
324
325    fn reset(&mut self) {
326        self.partial = None;
327    }
328
329    fn partial_scalar(&self) -> VortexResult<Scalar> {
330        let args = self.dtypes.args(&self.options);
331        let partial = match &self.partial {
332            Some(partial) => self.vtable.to_scalar(args, partial)?,
333            None => self.vtable.to_scalar(args, &self.empty_partial()?)?,
334        };
335
336        #[cfg(debug_assertions)]
337        {
338            vortex_ensure_eq!(
339                partial.dtype(),
340                args.partial_dtype,
341                "Aggregate returned incorrect DType on partial_scalar",
342            );
343        }
344
345        Ok(partial)
346    }
347
348    fn final_scalar(&self) -> VortexResult<Scalar> {
349        let args = self.dtypes.args(&self.options);
350        let result = match &self.partial {
351            Some(partial) => self.vtable.finalize_scalar(args, partial)?,
352            None => self.vtable.finalize_scalar(args, &self.empty_partial()?)?,
353        };
354
355        vortex_ensure_eq!(
356            result.dtype(),
357            args.return_dtype,
358            "Aggregate returned incorrect DType on final_scalar",
359        );
360
361        Ok(result)
362    }
363
364    fn flush(&mut self) -> VortexResult<Scalar> {
365        let partial = self.partial_scalar()?;
366        self.reset();
367        Ok(partial)
368    }
369
370    fn finish(&mut self) -> VortexResult<Scalar> {
371        let result = self.final_scalar()?;
372        self.reset();
373        Ok(result)
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use vortex_buffer::buffer;
380    use vortex_error::VortexResult;
381    use vortex_session::SessionExt;
382    use vortex_session::VortexSession;
383
384    use crate::ArrayRef;
385    use crate::ExecutionCtx;
386    use crate::IntoArray;
387    use crate::VortexSessionExecute;
388    use crate::aggregate_fn::Accumulator;
389    use crate::aggregate_fn::AccumulatorRef;
390    use crate::aggregate_fn::AggregateFnRef;
391    use crate::aggregate_fn::AggregateFnVTable;
392    use crate::aggregate_fn::DynAccumulator;
393    use crate::aggregate_fn::NumericalAggregateOpts;
394    use crate::aggregate_fn::combined::Combined;
395    use crate::aggregate_fn::combined::PairOptions;
396    use crate::aggregate_fn::fns::mean::Mean;
397    use crate::aggregate_fn::fns::min::Min;
398    use crate::aggregate_fn::fns::sum::Sum;
399    use crate::aggregate_fn::kernels::DynAggregateKernel;
400    use crate::aggregate_fn::session::AggregateFnSession;
401    use crate::array::VTable;
402    use crate::arrays::Dict;
403    use crate::arrays::DictArray;
404    use crate::dtype::DType;
405    use crate::dtype::Nullability;
406    use crate::dtype::PType;
407    use crate::expr::stats::Precision;
408    use crate::expr::stats::Stat;
409    use crate::expr::stats::StatsProvider;
410    use crate::expr::stats::StatsProviderExt;
411    use crate::scalar::Scalar;
412    use crate::scalar::ScalarValue;
413
414    /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the
415    /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate`
416    /// would produce for `dict_of_seven()`.
417    #[derive(Debug)]
418    struct SentinelMeanPartialKernel;
419    impl DynAggregateKernel for SentinelMeanPartialKernel {
420        fn aggregate(
421            &self,
422            _aggregate_fn: &AggregateFnRef,
423            _batch: &ArrayRef,
424            _ctx: &mut ExecutionCtx,
425        ) -> VortexResult<Option<Scalar>> {
426            Ok(Some(sentinel_partial()))
427        }
428    }
429
430    /// Returns `Ok(None)` => kernel declined, dispatch falls through.
431    #[derive(Debug)]
432    struct DeclineKernel;
433    impl DynAggregateKernel for DeclineKernel {
434        fn aggregate(
435            &self,
436            _aggregate_fn: &AggregateFnRef,
437            _batch: &ArrayRef,
438            _ctx: &mut ExecutionCtx,
439        ) -> VortexResult<Option<Scalar>> {
440            Ok(None)
441        }
442    }
443
444    /// A sentinel sum of `42.0`, distinguishable from `dict_of_seven()`'s natural sum of `7.0`.
445    #[derive(Debug)]
446    struct SentinelSumPartialKernel;
447    impl DynAggregateKernel for SentinelSumPartialKernel {
448        fn aggregate(
449            &self,
450            _aggregate_fn: &AggregateFnRef,
451            _batch: &ArrayRef,
452            _ctx: &mut ExecutionCtx,
453        ) -> VortexResult<Option<Scalar>> {
454            Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable)))
455        }
456    }
457
458    fn fresh_session() -> VortexSession {
459        crate::array_session()
460    }
461
462    fn dict_of_seven() -> ArrayRef {
463        DictArray::try_new(buffer![0u32].into_array(), buffer![7.0f64].into_array())
464            .expect("valid dictionary")
465            .into_array()
466    }
467
468    fn mean_f64_accumulator() -> VortexResult<Accumulator<Combined<Mean>>> {
469        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
470        Accumulator::try_new(
471            Mean::combined(),
472            PairOptions(
473                NumericalAggregateOpts::default(),
474                NumericalAggregateOpts::default(),
475            ),
476            dtype,
477        )
478    }
479
480    fn sentinel_partial() -> Scalar {
481        let acc = mean_f64_accumulator().expect("build accumulator");
482        let sum = Scalar::primitive(42.0f64, Nullability::Nullable);
483        let count = Scalar::primitive(1u64, Nullability::NonNullable);
484        Scalar::struct_(acc.dtypes.partial_dtype, vec![sum, count])
485    }
486
487    /// Kernel registered for `(Dict, Combined<Mean>)` fires in preference to
488    /// `Combined::try_accumulate`'s fan-out path — proves the dispatch reorder.
489    #[test]
490    fn combined_kernel_fires() -> VortexResult<()> {
491        static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel;
492        let session = fresh_session();
493        session
494            .get::<AggregateFnSession>()
495            .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL);
496        let mut ctx = session.create_execution_ctx();
497
498        let mut acc = mean_f64_accumulator()?;
499        acc.accumulate(&dict_of_seven(), &mut ctx)?;
500        let partial = acc.flush()?;
501
502        let s = partial.as_struct();
503        assert_eq!(
504            s.field("sum").unwrap().as_primitive().as_::<f64>(),
505            Some(42.0)
506        );
507        assert_eq!(
508            s.field("count").unwrap().as_primitive().as_::<u64>(),
509            Some(1)
510        );
511        Ok(())
512    }
513
514    /// An execution context snapshots the aggregate kernels on first use, so a kernel registered
515    /// afterwards only fires through a new context.
516    #[test]
517    fn execution_ctx_snapshots_aggregate_kernels_on_first_use() -> VortexResult<()> {
518        static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel;
519        let session = fresh_session();
520        let mut ctx = session.create_execution_ctx();
521
522        let mut acc = mean_f64_accumulator()?;
523        acc.accumulate(&dict_of_seven(), &mut ctx)?;
524        assert_eq!(
525            acc.flush()?
526                .as_struct()
527                .field("sum")
528                .unwrap()
529                .as_primitive()
530                .as_::<f64>(),
531            Some(7.0)
532        );
533
534        session
535            .get::<AggregateFnSession>()
536            .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL);
537
538        acc.accumulate(&dict_of_seven(), &mut ctx)?;
539        assert_eq!(
540            acc.flush()?
541                .as_struct()
542                .field("sum")
543                .unwrap()
544                .as_primitive()
545                .as_::<f64>(),
546            Some(7.0)
547        );
548
549        acc.accumulate(&dict_of_seven(), &mut session.create_execution_ctx())?;
550        assert_eq!(acc.flush()?, sentinel_partial());
551        Ok(())
552    }
553
554    /// Kernel returns `Ok(None)` => dispatch falls through to `Combined::try_accumulate`'s
555    /// natural fan-out. The natural partial is `{sum: 7.0, count: 1}`.
556    #[test]
557    fn fallback_when_kernel_declines() -> VortexResult<()> {
558        static KERNEL: DeclineKernel = DeclineKernel;
559        let session = fresh_session();
560        session
561            .get::<AggregateFnSession>()
562            .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL);
563        let mut ctx = session.create_execution_ctx();
564
565        let mut acc = mean_f64_accumulator()?;
566        acc.accumulate(&dict_of_seven(), &mut ctx)?;
567        let partial = acc.flush()?;
568
569        let s = partial.as_struct();
570        assert_eq!(
571            s.field("sum").unwrap().as_primitive().as_::<f64>(),
572            Some(7.0)
573        );
574        assert_eq!(
575            s.field("count").unwrap().as_primitive().as_::<u64>(),
576            Some(1)
577        );
578        Ok(())
579    }
580
581    /// A kernel registered for the inner `(Dict, Sum)` child fires when accumulating a
582    /// Dict batch through `Combined<Mean>`. This is the reusable-primitive case the
583    /// refactor enables: no `(Dict, Combined<Mean>)` kernel is needed.
584    #[test]
585    fn child_kernel_fires_through_combined() -> VortexResult<()> {
586        static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel;
587        let session = fresh_session();
588        session
589            .get::<AggregateFnSession>()
590            .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL);
591        let mut ctx = session.create_execution_ctx();
592
593        let mut acc = mean_f64_accumulator()?;
594        acc.accumulate(&dict_of_seven(), &mut ctx)?;
595        let partial = acc.flush()?;
596
597        let s = partial.as_struct();
598        // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired
599        // via `Combined<Mean>`'s fan-out. `Count`'s native `try_accumulate` reads the
600        // batch's valid_count, so count is the real 1.
601        assert_eq!(
602            s.field("sum").unwrap().as_primitive().as_::<f64>(),
603            Some(42.0)
604        );
605        assert_eq!(
606            s.field("count").unwrap().as_primitive().as_::<u64>(),
607            Some(1)
608        );
609        Ok(())
610    }
611
612    #[test]
613    fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> {
614        static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel;
615        let session = fresh_session();
616        session
617            .get::<AggregateFnSession>()
618            .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL);
619        let mut ctx = session.create_execution_ctx();
620
621        let batch = dict_of_seven();
622        batch
623            .statistics()
624            .set(Stat::Sum, Precision::Exact(ScalarValue::from(7.0f64)));
625
626        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
627        let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?;
628        acc.accumulate(&batch, &mut ctx)?;
629
630        assert_eq!(acc.finish()?.as_primitive().as_::<f64>(), Some(7.0));
631        Ok(())
632    }
633
634    /// Each batch caches its own result, not the running result and not on the canonical array
635    /// it executes into.
636    #[test]
637    fn caches_result_on_requested_batch() -> VortexResult<()> {
638        let mut ctx = fresh_session().create_execution_ctx();
639        let dict = |codes: ArrayRef| -> VortexResult<ArrayRef> {
640            Ok(DictArray::try_new(codes, buffer![5i32, 1, 9].into_array())?.into_array())
641        };
642        let batch1 = dict(buffer![0u32, 1, 2].into_array())?;
643        let batch2 = dict(buffer![2u32, 0].into_array())?;
644
645        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
646        let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?;
647        acc.accumulate(&batch1, &mut ctx)?;
648        acc.accumulate(&batch2, &mut ctx)?;
649
650        assert_eq!(
651            acc.finish()?,
652            Scalar::primitive(1i32, Nullability::Nullable)
653        );
654        assert_eq!(
655            batch1.statistics().get_as::<i32>(Stat::Min),
656            Precision::exact(1i32)
657        );
658        assert_eq!(
659            batch2.statistics().get_as::<i32>(Stat::Min),
660            Precision::exact(5i32)
661        );
662        Ok(())
663    }
664
665    #[test]
666    fn overflowed_sum_is_not_cached() -> VortexResult<()> {
667        let mut ctx = fresh_session().create_execution_ctx();
668        let batch = buffer![i64::MAX, i64::MAX].into_array();
669
670        let dtype = DType::Primitive(PType::I64, Nullability::NonNullable);
671        let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?;
672        acc.accumulate(&batch, &mut ctx)?;
673
674        assert!(acc.finish()?.is_null());
675        assert!(batch.statistics().get(Stat::Sum).as_exact().is_none());
676        Ok(())
677    }
678
679    fn sum_i32_accumulator(options: NumericalAggregateOpts) -> VortexResult<AccumulatorRef> {
680        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
681        Ok(Box::new(Accumulator::try_new(Sum, options, dtype)?))
682    }
683
684    #[test]
685    fn merge_from_drains_other_accumulator() -> VortexResult<()> {
686        let mut ctx = fresh_session().create_execution_ctx();
687        let mut global = sum_i32_accumulator(NumericalAggregateOpts::default())?;
688        let mut local = sum_i32_accumulator(NumericalAggregateOpts::default())?;
689
690        global.accumulate(&buffer![10i32, 20].into_array(), &mut ctx)?;
691        local.accumulate(&buffer![5i32].into_array(), &mut ctx)?;
692        global.merge_from(local.as_mut())?;
693
694        assert_eq!(
695            global.finish()?,
696            Scalar::primitive(35i64, Nullability::Nullable)
697        );
698        // The merged-from accumulator is reset back to the empty state.
699        assert_eq!(
700            local.finish()?,
701            Scalar::primitive(0i64, Nullability::Nullable)
702        );
703        Ok(())
704    }
705
706    #[test]
707    fn merge_from_rejects_a_different_aggregate() -> VortexResult<()> {
708        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
709        let mut sum = sum_i32_accumulator(NumericalAggregateOpts::default())?;
710        let mut min: AccumulatorRef = Box::new(Accumulator::try_new(
711            Min,
712            NumericalAggregateOpts::default(),
713            dtype,
714        )?);
715
716        assert!(sum.merge_from(min.as_mut()).is_err());
717        Ok(())
718    }
719
720    #[test]
721    fn merge_from_rejects_mismatched_options() -> VortexResult<()> {
722        let mut skipping = sum_i32_accumulator(NumericalAggregateOpts::skip_nans())?;
723        let mut including = sum_i32_accumulator(NumericalAggregateOpts::include_nans())?;
724
725        assert!(skipping.merge_from(including.as_mut()).is_err());
726        Ok(())
727    }
728
729    #[test]
730    fn merge_from_combines_child_accumulators() -> VortexResult<()> {
731        let mut ctx = fresh_session().create_execution_ctx();
732        let mut global: AccumulatorRef = Box::new(mean_f64_accumulator()?);
733        let mut local: AccumulatorRef = Box::new(mean_f64_accumulator()?);
734
735        global.accumulate(&buffer![1.0f64, 2.0].into_array(), &mut ctx)?;
736        local.accumulate(&buffer![6.0f64].into_array(), &mut ctx)?;
737        global.merge_from(local.as_mut())?;
738
739        assert_eq!(global.finish()?.as_primitive().as_::<f64>(), Some(3.0));
740        Ok(())
741    }
742}
vortex_array/aggregate_fn/accumulator.rs

vortex_array/aggregate_fn/
accumulator.rs
