vortex_array/arrays/primitive/compute/
sum.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_buffer::BooleanBuffer;
5use itertools::Itertools;
6use num_traits::{CheckedAdd, Float, ToPrimitive};
7use vortex_dtype::{NativePType, match_each_native_ptype};
8use vortex_error::{VortexExpect, VortexResult};
9use vortex_mask::AllOr;
10use vortex_scalar::Scalar;
11
12use crate::arrays::{PrimitiveArray, PrimitiveVTable};
13use crate::compute::{SumKernel, SumKernelAdapter};
14use crate::register_kernel;
15use crate::stats::Stat;
16
17impl SumKernel for PrimitiveVTable {
18    fn sum(&self, array: &PrimitiveArray) -> VortexResult<Scalar> {
19        Ok(match array.validity_mask()?.boolean_buffer() {
20            AllOr::All => {
21                // All-valid
22                match_each_native_ptype!(
23                    array.ptype(),
24                    unsigned: |T| { sum_integer::<_, u64>(array.as_slice::<T>()).into() },
25                    signed: |T| { sum_integer::<_, i64>(array.as_slice::<T>()).into() },
26                    floating: |T| { Some(sum_float(array.as_slice::<T>())).into() }
27                )
28            }
29            AllOr::None => {
30                // All-invalid
31                return Ok(Scalar::null(
32                    Stat::Sum
33                        .dtype(array.dtype())
34                        .vortex_expect("Sum dtype must be defined for primitive type"),
35                ));
36            }
37            AllOr::Some(validity_mask) => {
38                // Some-valid
39                match_each_native_ptype!(
40                    array.ptype(),
41                    unsigned: |T| {
42                        sum_integer_with_validity::<_, u64>(array.as_slice::<T>(), validity_mask).into()
43                    },
44                    signed: |T| {
45                        sum_integer_with_validity::<_, i64>(array.as_slice::<T>(), validity_mask).into()
46                    },
47                    floating: |T| {
48                        Some(sum_float_with_validity(array.as_slice::<T>(), validity_mask)).into()
49                    }
50                )
51            }
52        })
53    }
54}
55
56register_kernel!(SumKernelAdapter(PrimitiveVTable).lift());
57
58fn sum_integer<T: NativePType + ToPrimitive, R: NativePType + CheckedAdd>(
59    values: &[T],
60) -> Option<R> {
61    let mut sum = R::zero();
62    for &x in values {
63        sum = sum.checked_add(&R::from(x)?)?;
64    }
65    Some(sum)
66}
67
68fn sum_integer_with_validity<T: NativePType + ToPrimitive, R: NativePType + CheckedAdd>(
69    values: &[T],
70    validity: &BooleanBuffer,
71) -> Option<R> {
72    let mut sum = R::zero();
73    for (&x, valid) in values.iter().zip_eq(validity.iter()) {
74        if valid {
75            sum = sum.checked_add(&R::from(x)?)?;
76        }
77    }
78    Some(sum)
79}
80
81fn sum_float<T: NativePType + Float>(values: &[T]) -> f64 {
82    let mut sum = 0.0;
83    for &x in values {
84        sum += x.to_f64().vortex_expect("Failed to cast value to f64");
85    }
86    sum
87}
88
89fn sum_float_with_validity<T: NativePType + Float>(array: &[T], validity: &BooleanBuffer) -> f64 {
90    let mut sum = 0.0;
91    for (&x, valid) in array.iter().zip_eq(validity.iter()) {
92        if valid {
93            sum += x.to_f64().vortex_expect("Failed to cast value to f64");
94        }
95    }
96    sum
97}
vortex_array/arrays/primitive/compute/sum.rs

vortex_array/arrays/primitive/compute/
sum.rs
