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..e9ecdde6859 --- /dev/null +++ b/encodings/runend/benches/run_end_sum.rs @@ -0,0 +1,181 @@ +// 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::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +const LEN: usize = 16_384; + +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); + for encoding in [RunEnd.id(), Constant.id()] { + session + .aggregate_fns() + .register_grouped_encoding_kernel(encoding, 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 = [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); +} + +#[divan::bench(consts = [1, 2, 8, 128])] +fn grouped_constant(bencher: Bencher) { + bench_grouped( + bencher, + ConstantArray::new(3i32, LEN).into_array(), + GROUP_SIZE, + &SESSION, + ); +} + +#[divan::bench(consts = [1, 2, 8, 128])] +fn grouped_constant_fallback(bencher: Bencher) { + bench_grouped( + bencher, + ConstantArray::new(3i32, LEN).into_array(), + GROUP_SIZE, + &FALLBACK_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/mod.rs b/encodings/runend/src/compute/sum/mod.rs new file mode 100644 index 00000000000..055a3cab096 --- /dev/null +++ b/encodings/runend/src/compute/sum/mod.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive sums over run-end encoded arrays. +//! +//! The kernels decode the run ends and values once, then sum each value weighted by its run +//! length. Grouped sums intersect runs with each group's range. Decimal inputs use the fallback. +//! Float multiplication can round differently from repeated addition, as with constant sums. + +mod primitive; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +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::DynAggregateKernel; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use self::primitive::RunEndSums; +use crate::RunEnd; + +/// Whole-array and grouped primitive sum kernels for [`RunEnd`]. +#[derive(Debug)] +pub(crate) struct RunEndSumKernel; + +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 sums = RunEndSums::new(array, ctx, options.skip_nans)?; + let (sum, is_empty) = sums.sum(0..batch.len()); + Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?)) + } +} + +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 ranges = groups.group_ranges(ctx)?; + let validity = groups.group_validity(ctx)?; + let sums = RunEndSums::new(elements, ctx, options.skip_nans)?; + let (results, empty_groups) = sums.grouped_sum(&ranges, &validity); + + 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 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/primitive.rs b/encodings/runend/src/compute/sum/primitive.rs new file mode 100644 index 00000000000..d9eba21d649 --- /dev/null +++ b/encodings/runend/src/compute/sum/primitive.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Weighted primitive reductions over materialized run ends, values, and validity. +//! +//! Each range visits only its valid runs. Signed products use wider arithmetic so that a run +//! can cancel a preceding sum even when its product alone does not fit in the result type. +//! Fixed-size groups share a forward run cursor; list-view groups locate their runs independently. + +use std::iter::Peekable; +use std::ops::Range; + +use itertools::Either; +use num_traits::AsPrimitive; +use num_traits::ToPrimitive; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::GroupRanges; +use vortex_array::arrays::PrimitiveArray; +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_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::RunEnd; +use crate::RunEndArrayExt; +use crate::RunEndArraySlotsExt; + +pub(super) struct RunEndSums { + ends: PrimitiveArray, + values: PrimitiveArray, + validity: Mask, + offset: usize, + skip_nans: bool, +} + +impl RunEndSums { + pub(super) fn new( + array: ArrayView<'_, RunEnd>, + ctx: &mut ExecutionCtx, + skip_nans: bool, + ) -> VortexResult { + let ends = array.ends().clone().execute::(ctx)?; + let values = array.values().clone().execute::(ctx)?; + let validity = values.validity()?.execute_mask(values.len(), ctx)?; + + Ok(Self { + ends, + values, + validity, + offset: array.offset(), + skip_nans, + }) + } + + /// Return the widened sum and whether the range contains no valid values. + pub(super) fn sum(&self, range: Range) -> (Scalar, bool) { + let range = self.offset + range.start..self.offset + range.end; + let valid_runs = self.validity.indices(); + + match_each_unsigned_integer_ptype!(self.ends.ptype(), |E| { + let ends = self.ends.as_slice::(); + match_each_native_ptype!(self.values.ptype(), + unsigned: |T| { + sum_scalar(ends, self.values.as_slice::(), &valid_runs, range, add_unsigned_run) + }, + signed: |T| { + sum_scalar(ends, self.values.as_slice::(), &valid_runs, range, add_signed_run) + }, + floating: |T| { + sum_scalar(ends, self.values.as_slice::(), &valid_runs, range, + |sum, value, len| add_float_run(sum, value, len, self.skip_nans)) + } + ) + }) + } + + /// Sum all groups with one native type dispatch and return their sums and empty flags. + pub(super) fn grouped_sum( + &self, + ranges: &GroupRanges, + group_validity: &Mask, + ) -> (PrimitiveArray, BitBuffer) { + let valid_runs = self.validity.indices(); + + match_each_unsigned_integer_ptype!(self.ends.ptype(), |E| { + let ends = self.ends.as_slice::(); + match_each_native_ptype!(self.values.ptype(), + unsigned: |T| { + collect_sums(ends, self.values.as_slice::(), valid_runs, ranges, + group_validity, self.offset, add_unsigned_run) + }, + signed: |T| { + collect_sums(ends, self.values.as_slice::(), valid_runs, ranges, + group_validity, self.offset, add_signed_run) + }, + floating: |T| { + collect_sums(ends, self.values.as_slice::(), valid_runs, ranges, + group_validity, self.offset, + |sum, value, len| add_float_run(sum, value, len, self.skip_nans)) + } + ) + }) + } +} + +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)) +} + +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() +} + +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)) +} + +fn sum_scalar>( + ends: &[E], + values: &[T], + validity: &AllOr<&[usize]>, + range: Range, + add_run: impl Fn(A, T, usize) -> Option, +) -> (Scalar, bool) { + let (sum, is_empty) = sum_runs(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 collect_sums( + ends: &[E], + values: &[T], + validity: AllOr<&[usize]>, + ranges: &GroupRanges, + group_validity: &Mask, + offset: usize, + add_run: impl Fn(A, T, usize) -> Option, +) -> (PrimitiveArray, BitBuffer) { + match ranges { + GroupRanges::FixedSizeList { .. } => { + let indices = match validity { + AllOr::All => Either::Left(0..ends.len()), + AllOr::None => Either::Left(0..0), + AllOr::Some(indices) => Either::Right(indices.iter().copied()), + }; + let mut indices = indices.peekable(); + + collect_group_sums(ranges, group_validity, offset, |range| { + sum_consecutive_runs(ends, values, &mut indices, range, &add_run) + }) + } + GroupRanges::ListView { .. } => { + collect_group_sums(ranges, group_validity, offset, |range| { + sum_runs(ends, values, &validity, 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()) +} + +/// Sum consecutive, non-overlapping ranges while retaining runs that cross a group boundary. +fn sum_consecutive_runs( + 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_(); + // Null groups and overflow can leave the cursor behind the next group's start. + if end <= range.start { + indices.next(); + continue; + } + + let start = if index == 0 { + range.start + } else { + ends[index - 1].as_().max(range.start) + }; + if start >= range.end { + break; + } + + is_empty = false; + let Some(next) = add_run(sum, values[index], end.min(range.end) - start) else { + return (None, false); + }; + sum = next; + if end > range.end { + break; + } + indices.next(); + } + + (Some(sum), is_empty) +} + +fn sum_runs( + ends: &[E], + values: &[T], + validity: &AllOr<&[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); + } + + let first = ends.partition_point(|end| end.as_() <= range.start); + let last = (ends.partition_point(|end| end.as_() < range.end) + 1).min(ends.len()); + let indices = match validity { + AllOr::All => Either::Left(first..last), + AllOr::None => return (Some(sum), true), + AllOr::Some(indices) => { + let start = indices.partition_point(|&index| index < first); + let end = indices.partition_point(|&index| index < last); + Either::Right(indices[start..end].iter().copied()) + } + }; + let mut is_empty = true; + + for index in indices { + let start = if index == 0 { + range.start + } else { + ends[index - 1].as_().max(range.start) + }; + let end = ends[index].as_().min(range.end); + is_empty = false; + let Some(next) = add_run(sum, values[index], end - start) else { + return (None, false); + }; + sum = 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..a05f32c7362 --- /dev/null +++ b/encodings/runend/src/compute/sum/tests.rs @@ -0,0 +1,373 @@ +// 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::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::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +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] +fn primitive_types( + #[values(PType::U8, PType::U16, PType::U32, PType::U64)] ends_type: PType, + #[values( + PType::U8, PType::U16, PType::U32, PType::U64, PType::I8, PType::I16, PType::I32, + PType::I64, PType::F16, PType::F32, PType::F64 + )] + values_type: PType, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ends = buffer![2u64, 5, 9] + .into_array() + .cast(DType::Primitive(ends_type, NonNullable))?; + let values = buffer![1u64, 3, 7] + .into_array() + .cast(DType::Primitive(values_type, NonNullable))?; + let array = RunEnd::try_new_offset_length(ends, values, 1, 7, &mut ctx)?.into_array(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::all_null(Validity::AllInvalid)] +#[case::partially_valid(Validity::from_iter([true, false, true]))] +fn nullable_runs(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::new(buffer![-3i32, 100, 7], validity).into_array(); + let array = RunEnd::try_new(buffer![2u16, 5, 9].into_array(), values, &mut ctx)?.into_array(); + + check_sum(array, NumericalAggregateOpts::default()) +} + +#[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, + #[values(false, true)] skip_nans: 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 { skip_nans }); + } + + 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 { skip_nans }), + SumV2.bind(NumericalAggregateOpts { skip_nans }), + ] { + 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()) +} + +#[test] +#[cfg(target_pointer_width = "64")] +fn huge_run_uses_registered_kernel() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let len = usize::try_from(i64::MAX)? + 1; + let array = RunEnd::try_new( + buffer![len as u64].into_array(), + buffer![-1i64].into_array(), + &mut ctx, + )? + .into_array(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + let mut acc = aggregate.accumulator(array.dtype())?; + acc.accumulate(&array, &mut ctx)?; + assert_eq!(acc.finish()?, Scalar::from(i64::MIN)); + } + + Ok(()) +} + +#[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)] +#[case::single_elements(0, 1, false)] +#[case::short_groups(0, 2, false)] +#[case::sliced_run(1, 2, false)] +#[case::run_boundary(3, 2, false)] +#[case::null_groups(1, 8, true)] +#[case::multiple_runs(0, 8, false)] +fn consecutive_groups( + #[case] offset: usize, + #[case] size: u32, + #[case] null_groups: bool, + #[values( + buffer![1i64, -2, 3, -4, 5].into_array(), + PrimitiveArray::from_option_iter([None, Some(2i64), None, Some(4), Some(5)]).into_array(), + PrimitiveArray::from_option_iter([None::; 5]).into_array(), + buffer![i64::MAX, 1, -2, i64::MIN, 3].into_array() + )] + 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/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/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index ced92af1d68..07aeb698bc5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -50,14 +50,15 @@ pub(crate) fn multiply_constant( } DType::Primitive(PType::I64, _) => { let val = pvalue.cast::()?; - match i64::try_from(len).ok().and_then(|l| val.checked_mul(l)) { - Some(product) => Scalar::primitive(product, Nullability::Nullable), - None => Scalar::null(return_dtype.as_nullable()), + match i64::try_from(val as i128 * len as i128) { + Ok(product) => Scalar::primitive(product, Nullability::Nullable), + Err(_) => Scalar::null(return_dtype.as_nullable()), } } DType::Primitive(PType::F64, _) => { let val = pvalue.cast::()?; - Scalar::primitive(val * len as f64, Nullability::Nullable) + // Sums start at positive zero, including groups of negative zeros. + Scalar::primitive(0.0 + val * len as f64, Nullability::Nullable) } _ => vortex_bail!( "Unexpected return dtype for primitive sum: {}", 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; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 65b00d81031..3b9ad5d6e03 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -41,9 +41,11 @@ use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; +use crate::arrays::Constant; use crate::arrays::Dict; use crate::arrays::Primitive; use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; +use crate::arrays::constant::compute::sum::ConstantGroupedSumKernel; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; @@ -134,6 +136,9 @@ impl Default for AggregateFnSession { SumV2.id(), &PrimitiveGroupedSumV2EncodingKernel, ); + for sum in [Sum.id(), SumV2.id()] { + this.register_grouped_encoding_kernel(Constant.id(), sum, &ConstantGroupedSumKernel); + } this } diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index a996bf79a0f..0d5f2449471 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -8,6 +8,7 @@ mod filter; mod not; pub(crate) mod rules; mod slice; +pub(crate) mod sum; mod take; pub(crate) mod uncompressed_size; diff --git a/vortex-array/src/arrays/constant/compute/sum.rs b/vortex-array/src/arrays/constant/compute/sum.rs new file mode 100644 index 00000000000..68512161d4a --- /dev/null +++ b/vortex-array/src/arrays/constant/compute/sum.rs @@ -0,0 +1,250 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive sums over constant elements. +//! +//! Whole-array sums already multiply the scalar by its length. Grouped sums apply the same +//! arithmetic to each group size without decoding the elements or slicing a constant per group. + +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::multiply_constant; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::Constant; +use crate::builders::builder_with_capacity_in; +use crate::scalar::Scalar; + +/// Grouped primitive sum kernel for constant elements. +#[derive(Debug)] +pub(crate) struct ConstantGroupedSumKernel; + +impl DynGroupedAggregateKernel for ConstantGroupedSumKernel { + 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 !elements.dtype().is_primitive() { + return Ok(None); + } + + let Some(sum_dtype) = aggregate_fn.return_dtype(elements.dtype()) else { + return Ok(None); + }; + let ranges = groups.group_ranges(ctx)?; + let validity = groups.group_validity(ctx)?; + let scalar = elements.scalar(); + let skip_nan = options.skip_nans && scalar.as_primitive().is_nan(); + let zero = Scalar::zero_value(&sum_dtype); + let mut sums = builder_with_capacity_in(&sum_dtype, groups.len(), ctx.allocator()); + let mut empty_groups = BitBufferMut::new_unset(groups.len()); + + for (index, ((_, size), valid)) in ranges.iter().zip(validity.iter()).enumerate() { + if !valid { + sums.append_null(); + continue; + } + + let is_empty = size == 0 || scalar.is_null(); + let sum = if is_empty || skip_nan { + zero.clone() + } else { + multiply_constant(scalar, size, &sum_dtype)?.unwrap_or_else(|| zero.clone()) + }; + sums.append_scalar(&sum)?; + empty_groups.set_to(index, is_empty); + } + + let sums = sums.finish(); + if aggregate_fn.is::() { + Ok(Some(SumV2::partials_from_sums( + sums, + empty_groups.freeze(), + validity, + ctx, + )?)) + } else { + Ok(Some(sums)) + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::ConstantGroupedSumKernel; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateFnVTableExt; + use crate::aggregate_fn::GroupedArray; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum_v2::SumV2; + use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; + use crate::array_session; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListViewArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::half::f16; + use crate::scalar::Scalar; + use crate::validity::Validity; + + fn check_groups( + scalar: Scalar, + fixed_size: bool, + options: NumericalAggregateOpts, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = ConstantArray::new(scalar, 8).into_array(); + let decoded = elements + .clone() + .execute::(&mut ctx)? + .into_array(); + let make_groups = |elements: ArrayRef| -> VortexResult { + if fixed_size { + Ok(FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 4)?.into()) + } else { + Ok(ListViewArray::try_new( + elements, + buffer![0u32, 4, 2, 8, 1].into_array(), + buffer![3u32, 2, 4, 0, 0].into_array(), + Validity::from_iter([true, false, true, true, true]), + )? + .into()) + } + }; + let groups = make_groups(elements.clone())?; + let reference = make_groups(decoded.clone())?; + 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!( + ConstantGroupedSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_some() + ); + let mut actual = aggregate.accumulator_grouped(elements.dtype())?; + actual.accumulate_list(&groups_array, &mut ctx)?; + let mut expected = aggregate.accumulator_grouped(decoded.dtype())?; + expected.accumulate_list(&reference_array, &mut ctx)?; + assert_arrays_eq!(actual.finish()?, expected.finish()?, &mut ctx); + } + + Ok(()) + } + + #[rstest] + #[case::u8(3u8.into())] + #[case::u16(3u16.into())] + #[case::u32(3u32.into())] + #[case::u64(3u64.into())] + #[case::i8((-3i8).into())] + #[case::i16((-3i16).into())] + #[case::i32((-3i32).into())] + #[case::i64((-3i64).into())] + #[case::f16(f16::from_f32(1.25).into())] + #[case::f32(1.25f32.into())] + #[case::f64(1.25f64.into())] + #[case::null(Scalar::null_native::())] + #[case::signed_overflow(i64::MAX.into())] + #[case::unsigned_overflow(u64::MAX.into())] + #[case::negative_zero((-0.0f64).into())] + fn primitive_groups( + #[case] scalar: Scalar, + #[values(false, true)] fixed_size: bool, + ) -> VortexResult<()> { + check_groups(scalar, fixed_size, NumericalAggregateOpts::default()) + } + + #[rstest] + fn nan_groups(#[values(true, false)] skip_nans: bool) -> VortexResult<()> { + check_groups(f64::NAN.into(), false, NumericalAggregateOpts { skip_nans }) + } + + #[rstest] + #[cfg(target_pointer_width = "64")] + #[case::zero(0i64, 0i64)] + #[case::negative_one(-1i64, i64::MIN)] + fn huge_constants(#[case] value: i64, #[case] expected: i64) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let len = usize::try_from(i64::MAX)? + 1; + let array = ConstantArray::new(value, len).into_array(); + let groups = ListViewArray::try_new( + array.clone(), + buffer![0u64].into_array(), + buffer![len as u64].into_array(), + Validity::NonNullable, + )? + .into_array(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + let mut acc = aggregate.accumulator(array.dtype())?; + acc.accumulate(&array, &mut ctx)?; + assert_eq!(acc.finish()?, Scalar::from(expected)); + let mut grouped = aggregate.accumulator_grouped(array.dtype())?; + grouped.accumulate_list(&groups, &mut ctx)?; + assert_arrays_eq!( + grouped.finish()?, + PrimitiveArray::from_option_iter([Some(expected)]).into_array(), + &mut ctx + ); + } + + Ok(()) + } + + #[test] + fn decimal_kernel_declines() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(10, 2), NonNullable); + let elements = ConstantArray::new(Scalar::zero_value(&dtype), 4).into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + assert!( + ConstantGroupedSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_none() + ); + } + + Ok(()) + } +}