Add StandardSum: a sum aggregate where zero valid values yields null

vortex-array/benches/aggregate_grouped.rs

@@ -11,13 +11,15 @@ 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;
 use vortex_array::aggregate_fn::DynGroupedAccumulator;
 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::standard_sum::StandardSum;
 use vortex_array::aggregate_fn::fns::sum::Sum;
 use vortex_array::arrays::ListViewArray;
 use vortex_array::arrays::PrimitiveArray;
@@ -199,6 +201,125 @@ fn sum_f64_clustered_nulls(bencher: Bencher) {
         .bench_refs(|input| grouped_accumulator(input, Sum));
 }
 
+#[divan::bench]
+fn standard_sum_i32_nullable_all_valid(bencher: Bencher) {
+    let input = i32_nullable_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, StandardSum));
+}
+
+#[divan::bench]
+fn standard_sum_i32_clustered_nulls(bencher: Bencher) {
+    let input = i32_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, StandardSum));
+}
+
+#[divan::bench]
+fn standard_sum_f64_all_valid(bencher: Bencher) {
+    let input = f64_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, StandardSum));
+}
+
+#[divan::bench]
+fn standard_sum_f64_clustered_nulls(bencher: Bencher) {
+    let input = f64_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator(input, StandardSum));
+}
+
+/// Like [`grouped_accumulator`], but executes the lazy finalize result to canonical so the
+/// bench measures 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_standard_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, StandardSum));
+}
+
+#[divan::bench]
+fn canonical_standard_sum_i32_clustered_nulls(bencher: Bencher) {
+    let input = i32_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum));
+}
+
+#[divan::bench]
+fn canonical_standard_sum_f64_all_valid(bencher: Bencher) {
+    let input = f64_all_valid_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum));
+}
+
+#[divan::bench]
+fn canonical_standard_sum_f64_clustered_nulls(bencher: Bencher) {
+    let input = f64_clustered_nulls_input();
+    bencher
+        .with_inputs(|| &input)
+        .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum));
+}
+
 #[divan::bench]
 fn count_i32_clustered_nulls(bencher: Bencher) {
     let input = i32_clustered_nulls_input();

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

@@ -19,5 +19,6 @@ pub mod min;
 pub mod min_max;
 pub mod nan_count;
 pub mod null_count;
+pub mod standard_sum;
 pub mod sum;
 pub mod uncompressed_size_in_bytes;