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 =
735            ListViewArray::try_new(elements, offsets, sizes, validity, &mut ctx)?.into_array();
736
737        let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
738        let result = run_grouped_sum(&groups, &elem_dtype)?;
739
740        // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200.
741        let expected =
742            PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array();
743        assert_arrays_eq!(&result, &expected, &mut ctx);
744        Ok(())
745    }
746
747    // Chunked array tests
748
749    #[test]
750    fn sum_chunked_floats_with_nulls() -> VortexResult<()> {
751        let chunk1 =
752            PrimitiveArray::from_option_iter(vec![Some(1.5f64), None, Some(3.2), Some(4.8)]);
753        let chunk2 = PrimitiveArray::from_option_iter(vec![Some(2.1f64), Some(5.7), None]);
754        let chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]);
755        let dtype = chunk1.dtype().clone();
756        let chunked = ChunkedArray::try_new(
757            vec![
758                chunk1.into_array(),
759                chunk2.into_array(),
760                chunk3.into_array(),
761            ],
762            dtype,
763        )?;
764
765        let result = sum(
766            &chunked.into_array(),
767            &mut array_session().create_execution_ctx(),
768        )?;
769        assert_eq!(result.as_primitive().as_::<f64>(), Some(20.8));
770        Ok(())
771    }
772
773    #[test]
774    fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> {
775        let chunk1 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None, None]);
776        let chunk2 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None]);
777        let dtype = chunk1.dtype().clone();
778        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
779        let result = sum(
780            &chunked.into_array(),
781            &mut array_session().create_execution_ctx(),
782        )?;
783        assert_eq!(result, Scalar::primitive(0f64, Nullable));
784        Ok(())
785    }
786
787    #[test]
788    fn sum_chunked_floats_empty_chunks() -> VortexResult<()> {
789        let chunk1 = PrimitiveArray::from_option_iter(vec![Some(10.5f64), Some(20.3)]);
790        let chunk2 = ConstantArray::new(Scalar::primitive(0f64, Nullable), 0);
791        let chunk3 = PrimitiveArray::from_option_iter(vec![Some(5.2f64)]);
792        let dtype = chunk1.dtype().clone();
793        let chunked = ChunkedArray::try_new(
794            vec![
795                chunk1.into_array(),
796                chunk2.into_array(),
797                chunk3.into_array(),
798            ],
799            dtype,
800        )?;
801
802        let result = sum(
803            &chunked.into_array(),
804            &mut array_session().create_execution_ctx(),
805        )?;
806        assert_eq!(result.as_primitive().as_::<f64>(), Some(36.0));
807        Ok(())
808    }
809
810    #[test]
811    fn sum_chunked_int_almost_all_null() -> VortexResult<()> {
812        let chunk1 = PrimitiveArray::from_option_iter::<u32, _>(vec![Some(1)]);
813        let chunk2 = PrimitiveArray::from_option_iter::<u32, _>(vec![None]);
814        let dtype = chunk1.dtype().clone();
815        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
816
817        let result = sum(
818            &chunked.into_array(),
819            &mut array_session().create_execution_ctx(),
820        )?;
821        assert_eq!(result.as_primitive().as_::<u64>(), Some(1));
822        Ok(())
823    }
824
825    #[test]
826    fn sum_chunked_decimals() -> VortexResult<()> {
827        let decimal_dtype = DecimalDType::new(10, 2);
828        let chunk1 = DecimalArray::new(
829            buffer![100i32, 100i32, 100i32, 100i32, 100i32],
830            decimal_dtype,
831            Validity::AllValid,
832        );
833        let chunk2 = DecimalArray::new(
834            buffer![200i32, 200i32, 200i32],
835            decimal_dtype,
836            Validity::AllValid,
837        );
838        let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid);
839        let dtype = chunk1.dtype().clone();
840        let chunked = ChunkedArray::try_new(
841            vec![
842                chunk1.into_array(),
843                chunk2.into_array(),
844                chunk3.into_array(),
845            ],
846            dtype,
847        )?;
848
849        let result = sum(
850            &chunked.into_array(),
851            &mut array_session().create_execution_ctx(),
852        )?;
853        let decimal_result = result.as_decimal();
854        assert_eq!(
855            decimal_result.decimal_value(),
856            Some(DecimalValue::I256(i256::from_i128(1700)))
857        );
858        Ok(())
859    }
860
861    #[test]
862    fn sum_chunked_decimals_with_nulls() -> VortexResult<()> {
863        let decimal_dtype = DecimalDType::new(10, 2);
864        let chunk1 = DecimalArray::new(
865            buffer![100i32, 100i32, 100i32],
866            decimal_dtype,
867            Validity::AllValid,
868        );
869        let chunk2 = DecimalArray::new(
870            buffer![0i32, 0i32],
871            decimal_dtype,
872            Validity::from_iter([false, false]),
873        );
874        let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid);
875        let dtype = chunk1.dtype().clone();
876        let chunked = ChunkedArray::try_new(
877            vec![
878                chunk1.into_array(),
879                chunk2.into_array(),
880                chunk3.into_array(),
881            ],
882            dtype,
883        )?;
884
885        let result = sum(
886            &chunked.into_array(),
887            &mut array_session().create_execution_ctx(),
888        )?;
889        let decimal_result = result.as_decimal();
890        assert_eq!(
891            decimal_result.decimal_value(),
892            Some(DecimalValue::I256(i256::from_i128(700)))
893        );
894        Ok(())
895    }
896
897    #[test]
898    fn sum_chunked_decimals_large() -> VortexResult<()> {
899        let decimal_dtype = DecimalDType::new(3, 0);
900        let chunk1 = ConstantArray::new(
901            Scalar::decimal(
902                DecimalValue::I16(500),
903                decimal_dtype,
904                Nullability::NonNullable,
905            ),
906            1,
907        );
908        let chunk2 = ConstantArray::new(
909            Scalar::decimal(
910                DecimalValue::I16(600),
911                decimal_dtype,
912                Nullability::NonNullable,
913            ),
914            1,
915        );
916        let dtype = chunk1.dtype().clone();
917        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
918
919        let result = sum(
920            &chunked.into_array(),
921            &mut array_session().create_execution_ctx(),
922        )?;
923        let decimal_result = result.as_decimal();
924        assert_eq!(
925            decimal_result.decimal_value(),
926            Some(DecimalValue::I256(i256::from_i128(1100)))
927        );
928        assert_eq!(
929            result.dtype(),
930            &DType::Decimal(DecimalDType::new(13, 0), Nullable)
931        );
932        Ok(())
933    }
934}
vortex_array/aggregate_fn/fns/sum/mod.rs

vortex_array/aggregate_fn/fns/sum/
mod.rs
