fix[array]: sum with initial value to fix op assoc (#5278)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/arrays/bool/compute/sum.rs

@@ -3,7 +3,7 @@
 
 use std::ops::BitAnd;
 
-use vortex_error::VortexResult;
+use vortex_error::{VortexExpect, VortexResult};
 use vortex_mask::AllOr;
 use vortex_scalar::Scalar;
 
@@ -12,7 +12,7 @@ use crate::compute::{SumKernel, SumKernelAdapter};
 use crate::register_kernel;
 
 impl SumKernel for BoolVTable {
-    fn sum(&self, array: &BoolArray) -> VortexResult<Scalar> {
+    fn sum(&self, array: &BoolArray, accumulator: &Scalar) -> VortexResult<Scalar> {
         let true_count: Option<u64> = match array.validity_mask().bit_buffer() {
             AllOr::All => {
                 // All-valid
@@ -26,7 +26,14 @@ impl SumKernel for BoolVTable {
                 Some(array.bit_buffer().bitand(validity_mask).true_count() as u64)
             }
         };
-        Ok(Scalar::from(true_count))
+
+        let accumulator = accumulator
+            .as_primitive()
+            .as_::<u64>()
+            .vortex_expect("cannot be null");
+        Ok(Scalar::from(
+            true_count.and_then(|tc| accumulator.checked_add(tc)),
+        ))
     }
 }
 

vortex-array/src/arrays/chunked/compute/sum.rs

@@ -1,100 +1,26 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use num_traits::PrimInt;
-use vortex_dtype::Nullability::Nullable;
-use vortex_dtype::{DType, DecimalDType, NativePType, i256, match_each_native_ptype};
-use vortex_error::{VortexResult, vortex_bail, vortex_err};
-use vortex_scalar::{DecimalScalar, DecimalValue, Scalar};
+use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
 
 use crate::arrays::{ChunkedArray, ChunkedVTable};
-use crate::compute::{SumKernel, SumKernelAdapter, sum};
-use crate::stats::Stat;
-use crate::{ArrayRef, register_kernel};
+use crate::compute::{SumKernel, SumKernelAdapter, sum_with_accumulator};
+use crate::register_kernel;
 
 impl SumKernel for ChunkedVTable {
-    fn sum(&self, array: &ChunkedArray) -> VortexResult<Scalar> {
-        let sum_dtype = Stat::Sum
-            .dtype(array.dtype())
-            .ok_or_else(|| vortex_err!("Sum not supported for dtype {}", array.dtype()))?;
-
-        match sum_dtype {
-            DType::Decimal(decimal_dtype, _) => sum_decimal(array.chunks(), decimal_dtype),
-            DType::Primitive(sum_ptype, _) => {
-                let scalar_value = match_each_native_ptype!(
-                    sum_ptype,
-                    unsigned: |T| { sum_int::<u64>(array.chunks())?.into() },
-                    signed: |T| { sum_int::<i64>(array.chunks())?.into() },
-                    floating: |T| { sum_float(array.chunks())?.into() }
-                );
-
-                Ok(Scalar::new(sum_dtype, scalar_value))
-            }
-            _ => {
-                vortex_bail!("Sum not supported for dtype {}", sum_dtype);
-            }
-        }
+    fn sum(&self, array: &ChunkedArray, accumulator: &Scalar) -> VortexResult<Scalar> {
+        array
+            .chunks
+            .iter()
+            .try_fold(accumulator.clone(), |result, chunk| {
+                sum_with_accumulator(chunk, &result)
+            })
     }
 }
 
 register_kernel!(SumKernelAdapter(ChunkedVTable).lift());
 
-fn sum_int<T: NativePType + PrimInt>(chunks: &[ArrayRef]) -> VortexResult<Option<T>> {
-    let mut result: T = T::zero();
-    for chunk in chunks {
-        let chunk_sum = sum(chunk)?;
-        let Some(chunk_sum) = chunk_sum
-            .as_primitive()
-            .as_::<T>()
-            .and_then(|chunk_sum| result.checked_add(&chunk_sum))
-        else {
-            // Bail out on null or overflow
-            return Ok(None);
-        };
-        result = chunk_sum;
-    }
-    Ok(Some(result))
-}
-
-fn sum_float(chunks: &[ArrayRef]) -> VortexResult<Option<f64>> {
-    let mut result = 0f64;
-    for chunk in chunks {
-        let Some(chunk_sum) = sum(chunk)?.as_primitive().as_::<f64>() else {
-            return Ok(None);
-        };
-        result += chunk_sum;
-    }
-    Ok(Some(result))
-}
-
-fn sum_decimal(chunks: &[ArrayRef], result_decimal_type: DecimalDType) -> VortexResult<Scalar> {
-    let mut result = DecimalValue::I256(i256::ZERO);
-
-    let null = || Scalar::null(DType::Decimal(result_decimal_type, Nullable));
-
-    for chunk in chunks {
-        let chunk_sum = sum(chunk)?;
-
-        let chunk_decimal = DecimalScalar::try_from(&chunk_sum)?;
-        let Some(r) = chunk_decimal
-            .decimal_value()
-            // TODO(joe): added a precision capped checked_add.
-            .and_then(|c_sum| result.checked_add(&c_sum))
-            .filter(|sum_value| {
-                sum_value
-                    .fits_in_precision(result_decimal_type)
-                    .unwrap_or(false)
-            })
-        else {
-            // null if any chunk is null or the sum overflows
-            return Ok(null());
-        };
-        result = r;
-    }
-
-    Ok(Scalar::decimal(result, result_decimal_type, Nullable))
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_buffer::buffer;

vortex-array/src/arrays/constant/compute/sum.rs

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use num_traits::{CheckedMul, ToPrimitive};
+use arrow_array::ArrowNativeTypeOp;
+use num_traits::{CheckedAdd, CheckedMul, ToPrimitive};
 use vortex_dtype::{DType, DecimalDType, NativePType, Nullability, i256, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_scalar::{DecimalScalar, DecimalValue, PrimitiveScalar, Scalar, ScalarValue};
@@ -12,32 +13,44 @@ use crate::register_kernel;
 use crate::stats::Stat;
 
 impl SumKernel for ConstantVTable {
-    fn sum(&self, array: &ConstantArray) -> VortexResult<Scalar> {
+    fn sum(&self, array: &ConstantArray, accumulator: &Scalar) -> VortexResult<Scalar> {
         // Compute the expected dtype of the sum.
         let sum_dtype = Stat::Sum
             .dtype(array.dtype())
             .ok_or_else(|| vortex_err!("Sum not supported for dtype {}", array.dtype()))?;
 
-        let sum_value = sum_scalar(array.scalar(), array.len())?;
+        let sum_value = sum_scalar(array.scalar(), array.len(), accumulator)?;
         Ok(Scalar::new(sum_dtype, sum_value))
     }
 }
 
-fn sum_scalar(scalar: &Scalar, len: usize) -> VortexResult<ScalarValue> {
+fn sum_scalar(scalar: &Scalar, len: usize, accumulator: &Scalar) -> VortexResult<ScalarValue> {
     match scalar.dtype() {
-        DType::Bool(_) => Ok(ScalarValue::from(match scalar.as_bool().value() {
-            None => unreachable!("Handled before reaching this point"),
-            Some(false) => 0u64,
-            Some(true) => len as u64,
-        })),
-        DType::Primitive(ptype, _) => Ok(match_each_native_ptype!(
-            ptype,
-            unsigned: |T| { sum_integral::<u64>(scalar.as_primitive(), len)?.into() },
-            signed: |T| { sum_integral::<i64>(scalar.as_primitive(), len)?.into() },
-            floating: |T| { sum_float(scalar.as_primitive(), len)?.into() }
-        )),
-        DType::Decimal(decimal_dtype, _) => sum_decimal(scalar.as_decimal(), len, *decimal_dtype),
-        DType::Extension(_) => sum_scalar(&scalar.as_extension().storage(), len),
+        DType::Bool(_) => {
+            let count = match scalar.as_bool().value() {
+                None => unreachable!("Handled before reaching this point"),
+                Some(false) => 0u64,
+                Some(true) => len as u64,
+            };
+            let accumulator = accumulator
+                .as_primitive()
+                .as_::<u64>()
+                .vortex_expect("cannot be null");
+            Ok(ScalarValue::from(accumulator.checked_add(count)))
+        }
+        DType::Primitive(ptype, _) => {
+            let result = match_each_native_ptype!(
+                ptype,
+                unsigned: |T| { sum_integral::<u64>(scalar.as_primitive(), len, accumulator)?.into() },
+                signed: |T| { sum_integral::<i64>(scalar.as_primitive(), len, accumulator)?.into() },
+                floating: |T| { sum_float(scalar.as_primitive(), len, accumulator)?.into() }
+            );
+            Ok(result)
+        }
+        DType::Decimal(decimal_dtype, _) => {
+            sum_decimal(scalar.as_decimal(), len, *decimal_dtype, accumulator)
+        }
+        DType::Extension(_) => sum_scalar(&scalar.as_extension().storage(), len, accumulator),
         dtype => vortex_bail!("Unsupported dtype for sum: {}", dtype),
     }
 }
@@ -46,6 +59,7 @@ fn sum_decimal(
     decimal_scalar: DecimalScalar,
     array_len: usize,
     decimal_dtype: DecimalDType,
+    accumulator: &Scalar,
 ) -> VortexResult<ScalarValue> {
     let result_dtype = Stat::Sum
         .dtype(&DType::Decimal(decimal_dtype, Nullability::Nullable))
@@ -63,43 +77,82 @@ fn sum_decimal(
     let len_value = DecimalValue::I256(i256::from_i128(array_len as i128));
 
     // Multiply value * len
-    let sum = value.checked_mul(&len_value).and_then(|result| {
+    let array_sum = value.checked_mul(&len_value).and_then(|result| {
         // Check if result fits in the precision
         result
             .fits_in_precision(*result_decimal_type)
             .unwrap_or(false)
             .then_some(result)
     });
 
-    match sum {
-        Some(result_value) => Ok(ScalarValue::from(result_value)),
+    // Add accumulator to array_sum
+    let initial_decimal = DecimalScalar::try_from(accumulator)?;
+    let initial_dec_value = initial_decimal
+        .decimal_value()
+        .unwrap_or(DecimalValue::I256(i256::ZERO));
+
+    match array_sum {
+        Some(array_sum_value) => {
+            let total = array_sum_value
+                .checked_add(&initial_dec_value)
+                .and_then(|result| {
+                    result
+                        .fits_in_precision(*result_decimal_type)
+                        .unwrap_or(false)
+                        .then_some(result)
+                });
+            match total {
+                Some(result_value) => Ok(ScalarValue::from(result_value)),
+                None => Ok(ScalarValue::null()), // Overflow
+            }
+        }
         None => Ok(ScalarValue::null()), // Overflow
     }
 }
 
 fn sum_integral<T>(
     primitive_scalar: PrimitiveScalar<'_>,
     array_len: usize,
+    accumulator: &Scalar,
 ) -> VortexResult<Option<T>>
 where
-    T: NativePType + CheckedMul,
+    T: NativePType + CheckedMul + CheckedAdd,
     Scalar: From<Option<T>>,
 {
     let v = primitive_scalar.as_::<T>();
     let array_len =
         T::from(array_len).ok_or_else(|| vortex_err!("array_len must fit the sum type"))?;
-    let sum = v.and_then(|v| v.checked_mul(&array_len));
+    let Some(array_sum) = v.and_then(|v| v.checked_mul(&array_len)) else {
+        return Ok(None);
+    };
 
-    Ok(sum)
+    let initial = accumulator
+        .as_primitive()
+        .as_::<T>()
+        .vortex_expect("cannot be null");
+    Ok(initial.checked_add(&array_sum))
 }
 
-fn sum_float(primitive_scalar: PrimitiveScalar<'_>, array_len: usize) -> VortexResult<Option<f64>> {
-    let v = primitive_scalar.as_::<f64>();
+fn sum_float(
+    primitive_scalar: PrimitiveScalar<'_>,
+    array_len: usize,
+    accumulator: &Scalar,
+) -> VortexResult<Option<f64>> {
+    let v = primitive_scalar
+        .as_::<f64>()
+        .vortex_expect("cannot be null");
     let array_len = array_len
         .to_f64()
         .ok_or_else(|| vortex_err!("array_len must fit the sum type"))?;
 
-    Ok(v.map(|v| v * array_len))
+    let Ok(array_sum) = v.mul_checked(array_len) else {
+        return Ok(None);
+    };
+    let initial = accumulator
+        .as_primitive()
+        .as_::<f64>()
+        .vortex_expect("cannot be null");
+    Ok(Some(initial + array_sum))
 }
 
 register_kernel!(SumKernelAdapter(ConstantVTable).lift());