vortex_array/aggregate_fn/fns/sum/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod bool;
5mod constant;
6mod decimal;
7mod grouped;
8mod primitive;
9pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel;
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_err;
14use vortex_error::vortex_panic;
15use vortex_session::VortexSession;
16use vortex_session::registry::CachedId;
17
18pub(crate) use self::bool::accumulate_bool;
19pub(crate) use self::constant::multiply_constant;
20pub(crate) use self::decimal::accumulate_decimal;
21pub(crate) use self::primitive::accumulate_primitive;
22pub(crate) use self::primitive::sum_float_all;
23pub(crate) use self::primitive::sum_signed_all;
24pub(crate) use self::primitive::sum_unsigned_all;
25use crate::ArrayRef;
26use crate::Canonical;
27use crate::Columnar;
28use crate::ExecutionCtx;
29use crate::aggregate_fn::Accumulator;
30use crate::aggregate_fn::AggregateFnId;
31use crate::aggregate_fn::AggregateFnVTable;
32use crate::aggregate_fn::DynAccumulator;
33use crate::aggregate_fn::NumericalAggregateOpts;
34use crate::dtype::DType;
35use crate::dtype::DecimalDType;
36use crate::dtype::MAX_PRECISION;
37use crate::dtype::Nullability;
38use crate::dtype::PType;
39use crate::expr::stats::Precision;
40use crate::expr::stats::Stat;
41use crate::expr::stats::StatsProvider;
42use crate::expr::stats::StatsProviderExt;
43use crate::scalar::DecimalValue;
44use crate::scalar::Scalar;
45
46/// Return the sum of an array.
47///
48/// See [`Sum`] for details.
49pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
50    // Short-circuit using cached array statistics.
51    if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) {
52        return Ok(sum_scalar);
53    }
54
55    // Compute using Accumulator<Sum>.
56    // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe.
57    let mut acc = Accumulator::try_new(
58        Sum,
59        NumericalAggregateOpts::default(),
60        array.dtype().clone(),
61    )?;
62    acc.accumulate(array, ctx)?;
63    let result = acc.finish()?;
64
65    // Cache the computed sum as a statistic (only if non-null, i.e. no overflow).
66    if let Some(val) = result.value().cloned() {
67        array.statistics().set(Stat::Sum, Precision::Exact(val));
68    }
69
70    Ok(result)
71}
72
73/// Sum an array, starting from zero.
74///
75/// If the sum overflows, a null scalar will be returned.
76/// If the array is all-invalid, the sum will be zero.
77///
78/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the
79/// default) NaN values contribute nothing, otherwise any NaN value poisons the sum to NaN.
80#[derive(Clone, Debug)]
81pub struct Sum;
82
83// Both Spark and DataFusion use this heuristic.
84// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66
85// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188
86pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType {
87    DecimalDType::new(
88        u8::min(MAX_PRECISION, input.precision() + 10),
89        input.scale(),
90    )
91}
92
93impl AggregateFnVTable for Sum {
94    type Options = NumericalAggregateOpts;
95    type Partial = SumPartial;
96
97    fn id(&self) -> AggregateFnId {
98        static ID: CachedId = CachedId::new("vortex.sum");
99        *ID
100    }
101
102    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
103        Ok(Some(options.serialize()))
104    }
105
106    fn deserialize(
107        &self,
108        metadata: &[u8],
109        _session: &VortexSession,
110    ) -> VortexResult<Self::Options> {
111        NumericalAggregateOpts::deserialize(metadata)
112    }
113
114    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
115        // When a sum overflows, we return a sum _value_ of null. Therefore, we all return dtypes
116        // are nullable.
117        use Nullability::Nullable;
118
119        Some(match input_dtype {
120            DType::Bool(_) => DType::Primitive(PType::U64, Nullable),
121            DType::Primitive(ptype, _) => match ptype {
122                PType::U8 | PType::U16 | PType::U32 | PType::U64 => {
123                    DType::Primitive(PType::U64, Nullable)
124                }
125                PType::I8 | PType::I16 | PType::I32 | PType::I64 => {
126                    DType::Primitive(PType::I64, Nullable)
127                }
128                PType::F16 | PType::F32 | PType::F64 => {
129                    // Float sums cannot overflow, but all null floats still end up as null
130                    DType::Primitive(PType::F64, Nullable)
131                }
132            },
133            DType::Decimal(decimal_dtype, _) => {
134                DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable)
135            }
136            // Unsupported types
137            _ => return None,
138        })
139    }
140
141    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
142        self.return_dtype(options, input_dtype)
143    }
144
145    fn empty_partial(
146        &self,
147        options: &Self::Options,
148        input_dtype: &DType,
149    ) -> VortexResult<Self::Partial> {
150        let return_dtype = self
151            .return_dtype(options, input_dtype)
152            .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?;
153        let initial = make_zero_state(&return_dtype);
154
155        Ok(SumPartial {
156            return_dtype,
157            current: Some(initial),
158            skip_nans: options.skip_nans,
159        })
160    }
161
162    fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
163        if other.is_null() {
164            // A null partial means the sub-accumulator saturated (overflow).
165            partial.current = None;
166            return Ok(());
167        }
168        let Some(ref mut inner) = partial.current else {
169            return Ok(());
170        };
171        let saturated = match inner {
172            SumState::Unsigned(acc) => {
173                let val = other
174                    .as_primitive()
175                    .typed_value::<u64>()
176                    .vortex_expect("checked non-null");
177                checked_add_u64(acc, val)
178            }
179            SumState::Signed(acc) => {
180                let val = other
181                    .as_primitive()
182                    .typed_value::<i64>()
183                    .vortex_expect("checked non-null");
184                checked_add_i64(acc, val)
185            }
186            SumState::Float(acc) => {
187                let val = other
188                    .as_primitive()
189                    .typed_value::<f64>()
190                    .vortex_expect("checked non-null");
191                *acc += val;
192                false
193            }
194            SumState::Decimal { value, dtype } => {
195                let val = other
196                    .as_decimal()
197                    .decimal_value()
198                    .vortex_expect("checked non-null");
199                match value.checked_add(&val) {
200                    Some(r) => {
201                        *value = r;
202                        !value.fits_in_precision(*dtype)
203                    }
204                    None => true,
205                }
206            }
207        };
208        if saturated {
209            partial.current = None;
210        }
211        Ok(())
212    }
213
214    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
215        Ok(match &partial.current {
216            None => Scalar::null(partial.return_dtype.as_nullable()),
217            Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable),
218            Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable),
219            Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable),
220            Some(SumState::Decimal { value, .. }) => {
221                let decimal_dtype = *partial
222                    .return_dtype
223                    .as_decimal_opt()
224                    .vortex_expect("return dtype must be decimal");
225                Scalar::decimal(*value, decimal_dtype, Nullability::Nullable)
226            }
227        })
228    }
229
230    fn reset(&self, partial: &mut Self::Partial) {
231        partial.current = Some(make_zero_state(&partial.return_dtype));
232    }
233
234    #[inline]
235    fn is_saturated(&self, partial: &Self::Partial) -> bool {
236        match partial.current.as_ref() {
237            None => true,
238            Some(SumState::Float(v)) => v.is_nan(),
239            Some(_) => false,
240        }
241    }
242
243    fn try_accumulate(
244        &self,
245        partial: &mut Self::Partial,
246        batch: &ArrayRef,
247        _ctx: &mut ExecutionCtx,
248    ) -> VortexResult<bool> {
249        // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes
250        // the default dispatch path.
251        if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) {
252            return Ok(false);
253        }
254        match batch.statistics().get_as::<u64>(Stat::NaNCount) {
255            Precision::Exact(0) => {
256                // NaN-free batch: the cached NaN-skipping sum (if any) equals the
257                // NaN-including sum.
258                if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) {
259                    let sum = if sum.dtype() == &partial.return_dtype {
260                        sum
261                    } else {
262                        sum.cast(&partial.return_dtype)?
263                    };
264                    self.combine_partials(partial, sum)?;
265                    return Ok(true);
266                }
267                Ok(false)
268            }
269            Precision::Exact(_) => {
270                // At least one NaN value: the sum is NaN without scanning the batch.
271                if let Some(SumState::Float(acc)) = partial.current.as_mut() {
272                    *acc = f64::NAN;
273                }
274                Ok(true)
275            }
276            _ => Ok(false),
277        }
278    }
279
280    fn accumulate(
281        &self,
282        partial: &mut Self::Partial,
283        batch: &Columnar,
284        ctx: &mut ExecutionCtx,
285    ) -> VortexResult<()> {
286        // Constants compute scalar * len and combine via combine_partials.
287        if let Columnar::Constant(c) = batch {
288            // NaN constants are treated as missing when skipping NaNs.
289            if partial.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) {
290                return Ok(());
291            }
292            if let Some(product) = multiply_constant(c.scalar(), c.len(), &partial.return_dtype)? {
293                self.combine_partials(partial, product)?;
294            }
295            return Ok(());
296        }
297
298        let skip_nans = partial.skip_nans;
299        let mut inner = match partial.current.take() {
300            Some(inner) => inner,
301            None => return Ok(()),
302        };
303
304        let result = match batch {
305            Columnar::Canonical(c) => match c {
306                Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans),
307                Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx),
308                Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx),
309                _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()),
310            },
311            Columnar::Constant(_) => unreachable!(),
312        };
313
314        match result {
315            Ok(false) => partial.current = Some(inner),
316            Ok(true) => {} // saturated: current stays None
317            Err(e) => {
318                partial.current = Some(inner);
319                return Err(e);
320            }
321        }
322        Ok(())
323    }
324
325    fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
326        Ok(partials)
327    }
328
329    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
330        self.to_scalar(partial)
331    }
332}
333
334/// The group state for a sum aggregate, containing the accumulated value and configuration
335/// needed for reset/result without external context.
336pub struct SumPartial {
337    return_dtype: DType,
338    /// The current accumulated state, or `None` if saturated (checked overflow).
339    current: Option<SumState>,
340    /// Whether NaN values in float inputs are skipped.
341    skip_nans: bool,
342}
343
344/// The accumulated sum value.
345// TODO(ngates): instead of an enum, we should use a Box<dyn State> to avoid dispatcher over the
346//  input type every time? Perhaps?
347pub enum SumState {
348    Unsigned(u64),
349    Signed(i64),
350    Float(f64),
351    Decimal {
352        value: DecimalValue,
353        dtype: DecimalDType,
354    },
355}
356
357pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState {
358    match return_dtype {
359        DType::Primitive(ptype, _) => match ptype {
360            PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0),
361            PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0),
362            PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0),
363        },
364        DType::Decimal(decimal, _) => SumState::Decimal {
365            value: DecimalValue::zero(decimal),
366            dtype: *decimal,
367        },
368        _ => vortex_panic!("Unsupported sum type"),
369    }
370}
371
372/// Checked add for u64, returning true if overflow occurred.
373#[allow(clippy::inline_always)]
374#[inline(always)]
375fn checked_add_u64(acc: &mut u64, val: u64) -> bool {
376    match acc.checked_add(val) {
377        Some(r) => {
378            *acc = r;
379            false
380        }
381        None => true,
382    }
383}
384
385/// Checked add for i64, returning true if overflow occurred.
386#[allow(clippy::inline_always)]
387#[inline(always)]
388fn checked_add_i64(acc: &mut i64, val: i64) -> bool {
389    match acc.checked_add(val) {
390        Some(r) => {
391            *acc = r;
392            false
393        }
394        None => true,
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use num_traits::CheckedAdd;
401    use vortex_buffer::buffer;
402    use vortex_error::VortexExpect;
403    use vortex_error::VortexResult;
404
405    use crate::ArrayRef;
406    use crate::IntoArray;
407    use crate::VortexSessionExecute;
408    use crate::aggregate_fn::Accumulator;
409    use crate::aggregate_fn::AggregateFnVTable;
410    use crate::aggregate_fn::DynAccumulator;
411    use crate::aggregate_fn::DynGroupedAccumulator;
412    use crate::aggregate_fn::GroupedAccumulator;
413    use crate::aggregate_fn::NumericalAggregateOpts;
414    use crate::aggregate_fn::fns::sum::Sum;
415    use crate::aggregate_fn::fns::sum::sum;
416    use crate::array_session;
417    use crate::arrays::BoolArray;
418    use crate::arrays::ChunkedArray;
419    use crate::arrays::ConstantArray;
420    use crate::arrays::DecimalArray;
421    use crate::arrays::FixedSizeListArray;
422    use crate::arrays::ListViewArray;
423    use crate::arrays::PrimitiveArray;
424    use crate::assert_arrays_eq;
425    use crate::dtype::DType;
426    use crate::dtype::DecimalDType;
427    use crate::dtype::Nullability;
428    use crate::dtype::Nullability::Nullable;
429    use crate::dtype::PType;
430    use crate::dtype::i256;
431    use crate::expr::stats::Precision;
432    use crate::expr::stats::Stat;
433    use crate::expr::stats::StatsProvider;
434    use crate::scalar::DecimalValue;
435    use crate::scalar::NumericOperator;
436    use crate::scalar::Scalar;
437    use crate::validity::Validity;
438
439    /// Sum an array with an initial value (test-only helper).
440    fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult<Scalar> {
441        let mut ctx = array_session().create_execution_ctx();
442        if accumulator.is_null() {
443            return Ok(accumulator.clone());
444        }
445        if accumulator.is_zero() == Some(true) {
446            return sum(array, &mut ctx);
447        }
448
449        let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| {
450            vortex_error::vortex_err!("Sum not supported for dtype: {}", array.dtype())
451        })?;
452
453        // For non-float types, try statistics short-circuit with accumulator.
454        if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float())
455            && let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum)
456        {
457            return add_scalars(&sum_dtype, &sum_scalar, accumulator);
458        }
459
460        // Compute array sum from zero (also caches stats).
461        let array_sum = sum(array, &mut ctx)?;
462
463        // Combine with the accumulator.
464        add_scalars(&sum_dtype, &array_sum, accumulator)
465    }
466
467    /// Add two sum scalars with overflow checking.
468    fn add_scalars(sum_dtype: &DType, lhs: &Scalar, rhs: &Scalar) -> VortexResult<Scalar> {
469        if lhs.is_null() || rhs.is_null() {
470            return Ok(Scalar::null(sum_dtype.as_nullable()));
471        }
472
473        Ok(match sum_dtype {
474            DType::Primitive(ptype, _) if ptype.is_float() => {
475                let lhs_val = f64::try_from(lhs)?;
476                let rhs_val = f64::try_from(rhs)?;
477                Scalar::primitive(lhs_val + rhs_val, Nullable)
478            }
479            DType::Primitive(..) => lhs
480                .as_primitive()
481                .checked_add(&rhs.as_primitive())
482                .map(Scalar::from)
483                .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())),
484            // Add widens the result precision, so restate the sum in the accumulator's own
485            // decimal type, treating a value that no longer fits as an overflow.
486            DType::Decimal(decimal_dtype, _) => lhs
487                .as_decimal()
488                .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add)?
489                .and_then(|scalar| scalar.as_decimal().decimal_value())
490                .filter(|value| value.fits_in_precision(*decimal_dtype))
491                .map(|value| Scalar::decimal(value, *decimal_dtype, Nullable))
492                .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())),
493            _ => unreachable!("Sum will always be a decimal or a primitive dtype"),
494        })
495    }
496
497    // Multi-batch and reset tests
498
499    #[test]
500    fn sum_multi_batch() -> VortexResult<()> {
501        let mut ctx = array_session().create_execution_ctx();
502        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
503        let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?;
504
505        let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
506        acc.accumulate(&batch1, &mut ctx)?;
507
508        let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array();
509        acc.accumulate(&batch2, &mut ctx)?;
510
511        let result = acc.finish()?;
512        assert_eq!(result.as_primitive().typed_value::<i64>(), Some(48));
513        Ok(())
514    }
515
516    #[test]
517    fn sum_finish_resets_state() -> VortexResult<()> {
518        let mut ctx = array_session().create_execution_ctx();
519        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
520        let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?;
521
522        let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
523        acc.accumulate(&batch1, &mut ctx)?;
524        let result1 = acc.finish()?;
525        assert_eq!(result1.as_primitive().typed_value::<i64>(), Some(30));
526
527        let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array();
528        acc.accumulate(&batch2, &mut ctx)?;
529        let result2 = acc.finish()?;
530        assert_eq!(result2.as_primitive().typed_value::<i64>(), Some(18));
531        Ok(())
532    }
533
534    // State merge tests (vtable-level)
535
536    #[test]
537    fn sum_state_merge() -> VortexResult<()> {
538        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
539        let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?;
540
541        let scalar1 = Scalar::primitive(100i64, Nullable);
542        Sum.combine_partials(&mut state, scalar1)?;
543
544        let scalar2 = Scalar::primitive(50i64, Nullable);
545        Sum.combine_partials(&mut state, scalar2)?;
546
547        let result = Sum.to_scalar(&state)?;
548        Sum.reset(&mut state);
549        assert_eq!(result.as_primitive().typed_value::<i64>(), Some(150));
550        Ok(())
551    }
552
553    // Stats caching test
554
555    #[test]
556    fn sum_stats() -> VortexResult<()> {
557        let array = ChunkedArray::try_new(
558            vec![
559                PrimitiveArray::from_iter([1, 1, 1]).into_array(),
560                PrimitiveArray::from_iter([2, 2, 2]).into_array(),
561            ],
562            DType::Primitive(PType::I32, Nullability::NonNullable),
563        )
564        .vortex_expect("operation should succeed in test");
565        let array = array.into_array();
566        // compute sum with accumulator to populate stats
567        sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?;
568
569        let sum_without_acc = sum(&array, &mut array_session().create_execution_ctx())?;
570        assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable));
571        Ok(())
572    }
573
574    // Constant float non-multiply test
575
576    #[test]
577    fn sum_constant_float_non_multiply() -> VortexResult<()> {
578        let acc = -2048669276050936500000000000f64;
579        let array = ConstantArray::new(6.1811675e16f64, 25);
580        let result = sum_with_accumulator(&array.into_array(), &Scalar::primitive(acc, Nullable))
581            .vortex_expect("operation should succeed in test");
582        assert_eq!(
583            f64::try_from(&result).vortex_expect("operation should succeed in test"),
584            -2048669274505644600000000000f64
585        );
586        Ok(())
587    }
588
589    // Grouped sum tests
590
591    fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult<ArrayRef> {
592        let mut acc = GroupedAccumulator::try_new(
593            Sum,
594            NumericalAggregateOpts::default(),
595            elem_dtype.clone(),
596        )?;
597        acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?;
598        acc.finish()
599    }
600
601    #[test]
602    fn grouped_sum_fixed_size_list() -> VortexResult<()> {
603        let mut ctx = array_session().create_execution_ctx();
604        let elements =
605            PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array();
606        let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?;
607
608        let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
609        let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?;
610
611        let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array();
612        assert_arrays_eq!(&result, &expected, &mut ctx);
613        Ok(())
614    }
615
616    #[test]
617    fn grouped_sum_with_null_elements() -> VortexResult<()> {
618        let mut ctx = array_session().create_execution_ctx();
619        let elements =
620            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)])
621                .into_array();
622        let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?;
623
624        let elem_dtype = DType::Primitive(PType::I32, Nullable);
625        let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?;
626
627        let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array();
628        assert_arrays_eq!(&result, &expected, &mut ctx);
629        Ok(())
630    }
631
632    #[test]
633    fn grouped_sum_with_null_group() -> VortexResult<()> {
634        let mut ctx = array_session().create_execution_ctx();
635        let elements =
636            PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable)
637                .into_array();
638        let validity = Validity::from_iter([true, false, true]);
639        let groups = FixedSizeListArray::try_new(elements, 3, validity, 3)?;
640
641        let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
642        let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?;
643
644        let expected =
645            PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array();
646        assert_arrays_eq!(&result, &expected, &mut ctx);
647        Ok(())
648    }
649
650    #[test]
651    fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> {
652        let mut ctx = array_session().create_execution_ctx();
653        let elements =
654            PrimitiveArray::from_option_iter([None::<i32>, None, Some(3), Some(4)]).into_array();
655        let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?;
656
657        let elem_dtype = DType::Primitive(PType::I32, Nullable);
658        let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?;
659
660        let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array();
661        assert_arrays_eq!(&result, &expected, &mut ctx);
662        Ok(())
663    }
664
665    #[test]
666    fn grouped_sum_bool() -> VortexResult<()> {
667        let mut ctx = array_session().create_execution_ctx();
668        let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect();
669        let groups =
670            FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?;
671
672        let elem_dtype = DType::Bool(Nullability::NonNullable);
673        let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?;
674
675        let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).into_array();
676        assert_arrays_eq!(&result, &expected, &mut ctx);
677        Ok(())
678    }
679
680    #[test]
681    fn grouped_sum_finish_resets() -> VortexResult<()> {
682        let mut ctx = array_session().create_execution_ctx();
683        let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
684        let mut acc =
685            GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype)?;
686
687        let elements1 =
688            PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array();
689        let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?;
690        acc.accumulate_list(&groups1.into_array(), &mut ctx)?;
691        let result1 = acc.finish()?;
692
693        let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array();
694        assert_arrays_eq!(&result1, &expected1, &mut ctx);
695
696        let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
697        let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?;
698        acc.accumulate_list(&groups2.into_array(), &mut ctx)?;
699        let result2 = acc.finish()?;
700
701        let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array();
702        assert_arrays_eq!(&result2, &expected2, &mut ctx);
703        Ok(())
704    }
705
706    #[test]
707    fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> {
708        let mut ctx = array_session().create_execution_ctx();
709        let elements =
710            PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array();
711        let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array();
712        let sizes = PrimitiveArray::new(buffer![1i32, 1, 1], Validity::NonNullable).into_array();
713        let validity = Validity::from_iter([true, false, true]);
714        let groups = ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array();
715
716        let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
717        let result = run_grouped_sum(&groups, &elem_dtype)?;
718
719        // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200.
720        let expected =
721            PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array();
722        assert_arrays_eq!(&result, &expected, &mut ctx);
723        Ok(())
724    }
725
726    // Chunked array tests
727
728    #[test]
729    fn sum_chunked_floats_with_nulls() -> VortexResult<()> {
730        let chunk1 =
731            PrimitiveArray::from_option_iter(vec![Some(1.5f64), None, Some(3.2), Some(4.8)]);
732        let chunk2 = PrimitiveArray::from_option_iter(vec![Some(2.1f64), Some(5.7), None]);
733        let chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]);
734        let dtype = chunk1.dtype().clone();
735        let chunked = ChunkedArray::try_new(
736            vec![
737                chunk1.into_array(),
738                chunk2.into_array(),
739                chunk3.into_array(),
740            ],
741            dtype,
742        )?;
743
744        let result = sum(
745            &chunked.into_array(),
746            &mut array_session().create_execution_ctx(),
747        )?;
748        assert_eq!(result.as_primitive().as_::<f64>(), Some(20.8));
749        Ok(())
750    }
751
752    #[test]
753    fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> {
754        let chunk1 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None, None]);
755        let chunk2 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None]);
756        let dtype = chunk1.dtype().clone();
757        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
758        let result = sum(
759            &chunked.into_array(),
760            &mut array_session().create_execution_ctx(),
761        )?;
762        assert_eq!(result, Scalar::primitive(0f64, Nullable));
763        Ok(())
764    }
765
766    #[test]
767    fn sum_chunked_floats_empty_chunks() -> VortexResult<()> {
768        let chunk1 = PrimitiveArray::from_option_iter(vec![Some(10.5f64), Some(20.3)]);
769        let chunk2 = ConstantArray::new(Scalar::primitive(0f64, Nullable), 0);
770        let chunk3 = PrimitiveArray::from_option_iter(vec![Some(5.2f64)]);
771        let dtype = chunk1.dtype().clone();
772        let chunked = ChunkedArray::try_new(
773            vec![
774                chunk1.into_array(),
775                chunk2.into_array(),
776                chunk3.into_array(),
777            ],
778            dtype,
779        )?;
780
781        let result = sum(
782            &chunked.into_array(),
783            &mut array_session().create_execution_ctx(),
784        )?;
785        assert_eq!(result.as_primitive().as_::<f64>(), Some(36.0));
786        Ok(())
787    }
788
789    #[test]
790    fn sum_chunked_int_almost_all_null() -> VortexResult<()> {
791        let chunk1 = PrimitiveArray::from_option_iter::<u32, _>(vec![Some(1)]);
792        let chunk2 = PrimitiveArray::from_option_iter::<u32, _>(vec![None]);
793        let dtype = chunk1.dtype().clone();
794        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
795
796        let result = sum(
797            &chunked.into_array(),
798            &mut array_session().create_execution_ctx(),
799        )?;
800        assert_eq!(result.as_primitive().as_::<u64>(), Some(1));
801        Ok(())
802    }
803
804    #[test]
805    fn sum_chunked_decimals() -> VortexResult<()> {
806        let decimal_dtype = DecimalDType::new(10, 2);
807        let chunk1 = DecimalArray::new(
808            buffer![100i32, 100i32, 100i32, 100i32, 100i32],
809            decimal_dtype,
810            Validity::AllValid,
811        );
812        let chunk2 = DecimalArray::new(
813            buffer![200i32, 200i32, 200i32],
814            decimal_dtype,
815            Validity::AllValid,
816        );
817        let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid);
818        let dtype = chunk1.dtype().clone();
819        let chunked = ChunkedArray::try_new(
820            vec![
821                chunk1.into_array(),
822                chunk2.into_array(),
823                chunk3.into_array(),
824            ],
825            dtype,
826        )?;
827
828        let result = sum(
829            &chunked.into_array(),
830            &mut array_session().create_execution_ctx(),
831        )?;
832        let decimal_result = result.as_decimal();
833        assert_eq!(
834            decimal_result.decimal_value(),
835            Some(DecimalValue::I256(i256::from_i128(1700)))
836        );
837        Ok(())
838    }
839
840    #[test]
841    fn sum_chunked_decimals_with_nulls() -> VortexResult<()> {
842        let decimal_dtype = DecimalDType::new(10, 2);
843        let chunk1 = DecimalArray::new(
844            buffer![100i32, 100i32, 100i32],
845            decimal_dtype,
846            Validity::AllValid,
847        );
848        let chunk2 = DecimalArray::new(
849            buffer![0i32, 0i32],
850            decimal_dtype,
851            Validity::from_iter([false, false]),
852        );
853        let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid);
854        let dtype = chunk1.dtype().clone();
855        let chunked = ChunkedArray::try_new(
856            vec![
857                chunk1.into_array(),
858                chunk2.into_array(),
859                chunk3.into_array(),
860            ],
861            dtype,
862        )?;
863
864        let result = sum(
865            &chunked.into_array(),
866            &mut array_session().create_execution_ctx(),
867        )?;
868        let decimal_result = result.as_decimal();
869        assert_eq!(
870            decimal_result.decimal_value(),
871            Some(DecimalValue::I256(i256::from_i128(700)))
872        );
873        Ok(())
874    }
875
876    #[test]
877    fn sum_chunked_decimals_large() -> VortexResult<()> {
878        let decimal_dtype = DecimalDType::new(3, 0);
879        let chunk1 = ConstantArray::new(
880            Scalar::decimal(
881                DecimalValue::I16(500),
882                decimal_dtype,
883                Nullability::NonNullable,
884            ),
885            1,
886        );
887        let chunk2 = ConstantArray::new(
888            Scalar::decimal(
889                DecimalValue::I16(600),
890                decimal_dtype,
891                Nullability::NonNullable,
892            ),
893            1,
894        );
895        let dtype = chunk1.dtype().clone();
896        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
897
898        let result = sum(
899            &chunked.into_array(),
900            &mut array_session().create_execution_ctx(),
901        )?;
902        let decimal_result = result.as_decimal();
903        assert_eq!(
904            decimal_result.decimal_value(),
905            Some(DecimalValue::I256(i256::from_i128(1100)))
906        );
907        assert_eq!(
908            result.dtype(),
909            &DType::Decimal(DecimalDType::new(13, 0), Nullable)
910        );
911        Ok(())
912    }
913}
vortex_array/aggregate_fn/fns/sum/mod.rs

vortex_array/aggregate_fn/fns/sum/
mod.rs
