perf: scan all-valid run-end sums directly

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

encodings/runend/benches/run_end_sum.rs

@@ -150,20 +150,7 @@ fn sum_runend_fallback(bencher: Bencher, run_length: usize) {
 
 #[divan::bench(args = [Validity::NonNullable, Validity::AllValid, Validity::AllInvalid])]
 fn sum_runend_validity(bencher: Bencher, validity: &Validity) {
-    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(),
-    );
-    let array = RunEnd::try_new(
-        ends.into_array(),
-        values.into_array(),
-        &mut SESSION.create_execution_ctx(),
-    )
-    .unwrap()
-    .into_array();
+    let array = runend_with_validity(validity);
 
     bencher
         .with_inputs(|| SESSION.create_execution_ctx())
@@ -179,3 +166,30 @@ fn grouped_runend<const GROUP_SIZE: u32>(bencher: Bencher, run_length: usize) {
 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/sum/grouped.rs

@@ -9,7 +9,6 @@
 
 use std::ops::Range;
 
-use itertools::Either;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
@@ -31,7 +30,6 @@ use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
 use vortex_error::VortexResult;
-use vortex_mask::AllOr;
 use vortex_mask::Mask;
 
 use super::RunEndInputs;
@@ -40,8 +38,9 @@ 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_next_range;
-use super::runs::sum_range;
+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 {
@@ -78,21 +77,20 @@ impl DynGroupedAggregateKernel for RunEndSumKernel {
         };
 
         let ranges = groups.group_ranges(ctx)?;
-        let valid_runs = runs.validity.indices();
 
         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>(), valid_runs, &ranges,
+                    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>(), valid_runs, &ranges,
+                    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>(), valid_runs, &ranges,
+                    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))
                 }
@@ -116,30 +114,38 @@ impl DynGroupedAggregateKernel for RunEndSumKernel {
 fn sum_groups<E: IntegerPType, T: NativePType, A: NativePType>(
     ends: &[E],
     values: &[T],
-    validity: AllOr<&[usize]>,
+    validity: &Mask,
     ranges: &GroupRanges,
     group_validity: &Mask,
     offset: usize,
     add_run: impl Fn(A, T, usize) -> Option<A>,
 ) -> (PrimitiveArray, BitBuffer) {
-    match ranges {
-        // Consecutive groups reuse the cursor, including a run split across group boundaries.
-        GroupRanges::FixedSizeList { .. } => {
-            let indices = match validity {
-                AllOr::All => Either::Left(0..ends.len()),
-                AllOr::None => Either::Left(0..0),
-                AllOr::Some(indices) => Either::Right(indices.iter().copied()),
-            };
-            let mut indices = indices.peekable();
-
+    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_range(ends, values, &mut indices, range, &add_run)
+                sum_next_valid_range(ends, values, &mut indices, range, &add_run)
             })
         }
-        // These ranges can overlap or go backwards, so each group seeks independently.
-        GroupRanges::ListView { .. } => {
+        (Mask::Values(validity), GroupRanges::ListView { .. }) => {
+            let indices = validity.indices();
             collect_group_sums(ranges, group_validity, offset, |range| {
-                sum_range(ends, values, &validity, range, &add_run)
+                sum_valid_range(ends, values, indices, range, &add_run)
             })
         }
     }

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

@@ -5,6 +5,7 @@
 //!
 //! Empty arrays and all-null inputs return before decoding the children. Otherwise, each valid
 //! run contributes its value multiplied by the length included in the input.
+//! All-valid inputs scan the end and value slices directly. Only partially valid inputs use indices.
 //!
 //! Whole-array sums visit one range, clipped at the array's slice boundaries. Fixed-size groups
 //! share a forward cursor. List-view groups can overlap or arrive out of order, so each group

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

@@ -3,20 +3,19 @@
 
 //! Weighted reduction of the valid runs that intersect a logical range.
 //!
-//! Seeking and forward traversal share one loop that clips each run to the requested range.
+//! All-valid inputs traverse the end and value slices directly. Partially valid inputs visit only
+//! their valid run indices. Both paths clip the boundary runs to the requested range.
 //! Signed arithmetic widens the product, and floating-point arithmetic uses fused multiply-add,
 //! so a run can cancel a preceding sum even when its product alone exceeds the result type.
 
 use std::iter::Peekable;
 use std::ops::Range;
 
-use itertools::Either;
 use num_traits::AsPrimitive;
 use num_traits::ToPrimitive;
 use vortex_array::dtype::IntegerPType;
 use vortex_array::dtype::NativePType;
 use vortex_error::VortexExpect;
-use vortex_mask::AllOr;
 
 pub(super) fn add_unsigned_run<T: AsPrimitive<u64>>(sum: u64, value: T, len: usize) -> Option<u64> {
     value
@@ -44,32 +43,66 @@ pub(super) fn add_float_run<T: NativePType>(
     Some(value.mul_add(len as f64, sum))
 }
 
-/// Locate the first intersecting valid run before summing a range.
-pub(super) fn sum_range<E: IntegerPType, T: NativePType, A: NativePType>(
+/// Sum an all-valid range directly from the end and value slices.
+///
+/// The cursor is a position in the slices, retained for consecutive groups. Ranges must be ordered
+/// and non-overlapping. Arbitrary ranges must first position the cursor with a binary search.
+pub(super) fn sum_all_valid<E: IntegerPType, T: NativePType, A: NativePType>(
     ends: &[E],
     values: &[T],
-    validity: &AllOr<&[usize]>,
+    cursor: &mut usize,
     range: Range<usize>,
     add_run: impl Fn(A, T, usize) -> Option<A>,
 ) -> (Option<A>, bool) {
-    let first = ends.partition_point(|end| end.as_() <= range.start);
-    let indices = match validity {
-        AllOr::All => Either::Left(first..ends.len()),
-        AllOr::None => return (Some(A::default()), true),
-        AllOr::Some(indices) => {
-            let start = indices.partition_point(|&index| index < first);
-            Either::Right(indices[start..].iter().copied())
+    let mut sum = A::default();
+    if range.is_empty() {
+        return (Some(sum), true);
+    }
+
+    // Skipped null groups or an earlier overflow can leave the cursor behind this range.
+    while ends[*cursor].as_() <= range.start {
+        *cursor += 1;
+    }
+
+    let mut start = range.start;
+    for (&end, &value) in ends[*cursor..].iter().zip(&values[*cursor..]) {
+        let end = end.as_();
+        if end >= range.end {
+            *cursor += usize::from(end == range.end);
+            return (add_run(sum, value, range.end - start), false);
         }
-    };
 
-    sum_next_range(ends, values, &mut indices.peekable(), range, add_run)
+        *cursor += 1;
+        let Some(next) = add_run(sum, value, end - start) else {
+            return (None, false);
+        };
+        sum = next;
+        start = end;
+    }
+
+    (Some(sum), false)
+}
+
+/// Locate the first intersecting valid run before summing an arbitrary range.
+pub(super) fn sum_valid_range<E: IntegerPType, T: NativePType, A: NativePType>(
+    ends: &[E],
+    values: &[T],
+    indices: &[usize],
+    range: Range<usize>,
+    add_run: impl Fn(A, T, usize) -> Option<A>,
+) -> (Option<A>, bool) {
+    let first = ends.partition_point(|end| end.as_() <= range.start);
+    let start = indices.partition_point(|&index| index < first);
+    let mut indices = indices[start..].iter().copied().peekable();
+
+    sum_next_valid_range(ends, values, &mut indices, range, add_run)
 }
 
 /// Sum the next range while retaining a run that crosses its end.
 ///
 /// The caller must supply non-overlapping ranges in increasing order. Skipped null groups and
 /// early overflow returns can leave the cursor behind the next range's start.
-pub(super) fn sum_next_range<E: IntegerPType, T: NativePType, A: NativePType>(
+pub(super) fn sum_next_valid_range<E: IntegerPType, T: NativePType, A: NativePType>(
     ends: &[E],
     values: &[T],
     indices: &mut Peekable<impl Iterator<Item = usize>>,

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

@@ -306,6 +306,7 @@ fn grouped_sums(
 #[rstest]
 #[case::empty(0, 0, false, buffer![1i64, -2, 3, -4, 5].into_array())]
 #[case::short_groups(0, 2, false, buffer![1i64, -2, 3, -4, 5].into_array())]
+#[case::all_valid_with_null_groups(1, 8, true, buffer![1i64, -2, 3, -4, 5].into_array())]
 #[case::sliced_null_groups(1, 8, true,
     PrimitiveArray::from_option_iter([None, Some(2i64), None, Some(4), Some(5)]).into_array())]
 #[case::all_null(1, 8, true, PrimitiveArray::from_option_iter([None::<i64>; 5]).into_array())]

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

@@ -22,7 +22,7 @@ use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::scalar::PValue;
 use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
-use vortex_mask::AllOr;
+use vortex_mask::Mask;
 
 use super::RunEndInputs;
 use super::RunEndSumKernel;
@@ -31,7 +31,8 @@ use super::partial_scalar;
 use super::runs::add_float_run;
 use super::runs::add_signed_run;
 use super::runs::add_unsigned_run;
-use super::runs::sum_range;
+use super::runs::sum_all_valid;
+use super::runs::sum_valid_range;
 use crate::RunEnd;
 
 impl DynAggregateKernel for RunEndSumKernel {
@@ -58,19 +59,18 @@ impl DynAggregateKernel for RunEndSumKernel {
             return Ok(Some(empty_partial(aggregate_fn, batch.dtype())?));
         };
         let range = runs.offset..runs.offset + batch.len();
-        let valid_runs = runs.validity.indices();
 
         let (sum, is_empty) = 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_scalar(ends, runs.values.as_slice::<T>(), &valid_runs, range, add_unsigned_run)
+                    sum_scalar(ends, runs.values.as_slice::<T>(), &runs.validity, range, add_unsigned_run)
                 },
                 signed: |T| {
-                    sum_scalar(ends, runs.values.as_slice::<T>(), &valid_runs, range, add_signed_run)
+                    sum_scalar(ends, runs.values.as_slice::<T>(), &runs.validity, range, add_signed_run)
                 },
                 floating: |T| {
-                    sum_scalar(ends, runs.values.as_slice::<T>(), &valid_runs, range,
+                    sum_scalar(ends, runs.values.as_slice::<T>(), &runs.validity, range,
                         |sum, value, len| add_float_run(sum, value, len, options.skip_nans))
                 }
             )
@@ -83,11 +83,18 @@ impl DynAggregateKernel for RunEndSumKernel {
 fn sum_scalar<E: IntegerPType, T: NativePType, A: NativePType + Into<PValue>>(
     ends: &[E],
     values: &[T],
-    validity: &AllOr<&[usize]>,
+    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, is_empty) = match validity {
+        Mask::AllTrue(_) => {
+            let mut cursor = ends.partition_point(|end| end.as_() <= range.start);
+            sum_all_valid(ends, values, &mut cursor, range, add_run)
+        }
+        Mask::AllFalse(_) => (Some(A::default()), true),
+        Mask::Values(validity) => sum_valid_range(ends, values, validity.indices(), range, add_run),
+    };
     let sum = match sum {
         Some(sum) => Scalar::primitive(sum, Nullable),
         None => Scalar::null(DType::Primitive(A::PTYPE, Nullable)),