fix[array]: sum with initial value to fix op assoc

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

Author: joseph-isaacs

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

@@ -1,9 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::ops::AddAssign;
-
-use num_traits::PrimInt;
+use num_traits::CheckedAdd;
 use vortex_dtype::{DType, NativePType, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_scalar::Scalar;
@@ -26,7 +24,7 @@ impl SumKernel for ChunkedVTable {
                     sum_ptype,
                     unsigned: |T| { sum_int::<u64>(array.chunks(), initial_value.as_primitive().as_::<u64>().vortex_expect("cannot be null"))?.into() },
                     signed: |T| { sum_int::<i64>(array.chunks(), initial_value.as_primitive().as_::<i64>().vortex_expect("cannot be null"))?.into() },
-                    floating: |T| { sum_float::<f64>(array.chunks(), initial_value.as_primitive().as_::<f64>().vortex_expect("cannot be null"))?.into() }
+                    floating: |T| { sum_float(array.chunks(), initial_value.as_primitive().as_::<f64>().vortex_expect("cannot be null"))?.into() }
                 );
 
                 Ok(Scalar::new(sum_dtype, scalar_value))
@@ -40,7 +38,7 @@ impl SumKernel for ChunkedVTable {
 
 register_kernel!(SumKernelAdapter(ChunkedVTable).lift());
 
-fn sum_int<T: NativePType + PrimInt>(
+fn sum_int<T: NativePType + CheckedAdd>(
     chunks: &[ArrayRef],
     initial_value: T,
 ) -> VortexResult<Option<T>> {
@@ -60,13 +58,10 @@ fn sum_int<T: NativePType + PrimInt>(
     Ok(Some(result))
 }
 
-fn sum_float<T: NativePType + AddAssign>(
-    chunks: &[ArrayRef],
-    initial_value: T,
-) -> VortexResult<Option<T>> {
+fn sum_float(chunks: &[ArrayRef], initial_value: f64) -> VortexResult<Option<f64>> {
     let mut result = initial_value;
     for chunk in chunks {
-        let Some(chunk_sum) = sum(chunk)?.as_primitive().as_::<T>() else {
+        let Some(chunk_sum) = sum(chunk)?.as_primitive().as_::<f64>() else {
             return Ok(None);
         };
         result += chunk_sum;
@@ -78,7 +73,7 @@ fn sum_decimal(chunks: &[ArrayRef], initial_value: &Scalar) -> VortexResult<Scal
     let mut result = initial_value.clone();
 
     for chunk in chunks {
-        result = sum_with_initial(chunk, result)?;
+        result = sum_with_initial(chunk, &result)?;
     }
 
     Ok(result)

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

@@ -10,7 +10,7 @@ use crate::register_kernel;
 
 impl SumKernel for ExtensionVTable {
     fn sum(&self, array: &ExtensionArray, initial_value: &Scalar) -> VortexResult<Scalar> {
-        compute::sum_with_initial(array.storage(), initial_value.clone())
+        compute::sum_with_initial(array.storage(), initial_value)
     }
 }
 

vortex-array/src/compute/sum.rs

@@ -6,13 +6,11 @@ use std::sync::LazyLock;
 
 use arcref::ArcRef;
 use vortex_dtype::DType;
-use vortex_error::{VortexResult, vortex_err, vortex_panic};
+use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err, vortex_panic};
 use vortex_scalar::Scalar;
 
 use crate::Array;
-use crate::compute::{
-    ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output, UnaryArgs,
-};
+use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
 use crate::stats::{Precision, Stat, StatsProvider};
 use crate::vtable::VTable;
 
@@ -45,11 +43,11 @@ pub(crate) fn warm_up_vtable() -> usize {
 /// If the sum is not supported for the array's dtype, an error will be raised.
 /// If the array is all-invalid, the sum will be the initial_value.
 /// The initial_value must have a dtype compatible with the sum result dtype.
-pub(crate) fn sum_with_initial(array: &dyn Array, initial_value: Scalar) -> VortexResult<Scalar> {
+pub(crate) fn sum_with_initial(array: &dyn Array, initial_value: &Scalar) -> VortexResult<Scalar> {
     SUM_FN
         .invoke(&InvocationArgs {
-            inputs: &[array.into()],
-            options: &SumOptions { initial_value },
+            inputs: &[array.into(), initial_value.into()],
+            options: &(),
         })?
         .unwrap_scalar()
 }
@@ -64,7 +62,30 @@ pub fn sum(array: &dyn Array) -> VortexResult<Scalar> {
         .dtype(array.dtype())
         .ok_or_else(|| vortex_err!("Sum not supported for dtype: {}", array.dtype()))?;
     let zero = Scalar::zero_value(sum_dtype);
-    sum_with_initial(array, zero)
+    sum_with_initial(array, &zero)
+}
+
+/// For unary compute functions, it's useful to just have this short-cut.
+pub struct SumArgs<'a> {
+    pub array: &'a dyn Array,
+    pub accumulator: &'a Scalar,
+}
+
+impl<'a> TryFrom<&InvocationArgs<'a>> for SumArgs<'a> {
+    type Error = VortexError;
+
+    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
+        if value.inputs.len() != 2 {
+            vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());
+        }
+        let array = value.inputs[0]
+            .array()
+            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
+        let accumulator = value.inputs[1]
+            .scalar()
+            .ok_or_else(|| vortex_err!("Expected input 1 to be a scalar"))?;
+        Ok(SumArgs { array, accumulator })
+    }
 }
 
 struct Sum;
@@ -75,8 +96,7 @@ impl ComputeFnVTable for Sum {
         args: &InvocationArgs,
         kernels: &[ArcRef<dyn Kernel>],
     ) -> VortexResult<Output> {
-        let UnaryArgs { array, options } = UnaryArgs::<SumOptions>::try_from(args)?;
-        let initial_value = &options.initial_value;
+        let SumArgs { array, accumulator } = args.try_into()?;
 
         // Compute the expected dtype of the sum.
         let sum_dtype = self.return_dtype(args)?;
@@ -86,7 +106,7 @@ impl ComputeFnVTable for Sum {
             return Ok(sum.into());
         }
 
-        let sum_scalar = sum_impl(array, sum_dtype, initial_value, kernels)?;
+        let sum_scalar = sum_impl(array, sum_dtype, accumulator, kernels)?;
 
         // Update the statistics with the computed sum.
         array
@@ -97,7 +117,7 @@ impl ComputeFnVTable for Sum {
     }
 
     fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
-        let UnaryArgs { array, .. } = UnaryArgs::<SumOptions>::try_from(args)?;
+        let SumArgs { array, .. } = args.try_into()?;
         Stat::Sum
             .dtype(array.dtype())
             .ok_or_else(|| vortex_err!("Sum not supported for dtype: {}", array.dtype()))
@@ -136,11 +156,14 @@ impl<V: VTable + SumKernel> SumKernelAdapter<V> {
 
 impl<V: VTable + SumKernel> Kernel for SumKernelAdapter<V> {
     fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
-        let UnaryArgs { array, options } = UnaryArgs::<SumOptions>::try_from(args)?;
+        let SumArgs {
+            array,
+            accumulator: initial_value,
+        } = args.try_into()?;
         let Some(array) = array.as_opt::<V>() else {
             return Ok(None);
         };
-        Ok(Some(V::sum(&self.0, array, &options.initial_value)?.into()))
+        Ok(Some(V::sum(&self.0, array, initial_value)?.into()))
     }
 }
 
@@ -161,10 +184,8 @@ pub fn sum_impl(
 
     // Try to find a sum kernel
     let args = InvocationArgs {
-        inputs: &[array.into()],
-        options: &SumOptions {
-            initial_value: initial_value.clone(),
-        },
+        inputs: &[array.into(), initial_value.into()],
+        options: &(),
     };
     for kernel in kernels {
         if let Some(output) = kernel.invoke(&args)? {
@@ -184,7 +205,7 @@ pub fn sum_impl(
             array.encoding_id()
         );
     }
-    sum_with_initial(array.to_canonical().as_ref(), initial_value.clone())
+    sum_with_initial(array.to_canonical().as_ref(), initial_value)
 }
 
 #[cfg(test)]