fix[array]: handle empty/all_invalid sum correctly (#5180)

empty/all_invalid arrays have a sum of 0.

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

Author: joseph-isaacs

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

@@ -42,37 +42,31 @@ register_kernel!(SumKernelAdapter(ChunkedVTable).lift());
 fn sum_int<T: NativePType + PrimInt + FromPrimitiveOrF16>(
     chunks: &[ArrayRef],
 ) -> VortexResult<Option<T>> {
-    let mut result: Option<T> = None;
+    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>() else {
-            // Skip missing null chunk
-            continue;
+        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 = Some(match result {
-            None => chunk_sum,
-            Some(result) => {
-                let Some(chunk_result) = result.checked_add(&chunk_sum) else {
-                    // Bail out on overflow
-                    return Ok(None);
-                };
-                chunk_result
-            }
-        });
+        result = chunk_sum;
     }
-    Ok(result)
+    Ok(Some(result))
 }
 
-fn sum_float(chunks: &[ArrayRef]) -> VortexResult<f64> {
+fn sum_float(chunks: &[ArrayRef]) -> VortexResult<Option<f64>> {
     let mut result = 0f64;
     for chunk in chunks {
-        if let Some(chunk_sum) = sum(chunk)?.as_primitive().as_::<f64>() {
-            result += chunk_sum;
+        let Some(chunk_sum) = sum(chunk)?.as_primitive().as_::<f64>() else {
+            return Ok(None);
         };
+        result += chunk_sum;
     }
-    Ok(result)
+    Ok(Some(result))
 }
 
 fn sum_decimal(chunks: &[ArrayRef], result_decimal_type: DecimalDType) -> VortexResult<Scalar> {
@@ -84,21 +78,19 @@ fn sum_decimal(chunks: &[ArrayRef], result_decimal_type: DecimalDType) -> Vortex
         let chunk_sum = sum(chunk)?;
 
         let chunk_decimal = DecimalScalar::try_from(&chunk_sum)?;
-        let Some(chunk_value) = chunk_decimal.decimal_value() else {
-            // skips all null chunks
-            continue;
-        };
-
-        // Perform checked addition with current result
-        let Some(r) = result.checked_add(&chunk_value).filter(|sum_value| {
-            sum_value
-                .fits_in_precision(result_decimal_type)
-                .unwrap_or(false)
-        }) else {
-            // Overflow
+        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;
     }
 
@@ -146,18 +138,17 @@ mod tests {
     }
 
     #[test]
-    fn test_sum_chunked_floats_all_nulls() {
+    fn test_sum_chunked_floats_all_nulls_is_zero() {
         // Create chunks with all nulls
         let chunk1 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None, None]);
         let chunk2 = PrimitiveArray::from_option_iter::<f32, _>(vec![None, None]);
 
         let dtype = chunk1.dtype().clone();
         let chunked =
             ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype).unwrap();
-
         // Compute sum - should return null for all nulls
         let result = sum(chunked.as_ref()).unwrap();
-        assert!(result.as_primitive().as_::<f64>().is_none());
+        assert_eq!(result, Scalar::primitive(0f64, Nullability::Nullable));
     }
 
     #[test]

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

@@ -108,7 +108,8 @@ register_kernel!(SumKernelAdapter(ConstantVTable).lift());
 
 #[cfg(test)]
 mod tests {
-    use vortex_dtype::{DType, DecimalDType, Nullability, PType};
+    use vortex_dtype::Nullability::Nullable;
+    use vortex_dtype::{DType, DecimalDType, Nullability, PType, i256};
     use vortex_scalar::{DecimalValue, Scalar};
 
     use crate::arrays::ConstantArray;
@@ -132,13 +133,10 @@ mod tests {
 
     #[test]
     fn test_sum_nullable_value() {
-        let array = ConstantArray::new(
-            Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)),
-            10,
-        )
-        .into_array();
+        let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10)
+            .into_array();
         let result = sum(&array).unwrap();
-        assert!(result.is_null());
+        assert_eq!(result, Scalar::primitive(0u64, Nullable));
     }
 
     #[test]
@@ -157,10 +155,9 @@ mod tests {
 
     #[test]
     fn test_sum_bool_null() {
-        let array =
-            ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 10).into_array();
+        let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array();
         let result = sum(&array).unwrap();
-        assert!(result.is_null());
+        assert_eq!(result, Scalar::primitive(0u64, Nullable));
     }
 
     #[test]
@@ -180,22 +177,26 @@ mod tests {
 
         assert_eq!(
             result.as_decimal().decimal_value(),
-            Some(DecimalValue::I256(vortex_scalar::i256::from_i128(500)))
+            Some(DecimalValue::I256(i256::from_i128(500)))
         );
         assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap());
     }
 
     #[test]
     fn test_sum_decimal_null() {
         let decimal_dtype = DecimalDType::new(10, 2);
-        let array = ConstantArray::new(
-            Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)),
-            10,
-        )
-        .into_array();
+        let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10)
+            .into_array();
 
         let result = sum(&array).unwrap();
-        assert!(result.is_null());
+        assert_eq!(
+            result,
+            Scalar::decimal(
+                DecimalValue::I256(i256::ZERO),
+                DecimalDType::new(20, 2),
+                Nullable
+            )
+        );
     }
 
     #[test]
@@ -214,9 +215,7 @@ mod tests {
         let result = sum(&array).unwrap();
         assert_eq!(
             result.as_decimal().decimal_value(),
-            Some(DecimalValue::I256(vortex_scalar::i256::from_i128(
-                99_999_999_900
-            )))
+            Some(DecimalValue::I256(i256::from_i128(99_999_999_900)))
         );
     }
 }

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

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools;
-use num_traits::{CheckedAdd, Float, ToPrimitive};
+use num_traits::{CheckedAdd, Float, ToPrimitive, Zero};
 use vortex_buffer::BitBuffer;
 use vortex_dtype::{NativePType, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult};
@@ -28,11 +28,12 @@ impl SumKernel for PrimitiveVTable {
             }
             AllOr::None => {
                 // All-invalid
-                return Ok(Scalar::null(
-                    Stat::Sum
-                        .dtype(array.dtype())
-                        .vortex_expect("Sum dtype must be defined for primitive type"),
-                ));
+                let sum_dtype = Stat::Sum
+                    .dtype(array.dtype())
+                    .vortex_expect("Sum dtype must be defined for primitive type");
+                return Ok(match_each_native_ptype!(sum_dtype.as_ptype(), |P| {
+                    Scalar::primitive(P::zero(), sum_dtype.nullability())
+                }));
             }
             AllOr::Some(validity_mask) => {
                 // Some-valid

vortex-array/src/compute/sum.rs

@@ -5,6 +5,7 @@ use std::sync::LazyLock;
 
 use arcref::ArcRef;
 use vortex_dtype::DType;
+use vortex_dtype::Nullability::NonNullable;
 use vortex_error::{VortexResult, vortex_err, vortex_panic};
 use vortex_scalar::Scalar;
 
@@ -124,13 +125,8 @@ pub fn sum_impl(
     sum_dtype: DType,
     kernels: &[ArcRef<dyn Kernel>],
 ) -> VortexResult<Scalar> {
-    if array.is_empty() {
-        return Ok(Scalar::default_value(sum_dtype));
-    }
-
-    // Sum of all null is null.
-    if array.all_invalid() {
-        return Ok(Scalar::null(sum_dtype));
+    if array.is_empty() || array.all_invalid() {
+        return Scalar::default_value(sum_dtype.with_nullability(NonNullable)).cast(&sum_dtype);
     }
 
     // Try to find a sum kernel
@@ -162,7 +158,7 @@ pub fn sum_impl(
 #[cfg(test)]
 mod test {
     use vortex_buffer::buffer;
-    use vortex_dtype::{DType, Nullability, PType};
+    use vortex_dtype::Nullability;
     use vortex_scalar::Scalar;
 
     use crate::IntoArray as _;
@@ -173,20 +169,14 @@ mod test {
     fn sum_all_invalid() {
         let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]);
         let result = sum(array.as_ref()).unwrap();
-        assert_eq!(
-            result,
-            Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))
-        );
+        assert_eq!(result, Scalar::primitive(0i64, Nullability::Nullable));
     }
 
     #[test]
     fn sum_all_invalid_float() {
         let array = PrimitiveArray::from_option_iter::<f32, _>([None, None, None]);
         let result = sum(array.as_ref()).unwrap();
-        assert_eq!(
-            result,
-            Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
-        );
+        assert_eq!(result, Scalar::primitive(0f64, Nullability::Nullable));
     }
 
     #[test]