vortex_array/scalar_fn/fns/
list_sum.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7use vortex_error::vortex_err;
8use vortex_mask::AllOr;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayRef;
13use crate::Canonical;
14use crate::Columnar;
15use crate::ExecutionCtx;
16use crate::IntoArray;
17use crate::aggregate_fn::AggregateFnVTable;
18use crate::aggregate_fn::DynGroupedAccumulator;
19use crate::aggregate_fn::GroupRanges;
20use crate::aggregate_fn::GroupedAccumulator;
21use crate::aggregate_fn::GroupedArray;
22use crate::aggregate_fn::NumericalAggregateOpts;
23use crate::aggregate_fn::fns::sum::Sum;
24use crate::arrays::BoolArray;
25use crate::arrays::ConstantArray;
26use crate::arrays::FixedSizeList;
27use crate::arrays::ListView;
28use crate::builtins::ArrayBuiltins;
29use crate::dtype::DType;
30use crate::scalar_fn::Arity;
31use crate::scalar_fn::ChildName;
32use crate::scalar_fn::ExecutionArgs;
33use crate::scalar_fn::ScalarFnId;
34use crate::scalar_fn::ScalarFnVTable;
35use crate::validity::Validity;
36
37/// Sum of the elements in each list of a `List` or `FixedSizeList` typed array.
38///
39/// Follows SQL `SUM` semantics per list, matching DuckDB's `list_sum`: null lists, empty
40/// lists, and lists whose elements are all null yield a null sum; null elements are skipped.
41/// Integer and decimal overflow yields a null sum value, matching [`Sum`]. The result dtype
42/// follows [`Sum`]'s widening rules and is always nullable.
43///
44/// NaN handling for float elements is controlled by [`NumericalAggregateOpts`]: with
45/// `skip_nans` (the default) NaN values contribute nothing, otherwise any NaN poisons the
46/// list's sum to NaN.
47#[derive(Clone)]
48pub struct ListSum;
49
50impl ScalarFnVTable for ListSum {
51    type Options = NumericalAggregateOpts;
52
53    fn id(&self) -> ScalarFnId {
54        static ID: CachedId = CachedId::new("vortex.list.sum");
55        *ID
56    }
57
58    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
59        Ok(Some(options.serialize()))
60    }
61
62    fn deserialize(
63        &self,
64        metadata: &[u8],
65        _session: &VortexSession,
66    ) -> VortexResult<Self::Options> {
67        NumericalAggregateOpts::deserialize(metadata)
68    }
69
70    fn arity(&self, _options: &Self::Options) -> Arity {
71        Arity::Exact(1)
72    }
73
74    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
75        match child_idx {
76            0 => ChildName::from("input"),
77            _ => unreachable!("Invalid child index {child_idx} for list_sum()"),
78        }
79    }
80
81    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
82        let elem_dtype = match &arg_dtypes[0] {
83            DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref(),
84            other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"),
85        };
86        Sum.return_dtype(options, elem_dtype)
87            .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}"))
88    }
89
90    fn execute(
91        &self,
92        options: &Self::Options,
93        args: &dyn ExecutionArgs,
94        ctx: &mut ExecutionCtx,
95    ) -> VortexResult<ArrayRef> {
96        let input = args.get(0)?;
97
98        let elem_dtype = match input.dtype() {
99            DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref().clone(),
100            other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"),
101        };
102
103        // `mask_empty_lists` needs access to list elements validity and sizes
104        let columnar = input.execute::<Columnar>(ctx)?;
105
106        match columnar {
107            Columnar::Constant(constant) => {
108                // Canonicalize one row of the constant and broadcast its sum.
109                let one_row = ConstantArray::new(constant.scalar().clone(), 1)
110                    .into_array()
111                    .execute::<Canonical>(ctx)?
112                    .into_array();
113                let sum =
114                    list_sum_impl(one_row, elem_dtype, options, ctx)?.execute_scalar(0, ctx)?;
115                Ok(ConstantArray::new(sum, constant.len()).into_array())
116            }
117            Columnar::Canonical(canonical) => {
118                list_sum_impl(canonical.into_array(), elem_dtype, options, ctx)
119            }
120        }
121    }
122
123    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
124        false
125    }
126
127    fn is_fallible(&self, _options: &Self::Options) -> bool {
128        false
129    }
130}
131
132/// Sum each list of a canonical `array` into one value per list.
133///
134/// Note that we need to nullify sums produced by empty or all-null lists,
135/// since grouped sum kernels default to 0 for these.
136fn list_sum_impl(
137    canonical: ArrayRef,
138    elem_dtype: DType,
139    options: &NumericalAggregateOpts,
140    ctx: &mut ExecutionCtx,
141) -> VortexResult<ArrayRef> {
142    let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?;
143    acc.accumulate_list(&canonical, ctx)?;
144    let sums = acc.finish()?;
145
146    let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::<FixedSizeList>() {
147        fsl.into_owned().into()
148    } else if let Some(lv) = canonical.as_opt::<ListView>() {
149        lv.into_owned().into()
150    } else {
151        let dtype = canonical.dtype();
152        vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}")
153    };
154
155    mask_empty_lists(grouped, sums, ctx)
156}
157
158/// Applies a mask to `sums` that nullifies entries produced by lists without at least
159/// one valid element. This is necessary because the grouped `Sum` aggregate only produces
160/// nulls for null lists and sums that overflow.
161fn mask_empty_lists(
162    grouped: GroupedArray,
163    sums: ArrayRef,
164    ctx: &mut ExecutionCtx,
165) -> VortexResult<ArrayRef> {
166    let elements = grouped.elements();
167    let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?;
168    let ranges = grouped.group_ranges(ctx)?;
169
170    let has_valid_element: BitBuffer = match elem_mask.bit_buffer() {
171        AllOr::All => match &ranges {
172            // fixed-size lists of non-zero width cannot have empty lists.
173            GroupRanges::FixedSizeList { size, .. } if *size > 0 => return Ok(sums),
174            GroupRanges::FixedSizeList { len, .. } => BitBuffer::full(false, *len),
175            GroupRanges::ListView { ranges } => ranges.iter().map(|&(_, size)| size > 0).collect(),
176        },
177        AllOr::None => BitBuffer::full(false, ranges.len()),
178        AllOr::Some(bits) => ranges
179            .iter()
180            .map(|(offset, size)| size > 0 && bits.count_range(offset, offset + size) > 0)
181            .collect(),
182    };
183    if has_valid_element.true_count() == has_valid_element.len() {
184        return Ok(sums);
185    }
186
187    sums.mask(BoolArray::new(has_valid_element, Validity::NonNullable).into_array())
188}
189
190#[cfg(test)]
191mod tests {
192    use std::sync::Arc;
193
194    use prost::Message;
195    use rstest::rstest;
196    use vortex_buffer::buffer;
197    use vortex_error::VortexResult;
198    use vortex_proto::expr as pb;
199
200    use crate::ArrayRef;
201    use crate::IntoArray;
202    use crate::VortexSessionExecute;
203    use crate::aggregate_fn::NumericalAggregateOpts;
204    use crate::array_session;
205    use crate::arrays::BoolArray;
206    use crate::arrays::ConstantArray;
207    use crate::arrays::FixedSizeListArray;
208    use crate::arrays::ListArray;
209    use crate::arrays::ListViewArray;
210    use crate::arrays::PrimitiveArray;
211    use crate::assert_arrays_eq;
212    use crate::dtype::DType;
213    use crate::dtype::Nullability;
214    use crate::dtype::PType;
215    use crate::expr::Expression;
216    use crate::expr::list_sum;
217    use crate::expr::list_sum_opts;
218    use crate::expr::proto::ExprSerializeProtoExt;
219    use crate::expr::root;
220    use crate::scalar::Scalar;
221    use crate::scalar_fn::ScalarFnVTable;
222    use crate::scalar_fn::fns::list_sum::ListSum;
223    use crate::validity::Validity;
224
225    fn create_list_elements() -> ArrayRef {
226        PrimitiveArray::from_option_iter::<i32, _>([
227            Some(1),
228            Some(2),
229            Some(3),
230            Some(4),
231            Some(5),
232            Some(6),
233            None,
234        ])
235        .into_array()
236    }
237
238    #[rstest]
239    #[case(buffer![0u32, 2, 5, 5, 7].into_array())]
240    #[case(buffer![0u64, 2, 5, 5, 7].into_array())]
241    fn test_list_sum(#[case] offsets: ArrayRef) -> VortexResult<()> {
242        let elements = create_list_elements();
243        let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array();
244        let result = list.apply(&list_sum(root()))?;
245        let mut ctx = array_session().create_execution_ctx();
246        // [1, 2] = 3; [3, 4, 5] = 12; [] = null; [6, null] = 6 (nulls skipped).
247        let expected =
248            PrimitiveArray::from_option_iter::<i64, _>([Some(3), Some(12), None, Some(6)]);
249        assert_arrays_eq!(result, expected, &mut ctx);
250        Ok(())
251    }
252
253    #[test]
254    fn test_nullable_list_sum() -> VortexResult<()> {
255        let elements = create_list_elements();
256        let list = ListArray::try_new(
257            elements,
258            buffer![0u32, 2, 5, 5, 7].into_array(),
259            Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()),
260        )?
261        .into_array();
262        let result = list.apply(&list_sum(root()))?;
263
264        let mut ctx = array_session().create_execution_ctx();
265        // Row 1 and 3 are null lists; row 2 is a valid but empty list.
266        let expected = PrimitiveArray::from_option_iter::<i64, _>([Some(3), None, None, None]);
267        assert_arrays_eq!(result, expected, &mut ctx);
268        Ok(())
269    }
270
271    #[test]
272    fn test_all_null_elements_sum_to_null() -> VortexResult<()> {
273        let elements = PrimitiveArray::from_option_iter::<i32, _>([None, None, Some(1)]);
274        let list = ListArray::try_new(
275            elements.into_array(),
276            buffer![0u32, 2, 3].into_array(),
277            Validity::NonNullable,
278        )?
279        .into_array();
280        let result = list.apply(&list_sum(root()))?;
281
282        let mut ctx = array_session().create_execution_ctx();
283        let expected = PrimitiveArray::from_option_iter::<i64, _>([None, Some(1)]);
284        assert_arrays_eq!(result, expected, &mut ctx);
285        Ok(())
286    }
287
288    #[test]
289    fn test_overflow_sums_to_null() -> VortexResult<()> {
290        let elements = PrimitiveArray::from_iter([i64::MAX, 1, 1, 2]);
291        let list = ListArray::try_new(
292            elements.into_array(),
293            buffer![0u32, 2, 4].into_array(),
294            Validity::NonNullable,
295        )?
296        .into_array();
297        let result = list.apply(&list_sum(root()))?;
298
299        let mut ctx = array_session().create_execution_ctx();
300        let expected = PrimitiveArray::from_option_iter::<i64, _>([None, Some(3)]);
301        assert_arrays_eq!(result, expected, &mut ctx);
302        Ok(())
303    }
304
305    #[test]
306    fn test_unsigned_widening() -> VortexResult<()> {
307        let elements = PrimitiveArray::from_iter([1u8, 2, 3]);
308        let list = ListArray::try_new(
309            elements.into_array(),
310            buffer![0u32, 3].into_array(),
311            Validity::NonNullable,
312        )?
313        .into_array();
314        let result = list.apply(&list_sum(root()))?;
315
316        let mut ctx = array_session().create_execution_ctx();
317        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(6)]);
318        assert_arrays_eq!(result, expected, &mut ctx);
319        Ok(())
320    }
321
322    #[test]
323    fn test_bool_elements() -> VortexResult<()> {
324        let elements = BoolArray::from_iter([true, true, false, true]);
325        let list = ListArray::try_new(
326            elements.into_array(),
327            buffer![0u32, 3, 4].into_array(),
328            Validity::NonNullable,
329        )?
330        .into_array();
331        let result = list.apply(&list_sum(root()))?;
332
333        let mut ctx = array_session().create_execution_ctx();
334        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), Some(1)]);
335        assert_arrays_eq!(result, expected, &mut ctx);
336        Ok(())
337    }
338
339    #[test]
340    fn test_nan_skipped_by_default() -> VortexResult<()> {
341        let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]);
342        let list = ListArray::try_new(
343            elements.into_array(),
344            buffer![0u32, 3].into_array(),
345            Validity::NonNullable,
346        )?
347        .into_array();
348        let result = list.apply(&list_sum(root()))?;
349
350        let mut ctx = array_session().create_execution_ctx();
351        let expected = PrimitiveArray::from_option_iter::<f64, _>([Some(3.0)]);
352        assert_arrays_eq!(result, expected, &mut ctx);
353        Ok(())
354    }
355
356    #[test]
357    fn test_all_nan_list_sums_to_zero() -> VortexResult<()> {
358        // NaN elements are *valid*: with the default `skip_nans` they contribute nothing, but
359        // the list still has valid elements, so the sum is `0.0` rather than null (unlike an
360        // all-null list).
361        let elements = PrimitiveArray::from_iter([f64::NAN, f64::NAN]);
362        let list = ListArray::try_new(
363            elements.into_array(),
364            buffer![0u32, 2].into_array(),
365            Validity::NonNullable,
366        )?
367        .into_array();
368        let result = list.apply(&list_sum(root()))?;
369
370        let mut ctx = array_session().create_execution_ctx();
371        let expected = PrimitiveArray::from_option_iter::<f64, _>([Some(0.0)]);
372        assert_arrays_eq!(result, expected, &mut ctx);
373        Ok(())
374    }
375
376    #[test]
377    fn test_nan_poisons_with_include_nans() -> VortexResult<()> {
378        let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]);
379        let list = ListArray::try_new(
380            elements.into_array(),
381            buffer![0u32, 3].into_array(),
382            Validity::NonNullable,
383        )?
384        .into_array();
385        let result = list.apply(&list_sum_opts(
386            root(),
387            NumericalAggregateOpts::include_nans(),
388        ))?;
389
390        let mut ctx = array_session().create_execution_ctx();
391        assert!(result.is_valid(0, &mut ctx)?);
392        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
393        assert!(result.as_slice::<f64>()[0].is_nan());
394        Ok(())
395    }
396
397    #[test]
398    fn test_listview_sum() -> VortexResult<()> {
399        let elements = create_list_elements();
400        // Overlapping and out-of-order views over the shared elements.
401        let lv = ListViewArray::new(
402            elements,
403            buffer![5u32, 0, 4, 1].into_array(),
404            buffer![2u32, 3, 0, 2].into_array(),
405            Validity::NonNullable,
406        )
407        .into_array();
408        let result = lv.apply(&list_sum(root()))?;
409
410        let mut ctx = array_session().create_execution_ctx();
411        // [6, null] = 6; [1, 2, 3] = 6; [] = null; [2, 3] = 5.
412        let expected =
413            PrimitiveArray::from_option_iter::<i64, _>([Some(6), Some(6), None, Some(5)]);
414        assert_arrays_eq!(result, expected, &mut ctx);
415        Ok(())
416    }
417
418    fn create_fixed_size_list(validity: Validity) -> ArrayRef {
419        // 4 lists of size 2 over 8 primitive elements.
420        let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array();
421        FixedSizeListArray::new(elements, 2, validity, 4).into_array()
422    }
423
424    #[test]
425    fn test_fixed_size_list_sum() -> VortexResult<()> {
426        let fsl = create_fixed_size_list(Validity::NonNullable);
427        let result = fsl.apply(&list_sum(root()))?;
428
429        let mut ctx = array_session().create_execution_ctx();
430        let expected =
431            PrimitiveArray::from_option_iter::<i64, _>([Some(3), Some(7), Some(11), Some(15)]);
432        assert_arrays_eq!(result, expected, &mut ctx);
433        Ok(())
434    }
435
436    #[test]
437    fn test_fixed_size_list_sum_nullable() -> VortexResult<()> {
438        let fsl = create_fixed_size_list(Validity::Array(
439            BoolArray::from_iter([true, false, true, false]).into_array(),
440        ));
441        let result = fsl.apply(&list_sum(root()))?;
442
443        let mut ctx = array_session().create_execution_ctx();
444        let expected = PrimitiveArray::from_option_iter::<i64, _>([Some(3), None, Some(11), None]);
445        assert_arrays_eq!(result, expected, &mut ctx);
446        Ok(())
447    }
448
449    #[test]
450    fn test_list_sum_take() -> VortexResult<()> {
451        let elements = create_list_elements();
452        let list = ListArray::try_new(
453            elements,
454            buffer![0u32, 2, 5, 5, 7].into_array(),
455            Validity::NonNullable,
456        )?
457        .into_array();
458        let taken = list.take(buffer![3u64, 0, 2].into_array())?;
459
460        let result = taken.apply(&list_sum(root()))?;
461        let mut ctx = array_session().create_execution_ctx();
462        let expected = PrimitiveArray::from_option_iter::<i64, _>([Some(6), Some(3), None]);
463        assert_arrays_eq!(result, expected, &mut ctx);
464        Ok(())
465    }
466
467    #[test]
468    fn test_list_sum_slice() -> VortexResult<()> {
469        let elements = create_list_elements();
470        let list = ListArray::try_new(
471            elements,
472            buffer![0u32, 2, 5, 5, 7].into_array(),
473            Validity::NonNullable,
474        )?
475        .into_array();
476        let sliced = list.slice(1..4)?;
477
478        let result = sliced.apply(&list_sum(root()))?;
479        let mut ctx = array_session().create_execution_ctx();
480        let expected = PrimitiveArray::from_option_iter::<i64, _>([Some(12), None, Some(6)]);
481        assert_arrays_eq!(result, expected, &mut ctx);
482        Ok(())
483    }
484
485    #[test]
486    fn test_empty_array() -> VortexResult<()> {
487        let elements = PrimitiveArray::from_iter([0i32; 0]);
488        let list = ListArray::try_new(
489            elements.into_array(),
490            buffer![0u32].into_array(),
491            Validity::NonNullable,
492        )?
493        .into_array();
494        let result = list.apply(&list_sum(root()))?;
495
496        let mut ctx = array_session().create_execution_ctx();
497        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
498        assert_eq!(result.len(), 0);
499        Ok(())
500    }
501
502    #[test]
503    fn test_constant_list_sum() -> VortexResult<()> {
504        let elements = create_list_elements();
505        let list = ListArray::try_new(
506            elements,
507            buffer![0u32, 2, 5, 5, 7].into_array(),
508            Validity::NonNullable,
509        )?
510        .into_array();
511        let mut ctx = array_session().create_execution_ctx();
512        let scalar = list.execute_scalar(1, &mut ctx)?;
513
514        let constant = ConstantArray::new(scalar, 3).into_array();
515        let result = constant.apply(&list_sum(root()))?;
516        let expected = PrimitiveArray::from_option_iter::<i64, _>([Some(12), Some(12), Some(12)]);
517        assert_arrays_eq!(result, expected, &mut ctx);
518        Ok(())
519    }
520
521    #[test]
522    fn test_null_scalar_list_sum() -> VortexResult<()> {
523        let null_scalar = Scalar::null(DType::List(
524            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
525            Nullability::Nullable,
526        ));
527        let array = ConstantArray::new(null_scalar, 2).into_array();
528        let result = array.apply(&list_sum(root()))?;
529
530        let mut ctx = array_session().create_execution_ctx();
531        assert!(!result.is_valid(0, &mut ctx)?);
532        assert!(!result.is_valid(1, &mut ctx)?);
533        Ok(())
534    }
535
536    #[test]
537    fn test_unsupported_dtypes() {
538        let opts = NumericalAggregateOpts::default();
539        // Non-list inputs are rejected.
540        assert!(
541            ListSum
542                .return_dtype(
543                    &opts,
544                    &[DType::Primitive(PType::I32, Nullability::NonNullable)]
545                )
546                .is_err()
547        );
548        // Non-numeric element types are rejected.
549        assert!(
550            ListSum
551                .return_dtype(
552                    &opts,
553                    &[DType::List(
554                        Arc::new(DType::Utf8(Nullability::NonNullable)),
555                        Nullability::NonNullable,
556                    )],
557                )
558                .is_err()
559        );
560    }
561
562    #[test]
563    fn test_display() {
564        assert_eq!(list_sum(root()).to_string(), "vortex.list.sum($)");
565        assert_eq!(
566            list_sum_opts(root(), NumericalAggregateOpts::include_nans()).to_string(),
567            "vortex.list.sum($, opts=skip_nans=false)"
568        );
569    }
570
571    #[test]
572    fn test_proto_round_trip() -> VortexResult<()> {
573        for expr in [
574            list_sum(root()),
575            list_sum_opts(root(), NumericalAggregateOpts::include_nans()),
576        ] {
577            let proto = expr.serialize_proto()?;
578            let buf = proto.encode_to_vec();
579            let decoded = pb::Expr::decode(buf.as_slice())?;
580            let deser = Expression::from_proto(&decoded, &array_session())?;
581            assert_eq!(expr, deser);
582        }
583        Ok(())
584    }
585}
vortex_array/scalar_fn/fns/list_sum.rs

vortex_array/scalar_fn/fns/
list_sum.rs
