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

vortex_array/aggregate_fn/fns/sum/
mod.rs
