fix: Sum from stats has to account for accumulator (#5445)

fix: #5434 

Signed-off-by: Robert Kruszewski <[email protected]>

---------

Signed-off-by: Robert Kruszewski <[email protected]>
Co-authored-by: Joe Isaacs <[email protected]>

Author: robert3005

vortex-array/src/compute/sum.rs

@@ -4,11 +4,12 @@
 use std::sync::LazyLock;
 
 use arcref::ArcRef;
+use num_traits::{CheckedAdd, CheckedSub};
 use vortex_dtype::DType;
 use vortex_error::{
     VortexError, VortexResult, vortex_bail, vortex_ensure, vortex_err, vortex_panic,
 };
-use vortex_scalar::Scalar;
+use vortex_scalar::{NumericOperator, Scalar};
 
 use crate::Array;
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output};
@@ -102,15 +103,66 @@ impl ComputeFnVTable for Sum {
 
         // Short-circuit using array statistics.
         if let Some(Precision::Exact(sum)) = array.statistics().get(Stat::Sum) {
-            return Ok(sum.into());
+            // For floats only use stats if accumulator is zero. otherwise we might have numerical stability issues.
+            match sum_dtype {
+                DType::Primitive(p, _) => {
+                    if p.is_float() && accumulator.is_zero() {
+                        return Ok(sum.into());
+                    }
+                    let sum_from_stat = accumulator
+                        .as_primitive()
+                        .checked_add(&sum.as_primitive())
+                        .map(Scalar::from);
+                    return Ok(sum_from_stat
+                        .unwrap_or_else(|| Scalar::null(sum_dtype))
+                        .into());
+                }
+                DType::Decimal(..) => {
+                    let sum_from_stat = accumulator
+                        .as_decimal()
+                        .checked_binary_numeric(&sum.as_decimal(), NumericOperator::Add)
+                        .map(Scalar::from);
+                    return Ok(sum_from_stat
+                        .unwrap_or_else(|| Scalar::null(sum_dtype))
+                        .into());
+                }
+                _ => unreachable!("Sum will always be a decimal or a primitive dtype"),
+            }
         }
 
         let sum_scalar = sum_impl(array, accumulator, kernels)?;
 
-        // Update the statistics with the computed sum.
-        array
-            .statistics()
-            .set(Stat::Sum, Precision::Exact(sum_scalar.value().clone()));
+        // Update the statistics with the computed sum. Stored statistic shouldn't include the accumulator.
+        match sum_dtype {
+            DType::Primitive(p, _) => {
+                if p.is_float() && accumulator.is_zero() {
+                    array
+                        .statistics()
+                        .set(Stat::Sum, Precision::Exact(sum_scalar.value().clone()));
+                } else if p.is_int()
+                    && let Some(less_accumulator) = sum_scalar
+                        .as_primitive()
+                        .checked_sub(&accumulator.as_primitive())
+                {
+                    array.statistics().set(
+                        Stat::Sum,
+                        Precision::Exact(Scalar::from(less_accumulator).value().clone()),
+                    );
+                }
+            }
+            DType::Decimal(..) => {
+                if let Some(less_accumulator) = sum_scalar
+                    .as_decimal()
+                    .checked_binary_numeric(&accumulator.as_decimal(), NumericOperator::Sub)
+                {
+                    array.statistics().set(
+                        Stat::Sum,
+                        Precision::Exact(Scalar::from(less_accumulator).value().clone()),
+                    )
+                }
+            }
+            _ => unreachable!("Sum will always be a decimal or a primitive dtype"),
+        }
 
         Ok(sum_scalar.into())
     }
@@ -206,12 +258,13 @@ pub fn sum_impl(
 #[cfg(test)]
 mod test {
     use vortex_buffer::buffer;
-    use vortex_dtype::Nullability;
+    use vortex_dtype::{DType, Nullability, PType};
+    use vortex_error::VortexUnwrap;
     use vortex_scalar::Scalar;
 
     use crate::IntoArray as _;
-    use crate::arrays::{BoolArray, PrimitiveArray};
-    use crate::compute::sum;
+    use crate::arrays::{BoolArray, ChunkedArray, PrimitiveArray};
+    use crate::compute::{sum, sum_with_accumulator};
 
     #[test]
     fn sum_all_invalid() {
@@ -247,4 +300,28 @@ mod test {
         let result = sum(array.as_ref()).unwrap();
         assert_eq!(result.as_primitive().as_::<i32>(), Some(2));
     }
+
+    #[test]
+    fn sum_stats() {
+        let array = ChunkedArray::try_new(
+            vec![
+                PrimitiveArray::from_iter([1, 1, 1]).into_array(),
+                PrimitiveArray::from_iter([2, 2, 2]).into_array(),
+            ],
+            DType::Primitive(PType::I32, Nullability::NonNullable),
+        )
+        .vortex_unwrap();
+        // compute sum with accumulator to populate stats
+        sum_with_accumulator(
+            array.as_ref(),
+            &Scalar::primitive(2i64, Nullability::Nullable),
+        )
+        .unwrap();
+
+        let sum_without_acc = sum(array.as_ref()).unwrap();
+        assert_eq!(
+            sum_without_acc,
+            Scalar::primitive(9i64, Nullability::Nullable)
+        );
+    }
 }

vortex-scalar/src/scalar.rs

@@ -239,6 +239,44 @@ impl Scalar {
         }
     }
 
+    /// Returns true if the scalar is a zero value i.e., equal to a scalar returned from the ` zero_value ` method.
+    pub fn is_zero(&self) -> bool {
+        match self.dtype() {
+            DType::Null => true,
+            DType::Bool(_) => self.as_bool().value() == Some(false),
+            DType::Primitive(pt, _) => self.as_primitive().pvalue() == Some(PValue::zero(*pt)),
+            DType::Decimal(..) => {
+                self.as_decimal().decimal_value() == Some(DecimalValue::from(0i8))
+            }
+            DType::Utf8(_) => self
+                .as_utf8()
+                .value()
+                .map(|v| v.is_empty())
+                .unwrap_or(false),
+            DType::Binary(_) => self
+                .as_binary()
+                .value()
+                .map(|v| v.is_empty())
+                .unwrap_or(false),
+            DType::Struct(..) => self
+                .as_struct()
+                .fields()
+                .map(|mut sf| sf.all(|f| f.is_zero()))
+                .unwrap_or(false),
+            DType::List(..) => self
+                .as_list()
+                .elements()
+                .map(|vals| vals.is_empty())
+                .unwrap_or(false),
+            DType::FixedSizeList(..) => self
+                .as_list()
+                .elements()
+                .map(|vals| vals.iter().all(|f| f.is_zero()))
+                .unwrap_or(false),
+            DType::Extension(..) => self.as_extension().storage().is_zero(),
+        }
+    }
+
     /// Creates a "default" scalar value for the given data type.
     ///
     /// For nullable types, returns null. For non-nullable types, returns an appropriate zero/empty