Wire SumFn for chunked and constant (#2606)

Co-authored-by: Robert Kruszewski <[email protected]>

Author: AdamGS

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

@@ -5,8 +5,8 @@ use crate::arrays::ChunkedEncoding;
 use crate::arrays::chunked::ChunkedArray;
 use crate::compute::{
     BinaryBooleanFn, BinaryNumericFn, CastFn, CompareFn, FillNullFn, FilterFn, InvertFn,
-    IsConstantFn, IsSortedFn, MaskFn, MinMaxFn, ScalarAtFn, SliceFn, TakeFn, UncompressedSizeFn,
-    try_cast,
+    IsConstantFn, IsSortedFn, MaskFn, MinMaxFn, ScalarAtFn, SliceFn, SumFn, TakeFn,
+    UncompressedSizeFn, try_cast,
 };
 use crate::vtable::ComputeVTable;
 use crate::{Array, ArrayRef};
@@ -87,6 +87,10 @@ impl ComputeVTable for ChunkedEncoding {
     fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
         Some(self)
     }
+
+    fn sum_fn(&self) -> Option<&dyn SumFn<&dyn Array>> {
+        Some(self)
+    }
 }
 
 impl CastFn<&ChunkedArray> for ChunkedEncoding {

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

@@ -5,16 +5,19 @@ mod compare;
 mod invert;
 mod search_sorted;
 
-use vortex_error::VortexResult;
+use num_traits::{CheckedMul, ToPrimitive};
+use vortex_dtype::{NativePType, PType, match_each_native_ptype};
+use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_mask::Mask;
-use vortex_scalar::Scalar;
+use vortex_scalar::{FromPrimitiveOrF16, PrimitiveScalar, Scalar};
 
 use crate::arrays::ConstantEncoding;
 use crate::arrays::constant::ConstantArray;
 use crate::compute::{
     BinaryBooleanFn, BinaryNumericFn, CastFn, CompareFn, FilterFn, InvertFn, ScalarAtFn,
-    SearchSortedFn, SliceFn, TakeFn, UncompressedSizeFn,
+    SearchSortedFn, SliceFn, SumFn, TakeFn, UncompressedSizeFn,
 };
+use crate::stats::Stat;
 use crate::vtable::ComputeVTable;
 use crate::{Array, ArrayRef};
 
@@ -62,6 +65,10 @@ impl ComputeVTable for ConstantEncoding {
     fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
         Some(self)
     }
+
+    fn sum_fn(&self) -> Option<&dyn SumFn<&dyn Array>> {
+        Some(self)
+    }
 }
 
 impl ScalarAtFn<&ConstantArray> for ConstantEncoding {
@@ -100,6 +107,51 @@ impl UncompressedSizeFn<&ConstantArray> for ConstantEncoding {
     }
 }
 
+impl SumFn<&ConstantArray> for ConstantEncoding {
+    fn sum(&self, array: &ConstantArray) -> VortexResult<Scalar> {
+        let sum_dtype = Stat::Sum
+            .dtype(array.dtype())
+            .ok_or_else(|| vortex_err!("Sum not supported for dtype {}", array.dtype()))?;
+        let sum_ptype = PType::try_from(&sum_dtype).vortex_expect("sum dtype must be primitive");
+
+        let scalar = array.scalar();
+
+        let scalar_value = match_each_native_ptype!(
+            sum_ptype,
+            unsigned: |$T| { sum_integral::<u64>(scalar.as_primitive(), array.len())?.into() }
+            signed: |$T| { sum_integral::<i64>(scalar.as_primitive(), array.len())?.into() }
+            floating: |$T| { sum_float(scalar.as_primitive(), array.len())?.into() }
+        );
+
+        Ok(Scalar::new(sum_dtype, scalar_value))
+    }
+}
+
+fn sum_integral<T>(
+    primitive_scalar: PrimitiveScalar<'_>,
+    array_len: usize,
+) -> VortexResult<Option<T>>
+where
+    T: FromPrimitiveOrF16 + NativePType + CheckedMul,
+    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));
+
+    Ok(sum)
+}
+
+fn sum_float(primitive_scalar: PrimitiveScalar<'_>, array_len: usize) -> VortexResult<Option<f64>> {
+    let v = primitive_scalar.as_::<f64>()?;
+    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))
+}
+
 #[cfg(test)]
 mod test {
     use vortex_dtype::half::f16;