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

vortex_array/aggregate_fn/fns/sum/
mod.rs
