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..7b988b82b95 --- /dev/null +++ b/encodings/runend/benches/run_end_sum.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynGroupedAccumulator; +use vortex_array::aggregate_fn::GroupedAccumulator; +use vortex_array::aggregate_fn::GroupedArray; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +// Keep the one-element-group fallback below 1 ms in CodSpeed simulation. +const LEN: usize = 2_048; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); + +static FALLBACK_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session + .aggregate_fns() + .register_aggregate_kernel(RunEnd.id(), Some(SumV2.id()), &Decline); + session + .aggregate_fns() + .register_grouped_encoding_kernel(RunEnd.id(), SumV2.id(), &Decline); + session +}); + +/// Keep the pre-specialization dispatch paths available for benchmark comparisons. +#[derive(Debug)] +struct Decline; + +impl DynAggregateKernel for Decline { + fn aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +impl DynGroupedAggregateKernel for Decline { + fn grouped_aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _groups: &GroupedArray, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +fn main() { + LazyLock::force(&SESSION); + LazyLock::force(&FALLBACK_SESSION); + divan::main(); +} + +fn runend(run_length: usize) -> ArrayRef { + let ends = + PrimitiveArray::from_iter((run_length..=LEN).step_by(run_length).map(|end| end as u64)); + let values = PrimitiveArray::from_option_iter( + (0..ends.len()) + .map(|index| (index % 5 != 0).then_some(i32::try_from(index % 100).unwrap())), + ); + RunEnd::try_new( + ends.into_array(), + values.into_array(), + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array() +} + +fn bench_sum(bencher: Bencher, run_length: usize, session: &VortexSession) { + let array = runend(run_length); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| sum_v2(&array, ctx).unwrap()); +} + +fn bench_grouped(bencher: Bencher, elements: ArrayRef, group_size: u32, session: &VortexSession) { + let dtype = elements.dtype().clone(); + let groups = FixedSizeListArray::try_new( + elements, + group_size, + Validity::NonNullable, + LEN / group_size as usize, + ) + .unwrap() + .into_array(); + bencher + .with_inputs(|| { + ( + GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + dtype.clone(), + ) + .unwrap(), + session.create_execution_ctx(), + ) + }) + .bench_refs(|(acc, ctx)| { + acc.accumulate_list(&groups, ctx).unwrap(); + acc.finish() + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = [4, 64, 1024])] +fn sum_runend(bencher: Bencher, run_length: usize) { + bench_sum(bencher, run_length, &SESSION); +} + +#[divan::bench(args = [4, 64, 1024])] +fn sum_runend_fallback(bencher: Bencher, run_length: usize) { + bench_sum(bencher, run_length, &FALLBACK_SESSION); +} + +#[divan::bench(args = [Validity::NonNullable, Validity::AllValid, Validity::AllInvalid])] +fn sum_runend_validity(bencher: Bencher, validity: &Validity) { + let array = runend_with_validity(validity); + + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_refs(|ctx| sum_v2(&array, ctx).unwrap()); +} + +#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])] +fn grouped_runend(bencher: Bencher, run_length: usize) { + bench_grouped(bencher, runend(run_length), GROUP_SIZE, &SESSION); +} + +#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])] +fn grouped_runend_fallback(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::>(), + 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(bencher: Bencher) { + bench_grouped( + bencher, + runend_with_validity(&Validity::AllValid), + GROUP_SIZE, + &SESSION, + ); +} 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..7d4ae8580d9 --- /dev/null +++ b/encodings/runend/src/compute/sum/grouped.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Grouped aggregation with traversal selected by the group layout. +//! +//! Fixed-size groups share a forward run cursor. List-view ranges can overlap or arrive out of +//! order, so they locate their runs independently. Both paths weight runs by their intersection +//! with the group, and skip null groups before visiting any runs. + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::GroupRanges; +use vortex_array::aggregate_fn::GroupedArray; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::match_each_native_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::RunEndInputs; +use super::RunEndSumKernel; +use super::empty_partial; +use super::runs::add_float_run; +use super::runs::add_signed_run; +use super::runs::add_unsigned_run; +use super::runs::sum_all_valid; +use super::runs::sum_next_valid_range; +use super::runs::sum_valid_range; +use crate::RunEnd; + +impl DynGroupedAggregateKernel for RunEndSumKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = aggregate_fn + .as_opt::() + .or_else(|| aggregate_fn.as_opt::()) + 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| { + 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, + &validity, runs.offset, add_unsigned_run) + }, + signed: |T| { + sum_groups(ends, runs.values.as_slice::(), &runs.validity, &ranges, + &validity, runs.offset, add_signed_run) + }, + floating: |T| { + sum_groups(ends, runs.values.as_slice::(), &runs.validity, &ranges, + &validity, runs.offset, + |sum, value, len| add_float_run(sum, value, len, options.skip_nans)) + } + ) + }); + + let results = results.into_array(); + if aggregate_fn.is::() { + Ok(Some(SumV2::partials_from_sums( + results, + empty_groups, + validity, + ctx, + )?)) + } else { + Ok(Some(results)) + } + } +} + +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..ec1c68d5982 --- /dev/null +++ b/encodings/runend/src/compute/sum/mod.rs @@ -0,0 +1,94 @@ +// 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. +//! +//! 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::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(), + })) + } +} + +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..00b8e1a7246 --- /dev/null +++ b/encodings/runend/src/compute/sum/whole.rs @@ -0,0 +1,104 @@ +// 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::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +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 crate::RunEnd; + +impl DynAggregateKernel for RunEndSumKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = aggregate_fn + .as_opt::() + .or_else(|| aggregate_fn.as_opt::()) + 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| { + 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)) + } + ) + }); + + Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?)) + } +} + +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;