From 64a6336859b122fe66c8bef43c6b971a8015eb26 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 14 Sep 2026 10:16:13 -0400 Subject: [PATCH 1/3] perf: specialize primitive sums for run-end arrays Signed-off-by: Connor Tsui --- encodings/runend/src/compute/mod.rs | 1 + encodings/runend/src/compute/sum/grouped.rs | 183 +++++++++ encodings/runend/src/compute/sum/mod.rs | 108 +++++ encodings/runend/src/compute/sum/runs.rs | 144 +++++++ encodings/runend/src/compute/sum/tests.rs | 375 ++++++++++++++++++ encodings/runend/src/compute/sum/whole.rs | 110 +++++ encodings/runend/src/lib.rs | 14 + .../src/aggregate_fn/fns/sum_v2/mod.rs | 85 ++++ 8 files changed, 1020 insertions(+) create mode 100644 encodings/runend/src/compute/sum/grouped.rs create mode 100644 encodings/runend/src/compute/sum/mod.rs create mode 100644 encodings/runend/src/compute/sum/runs.rs create mode 100644 encodings/runend/src/compute/sum/tests.rs create mode 100644 encodings/runend/src/compute/sum/whole.rs diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs index fc7fc8804ec..296ab525952 100644 --- a/encodings/runend/src/compute/mod.rs +++ b/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; diff --git a/encodings/runend/src/compute/sum/grouped.rs b/encodings/runend/src/compute/sum/grouped.rs new file mode 100644 index 00000000000..04209171e21 --- /dev/null +++ b/encodings/runend/src/compute/sum/grouped.rs @@ -0,0 +1,183 @@ +// 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::NumericalAggregateOpts; +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 super::sum_options; +use crate::RunEnd; + +impl DynGroupedAggregateKernel for RunEndSumKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = sum_options(aggregate_fn) else { + return Ok(None); + }; + let Some(elements) = groups.elements().as_opt::() 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| { + sum_primitive_groups::(&runs, &ranges, &validity, options) + }); + + let results = results.into_array(); + if aggregate_fn.is::() { + Ok(Some(SumV2::partials_from_sums( + results, + empty_groups, + validity, + ctx, + )?)) + } else { + Ok(Some(results)) + } + } +} + +fn sum_primitive_groups( + runs: &RunEndInputs, + ranges: &GroupRanges, + group_validity: &Mask, + options: &NumericalAggregateOpts, +) -> (PrimitiveArray, BitBuffer) { + let ends = runs.ends.as_slice::(); + + match_each_native_ptype!(runs.values.ptype(), + unsigned: |T| { + sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, + group_validity, runs.offset, add_unsigned_run) + }, + signed: |T| { + sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, + group_validity, runs.offset, add_signed_run) + }, + floating: |T| { + sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, + group_validity, runs.offset, + |sum, value, len| add_float_run(sum, value, len, options.skip_nans)) + } + ) +} + +fn sum_groups( + ends: &[E], + values: &[T], + validity: &Mask, + ranges: &GroupRanges, + group_validity: &Mask, + offset: usize, + add_run: impl Fn(A, T, usize) -> Option, +) -> (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( + ranges: &GroupRanges, + group_validity: &Mask, + offset: usize, + mut sum_group: impl FnMut(Range) -> (Option, 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()) +} diff --git a/encodings/runend/src/compute/sum/mod.rs b/encodings/runend/src/compute/sum/mod.rs new file mode 100644 index 00000000000..91ccfa3a27c --- /dev/null +++ b/encodings/runend/src/compute/sum/mod.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive sums over run-end encoded arrays. +//! +//! 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 +//! locates its first run independently. The shared reduction in [`runs`] clips intersecting runs. +//! The entry points stay separate because grouped aggregation also handles group validity and +//! produces one partial per group, while whole-array aggregation produces a scalar partial. +//! +//! Both [`Sum`] and [`SumV2`] use this reduction. They retain their own partial representations and +//! empty-input semantics, as described in [`SumV2`]. +//! +//! Decimal inputs use the fallback. Floating-point multiplication can round differently from +//! repeated addition, as with constant sums. + +mod grouped; +mod runs; +mod whole; + +use vortex_array::ArrayView; +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::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::RunEnd; +use crate::RunEndArrayExt; +use crate::RunEndArraySlotsExt; + +/// Whole-array and grouped primitive sum kernels for [`RunEnd`]. +#[derive(Debug)] +pub(crate) struct RunEndSumKernel; + +struct RunEndInputs { + ends: PrimitiveArray, + values: PrimitiveArray, + validity: Mask, + offset: usize, +} + +impl RunEndInputs { + /// Skip materializing the children when the array is empty or every run is null. + fn new(array: ArrayView<'_, RunEnd>, ctx: &mut ExecutionCtx) -> VortexResult> { + if array.is_empty() { + return Ok(None); + } + + let validity = array + .values() + .validity()? + .execute_mask(array.values().len(), ctx)?; + if validity.all_false() { + return Ok(None); + } + + let ends = array.ends().clone().execute::(ctx)?; + let values = array.values().clone().execute::(ctx)?; + + Ok(Some(Self { + ends, + values, + validity, + offset: array.offset(), + })) + } +} + +/// Accept both sum IDs so legacy scalar partials and SumV2 struct partials remain supported. +fn sum_options(aggregate_fn: &AggregateFnRef) -> Option<&NumericalAggregateOpts> { + aggregate_fn + .as_opt::() + .or_else(|| aggregate_fn.as_opt::()) +} + +fn empty_partial(aggregate_fn: &AggregateFnRef, dtype: &DType) -> VortexResult { + 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) +} + +fn partial_scalar( + aggregate_fn: &AggregateFnRef, + sum: Scalar, + is_empty: bool, +) -> VortexResult { + if aggregate_fn.is::() { + SumV2::partial_from_sum(sum, is_empty) + } else { + Ok(sum) + } +} + +#[cfg(test)] +mod tests; diff --git a/encodings/runend/src/compute/sum/runs.rs b/encodings/runend/src/compute/sum/runs.rs new file mode 100644 index 00000000000..0e78174bfe4 --- /dev/null +++ b/encodings/runend/src/compute/sum/runs.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Weighted reduction of the valid runs that intersect a logical 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 num_traits::AsPrimitive; +use num_traits::ToPrimitive; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_error::VortexExpect; + +pub(super) fn add_unsigned_run>(sum: u64, value: T, len: usize) -> Option { + value + .as_() + .checked_mul(len as u64) + .and_then(|product| sum.checked_add(product)) +} + +pub(super) fn add_signed_run>(sum: i64, value: T, len: usize) -> Option { + i64::try_from(i128::from(sum) + i128::from(value.as_()) * len as i128).ok() +} + +pub(super) fn add_float_run( + sum: f64, + value: T, + len: usize, + skip_nans: bool, +) -> Option { + if skip_nans && value.is_nan() { + return Some(sum); + } + + let value = ToPrimitive::to_f64(&value).vortex_expect("Float values fit in f64"); + // Fuse the operations so a finite sum can cancel a product that exceeds f64::MAX. + Some(value.mul_add(len as f64, sum)) +} + +/// 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( + ends: &[E], + values: &[T], + cursor: &mut usize, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + 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); + } + + *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( + ends: &[E], + values: &[T], + indices: &[usize], + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, 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_valid_range( + ends: &[E], + values: &[T], + indices: &mut Peekable>, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + let mut sum = A::default(); + if range.is_empty() { + return (Some(sum), true); + } + + let mut is_empty = true; + while let Some(&index) = indices.peek() { + let end = ends[index].as_(); + if end <= range.start { + indices.next(); + continue; + } + + let run_start = if index == 0 { 0 } else { ends[index - 1].as_() }; + let start = run_start.max(range.start); + if start >= range.end { + break; + } + + let overlap_len = end.min(range.end) - start; + is_empty = false; + let Some(next) = add_run(sum, values[index], overlap_len) else { + return (None, false); + }; + sum = next; + if end > range.end { + break; + } + indices.next(); + } + + (Some(sum), is_empty) +} diff --git a/encodings/runend/src/compute/sum/tests.rs b/encodings/runend/src/compute/sum/tests.rs new file mode 100644 index 00000000000..5fb1177546b --- /dev/null +++ b/encodings/runend/src/compute/sum/tests.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::GroupedArray; +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::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +#[cfg(not(codspeed))] +use vortex_array::test_harness::trace::trace_op; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::RunEndSumKernel; +use crate::RunEnd; +use crate::tests::SESSION; + +/// Compare registered dispatch and the direct kernel with a decoded primitive reference. +fn check_sum(array: ArrayRef, options: NumericalAggregateOpts) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decoded = array + .clone() + .execute::(&mut ctx)? + .into_array(); + + for aggregate in [Sum.bind(options), SumV2.bind(options)] { + let mut reference = aggregate.accumulator(array.dtype())?; + reference.accumulate(&decoded, &mut ctx)?; + let expected = reference.finish()?; + let partial = RunEndSumKernel + .aggregate(&aggregate, &array, &mut ctx)? + .ok_or_else(|| vortex_err!("Primitive run-end kernel declined"))?; + let mut direct = aggregate.accumulator(array.dtype())?; + direct.combine_partials(partial)?; + let mut dispatched = aggregate.accumulator(array.dtype())?; + dispatched.accumulate(&array, &mut ctx)?; + + for actual in [direct.finish()?, dispatched.finish()?] { + if expected.as_primitive().is_nan() { + assert!(actual.as_primitive().is_nan()); + } else { + assert_eq!(actual, expected); + } + } + } + + Ok(()) +} + +/// Compare the registered grouped kernels with groups over decoded primitive elements. +fn check_groups( + groups: GroupedArray, + reference: GroupedArray, + options: NumericalAggregateOpts, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let as_array = |groups: &GroupedArray| match groups { + GroupedArray::ListView(array) => array.clone().into_array(), + GroupedArray::FixedSizeList(array) => array.clone().into_array(), + }; + let groups_array = as_array(&groups); + let reference_array = as_array(&reference); + + for aggregate in [Sum.bind(options), SumV2.bind(options)] { + assert!( + RunEndSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_some() + ); + let mut expected = aggregate.accumulator_grouped(reference.elements().dtype())?; + expected.accumulate_list(&reference_array, &mut ctx)?; + let mut actual = aggregate.accumulator_grouped(groups.elements().dtype())?; + actual.accumulate_list(&groups_array, &mut ctx)?; + assert_arrays_eq!(actual.finish()?, expected.finish()?, &mut ctx); + } + + Ok(()) +} + +#[rstest] +#[case::unsigned(buffer![1u64, 3, 7].into_array())] +#[case::signed(buffer![-1i32, 3, -7].into_array())] +#[case::float(buffer![1.25f64, 3.5, 7.75].into_array())] +#[case::nullable(PrimitiveArray::from_option_iter([Some(-3i32), None, Some(7)]).into_array())] +#[case::nulls(PrimitiveArray::from_option_iter([None::; 3]).into_array())] +fn sliced_sums(#[case] values: ArrayRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = + RunEnd::try_new_offset_length(buffer![2u32, 5, 9].into_array(), values, 1, 7, &mut ctx)? + .into_array(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[cfg(not(codspeed))] +#[rstest] +#[case::whole_array(false, Validity::NonNullable)] +#[case::null_runs(true, Validity::NonNullable)] +#[case::null_groups(true, Validity::AllInvalid)] +fn all_invalid_skips_decoding( + #[case] grouped: bool, + #[case] group_validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let value = if matches!(group_validity, Validity::AllInvalid) { + Scalar::from(3i32) + } else { + Scalar::null(DType::Primitive(PType::I32, Nullable)) + }; + let array = RunEnd::try_new( + ConstantArray::new(8u32, 1).into_array(), + ConstantArray::new(value, 1).into_array(), + &mut ctx, + )? + .into_array(); + + let groups = FixedSizeListArray::try_new(array.clone(), 2, group_validity, 4)?.into(); + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + let traced = trace_op(|| -> VortexResult<()> { + if grouped { + assert!( + RunEndSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_some() + ); + } else { + assert!( + RunEndSumKernel + .aggregate(&aggregate, &array, &mut ctx)? + .is_some() + ); + } + Ok(()) + })?; + assert!(!traced.trace.to_string().contains("execute_until")); + } + + Ok(()) +} + +#[rstest] +#[case::overflow(vec![i64::MAX, 1, -1], vec![1u64, 3, 4])] +#[case::underflow(vec![i64::MIN, -1, 1], vec![1u64, 3, 4])] +#[case::positive_cancellation(vec![-i64::MAX, i64::MAX], vec![1u64, 3])] +#[case::negative_cancellation(vec![i64::MAX, -i64::MAX], vec![1u64, 3])] +fn signed_overflow(#[case] values: Vec, #[case] ends: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + PrimitiveArray::from_iter(ends).into_array(), + PrimitiveArray::from_iter(values).into_array(), + &mut ctx, + )? + .into_array(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[rstest] +#[case::product(buffer![u64::MAX].into_array(), buffer![2u64].into_array())] +#[case::addition(buffer![u64::MAX, 1].into_array(), buffer![1u64, 2].into_array())] +fn unsigned_overflow(#[case] values: ArrayRef, #[case] ends: ArrayRef) -> VortexResult<()> { + let array = RunEnd::try_new(ends, values, &mut SESSION.create_execution_ctx())?.into_array(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[rstest] +#[case::nan(buffer![f64::NAN, 1.25, 2.5].into_array())] +#[case::all_nan(buffer![f64::NAN, f64::NAN, f64::NAN].into_array())] +#[case::infinities(buffer![f64::INFINITY, f64::NEG_INFINITY, 2.5].into_array())] +fn floats(#[case] values: ArrayRef, #[values(true, false)] skip_nans: bool) -> VortexResult<()> { + let array = RunEnd::try_new( + buffer![2u64, 4, 7].into_array(), + values, + &mut SESSION.create_execution_ctx(), + )? + .into_array(); + + check_sum(array, NumericalAggregateOpts { skip_nans }) +} + +#[rstest] +fn float_run_product_cancellation( + #[values(1e308, -1e308)] value: f64, + #[values(false, true)] grouped: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + buffer![1u64, 3].into_array(), + buffer![-value, value].into_array(), + &mut ctx, + )? + .into_array(); + + if !grouped { + return check_sum(array, NumericalAggregateOpts::default()); + } + + let groups = + FixedSizeListArray::try_new(array.clone(), 3, Validity::NonNullable, 1)?.into_array(); + let expected = PrimitiveArray::from_option_iter([Some(value)]).into_array(); + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + let mut acc = aggregate.accumulator_grouped(array.dtype())?; + acc.accumulate_list(&groups, &mut ctx)?; + assert_arrays_eq!(acc.finish()?, expected, &mut ctx); + } + + Ok(()) +} + +#[test] +fn empty_and_zero_length_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Zero-length boundary runs must not contribute a NaN or an overflow. + let array = RunEnd::try_new_offset_length( + buffer![2u64, 5, 8].into_array(), + buffer![f64::NAN, 3.0, f64::INFINITY].into_array(), + 2, + 3, + &mut ctx, + )? + .into_array(); + check_sum(array, NumericalAggregateOpts::include_nans())?; + + let empty = RunEnd::try_new( + PrimitiveArray::from_iter(Vec::::new()).into_array(), + PrimitiveArray::from_iter(Vec::::new()).into_array(), + &mut ctx, + )? + .into_array(); + check_sum(empty, NumericalAggregateOpts::default())?; + + let retained = RunEnd::try_new_offset_length( + buffer![2u64].into_array(), + buffer![f64::NAN].into_array(), + 0, + 0, + &mut ctx, + )? + .into_array(); + check_sum(retained, NumericalAggregateOpts::include_nans()) +} + +#[rstest] +#[case::nullable(PrimitiveArray::from_option_iter([Some(3i32), None, Some(5)]).into_array())] +#[case::overflow(buffer![u64::MAX, 1, 2].into_array())] +#[case::floats(buffer![f64::INFINITY, f64::NEG_INFINITY, f64::NAN].into_array())] +fn grouped_sums( + #[case] values: ArrayRef, + #[values(false, true)] fixed_size: bool, + #[values(false, true)] skip_nans: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let elements = + RunEnd::try_new_offset_length(buffer![3u32, 7, 12].into_array(), values, 1, 10, &mut ctx)? + .into_array(); + let decoded = elements + .clone() + .execute::(&mut ctx)? + .into_array(); + let make_groups = |values| -> VortexResult { + if fixed_size { + Ok(FixedSizeListArray::try_new(values, 2, Validity::NonNullable, 5)?.into()) + } else { + Ok(ListViewArray::try_new( + values, + buffer![6u32, 0, 3, 2, 10].into_array(), + buffer![3u32, 5, 4, 0, 0].into_array(), + Validity::from_iter([true, false, true, true, true]), + )? + .into()) + } + }; + check_groups( + make_groups(elements)?, + make_groups(decoded)?, + NumericalAggregateOpts { skip_nans }, + ) +} + +#[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::; 5]).into_array())] +#[case::overflow(1, 8, false, buffer![i64::MAX, 1, -2, i64::MIN, 3].into_array())] +fn consecutive_groups( + #[case] offset: usize, + #[case] size: u32, + #[case] null_groups: bool, + #[case] values: ArrayRef, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let elements = RunEnd::try_new_offset_length( + buffer![3u32, 7, 10, 17, 64].into_array(), + values, + offset, + size as usize * 6, + &mut ctx, + )? + .into_array(); + let decoded = elements + .clone() + .execute::(&mut ctx)? + .into_array(); + let validity = if null_groups { + Validity::from_iter([true, false, true, false, true, true]) + } else { + Validity::NonNullable + }; + + check_groups( + FixedSizeListArray::try_new(elements, size, validity.clone(), 6)?.into(), + FixedSizeListArray::try_new(decoded, size, validity, 6)?.into(), + NumericalAggregateOpts::default(), + ) +} + +#[test] +fn decimal_kernels_decline() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![100i64, 200], + DecimalDType::new(10, 2), + Validity::NonNullable, + ) + .into_array(); + let array = RunEnd::try_new(buffer![2u64, 4].into_array(), values, &mut ctx)?.into_array(); + let groups = FixedSizeListArray::try_new(array.clone(), 2, Validity::NonNullable, 2)?.into(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + assert!( + RunEndSumKernel + .aggregate(&aggregate, &array, &mut ctx)? + .is_none() + ); + assert!( + RunEndSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_none() + ); + } + + Ok(()) +} diff --git a/encodings/runend/src/compute/sum/whole.rs b/encodings/runend/src/compute/sum/whole.rs new file mode 100644 index 00000000000..46ed81ad614 --- /dev/null +++ b/encodings/runend/src/compute/sum/whole.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Whole-array aggregation over a single logical range. +//! +//! The first and last runs can be clipped by a slice. No group traversal or shared cursor is needed. + +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::kernels::DynAggregateKernel; +use vortex_array::dtype::DType; +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::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::RunEndInputs; +use super::RunEndSumKernel; +use super::empty_partial; +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_all_valid; +use super::runs::sum_valid_range; +use super::sum_options; +use crate::RunEnd; + +impl DynAggregateKernel for RunEndSumKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = sum_options(aggregate_fn) else { + return Ok(None); + }; + let Some(array) = batch.as_opt::() else { + return Ok(None); + }; + if !batch.dtype().is_primitive() { + return Ok(None); + } + + let Some(runs) = RunEndInputs::new(array, ctx)? else { + return Ok(Some(empty_partial(aggregate_fn, batch.dtype())?)); + }; + let range = runs.offset..runs.offset + batch.len(); + + let (sum, is_empty) = match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| { + sum_primitive::(&runs, range, options) + }); + + Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?)) + } +} + +fn sum_primitive( + runs: &RunEndInputs, + range: Range, + options: &NumericalAggregateOpts, +) -> (Scalar, bool) { + let ends = runs.ends.as_slice::(); + + match_each_native_ptype!(runs.values.ptype(), + unsigned: |T| { + sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, add_unsigned_run) + }, + signed: |T| { + sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, add_signed_run) + }, + floating: |T| { + sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, + |sum, value, len| add_float_run(sum, value, len, options.skip_nans)) + } + ) +} + +fn sum_scalar>( + ends: &[E], + values: &[T], + validity: &Mask, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Scalar, bool) { + 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)), + }; + + (sum, is_empty) +} diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index b991609c19c..d704e473cf9 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -33,6 +33,8 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::fns::min_max::MinMax; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; @@ -58,6 +60,18 @@ pub fn initialize(session: &VortexSession) { Some(IsSorted.id()), &compute::is_sorted::RunEndIsSortedKernel, ); + for sum in [Sum.id(), SumV2.id()] { + session.aggregate_fns().register_aggregate_kernel( + RunEnd.id(), + Some(sum), + &compute::sum::RunEndSumKernel, + ); + session.aggregate_fns().register_grouped_encoding_kernel( + RunEnd.id(), + sum, + &compute::sum::RunEndSumKernel, + ); + } } #[cfg(test)] diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index e1287448a8f..9bee8a9718c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -4,10 +4,12 @@ mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; +use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -16,6 +18,7 @@ use crate::ArrayView; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; @@ -28,13 +31,17 @@ use crate::aggregate_fn::fns::sum::accumulate_decimal; use crate::aggregate_fn::fns::sum::accumulate_primitive; use crate::aggregate_fn::fns::sum::make_zero_state; use crate::aggregate_fn::fns::sum::multiply_constant; +use crate::arrays::BoolArray; +use crate::arrays::PrimitiveArray; use crate::arrays::Struct; +use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -74,6 +81,84 @@ pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult #[derive(Clone, Copy, Debug)] pub struct SumV2; +impl SumV2 { + /// Build an encoding kernel's partial from a widened primitive sum. + /// + /// `sum` must have dtype `u64`, `i64`, or `f64`. A null sum records overflow. Set `is_empty` + /// only when there were no valid inputs. Valid NaNs make the input non-empty even when skipped. + pub fn partial_from_sum(sum: Scalar, is_empty: bool) -> VortexResult { + validate_primitive_sum_dtype(sum.dtype())?; + let sum_dtype = sum.dtype().as_nonnullable(); + let is_overflow = sum.is_null(); + let sum = if is_overflow { + Scalar::zero_value(&sum_dtype) + } else { + sum.cast(&sum_dtype)? + }; + + Ok(Scalar::struct_( + sum_v2_partial_dtype(sum_dtype), + vec![ + sum, + Scalar::bool(is_overflow, Nullability::NonNullable), + Scalar::bool(is_empty && !is_overflow, Nullability::NonNullable), + ], + )) + } + + /// Build grouped encoding-kernel partials from widened primitive sums and empty flags. + /// + /// Each input has one entry per group. Sums follow the rules in [`Self::partial_from_sum`]. + /// `is_empty` records groups without valid inputs, including zero-length groups. + /// Null groups are represented by `group_validity` and their sum and empty flag are ignored. + pub fn partials_from_sums( + sums: ArrayRef, + is_empty: BitBuffer, + group_validity: Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + validate_primitive_sum_dtype(sums.dtype())?; + vortex_ensure!( + sums.len() == group_validity.len() && is_empty.len() == group_validity.len(), + "Expected one sum and empty flag per group ({}), got {} sums and {} flags", + group_validity.len(), + sums.len(), + is_empty.len(), + ); + let sum_dtype = sums.dtype().as_nonnullable(); + let sums = sums.execute::(ctx)?.into_data_parts(); + let is_overflow = !sums.validity.execute_mask(group_validity.len(), ctx)?; + + // Overflow and null groups ignore the sum payload, so its physical values can be retained. + let sums = + PrimitiveArray::from_buffer_handle(sums.buffer, sums.ptype, Validity::NonNullable); + + Ok(StructArray::try_new_with_dtype( + vec![ + sums.into_array(), + BoolArray::new(is_overflow.to_bit_buffer(), Validity::NonNullable).into_array(), + BoolArray::new(is_empty, Validity::NonNullable).into_array(), + ], + sum_v2_partial_fields(sum_dtype), + group_validity.len(), + Validity::from_mask(group_validity, Nullability::Nullable), + )? + .into_array()) + } +} + +fn validate_primitive_sum_dtype(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!( + dtype, + DType::Primitive(PType::U64 | PType::I64 | PType::F64, _) + ), + "Expected a widened primitive sum, got {}", + dtype, + ); + Ok(()) +} + impl AggregateFnVTable for SumV2 { type Options = NumericalAggregateOpts; type Partial = SumV2Partial; From 7bcd47b5a6aba40d28e831e50dc0d58686539db4 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 14 Sep 2026 10:16:29 -0400 Subject: [PATCH 2/3] refactor: share run-end sum dispatch and range reduction Signed-off-by: Connor Tsui --- encodings/runend/src/compute/sum/grouped.rs | 68 ++++++--------------- encodings/runend/src/compute/sum/mod.rs | 32 ++++++++++ encodings/runend/src/compute/sum/runs.rs | 23 ++++++- encodings/runend/src/compute/sum/whole.rs | 44 ++----------- 4 files changed, 79 insertions(+), 88 deletions(-) diff --git a/encodings/runend/src/compute/sum/grouped.rs b/encodings/runend/src/compute/sum/grouped.rs index 04209171e21..a20e3ad0513 100644 --- a/encodings/runend/src/compute/sum/grouped.rs +++ b/encodings/runend/src/compute/sum/grouped.rs @@ -15,7 +15,6 @@ 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::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; use vortex_array::arrays::ConstantArray; @@ -24,23 +23,20 @@ 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::AllOr; use vortex_mask::Mask; use super::RunEndInputs; use super::RunEndSumKernel; +use super::dispatch_sum; 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 super::runs::sum_range; use super::sum_options; use crate::RunEnd; @@ -76,8 +72,16 @@ impl DynGroupedAggregateKernel for RunEndSumKernel { let ranges = groups.group_ranges(ctx)?; - let (results, empty_groups) = match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| { - sum_primitive_groups::(&runs, &ranges, &validity, options) + let (results, empty_groups) = dispatch_sum!(&runs, options, |ends, values, add_run| { + sum_groups( + ends, + values, + &runs.validity, + &ranges, + &validity, + runs.offset, + add_run, + ) }); let results = results.into_array(); @@ -94,31 +98,6 @@ impl DynGroupedAggregateKernel for RunEndSumKernel { } } -fn sum_primitive_groups( - runs: &RunEndInputs, - ranges: &GroupRanges, - group_validity: &Mask, - options: &NumericalAggregateOpts, -) -> (PrimitiveArray, BitBuffer) { - let ends = runs.ends.as_slice::(); - - match_each_native_ptype!(runs.values.ptype(), - unsigned: |T| { - sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, - group_validity, runs.offset, add_unsigned_run) - }, - signed: |T| { - sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, - group_validity, runs.offset, add_signed_run) - }, - floating: |T| { - sum_groups(ends, runs.values.as_slice::(), &runs.validity, ranges, - group_validity, runs.offset, - |sum, value, len| add_float_run(sum, value, len, options.skip_nans)) - } - ) -} - fn sum_groups( ends: &[E], values: &[T], @@ -128,32 +107,25 @@ fn sum_groups( offset: usize, add_run: impl Fn(A, T, usize) -> Option, ) -> (PrimitiveArray, BitBuffer) { - match (validity, ranges) { - (Mask::AllTrue(_), GroupRanges::FixedSizeList { .. }) => { + match (validity.indices(), ranges) { + (AllOr::All, 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, |_| { + (AllOr::None, _) => collect_group_sums(ranges, group_validity, offset, |_| { (Some(A::default()), true) }), - (Mask::Values(validity), GroupRanges::FixedSizeList { .. }) => { - let mut indices = validity.indices().iter().copied().peekable(); + (AllOr::Some(indices), GroupRanges::FixedSizeList { .. }) => { + let mut indices = 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(); + (valid_runs, GroupRanges::ListView { .. }) => { collect_group_sums(ranges, group_validity, offset, |range| { - sum_valid_range(ends, values, indices, range, &add_run) + sum_range(ends, values, &valid_runs, range, &add_run) }) } } diff --git a/encodings/runend/src/compute/sum/mod.rs b/encodings/runend/src/compute/sum/mod.rs index 91ccfa3a27c..86f6ccd92f8 100644 --- a/encodings/runend/src/compute/sum/mod.rs +++ b/encodings/runend/src/compute/sum/mod.rs @@ -10,6 +10,7 @@ //! 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 //! locates its first run independently. The shared reduction in [`runs`] clips intersecting runs. +//! Whole-array sums and list-view groups use [`runs::sum_range`] for independent ranges. //! The entry points stay separate because grouped aggregation also handles group validity and //! produces one partial per group, while whole-array aggregation produces a scalar partial. //! @@ -44,6 +45,37 @@ use crate::RunEndArraySlotsExt; #[derive(Debug)] pub(crate) struct RunEndSumKernel; +/// Bind typed run slices and their arithmetic operation for either aggregation entry point. +macro_rules! dispatch_sum { + ($runs:expr, $options:expr, |$ends:ident, $values:ident, $add_run:ident| $body:block) => {{ + let runs = $runs; + let options = $options; + vortex_array::match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| { + let $ends = runs.ends.as_slice::(); + vortex_array::match_each_native_ptype!(runs.values.ptype(), + unsigned: |T| { + let $values = runs.values.as_slice::(); + let $add_run = $crate::compute::sum::runs::add_unsigned_run::; + $body + }, + signed: |T| { + let $values = runs.values.as_slice::(); + let $add_run = $crate::compute::sum::runs::add_signed_run::; + $body + }, + floating: |T| { + let $values = runs.values.as_slice::(); + let $add_run = |sum, value, len| { + $crate::compute::sum::runs::add_float_run(sum, value, len, options.skip_nans) + }; + $body + } + ) + }) + }}; +} +use dispatch_sum; + struct RunEndInputs { ends: PrimitiveArray, values: PrimitiveArray, diff --git a/encodings/runend/src/compute/sum/runs.rs b/encodings/runend/src/compute/sum/runs.rs index 0e78174bfe4..db6d49fe91b 100644 --- a/encodings/runend/src/compute/sum/runs.rs +++ b/encodings/runend/src/compute/sum/runs.rs @@ -16,6 +16,7 @@ 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>(sum: u64, value: T, len: usize) -> Option { value @@ -43,6 +44,26 @@ pub(super) fn add_float_run( Some(value.mul_add(len as f64, sum)) } +/// Sum an independent range using valid run indices prepared by the caller. +/// +/// A `None` sum records overflow. The flag is true when the range contains no valid values. +pub(super) fn sum_range( + ends: &[E], + values: &[T], + valid_runs: &AllOr<&[usize]>, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + match valid_runs { + AllOr::All => { + let mut cursor = ends.partition_point(|end| end.as_() <= range.start); + sum_all_valid(ends, values, &mut cursor, range, add_run) + } + AllOr::None => (Some(A::default()), true), + AllOr::Some(indices) => sum_valid_range(ends, values, indices, range, add_run), + } +} + /// 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 @@ -84,7 +105,7 @@ pub(super) fn sum_all_valid( } /// Locate the first intersecting valid run before summing an arbitrary range. -pub(super) fn sum_valid_range( +fn sum_valid_range( ends: &[E], values: &[T], indices: &[usize], diff --git a/encodings/runend/src/compute/sum/whole.rs b/encodings/runend/src/compute/sum/whole.rs index 46ed81ad614..cc3f526e045 100644 --- a/encodings/runend/src/compute/sum/whole.rs +++ b/encodings/runend/src/compute/sum/whole.rs @@ -10,14 +10,11 @@ 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::kernels::DynAggregateKernel; use vortex_array::dtype::DType; 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::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -25,13 +22,10 @@ use vortex_mask::Mask; use super::RunEndInputs; use super::RunEndSumKernel; +use super::dispatch_sum; use super::empty_partial; 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_all_valid; -use super::runs::sum_valid_range; +use super::runs::sum_range; use super::sum_options; use crate::RunEnd; @@ -57,35 +51,14 @@ impl DynAggregateKernel for RunEndSumKernel { }; let range = runs.offset..runs.offset + batch.len(); - let (sum, is_empty) = match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| { - sum_primitive::(&runs, range, options) + let (sum, is_empty) = dispatch_sum!(&runs, options, |ends, values, add_run| { + sum_scalar(ends, values, &runs.validity, range, add_run) }); Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?)) } } -fn sum_primitive( - runs: &RunEndInputs, - range: Range, - options: &NumericalAggregateOpts, -) -> (Scalar, bool) { - let ends = runs.ends.as_slice::(); - - match_each_native_ptype!(runs.values.ptype(), - unsigned: |T| { - sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, add_unsigned_run) - }, - signed: |T| { - sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, add_signed_run) - }, - floating: |T| { - sum_scalar(ends, runs.values.as_slice::(), &runs.validity, range, - |sum, value, len| add_float_run(sum, value, len, options.skip_nans)) - } - ) -} - fn sum_scalar>( ends: &[E], values: &[T], @@ -93,14 +66,7 @@ fn sum_scalar>( range: Range, add_run: impl Fn(A, T, usize) -> Option, ) -> (Scalar, bool) { - 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, is_empty) = sum_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)), From d526f93fb1adbb5c66c81c67cf253ebc8c49aea0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 14 Sep 2026 10:16:34 -0400 Subject: [PATCH 3/3] bench: cover representative run-end sums Signed-off-by: Connor Tsui --- encodings/runend/Cargo.toml | 4 + encodings/runend/benches/run_end_sum.rs | 132 ++++++++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 encodings/runend/benches/run_end_sum.rs diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 5e607b78cc6..72a4581e5a6 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -58,3 +58,7 @@ harness = false [[bench]] name = "run_end_filter" harness = false + +[[bench]] +name = "run_end_sum" +harness = false diff --git a/encodings/runend/benches/run_end_sum.rs b/encodings/runend/benches/run_end_sum.rs new file mode 100644 index 00000000000..0fc6a823f9f --- /dev/null +++ b/encodings/runend/benches/run_end_sum.rs @@ -0,0 +1,132 @@ +// 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::DynGroupedAccumulator; +use vortex_array::aggregate_fn::GroupedAccumulator; +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::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +const LEN: usize = 2_048; +const RUN_LENGTH: usize = 64; + +static SESSION: LazyLock = 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 % 100).unwrap())), + ); + RunEnd::try_new( + ends.into_array(), + values.into_array(), + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array() +} + +fn bench_fixed_size_group_sum(bencher: Bencher, elements: ArrayRef, group_size: u32) { + 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::(ctx) + .unwrap() + }); +} + +#[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()); +} + +#[divan::bench(args = [Validity::AllValid, Validity::AllInvalid])] +fn whole_array_sum_uniform_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(consts = [2, 128])] +fn fixed_size_group_sum_partially_valid(bencher: Bencher) { + bench_fixed_size_group_sum(bencher, runend_with_null_runs(), GROUP_SIZE); +} + +fn runend_with_validity(validity: &Validity) -> ArrayRef { + let ends = + PrimitiveArray::from_iter((RUN_LENGTH..=LEN).step_by(RUN_LENGTH).map(|end| end as u64)); + let values = PrimitiveArray::new( + (0..ends.len()) + .map(|index| i32::try_from(index).unwrap()) + .collect::>(), + 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 fixed_size_group_sum_all_valid(bencher: Bencher) { + bench_fixed_size_group_sum( + bencher, + runend_with_validity(&Validity::AllValid), + GROUP_SIZE, + ); +}