feat: specialize primitive sums for constant and 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,181 @@
+// 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::ArrayVTable;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::AggregateFnVTable;
+use vortex_array::aggregate_fn::DynGroupedAccumulator;
+use vortex_array::aggregate_fn::GroupedAccumulator;
+use vortex_array::aggregate_fn::GroupedArray;
+use vortex_array::aggregate_fn::NumericalAggregateOpts;
+use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
+use vortex_array::aggregate_fn::fns::sum_v2::sum_v2;
+use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
+use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel;
+use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
+use vortex_array::arrays::Constant;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::FixedSizeListArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::scalar::Scalar;
+use vortex_array::validity::Validity;
+use vortex_error::VortexResult;
+use vortex_runend::RunEnd;
+use vortex_session::VortexSession;
+
+const LEN: usize = 16_384;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_runend::initialize(&session);
+    session
+});
+
+static FALLBACK_SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_runend::initialize(&session);
+    session
+        .aggregate_fns()
+        .register_aggregate_kernel(RunEnd.id(), Some(SumV2.id()), &Decline);
+    for encoding in [RunEnd.id(), Constant.id()] {
+        session
+            .aggregate_fns()
+            .register_grouped_encoding_kernel(encoding, SumV2.id(), &Decline);
+    }
+    session
+});
+
+/// Keep the pre-specialization dispatch paths available for benchmark comparisons.
+#[derive(Debug)]
+struct Decline;
+
+impl DynAggregateKernel for Decline {
+    fn aggregate(
+        &self,
+        _aggregate_fn: &AggregateFnRef,
+        _batch: &ArrayRef,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Scalar>> {
+        Ok(None)
+    }
+}
+
+impl DynGroupedAggregateKernel for Decline {
+    fn grouped_aggregate(
+        &self,
+        _aggregate_fn: &AggregateFnRef,
+        _groups: &GroupedArray,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        Ok(None)
+    }
+}
+
+fn main() {
+    LazyLock::force(&SESSION);
+    LazyLock::force(&FALLBACK_SESSION);
+    divan::main();
+}
+
+fn runend(run_length: usize) -> 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 % 100).unwrap())),
+    );
+    RunEnd::try_new(
+        ends.into_array(),
+        values.into_array(),
+        &mut SESSION.create_execution_ctx(),
+    )
+    .unwrap()
+    .into_array()
+}
+
+fn bench_sum(bencher: Bencher, run_length: usize, session: &VortexSession) {
+    let array = runend(run_length);
+    bencher
+        .with_inputs(|| session.create_execution_ctx())
+        .bench_refs(|ctx| sum_v2(&array, ctx).unwrap());
+}
+
+fn bench_grouped(bencher: Bencher, elements: ArrayRef, group_size: u32, session: &VortexSession) {
+    let dtype = elements.dtype().clone();
+    let groups = FixedSizeListArray::try_new(
+        elements,
+        group_size,
+        Validity::NonNullable,
+        LEN / group_size as usize,
+    )
+    .unwrap()
+    .into_array();
+    bencher
+        .with_inputs(|| {
+            (
+                GroupedAccumulator::try_new(
+                    SumV2,
+                    NumericalAggregateOpts::default(),
+                    dtype.clone(),
+                )
+                .unwrap(),
+                session.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(acc, ctx)| {
+            acc.accumulate_list(&groups, ctx).unwrap();
+            acc.finish()
+                .unwrap()
+                .execute::<PrimitiveArray>(ctx)
+                .unwrap()
+        });
+}
+
+#[divan::bench(args = [4, 64, 1024])]
+fn sum_runend(bencher: Bencher, run_length: usize) {
+    bench_sum(bencher, run_length, &SESSION);
+}
+
+#[divan::bench(args = [4, 64, 1024])]
+fn sum_runend_fallback(bencher: Bencher, run_length: usize) {
+    bench_sum(bencher, run_length, &FALLBACK_SESSION);
+}
+
+#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])]
+fn grouped_runend<const GROUP_SIZE: u32>(bencher: Bencher, run_length: usize) {
+    bench_grouped(bencher, runend(run_length), GROUP_SIZE, &SESSION);
+}
+
+#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])]
+fn grouped_runend_fallback<const GROUP_SIZE: u32>(bencher: Bencher, run_length: usize) {
+    bench_grouped(bencher, runend(run_length), GROUP_SIZE, &FALLBACK_SESSION);
+}
+
+#[divan::bench(consts = [1, 2, 8, 128])]
+fn grouped_constant<const GROUP_SIZE: u32>(bencher: Bencher) {
+    bench_grouped(
+        bencher,
+        ConstantArray::new(3i32, LEN).into_array(),
+        GROUP_SIZE,
+        &SESSION,
+    );
+}
+
+#[divan::bench(consts = [1, 2, 8, 128])]
+fn grouped_constant_fallback<const GROUP_SIZE: u32>(bencher: Bencher) {
+    bench_grouped(
+        bencher,
+        ConstantArray::new(3i32, LEN).into_array(),
+        GROUP_SIZE,
+        &FALLBACK_SESSION,
+    );
+}

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/mod.rs

@@ -0,0 +1,109 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Primitive sums over run-end encoded arrays.
+//!
+//! The kernels decode the run ends and values once, then sum each value weighted by its run
+//! length. Grouped sums intersect runs with each group's range. Decimal inputs use the fallback.
+//! Float multiplication can round differently from repeated addition, as with constant sums.
+
+mod primitive;
+
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::GroupedArray;
+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::aggregate_fn::kernels::DynGroupedAggregateKernel;
+use vortex_array::scalar::Scalar;
+use vortex_error::VortexResult;
+
+use self::primitive::RunEndSums;
+use crate::RunEnd;
+
+/// Whole-array and grouped primitive sum kernels 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>> {
+        let Some(options) = aggregate_fn
+            .as_opt::<Sum>()
+            .or_else(|| aggregate_fn.as_opt::<SumV2>())
+        else {
+            return Ok(None);
+        };
+        let Some(array) = batch.as_opt::<RunEnd>() else {
+            return Ok(None);
+        };
+        if !batch.dtype().is_primitive() {
+            return Ok(None);
+        }
+
+        let sums = RunEndSums::new(array, ctx, options.skip_nans)?;
+        let (sum, is_empty) = sums.sum(0..batch.len());
+        Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?))
+    }
+}
+
+impl DynGroupedAggregateKernel for RunEndSumKernel {
+    fn grouped_aggregate(
+        &self,
+        aggregate_fn: &AggregateFnRef,
+        groups: &GroupedArray,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let Some(options) = aggregate_fn
+            .as_opt::<Sum>()
+            .or_else(|| aggregate_fn.as_opt::<SumV2>())
+        else {
+            return Ok(None);
+        };
+        let Some(elements) = groups.elements().as_opt::<RunEnd>() else {
+            return Ok(None);
+        };
+        if !groups.elements().dtype().is_primitive() {
+            return Ok(None);
+        }
+
+        let ranges = groups.group_ranges(ctx)?;
+        let validity = groups.group_validity(ctx)?;
+        let sums = RunEndSums::new(elements, ctx, options.skip_nans)?;
+        let (results, empty_groups) = sums.grouped_sum(&ranges, &validity);
+
+        let results = results.into_array();
+        if aggregate_fn.is::<SumV2>() {
+            Ok(Some(SumV2::partials_from_sums(
+                results,
+                empty_groups,
+                validity,
+                ctx,
+            )?))
+        } else {
+            Ok(Some(results))
+        }
+    }
+}
+
+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)
+    }
+}
+
+#[cfg(test)]
+mod tests;