Fold aggregates with masked and lane-parallel kernels

vortex-array/Cargo.toml

@@ -103,6 +103,10 @@ vortex-bench-support = { workspace = true }
 name = "aggregate_max"
 harness = false
 
+[[bench]]
+name = "aggregate_nullable"
+harness = false
+
 [[bench]]
 name = "aggregate_sum"
 harness = false

vortex-array/benches/aggregate_max.rs

@@ -179,6 +179,36 @@ fn max_i32_nulls_scattered(bencher: Bencher) {
         .bench_refs(|(a, ctx)| a.statistics().compute_max::<i32>(ctx));
 }
 
+#[divan::bench]
+fn max_f32(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(6);
+    let data: Vec<f32> = (0..N).map(|_| rng.random::<f32>()).collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| a.statistics().compute_max::<f32>(ctx));
+}
+
+#[divan::bench]
+fn max_f64_nulls_scattered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(7);
+    let data: Vec<Option<f64>> = (0..N)
+        .map(|_| rng.random_bool(0.5).then(|| rng.random::<f64>()))
+        .collect();
+    bencher
+        .with_inputs(|| {
+            (
+                PrimitiveArray::from_option_iter(data.iter().copied()).into_array(),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(a, ctx)| a.statistics().compute_max::<f64>(ctx));
+}
+
 // String extrema. Every view carries its value's first four bytes inline, so the shape that
 // matters is how far into a value its order is settled: `distinct_prefixes` and
 // `low_cardinality` never have to look past the view, while the log-shaped columns (every value

vortex-array/benches/aggregate_nullable.rs

@@ -0,0 +1,186 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Aggregates over columns whose nulls are scattered at random, so their valid values form runs
+//! of one or two rows rather than the long runs of the clustered benchmarks.
+//!
+//! Each iteration gets a fresh array, so no statistic cached by a previous iteration is reused.
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use mimalloc::MiMalloc;
+use rand::prelude::*;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::NumericalAggregateOpts;
+use vortex_array::aggregate_fn::fns::min_max::min_max;
+use vortex_array::aggregate_fn::fns::nan_count::nan_count;
+use vortex_array::aggregate_fn::fns::null_count::null_count;
+use vortex_array::aggregate_fn::fns::sum::sum;
+use vortex_array::aggregate_fn::session::AggregateFnSession;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::dtype::DecimalDType;
+use vortex_array::optimizer::kernels::KernelSession;
+use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_session::VortexSession;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+// Sized to keep the CodSpeed simulation under 1ms per benchmark.
+const N: usize = 15_000;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    VortexSession::empty()
+        .with::<KernelSession>()
+        .with::<AggregateFnSession>()
+});
+
+/// Half the rows, chosen at random, are valid.
+fn scattered_validity() -> Validity {
+    let mut rng = StdRng::seed_from_u64(1);
+    Validity::from_iter((0..N).map(|_| rng.random_bool(0.5)))
+}
+
+fn f64_values() -> Buffer<f64> {
+    let mut rng = StdRng::seed_from_u64(2);
+    (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect()
+}
+
+fn bench_with<O>(
+    bencher: Bencher,
+    input: impl Fn() -> ArrayRef + Sync,
+    f: impl Fn(&ArrayRef) -> O + Sync,
+) {
+    bencher.with_inputs(&input).bench_refs(|array| f(array));
+}
+
+fn bench_min_max(bencher: Bencher, input: impl Fn() -> ArrayRef + Sync) {
+    bench_with(bencher, input, |array| {
+        min_max(
+            array,
+            &mut SESSION.create_execution_ctx(),
+            NumericalAggregateOpts::default(),
+        )
+    });
+}
+
+fn bench_sum(bencher: Bencher, input: impl Fn() -> ArrayRef + Sync) {
+    bench_with(bencher, input, |array| {
+        sum(array, &mut SESSION.create_execution_ctx())
+    });
+}
+
+#[divan::bench]
+fn min_max_f64_all_valid(bencher: Bencher) {
+    let values = f64_values();
+    bench_min_max(bencher, || {
+        PrimitiveArray::new(values.clone(), Validity::NonNullable).into_array()
+    });
+}
+
+#[divan::bench]
+fn min_max_f64_nulls_scattered(bencher: Bencher) {
+    let (values, validity) = (f64_values(), scattered_validity());
+    bench_min_max(bencher, || {
+        PrimitiveArray::new(values.clone(), validity.clone()).into_array()
+    });
+}
+
+fn decimal_i64_values() -> Buffer<i64> {
+    let mut rng = StdRng::seed_from_u64(3);
+    (0..N)
+        .map(|_| rng.random_range(-100_000..100_000))
+        .collect()
+}
+
+#[divan::bench]
+fn min_max_decimal_i64_nulls_scattered(bencher: Bencher) {
+    let (values, validity) = (decimal_i64_values(), scattered_validity());
+    bench_min_max(bencher, || {
+        DecimalArray::new(values.clone(), DecimalDType::new(18, 2), validity.clone()).into_array()
+    });
+}
+
+#[divan::bench]
+fn sum_i32_nulls_scattered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(4);
+    let values: Buffer<i32> = (0..N).map(|_| rng.random_range(-1000..1000)).collect();
+    let validity = scattered_validity();
+    bench_sum(bencher, || {
+        PrimitiveArray::new(values.clone(), validity.clone()).into_array()
+    });
+}
+
+#[divan::bench]
+fn sum_f64_nulls_scattered(bencher: Bencher) {
+    let (values, validity) = (f64_values(), scattered_validity());
+    bench_sum(bencher, || {
+        PrimitiveArray::new(values.clone(), validity.clone()).into_array()
+    });
+}
+
+#[divan::bench]
+fn sum_decimal_i64_all_valid(bencher: Bencher) {
+    let values = decimal_i64_values();
+    bench_sum(bencher, || {
+        DecimalArray::new(
+            values.clone(),
+            DecimalDType::new(18, 2),
+            Validity::NonNullable,
+        )
+        .into_array()
+    });
+}
+
+#[divan::bench]
+fn sum_decimal_i64_nulls_scattered(bencher: Bencher) {
+    let (values, validity) = (decimal_i64_values(), scattered_validity());
+    bench_sum(bencher, || {
+        DecimalArray::new(values.clone(), DecimalDType::new(18, 2), validity.clone()).into_array()
+    });
+}
+
+#[divan::bench]
+fn nan_count_f64_nulls_scattered(bencher: Bencher) {
+    let (values, validity) = (f64_values(), scattered_validity());
+    bench_with(
+        bencher,
+        || PrimitiveArray::new(values.clone(), validity.clone()).into_array(),
+        |array| nan_count(array, &mut SESSION.create_execution_ctx()),
+    );
+}
+
+/// Mask a nullable column, as a parent's nulls mask a field, and count its nulls, which reads the
+/// conjunction of the two validities.
+#[divan::bench]
+fn null_count_i32_masked(bencher: Bencher) {
+    let values = i32_values();
+    let validity = scattered_validity();
+    let mut rng = StdRng::seed_from_u64(6);
+    let mask = BoolArray::from_iter((0..N).map(|_| rng.random_bool(0.5))).into_array();
+    bench_with(
+        bencher,
+        || PrimitiveArray::new(values.clone(), validity.clone()).into_array(),
+        |array| {
+            let masked = array.clone().mask(mask.clone())?;
+            null_count(&masked, &mut SESSION.create_execution_ctx())
+        },
+    );
+}
+
+fn i32_values() -> Buffer<i32> {
+    let mut rng = StdRng::seed_from_u64(7);
+    (0..N).map(|_| rng.random_range(-1000..1000)).collect()
+}

vortex-array/benches/aggregate_sum.rs

@@ -267,3 +267,52 @@ fn sum_v2_i32_nulls_scattered(bencher: Bencher) {
         })
         .bench_refs(|(a, ctx)| sum_v2(a, ctx));
 }
+
+#[divan::bench]
+fn sum_u64(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(8);
+    let data: Vec<u64> = (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)| a.statistics().compute_as::<u64>(Stat::Sum, ctx));
+}
+
+#[divan::bench]
+fn sum_i64_nulls_scattered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(9);
+    let data: Vec<Option<i64>> = (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)| a.statistics().compute_as::<i64>(Stat::Sum, ctx));
+}
+
+#[divan::bench]
+fn sum_f64_nulls_scattered(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(10);
+    let data: Vec<Option<f64>> = (0..N)
+        .map(|_| {
+            rng.random_bool(0.5)
+                .then(|| 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)| a.statistics().compute_as::<f64>(Stat::Sum, ctx));
+}

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

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use itertools::Itertools;
+use vortex_compute::lane_kernels::IndexedFoldExt;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
@@ -79,25 +79,17 @@ pub fn valid_range<T: NativeDecimalType>(values: &[T], validity: &Mask) -> Optio
         "valid_range requires one validity bit per value"
     );
 
+    let min_max = |(min, max): (T, T), v: T| (min.min(v), max.max(v));
     match validity {
-        Mask::AllTrue(_) => compute_min_max(values.iter()),
+        Mask::AllTrue(_) => values
+            .first()
+            .map(|&first| values.iter().copied().fold((first, first), min_max)),
         Mask::AllFalse(_) => None,
-        Mask::Values(validity) => compute_min_max(
-            values
-                .iter()
-                .zip(validity.bit_buffer().iter())
-                .filter_map(|(value, valid)| valid.then_some(value)),
-        ),
-    }
-}
-
-fn compute_min_max<'a, T>(iter: impl Iterator<Item = &'a T>) -> Option<(T, T)>
-where
-    T: NativeDecimalType,
-{
-    match iter.minmax_by(|a, b| a.cmp(b)) {
-        itertools::MinMaxResult::NoElements => None,
-        itertools::MinMaxResult::OneElement(&value) => Some((value, value)),
-        itertools::MinMaxResult::MinMax(&min, &max) => Some((min, max)),
+        // Null lanes read as a valid value, which leaves both extrema unchanged, so the reduction
+        // needs no branch per lane.
+        Mask::Values(validity) => validity.bit_buffer().set_indices().next().map(|first| {
+            let first = values[first];
+            values.fold_masked(validity.bit_buffer(), first, (first, first), min_max)
+        }),
     }
 }