From 65f708bd1f1edb52362887d7e6522e76d6bd5217 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 14 Sep 2026 14:14:49 -0400 Subject: [PATCH 1/2] perf: specialize whole-array sums for run-end arrays Compute primitive Sum and SumV2 from run values and lengths without fully decoding the array. Clip boundary runs to the array slice and reuse cached validity indices or scan the bounded bitmap. Grouped sums and decimals continue to use the existing fallback. Cover slice boundaries, nulls, overflow, and NaNs, and add a whole-array benchmark. Weighted floating-point sums can round differently from repeated addition; shared signed-overflow limitations remain unchanged. Signed-off-by: "Connor Tsui" --- encodings/runend/Cargo.toml | 4 + encodings/runend/benches/run_end_sum.rs | 53 ++++ encodings/runend/src/compute/mod.rs | 1 + encodings/runend/src/compute/sum/kernel.rs | 152 ++++++++++ encodings/runend/src/compute/sum/mod.rs | 21 ++ encodings/runend/src/compute/sum/runs.rs | 176 ++++++++++++ encodings/runend/src/compute/sum/tests.rs | 267 ++++++++++++++++++ encodings/runend/src/lib.rs | 9 + .../src/aggregate_fn/fns/sum_v2/mod.rs | 39 ++- 9 files changed, 721 insertions(+), 1 deletion(-) create mode 100644 encodings/runend/benches/run_end_sum.rs create mode 100644 encodings/runend/src/compute/sum/kernel.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 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..9132d425c48 --- /dev/null +++ b/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 = 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()); +} 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/kernel.rs b/encodings/runend/src/compute/sum/kernel.rs new file mode 100644 index 00000000000..6f65f96e3f1 --- /dev/null +++ b/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> { + if !batch.dtype().is_primitive() { + return Ok(None); + } + + let Some(options) = aggregate_fn + .as_opt::() + .or_else(|| aggregate_fn.as_opt::()) + .copied() + else { + return Ok(None); + }; + + let Some(array) = batch.as_opt::() 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::(ctx)?; + let values = array.values().clone().execute::(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, + options: NumericalAggregateOpts, +) -> (Scalar, bool) { + match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + let ends = ends.as_slice::(); + + match_each_native_ptype!(values.ptype(), + unsigned: |T| { + sum_typed_runs(ends, values.as_slice::(), validity, range, add_unsigned_run::) + }, + signed: |T| { + sum_typed_runs(ends, values.as_slice::(), validity, range, add_signed_run::) + }, + floating: |T| { + sum_typed_runs(ends, values.as_slice::(), 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>( + ends: &[E], + values: &[T], + validity: &Mask, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (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 { + 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 { + if aggregate_fn.is::() { + SumV2::partial_from_sum(sum, is_empty) + } else { + Ok(sum) + } +} diff --git a/encodings/runend/src/compute/sum/mod.rs b/encodings/runend/src/compute/sum/mod.rs new file mode 100644 index 00000000000..6e495ad81a9 --- /dev/null +++ b/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; diff --git a/encodings/runend/src/compute/sum/runs.rs b/encodings/runend/src/compute/sum/runs.rs new file mode 100644 index 00000000000..fd40da33e4b --- /dev/null +++ b/encodings/runend/src/compute/sum/runs.rs @@ -0,0 +1,176 @@ +// 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 reuse +//! cached valid run indices or iterate the validity bitmap without building indices. Both paths +//! clip the boundary runs to the requested range. +//! +//! Ranges use positions in the unsliced array and must be covered by strictly increasing run ends. +//! Ends, values, and validity describe the same number of runs. Each reduction returns +//! `(sum, is_empty)`, where a `None` sum records overflow. Valid NaNs make the input non-empty even +//! when skipped. +//! +//! 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::ops::Range; + +use num_traits::AsPrimitive; +use num_traits::ToPrimitive; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::UnsignedPType; +use vortex_error::VortexExpect; +use vortex_mask::Mask; +use vortex_mask::MaskValues; + +/// Return `None` if the run's product or the accumulated sum exceeds `u64`. +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)) +} + +/// Widen the product and addition to `i128` before checking whether the result fits in `i64`. +pub(super) fn add_signed_run>(sum: i64, value: T, len: usize) -> Option { + let product = i128::from(value.as_()) * len as i128; + + i64::try_from(i128::from(sum) + product).ok() +} + +/// Add a run with fused multiplication, leaving the sum unchanged for skipped NaNs. +/// +/// Always returns `Some`, including for infinite or NaN results, to share the integer callback type. +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 independent range without materializing validity indices. +pub(super) fn sum_range( + ends: &[E], + values: &[T], + validity: &Mask, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + if range.is_empty() { + return (Some(A::default()), true); + } + + match validity { + Mask::AllTrue(_) => sum_all_valid(ends, values, range, add_run), + Mask::AllFalse(_) => (Some(A::default()), true), + Mask::Values(validity) => sum_partially_valid(ends, values, validity, range, add_run), + } +} + +/// Sum an all-valid range directly from the end and value slices. +/// +/// The caller must supply a non-empty range. +fn sum_all_valid( + ends: &[E], + values: &[T], + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + let mut sum = A::default(); + let first = ends.partition_point(|end| end.as_() <= range.start); + let mut start = range.start; + + for (&end, &value) in ends[first..].iter().zip(&values[first..]) { + let end = end.as_(); + if end >= range.end { + return (add_run(sum, value, range.end - start), false); + } + + let Some(next) = add_run(sum, value, end - start) else { + return (None, false); + }; + + sum = next; + start = end; + } + + (Some(sum), false) +} + +/// Sum a non-empty range using cached indices or a bounded validity bitmap. +fn sum_partially_valid( + ends: &[E], + values: &[T], + validity: &MaskValues, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + if let Some(indices) = validity.cached_indices() { + let start = indices.partition_point(|&index| ends[index].as_() <= range.start); + + return sum_indexed_runs( + ends, + values, + indices[start..].iter().copied(), + range, + add_run, + ); + } + + let first = ends.partition_point(|end| end.as_() <= range.start); + let last = ends.partition_point(|end| end.as_() < range.end); + + // Bound the bitmap so finding the next valid run cannot scan beyond the array's slice. + let bits = validity.bit_buffer().slice(first..=last); + let indices = bits.set_indices().map(|index| first + index); + + sum_indexed_runs(ends, values, indices, range, add_run) +} + +/// Sum valid runs in increasing order, clipping them to the range. +/// +/// Every supplied run must end after the range's start. +fn sum_indexed_runs( + ends: &[E], + values: &[T], + indices: impl Iterator, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Option, bool) { + let mut sum = A::default(); + let mut is_empty = true; + + for index in indices { + let end = ends[index].as_(); + 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; + } + } + + (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..4e375613772 --- /dev/null +++ b/encodings/runend/src/compute/sum/tests.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +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::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::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability::NonNullable; +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::BitBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::RunEndSumKernel; +use super::runs::add_unsigned_run; +use super::runs::sum_range; +use crate::RunEnd; +use crate::tests::SESSION; + +/// Compare registered dispatch and the direct kernel with a decoded primitive reference. +#[track_caller] +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)? + .vortex_expect("The fixture has a primitive dtype supported by the run-end sum kernel"); + 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(()) +} + +#[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()) +} + +#[rstest] +#[case::whole(0..9, 24, false)] +#[case::ends_inside_run(0..1, 2, false)] +#[case::ends_at_run_boundary(0..2, 4, false)] +#[case::clipped(1..8, 17, false)] +#[case::starts_in_null_run(2..6, 5, false)] +#[case::starts_at_run_boundary(5..7, 10, false)] +#[case::only_nulls(2..5, 0, true)] +#[case::empty_inside_run(1..1, 0, true)] +#[case::empty_at_end(9..9, 0, true)] +fn range_sum_preserves_validity_representation( + #[case] range: Range, + #[case] expected_sum: u64, + #[case] expected_empty: bool, + #[values(false, true)] cached: bool, +) { + // Exercise a validity bitmap that starts inside a byte. + let bits = BitBuffer::from_iter([false, false, false, false, false, true, false, true]); + let validity = Mask::from_buffer(bits.slice(5..)); + let Mask::Values(mask) = &validity else { + unreachable!("The fixture contains both valid and null runs"); + }; + + if cached { + mask.indices(); + } + + let result = sum_range( + &[2u32, 5, 9], + &[2u64, u64::MAX, 5], + &validity, + range, + add_unsigned_run, + ); + + assert_eq!(result, (Some(expected_sum), expected_empty)); + assert_eq!(mask.cached_indices().is_some(), cached); +} + +#[cfg(not(codspeed))] +#[test] +fn all_invalid_skips_decoding() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + ConstantArray::new(8u32, 1).into_array(), + ConstantArray::new(Scalar::null(DType::Primitive(PType::I32, Nullable)), 1).into_array(), + &mut ctx, + )? + .into_array(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + let traced = trace_op(|| -> VortexResult<()> { + assert!( + RunEndSumKernel + .aggregate(&aggregate, &array, &mut ctx)? + .is_some() + ); + Ok(()) + })?; + + assert!(!traced.trace.to_string().contains("execute_until")); + } + + Ok(()) +} + +#[rstest] +#[case::signed_overflow( + buffer![i64::MAX, 1, -1].into_array(), + buffer![1u64, 3, 4].into_array(), +)] +#[case::signed_underflow( + buffer![i64::MIN, -1, 1].into_array(), + buffer![1u64, 3, 4].into_array(), +)] +#[case::positive_cancellation( + buffer![-i64::MAX, i64::MAX].into_array(), + buffer![1u64, 3].into_array(), +)] +#[case::negative_cancellation( + buffer![i64::MAX, -i64::MAX].into_array(), + buffer![1u64, 3].into_array(), +)] +#[case::unsigned_product(buffer![u64::MAX].into_array(), buffer![2u64].into_array())] +#[case::unsigned_addition(buffer![u64::MAX, 1].into_array(), buffer![1u64, 2].into_array())] +fn integer_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) -> 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(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[rstest] +#[case::clipped_non_finite_runs( + buffer![2u64, 5, 8].into_array(), + buffer![f64::NAN, 3.0, f64::INFINITY].into_array(), + 2..5, +)] +#[case::empty_children( + PrimitiveArray::empty::(NonNullable).into_array(), + PrimitiveArray::empty::(NonNullable).into_array(), + 0..0, +)] +#[case::empty_slice( + buffer![2u64].into_array(), + buffer![f64::NAN].into_array(), + 0..0, +)] +fn empty_arrays_and_clipped_runs( + #[case] ends: ArrayRef, + #[case] values: ArrayRef, + #[case] range: Range, +) -> VortexResult<()> { + let array = RunEnd::try_new_offset_length( + ends, + values, + range.start, + range.len(), + &mut SESSION.create_execution_ctx(), + )? + .into_array(); + + check_sum(array, NumericalAggregateOpts::include_nans()) +} + +#[test] +fn decimal_kernel_declines() -> 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(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + assert!( + RunEndSumKernel + .aggregate(&aggregate, &array, &mut ctx)? + .is_none() + ); + } + + Ok(()) +} diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index b991609c19c..b9afd20151c 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,13 @@ 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, + ); + } } #[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..fccbdcbe9d8 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod grouped; +//! SQL-style sums with explicit overflow and empty-input state. +mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -35,6 +36,7 @@ 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 +76,41 @@ 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 and overrides + /// `is_empty`. Set `is_empty` only when there were no valid inputs. Valid NaNs make the input + /// non-empty even when skipped. The returned struct and its fields are non-null. + pub fn partial_from_sum(sum: Scalar, is_empty: bool) -> VortexResult { + vortex_ensure!( + matches!( + sum.dtype(), + DType::Primitive(PType::U64 | PType::I64 | PType::F64, _) + ), + "Expected a widened primitive sum, got {}", + 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), + ], + )) + } +} + impl AggregateFnVTable for SumV2 { type Options = NumericalAggregateOpts; type Partial = SumV2Partial; From 363be18d2e30a67bdf6aa270b9fb98d3d93b7793 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 14 Sep 2026 15:55:39 -0400 Subject: [PATCH 2/2] docs: explain why SumV2 is a separate aggregate Signed-off-by: "Connor Tsui" --- vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs | 4 ++++ 1 file changed, 4 insertions(+) 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 fccbdcbe9d8..f9f524fa94f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -2,6 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! SQL-style sums with explicit overflow and empty-input state. +//! +//! [`SumV2`] returns null when there are no valid values, while [`Sum`] returns zero. Tracking empty +//! input separately from overflow requires a different partial representation. A separate aggregate +//! preserves compatibility with the scalar partials stored by older Vortex files. mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel;