perf: specialize whole-array sums for run-end arrays

encodings/runend/Cargo.toml

@@ -58,3 +58,7 @@ harness = false
 [[bench]]
 name = "run_end_filter"
 harness = false
+
+[[bench]]
+name = "run_end_sum"
+harness = false

encodings/runend/benches/run_end_sum.rs

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::fns::sum_v2::sum_v2;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_runend::RunEnd;
+use vortex_session::VortexSession;
+
+const LEN: usize = 2_048;
+const RUN_LENGTH: usize = 64;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_runend::initialize(&session);
+    session
+});
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+fn runend_with_null_runs() -> ArrayRef {
+    let ends =
+        PrimitiveArray::from_iter((RUN_LENGTH..=LEN).step_by(RUN_LENGTH).map(|end| end as u64));
+    let values = PrimitiveArray::from_option_iter(
+        (0..ends.len()).map(|index| (index % 5 != 0).then_some(i32::try_from(index).unwrap())),
+    );
+
+    RunEnd::try_new(
+        ends.into_array(),
+        values.into_array(),
+        &mut SESSION.create_execution_ctx(),
+    )
+    .unwrap()
+    .into_array()
+}
+
+#[divan::bench]
+fn whole_array_sum_partially_valid(bencher: Bencher) {
+    let array = runend_with_null_runs();
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| sum_v2(&array, ctx).unwrap());
+}

encodings/runend/src/compute/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod filter;
 pub(crate) mod is_constant;
 pub(crate) mod is_sorted;
 pub(crate) mod min_max;
+pub(crate) mod sum;
 pub(crate) mod take;
 pub(crate) mod take_from;
 

encodings/runend/src/compute/sum/kernel.rs

@@ -0,0 +1,152 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Whole-array sum dispatch and scalar partial construction.
+//!
+//! Run traversal and arithmetic are implemented in [`super::runs`].
+
+use std::ops::Range;
+
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::NumericalAggregateOpts;
+use vortex_array::aggregate_fn::fns::sum::Sum;
+use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
+use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::NativePType;
+use vortex_array::dtype::Nullability::Nullable;
+use vortex_array::dtype::UnsignedPType;
+use vortex_array::match_each_native_ptype;
+use vortex_array::match_each_unsigned_integer_ptype;
+use vortex_array::scalar::PValue;
+use vortex_array::scalar::Scalar;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::runs::add_float_run;
+use super::runs::add_signed_run;
+use super::runs::add_unsigned_run;
+use super::runs::sum_range;
+use crate::RunEnd;
+use crate::RunEndArrayExt;
+use crate::RunEndArraySlotsExt;
+
+/// Whole-array primitive sum kernel for [`RunEnd`].
+#[derive(Debug)]
+pub(crate) struct RunEndSumKernel;
+
+impl DynAggregateKernel for RunEndSumKernel {
+    fn aggregate(
+        &self,
+        aggregate_fn: &AggregateFnRef,
+        batch: &ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Scalar>> {
+        if !batch.dtype().is_primitive() {
+            return Ok(None);
+        }
+
+        let Some(options) = aggregate_fn
+            .as_opt::<Sum>()
+            .or_else(|| aggregate_fn.as_opt::<SumV2>())
+            .copied()
+        else {
+            return Ok(None);
+        };
+
+        let Some(array) = batch.as_opt::<RunEnd>() else {
+            return Ok(None);
+        };
+
+        if array.is_empty() {
+            return Ok(Some(empty_partial(aggregate_fn, batch.dtype())?));
+        }
+
+        let validity = array
+            .values()
+            .validity()?
+            .execute_mask(array.values().len(), ctx)?;
+        if validity.all_false() {
+            return Ok(Some(empty_partial(aggregate_fn, batch.dtype())?));
+        }
+
+        let ends = array.ends().clone().execute::<PrimitiveArray>(ctx)?;
+        let values = array.values().clone().execute::<PrimitiveArray>(ctx)?;
+        let range = array.offset()..array.offset() + batch.len();
+        let (sum, is_empty) = sum_primitive_runs(&ends, &values, &validity, range, options);
+
+        Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?))
+    }
+}
+
+/// Select run-end and value types, widening sums to `u64`, `i64`, or `f64`.
+///
+/// Returns `(sum, is_empty)` with overflow represented by a null scalar. Input requirements and
+/// empty-input semantics are defined in [`super::runs`].
+fn sum_primitive_runs(
+    ends: &PrimitiveArray,
+    values: &PrimitiveArray,
+    validity: &Mask,
+    range: Range<usize>,
+    options: NumericalAggregateOpts,
+) -> (Scalar, bool) {
+    match_each_unsigned_integer_ptype!(ends.ptype(), |E| {
+        let ends = ends.as_slice::<E>();
+
+        match_each_native_ptype!(values.ptype(),
+            unsigned: |T| {
+                sum_typed_runs(ends, values.as_slice::<T>(), validity, range, add_unsigned_run::<T>)
+            },
+            signed: |T| {
+                sum_typed_runs(ends, values.as_slice::<T>(), validity, range, add_signed_run::<T>)
+            },
+            floating: |T| {
+                sum_typed_runs(ends, values.as_slice::<T>(), validity, range, |sum, value, len| {
+                    add_float_run(sum, value, len, options.skip_nans)
+                })
+            }
+        )
+    })
+}
+
+/// Convert the reduction's sum to a nullable scalar, preserving its empty-input flag.
+fn sum_typed_runs<E: UnsignedPType, T: NativePType, A: NativePType + Into<PValue>>(
+    ends: &[E],
+    values: &[T],
+    validity: &Mask,
+    range: Range<usize>,
+    add_run: impl Fn(A, T, usize) -> Option<A>,
+) -> (Scalar, bool) {
+    let (sum, is_empty) = sum_range(ends, values, validity, range, add_run);
+    let sum = match sum {
+        Some(sum) => Scalar::primitive(sum, Nullable),
+        None => Scalar::null(DType::Primitive(A::PTYPE, Nullable)),
+    };
+
+    (sum, is_empty)
+}
+
+fn empty_partial(aggregate_fn: &AggregateFnRef, dtype: &DType) -> VortexResult<Scalar> {
+    let sum_dtype = aggregate_fn
+        .return_dtype(dtype)
+        .vortex_expect("The primitive sum kernel accepts only supported dtypes");
+
+    partial_scalar(aggregate_fn, Scalar::zero_value(&sum_dtype), true)
+}
+
+/// Wrap the sum in the aggregate's partial representation without finalizing empty inputs.
+fn partial_scalar(
+    aggregate_fn: &AggregateFnRef,
+    sum: Scalar,
+    is_empty: bool,
+) -> VortexResult<Scalar> {
+    if aggregate_fn.is::<SumV2>() {
+        SumV2::partial_from_sum(sum, is_empty)
+    } else {
+        Ok(sum)
+    }
+}

encodings/runend/src/compute/sum/mod.rs

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Whole-array primitive sums over run-end encoded arrays.
+//!
+//! Both [`Sum`] and [`SumV2`] reduce valid run values and lengths without expanding the runs. Each
+//! aggregate retains its own partial representation and empty-input semantics. Grouped sums and
+//! decimal inputs use the fallback.
+//!
+//! Floating-point multiplication can round differently from repeated addition, as with constant sums.
+//!
+//! [`Sum`]: vortex_array::aggregate_fn::fns::sum::Sum
+//! [`SumV2`]: vortex_array::aggregate_fn::fns::sum_v2::SumV2
+
+mod kernel;
+pub(crate) use kernel::RunEndSumKernel;
+
+mod runs;
+
+#[cfg(test)]
+mod tests;