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

vortex_array/aggregate_fn/fns/sum/
mod.rs
