Add sum_v2 with SQL empty-input semantics

vortex-array/benches/aggregate_grouped.rs

@@ -11,6 +11,7 @@ use rand::RngExt;
 use rand::SeedableRng;
 use rand::rngs::StdRng;
 use vortex_array::ArrayRef;
+use vortex_array::Canonical;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::aggregate_fn::AggregateFnVTable;
@@ -19,6 +20,7 @@ use vortex_array::aggregate_fn::GroupedAccumulator;
 use vortex_array::aggregate_fn::NumericalAggregateOpts;
 use vortex_array::aggregate_fn::fns::count::Count;
 use vortex_array::aggregate_fn::fns::sum::Sum;
+use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
 use vortex_array::arrays::ListViewArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::VarBinViewArray;
@@ -199,6 +201,125 @@ fn sum_f64_clustered_nulls(bencher: Bencher) {
         .bench_refs(|input| grouped_accumulator(input, Sum));
 }
 
+#[divan::bench]
+fn sum_v2_i32_nullable_all_valid(bencher: Bencher) {
+    let input = i32_nullable_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, SumV2));
+}
+
+#[divan::bench]
+fn sum_v2_i32_clustered_nulls(bencher: Bencher) {
+    let input = i32_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, SumV2));
+}
+
+#[divan::bench]
+fn sum_v2_f64_all_valid(bencher: Bencher) {
+    let input = f64_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, SumV2));
+}
+
+#[divan::bench]
+fn sum_v2_f64_clustered_nulls(bencher: Bencher) {
+    let input = f64_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, SumV2));
+}
+
+/// Execute the lazy finalize result to canonical form so the benchmark includes the full cost of
+/// producing usable sums.
+fn grouped_accumulator_canonical<V>(list_view: &ArrayRef, vtable: V) -> ArrayRef
+where
+    V: AggregateFnVTable<Options = NumericalAggregateOpts> + Clone,
+{
+    let mut acc = GroupedAccumulator::try_new(
+        vtable,
+        NumericalAggregateOpts::default(),
+        list_element_dtype(list_view),
+    )
+    .unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    acc.accumulate_list(list_view, &mut ctx).unwrap();
+    let result = acc
+        .finish()
+        .unwrap()
+        .execute::<Canonical>(&mut ctx)
+        .unwrap()
+        .into_array();
+    divan::black_box(result)
+}
+
+#[divan::bench]
+fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) {
+    let input = i32_nullable_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, Sum));
+}
+
+#[divan::bench]
+fn canonical_sum_i32_clustered_nulls(bencher: Bencher) {
+    let input = i32_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, Sum));
+}
+
+#[divan::bench]
+fn canonical_sum_f64_all_valid(bencher: Bencher) {
+    let input = f64_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, Sum));
+}
+
+#[divan::bench]
+fn canonical_sum_f64_clustered_nulls(bencher: Bencher) {
+    let input = f64_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, Sum));
+}
+
+#[divan::bench]
+fn canonical_sum_v2_i32_nullable_all_valid(bencher: Bencher) {
+    let input = i32_nullable_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
+}
+
+#[divan::bench]
+fn canonical_sum_v2_i32_clustered_nulls(bencher: Bencher) {
+    let input = i32_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
+}
+
+#[divan::bench]
+fn canonical_sum_v2_f64_all_valid(bencher: Bencher) {
+    let input = f64_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
+}
+
+#[divan::bench]
+fn canonical_sum_v2_f64_clustered_nulls(bencher: Bencher) {
+    let input = f64_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
+}
+
 #[divan::bench]
 fn count_i32_clustered_nulls(bencher: Bencher) {
     let input = i32_clustered_nulls_input();

vortex-array/benches/aggregate_sum.rs

@@ -7,6 +7,7 @@ use divan::Bencher;
 use rand::prelude::*;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::fns::sum_v2::sum_v2;
 use vortex_array::array_session;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::expr::stats::Stat;
@@ -36,6 +37,20 @@ fn sum_i32(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<i64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_i32(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(1);
+    let data: Vec<i32> = (0..N).map(|_| rng.random_range(-1000..1000)).collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 #[divan::bench]
 fn sum_u32(bencher: Bencher) {
     let mut rng = StdRng::seed_from_u64(2);
@@ -50,6 +65,20 @@ fn sum_u32(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<u64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_u32(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(2);
+    let data: Vec<u32> = (0..N).map(|_| rng.random_range(0..2000)).collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 #[divan::bench]
 fn sum_i64(bencher: Bencher) {
     let mut rng = StdRng::seed_from_u64(3);
@@ -64,6 +93,20 @@ fn sum_i64(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<i64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_i64(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(3);
+    let data: Vec<i64> = (0..N).map(|_| rng.random_range(-1000..1000)).collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 #[divan::bench]
 fn sum_f64(bencher: Bencher) {
     let mut rng = StdRng::seed_from_u64(6);
@@ -78,6 +121,20 @@ fn sum_f64(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<f64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_f64(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(6);
+    let data: Vec<f64> = (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 #[divan::bench]
 fn sum_f64_nulls_clustered(bencher: Bencher) {
     let mut rng = StdRng::seed_from_u64(7);
@@ -100,6 +157,28 @@ fn sum_f64_nulls_clustered(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<f64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_f64_nulls_clustered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(7);
+    let data: Vec<Option<f64>> = (0..N)
+        .map(|i| {
+            if (i / 64) % 10 == 0 {
+                None
+            } else {
+                Some(rng.random_range(-1000.0..1000.0))
+            }
+        })
+        .collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_option_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 // Clustered nulls: long runs of valid values broken up by occasional null blocks. This is the
 // case the run-based valid path is expected to accelerate.
 #[divan::bench]
@@ -124,6 +203,28 @@ fn sum_i32_nulls_clustered(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<i64>(Stat::Sum, ctx));
 }
 
+#[divan::bench]
+fn sum_v2_i32_nulls_clustered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(4);
+    let data: Vec<Option<i32>> = (0..N)
+        .map(|i| {
+            if (i / 64) % 10 == 0 {
+                None
+            } else {
+                Some(rng.random_range(-1000..1000))
+            }
+        })
+        .collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_option_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}
+
 // Scattered nulls: ~50% nulls placed at random, producing many short runs. This is the worst case
 // for a run-based valid path, used to guard against regressions versus a per-element loop.
 #[divan::bench]
@@ -141,3 +242,19 @@ fn sum_i32_nulls_scattered(bencher: Bencher) {
         })
         .bench_refs(|(a, ctx)| a.statistics().compute_as::<i64>(Stat::Sum, ctx));
 }
+
+#[divan::bench]
+fn sum_v2_i32_nulls_scattered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(5);
+    let data: Vec<Option<i32>> = (0..N)
+        .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000..1000)))
+        .collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_option_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| sum_v2(a, ctx));
+}

vortex-array/src/aggregate_fn/fns/mod.rs

@@ -20,4 +20,5 @@ pub mod min_max;
 pub mod nan_count;
 pub mod null_count;
 pub mod sum;
+pub mod sum_v2;
 pub mod uncompressed_size_in_bytes;

vortex-array/src/aggregate_fn/fns/sum/bool.rs

@@ -13,7 +13,7 @@ use crate::ExecutionCtx;
 use crate::arrays::BoolArray;
 use crate::arrays::bool::BoolArrayExt;
 
-pub(super) fn accumulate_bool(
+pub(crate) fn accumulate_bool(
     inner: &mut SumState,
     b: &BoolArray,
     ctx: &mut ExecutionCtx,

vortex-array/src/aggregate_fn/fns/sum/constant.rs

@@ -15,7 +15,7 @@ use crate::scalar::Scalar;
 ///
 /// Returns `Ok(None)` if the scalar is null (no contribution to the sum).
 /// Returns a null scalar on overflow (saturation).
-pub(super) fn multiply_constant(
+pub(crate) fn multiply_constant(
     scalar: &Scalar,
     len: usize,
     return_dtype: &DType,

vortex-array/src/aggregate_fn/fns/sum/decimal.rs

@@ -23,7 +23,7 @@ use crate::scalar::DecimalValue;
 
 /// Accumulate a decimal array into the sum state.
 /// Returns Ok(true) if saturated (overflow), Ok(false) if not.
-pub(super) fn accumulate_decimal(
+pub(crate) fn accumulate_decimal(
     inner: &mut SumState,
     d: &DecimalArray,
     ctx: &mut ExecutionCtx,

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -15,10 +15,13 @@ use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use self::bool::accumulate_bool;
-use self::constant::multiply_constant;
-use self::decimal::accumulate_decimal;
-use self::primitive::accumulate_primitive;
+pub(crate) use self::bool::accumulate_bool;
+pub(crate) use self::constant::multiply_constant;
+pub(crate) use self::decimal::accumulate_decimal;
+pub(crate) use self::primitive::accumulate_primitive;
+pub(crate) use self::primitive::sum_float_all;
+pub(crate) use self::primitive::sum_signed_all;
+pub(crate) use self::primitive::sum_unsigned_all;
 use crate::ArrayRef;
 use crate::Canonical;
 use crate::Columnar;
@@ -351,7 +354,7 @@ pub enum SumState {
     },
 }
 
-fn make_zero_state(return_dtype: &DType) -> SumState {
+pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState {
     match return_dtype {
         DType::Primitive(ptype, _) => match ptype {
             PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0),