perf: specialize primitive 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,195 @@
+// 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::FixedSizeListArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::scalar::Scalar;
+use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_error::VortexResult;
+use vortex_runend::RunEnd;
+use vortex_session::VortexSession;
+
+// Keep the one-element-group fallback below 1 ms in CodSpeed simulation.
+const LEN: usize = 2_048;
+
+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);
+    session
+        .aggregate_fns()
+        .register_grouped_encoding_kernel(RunEnd.id(), 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 = [Validity::NonNullable, Validity::AllValid, Validity::AllInvalid])]
+fn sum_runend_validity(bencher: Bencher, validity: &Validity) {
+    let array = runend_with_validity(validity);
+
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| sum_v2(&array, ctx).unwrap());
+}
+
+#[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);
+}
+
+fn runend_with_validity(validity: &Validity) -> ArrayRef {
+    let ends = PrimitiveArray::from_iter((64..=LEN).step_by(64).map(|end| end as u64));
+    let values = PrimitiveArray::new(
+        (0..ends.len())
+            .map(|index| i32::try_from(index).unwrap())
+            .collect::<Buffer<_>>(),
+        validity.clone(),
+    );
+    RunEnd::try_new(
+        ends.into_array(),
+        values.into_array(),
+        &mut SESSION.create_execution_ctx(),
+    )
+    .unwrap()
+    .into_array()
+}
+
+#[divan::bench(consts = [2, 128])]
+fn grouped_runend_all_valid<const GROUP_SIZE: u32>(bencher: Bencher) {
+    bench_grouped(
+        bencher,
+        runend_with_validity(&Validity::AllValid),
+        GROUP_SIZE,
+        &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/grouped.rs

@@ -0,0 +1,175 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Grouped aggregation with traversal selected by the group layout.
+//!
+//! Fixed-size groups share a forward run cursor. List-view ranges can overlap or arrive out of
+//! order, so they locate their runs independently. Both paths weight runs by their intersection
+//! with the group, and skip null groups before visiting any runs.
+
+use std::ops::Range;
+
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::GroupRanges;
+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::DynGroupedAggregateKernel;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::dtype::IntegerPType;
+use vortex_array::dtype::NativePType;
+use vortex_array::dtype::Nullability::Nullable;
+use vortex_array::match_each_native_ptype;
+use vortex_array::match_each_unsigned_integer_ptype;
+use vortex_array::validity::Validity;
+use vortex_buffer::BitBuffer;
+use vortex_buffer::BitBufferMut;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::RunEndInputs;
+use super::RunEndSumKernel;
+use super::empty_partial;
+use super::runs::add_float_run;
+use super::runs::add_signed_run;
+use super::runs::add_unsigned_run;
+use super::runs::sum_all_valid;
+use super::runs::sum_next_valid_range;
+use super::runs::sum_valid_range;
+use crate::RunEnd;
+
+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 validity = groups.group_validity(ctx)?;
+        let runs = if validity.all_false() {
+            None
+        } else {
+            RunEndInputs::new(elements, ctx)?
+        };
+        let Some(runs) = runs else {
+            let partial = empty_partial(aggregate_fn, groups.elements().dtype())?;
+            let partials = ConstantArray::new(partial, groups.len()).into_array();
+            let validity = Validity::from_mask(validity, Nullable).to_array(groups.len());
+            return Ok(Some(partials.mask(validity)?));
+        };
+
+        let ranges = groups.group_ranges(ctx)?;
+
+        let (results, empty_groups) = match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| {
+            let ends = runs.ends.as_slice::<E>();
+            match_each_native_ptype!(runs.values.ptype(),
+                unsigned: |T| {
+                    sum_groups(ends, runs.values.as_slice::<T>(), &runs.validity, &ranges,
+                        &validity, runs.offset, add_unsigned_run)
+                },
+                signed: |T| {
+                    sum_groups(ends, runs.values.as_slice::<T>(), &runs.validity, &ranges,
+                        &validity, runs.offset, add_signed_run)
+                },
+                floating: |T| {
+                    sum_groups(ends, runs.values.as_slice::<T>(), &runs.validity, &ranges,
+                        &validity, runs.offset,
+                        |sum, value, len| add_float_run(sum, value, len, options.skip_nans))
+                }
+            )
+        });
+
+        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 sum_groups<E: IntegerPType, T: NativePType, A: NativePType>(
+    ends: &[E],
+    values: &[T],
+    validity: &Mask,
+    ranges: &GroupRanges,
+    group_validity: &Mask,
+    offset: usize,
+    add_run: impl Fn(A, T, usize) -> Option<A>,
+) -> (PrimitiveArray, BitBuffer) {
+    match (validity, ranges) {
+        (Mask::AllTrue(_), GroupRanges::FixedSizeList { .. }) => {
+            let mut cursor = 0;
+            collect_group_sums(ranges, group_validity, offset, |range| {
+                sum_all_valid(ends, values, &mut cursor, range, &add_run)
+            })
+        }
+        (Mask::AllTrue(_), GroupRanges::ListView { .. }) => {
+            collect_group_sums(ranges, group_validity, offset, |range| {
+                let mut cursor = ends.partition_point(|end| end.as_() <= range.start);
+                sum_all_valid(ends, values, &mut cursor, range, &add_run)
+            })
+        }
+        (Mask::AllFalse(_), _) => collect_group_sums(ranges, group_validity, offset, |_| {
+            (Some(A::default()), true)
+        }),
+        (Mask::Values(validity), GroupRanges::FixedSizeList { .. }) => {
+            let mut indices = validity.indices().iter().copied().peekable();
+            collect_group_sums(ranges, group_validity, offset, |range| {
+                sum_next_valid_range(ends, values, &mut indices, range, &add_run)
+            })
+        }
+        (Mask::Values(validity), GroupRanges::ListView { .. }) => {
+            let indices = validity.indices();
+            collect_group_sums(ranges, group_validity, offset, |range| {
+                sum_valid_range(ends, values, indices, range, &add_run)
+            })
+        }
+    }
+}
+
+fn collect_group_sums<A: NativePType>(
+    ranges: &GroupRanges,
+    group_validity: &Mask,
+    offset: usize,
+    mut sum_group: impl FnMut(Range<usize>) -> (Option<A>, bool),
+) -> (PrimitiveArray, BitBuffer) {
+    let mut empty_groups = BitBufferMut::new_unset(ranges.len());
+    let sums =
+        PrimitiveArray::from_option_iter(ranges.iter().zip(group_validity.iter()).enumerate().map(
+            |(index, ((start, len), valid))| {
+                if !valid {
+                    return None;
+                }
+
+                let (sum, is_empty) = sum_group(offset + start..offset + start + len);
+                empty_groups.set_to(index, is_empty);
+                sum
+            },
+        ));
+
+    (sums, empty_groups.freeze())
+}