vortex_array/compute/
sum.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::LazyLock;
5
6use arcref::ArcRef;
7use num_traits::CheckedAdd;
8use num_traits::CheckedSub;
9use vortex_dtype::DType;
10use vortex_error::VortexError;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_ensure;
14use vortex_error::vortex_err;
15use vortex_error::vortex_panic;
16
17use crate::Array;
18use crate::compute::ComputeFn;
19use crate::compute::ComputeFnVTable;
20use crate::compute::InvocationArgs;
21use crate::compute::Kernel;
22use crate::compute::Output;
23use crate::expr::stats::Precision;
24use crate::expr::stats::Stat;
25use crate::expr::stats::StatsProvider;
26use crate::scalar::NumericOperator;
27use crate::scalar::Scalar;
28use crate::vtable::VTable;
29
30static SUM_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
31    let compute = ComputeFn::new("sum".into(), ArcRef::new_ref(&Sum));
32    for kernel in inventory::iter::<SumKernelRef> {
33        compute.register_kernel(kernel.0.clone());
34    }
35    compute
36});
37
38pub(crate) fn warm_up_vtable() -> usize {
39    SUM_FN.kernels().len()
40}
41
42/// Sum an array with an initial value.
43///
44/// If the sum overflows, a null scalar will be returned.
45/// If the sum is not supported for the array's dtype, an error will be raised.
46/// If the array is all-invalid, the sum will be the accumulator.
47/// The accumulator must have a dtype compatible with the sum result dtype.
48pub(crate) fn sum_with_accumulator(
49    array: &dyn Array,
50    accumulator: &Scalar,
51) -> VortexResult<Scalar> {
52    SUM_FN
53        .invoke(&InvocationArgs {
54            inputs: &[array.into(), accumulator.into()],
55            options: &(),
56        })?
57        .unwrap_scalar()
58}
59
60/// Sum an array, starting from zero.
61///
62/// If the sum overflows, a null scalar will be returned.
63/// If the sum is not supported for the array's dtype, an error will be raised.
64/// If the array is all-invalid, the sum will be zero.
65pub fn sum(array: &dyn Array) -> VortexResult<Scalar> {
66    let sum_dtype = Stat::Sum
67        .dtype(array.dtype())
68        .ok_or_else(|| vortex_err!("Sum not supported for dtype: {}", array.dtype()))?;
69    let zero = Scalar::zero_value(&sum_dtype);
70    sum_with_accumulator(array, &zero)
71}
72
73/// For unary compute functions, it's useful to just have this short-cut.
74pub struct SumArgs<'a> {
75    pub array: &'a dyn Array,
76    pub accumulator: &'a Scalar,
77}
78
79impl<'a> TryFrom<&InvocationArgs<'a>> for SumArgs<'a> {
80    type Error = VortexError;
81
82    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
83        if value.inputs.len() != 2 {
84            vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());
85        }
86        let array = value.inputs[0]
87            .array()
88            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
89        let accumulator = value.inputs[1]
90            .scalar()
91            .ok_or_else(|| vortex_err!("Expected input 1 to be a scalar"))?;
92        Ok(SumArgs { array, accumulator })
93    }
94}
95
96struct Sum;
97
98impl ComputeFnVTable for Sum {
99    fn invoke(
100        &self,
101        args: &InvocationArgs,
102        kernels: &[ArcRef<dyn Kernel>],
103    ) -> VortexResult<Output> {
104        let SumArgs { array, accumulator } = args.try_into()?;
105
106        // Compute the expected dtype of the sum.
107        let sum_dtype = self.return_dtype(args)?;
108
109        vortex_ensure!(
110            &sum_dtype == accumulator.dtype(),
111            "sum_dtype {sum_dtype} must match accumulator dtype {}",
112            accumulator.dtype()
113        );
114
115        // Short-circuit using array statistics.
116        if let Some(Precision::Exact(sum_scalar)) = array.statistics().get(Stat::Sum) {
117            // For floats only use stats if accumulator is zero. otherwise we might have numerical
118            // stability issues.
119            match &sum_dtype {
120                DType::Primitive(p, _) => {
121                    if p.is_float() && accumulator.is_zero() == Some(true) {
122                        return Ok(sum_scalar.into());
123                    } else if p.is_int() {
124                        let sum_from_stat = accumulator
125                            .as_primitive()
126                            .checked_add(&sum_scalar.as_primitive())
127                            .map(Scalar::from);
128                        return Ok(sum_from_stat
129                            .unwrap_or_else(|| Scalar::null(sum_dtype))
130                            .into());
131                    }
132                }
133                DType::Decimal(..) => {
134                    let sum_from_stat = accumulator
135                        .as_decimal()
136                        .checked_binary_numeric(&sum_scalar.as_decimal(), NumericOperator::Add)
137                        .map(Scalar::from);
138                    return Ok(sum_from_stat
139                        .unwrap_or_else(|| Scalar::null(sum_dtype))
140                        .into());
141                }
142                _ => unreachable!("Sum will always be a decimal or a primitive dtype"),
143            }
144        }
145
146        let sum_scalar = sum_impl(array, accumulator, kernels)?;
147
148        // Update the statistics with the computed sum. Stored statistic shouldn't include the accumulator.
149        match sum_dtype {
150            DType::Primitive(p, _) => {
151                if p.is_float()
152                    && accumulator.is_zero() == Some(true)
153                    && let Some(sum_value) = sum_scalar.value().cloned()
154                {
155                    array
156                        .statistics()
157                        .set(Stat::Sum, Precision::Exact(sum_value));
158                } else if p.is_int()
159                    && let Some(less_accumulator) = sum_scalar
160                        .as_primitive()
161                        .checked_sub(&accumulator.as_primitive())
162                    && let Some(val) = Scalar::from(less_accumulator).into_value()
163                {
164                    array.statistics().set(Stat::Sum, Precision::Exact(val));
165                }
166            }
167            DType::Decimal(..) => {
168                if let Some(less_accumulator) = sum_scalar
169                    .as_decimal()
170                    .checked_binary_numeric(&accumulator.as_decimal(), NumericOperator::Sub)
171                    && let Some(val) = Scalar::from(less_accumulator).into_value()
172                {
173                    array.statistics().set(Stat::Sum, Precision::Exact(val));
174                }
175            }
176            _ => unreachable!("Sum will always be a decimal or a primitive dtype"),
177        }
178
179        Ok(sum_scalar.into())
180    }
181
182    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
183        let SumArgs { array, .. } = args.try_into()?;
184        Stat::Sum
185            .dtype(array.dtype())
186            .ok_or_else(|| vortex_err!("Sum not supported for dtype: {}", array.dtype()))
187    }
188
189    fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {
190        // The sum function always returns a single scalar value.
191        Ok(1)
192    }
193
194    fn is_elementwise(&self) -> bool {
195        false
196    }
197}
198
199pub struct SumKernelRef(ArcRef<dyn Kernel>);
200inventory::collect!(SumKernelRef);
201
202pub trait SumKernel: VTable {
203    /// # Preconditions
204    ///
205    /// * The array's DType is summable
206    /// * The array is not all-null
207    /// * The accumulator must have a dtype compatible with the sum result dtype
208    fn sum(&self, array: &Self::Array, accumulator: &Scalar) -> VortexResult<Scalar>;
209}
210
211#[derive(Debug)]
212pub struct SumKernelAdapter<V: VTable>(pub V);
213
214impl<V: VTable + SumKernel> SumKernelAdapter<V> {
215    pub const fn lift(&'static self) -> SumKernelRef {
216        SumKernelRef(ArcRef::new_ref(self))
217    }
218}
219
220impl<V: VTable + SumKernel> Kernel for SumKernelAdapter<V> {
221    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
222        let SumArgs { array, accumulator } = args.try_into()?;
223        let Some(array) = array.as_opt::<V>() else {
224            return Ok(None);
225        };
226        Ok(Some(V::sum(&self.0, array, accumulator)?.into()))
227    }
228}
229
230/// Sum an array.
231///
232/// If the sum overflows, a null scalar will be returned.
233/// If the sum is not supported for the array's dtype, an error will be raised.
234/// If the array is all-invalid, the sum will be the accumulator.
235pub fn sum_impl(
236    array: &dyn Array,
237    accumulator: &Scalar,
238    kernels: &[ArcRef<dyn Kernel>],
239) -> VortexResult<Scalar> {
240    if array.is_empty() || array.all_invalid()? || accumulator.is_null() {
241        return Ok(accumulator.clone());
242    }
243
244    // Try to find a sum kernel
245    let args = InvocationArgs {
246        inputs: &[array.into(), accumulator.into()],
247        options: &(),
248    };
249    for kernel in kernels {
250        if let Some(output) = kernel.invoke(&args)? {
251            return output.unwrap_scalar();
252        }
253    }
254
255    // Otherwise, canonicalize and try again.
256    tracing::debug!("No sum implementation found for {}", array.encoding_id());
257    if array.is_canonical() {
258        // Panic to avoid recursion, but it should never be hit.
259        vortex_panic!(
260            "No sum implementation found for canonical array: {}",
261            array.encoding_id()
262        );
263    }
264    sum_with_accumulator(array.to_canonical()?.as_ref(), accumulator)
265}
266
267#[cfg(test)]
268mod test {
269    use vortex_buffer::buffer;
270    use vortex_dtype::DType;
271    use vortex_dtype::Nullability;
272    use vortex_dtype::PType;
273    use vortex_error::VortexExpect;
274
275    use crate::IntoArray as _;
276    use crate::arrays::BoolArray;
277    use crate::arrays::ChunkedArray;
278    use crate::arrays::PrimitiveArray;
279    use crate::compute::sum;
280    use crate::compute::sum_with_accumulator;
281    use crate::scalar::Scalar;
282
283    #[test]
284    fn sum_all_invalid() {
285        let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]);
286        let result = sum(array.as_ref()).unwrap();
287        assert_eq!(result, Scalar::primitive(0i64, Nullability::Nullable));
288    }
289
290    #[test]
291    fn sum_all_invalid_float() {
292        let array = PrimitiveArray::from_option_iter::<f32, _>([None, None, None]);
293        let result = sum(array.as_ref()).unwrap();
294        assert_eq!(result, Scalar::primitive(0f64, Nullability::Nullable));
295    }
296
297    #[test]
298    fn sum_constant() {
299        let array = buffer![1, 1, 1, 1].into_array();
300        let result = sum(array.as_ref()).unwrap();
301        assert_eq!(result.as_primitive().as_::<i32>(), Some(4));
302    }
303
304    #[test]
305    fn sum_constant_float() {
306        let array = buffer![1., 1., 1., 1.].into_array();
307        let result = sum(array.as_ref()).unwrap();
308        assert_eq!(result.as_primitive().as_::<f32>(), Some(4.));
309    }
310
311    #[test]
312    fn sum_boolean() {
313        let array = BoolArray::from_iter([true, false, false, true]);
314        let result = sum(array.as_ref()).unwrap();
315        assert_eq!(result.as_primitive().as_::<i32>(), Some(2));
316    }
317
318    #[test]
319    fn sum_stats() {
320        let array = ChunkedArray::try_new(
321            vec![
322                PrimitiveArray::from_iter([1, 1, 1]).into_array(),
323                PrimitiveArray::from_iter([2, 2, 2]).into_array(),
324            ],
325            DType::Primitive(PType::I32, Nullability::NonNullable),
326        )
327        .vortex_expect("operation should succeed in test");
328        // compute sum with accumulator to populate stats
329        sum_with_accumulator(
330            array.as_ref(),
331            &Scalar::primitive(2i64, Nullability::Nullable),
332        )
333        .unwrap();
334
335        let sum_without_acc = sum(array.as_ref()).unwrap();
336        assert_eq!(
337            sum_without_acc,
338            Scalar::primitive(9i64, Nullability::Nullable)
339        );
340    }
341}
vortex_array/compute/sum.rs

vortex_array/compute/
sum.rs
