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..f9f524fa94f 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,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod grouped; +//! 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; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -35,6 +40,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 +80,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;