vortex_array/arrays/chunked/compute/
sum.rs

Help
1use num_traits::PrimInt;
2use vortex_dtype::{NativePType, PType, match_each_native_ptype};
3use vortex_error::{VortexExpect, VortexResult, vortex_err};
4use vortex_scalar::{FromPrimitiveOrF16, Scalar};
5
6use crate::arrays::{ChunkedArray, ChunkedEncoding};
7use crate::compute::{SumFn, sum};
8use crate::stats::Stat;
9use crate::{Array, ArrayRef};
10
11impl SumFn<&ChunkedArray> for ChunkedEncoding {
12    fn sum(&self, array: &ChunkedArray) -> VortexResult<Scalar> {
13        let sum_dtype = Stat::Sum
14            .dtype(array.dtype())
15            .ok_or_else(|| vortex_err!("Sum not supported for dtype {}", array.dtype()))?;
16        let sum_ptype = PType::try_from(&sum_dtype).vortex_expect("sum dtype must be primitive");
17
18        let scalar_value = match_each_native_ptype!(
19            sum_ptype,
20            unsigned: |$T| { sum_int::<u64>(array.chunks())?.into() }
21            signed: |$T| { sum_int::<i64>(array.chunks())?.into() }
22            floating: |$T| { sum_float(array.chunks())?.into() }
23        );
24
25        Ok(Scalar::new(sum_dtype, scalar_value))
26    }
27}
28
29fn sum_int<T: NativePType + PrimInt + FromPrimitiveOrF16>(
30    chunks: &[ArrayRef],
31) -> VortexResult<Option<T>> {
32    let mut result = T::zero();
33    for chunk in chunks {
34        let chunk_sum = sum(chunk)?;
35
36        let Some(chunk_sum) = chunk_sum.as_primitive().as_::<T>()? else {
37            // Bail out on overflow
38            return Ok(None);
39        };
40
41        let Some(chunk_result) = result.checked_add(&chunk_sum) else {
42            // Bail out on overflow
43            return Ok(None);
44        };
45
46        result = chunk_result;
47    }
48    Ok(Some(result))
49}
50
51fn sum_float(chunks: &[ArrayRef]) -> VortexResult<f64> {
52    let mut result = 0f64;
53    for chunk in chunks {
54        let chunk_sum = sum(chunk)?;
55        let chunk_sum = chunk_sum
56            .as_primitive()
57            .as_::<f64>()?
58            .vortex_expect("Float sum should never be null");
59        result += chunk_sum;
60    }
61    Ok(result)
62}
vortex_array/arrays/chunked/compute/sum.rs

vortex_array/arrays/chunked/compute/
sum.rs
