diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..1dc57c431c2 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -11,6 +11,7 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTable; @@ -18,6 +19,7 @@ use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::count::Count; +use vortex_array::aggregate_fn::fns::standard_sum::StandardSum; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; @@ -199,6 +201,125 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Sum)); } +#[divan::bench] +fn standard_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, StandardSum)); +} + +#[divan::bench] +fn standard_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, StandardSum)); +} + +#[divan::bench] +fn standard_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, StandardSum)); +} + +#[divan::bench] +fn standard_sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, StandardSum)); +} + +/// Like [`grouped_accumulator`], but executes the lazy finalize result to canonical so the +/// bench measures the full cost of producing usable sums. +fn grouped_accumulator_canonical(list_view: &ArrayRef, vtable: V) -> ArrayRef +where + V: AggregateFnVTable + Clone, +{ + let mut acc = GroupedAccumulator::try_new( + vtable, + NumericalAggregateOpts::default(), + list_element_dtype(list_view), + ) + .unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + acc.accumulate_list(list_view, &mut ctx).unwrap(); + let result = acc + .finish() + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + divan::black_box(result) +} + +#[divan::bench] +fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_standard_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); +} + +#[divan::bench] +fn canonical_standard_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); +} + +#[divan::bench] +fn canonical_standard_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); +} + +#[divan::bench] +fn canonical_standard_sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); +} + #[divan::bench] fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index da81e67ecb9..f7a6b29953a 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -19,5 +19,6 @@ pub mod min; pub mod min_max; pub mod nan_count; pub mod null_count; +pub mod standard_sum; pub mod sum; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs new file mode 100644 index 00000000000..a4c0fb46da9 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -0,0 +1,436 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::StandardSum; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::GroupRanges; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::fns::sum::sum_float_all; +use crate::aggregate_fn::fns::sum::sum_signed_all; +use crate::aggregate_fn::fns::sum::sum_unsigned_all; +use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::BoolArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::match_each_native_ptype; +use crate::validity::Validity; + +/// Encoding-specific grouped [`StandardSum`] kernel for primitive element arrays. +#[derive(Debug)] +pub(crate) struct PrimitiveGroupedStandardSumEncodingKernel; + +impl DynGroupedAggregateKernel for PrimitiveGroupedStandardSumEncodingKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(None); + }; + try_grouped_sum(groups, ctx, options.skip_nans) + } +} + +/// Grouped [`StandardSum`] implementation for canonical primitive elements. +/// +/// Reuses the scalar primitive-sum reductions ([`sum_unsigned_all`]/[`sum_signed_all`]/ +/// [`sum_float_all`]) so the per-group semantics match scalar `sum` exactly (overflow saturates to +/// a null sum, NaNs are skipped). The element validity mask is materialized once and sliced per +/// group, rather than the per-group accumulator setup of the generic fallback path. +/// +/// Produces `{sum, overflowed}` partial rows (see `StandardSum`'s partial dtype): `sum` is null for +/// empty, all-null, and overflowed groups, while `overflowed` distinguishes saturation from the +/// empty identity during partial merges. Null groups are null struct rows. +pub(super) fn try_grouped_sum( + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + skip_nans: bool, +) -> VortexResult> { + if !groups.elements().is::() { + return Ok(None); + } + let elements = groups.elements().clone().downcast::(); + let group_ranges = groups.group_ranges(ctx)?; + let group_validity = groups.group_validity(ctx)?; + + Ok(Some(grouped_sum( + &elements, + &group_ranges, + &group_validity, + ctx, + skip_nans, + )?)) +} + +/// StandardSum each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. +fn grouped_sum( + elements: &PrimitiveArray, + group_ranges: &GroupRanges, + group_validity: &Mask, + ctx: &mut ExecutionCtx, + skip_nans: bool, +) -> VortexResult { + let elem_mask = elements + .as_ref() + .validity()? + .execute_mask(elements.as_ref().len(), ctx)?; + let all_valid = matches!(elem_mask.slices(), AllOr::All); + + let (sums, overflowed) = match_each_native_ptype!(elements.ptype(), + unsigned: |T| { + let values = elements.as_slice::(); + collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, + sum_unsigned_all) + }, + signed: |T| { + let values = elements.as_slice::(); + collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, + sum_signed_all) + }, + floating: |T| { + let values = elements.as_slice::(); + collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, + |acc, slice| { sum_float_all(acc, slice, skip_nans); false }) + } + ); + + Ok(StructArray::try_new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("overflowed")]), + vec![ + sums.into_array(), + BoolArray::new(overflowed, Validity::NonNullable).into_array(), + ], + group_validity.len(), + Validity::from_mask(group_validity.clone(), Nullability::Nullable), + )? + .into_array()) +} + +/// Reduce each group's element slice into a nullable sum plus a per-group overflow bit. A sum is +/// null when the group is invalid, has no valid elements, or summing it overflows. With +/// `skip_nans`, a valid all-NaN group has a real zero sum because NaNs remain valid elements. +fn collect_sums( + values: &[T], + group_ranges: &GroupRanges, + group_validity: &Mask, + elem_mask: &Mask, + all_valid: bool, + sum_run: impl Fn(&mut A, &[T]) -> bool, +) -> (PrimitiveArray, BitBuffer) { + let mut overflowed = BitBufferMut::with_capacity(group_ranges.len()); + let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { + if !group_validity.value(i) { + overflowed.append(false); + return None; + } + let mut acc = A::default(); + let (overflow, any_valid) = if all_valid { + (sum_run(&mut acc, &values[offset..offset + size]), size > 0) + } else { + sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) + }; + overflowed.append(overflow); + (!overflow && any_valid).then_some(acc) + }); + let sums = PrimitiveArray::from_option_iter(sums); + (sums, overflowed.freeze()) +} + +/// StandardSum the valid elements of a single group, using the contiguous valid runs of the element mask +/// intersected with the group's `[offset, offset + size)` range. +/// +/// Returns `(overflow, any_valid)`, where `any_valid` records whether the group contained at +/// least one valid element. +fn sum_masked_group( + acc: &mut A, + values: &[T], + offset: usize, + size: usize, + elem_mask: &Mask, + sum_run: &impl Fn(&mut A, &[T]) -> bool, +) -> (bool, bool) { + match elem_mask.slice(offset..offset + size).slices() { + AllOr::All => (sum_run(acc, &values[offset..offset + size]), size > 0), + AllOr::None => (false, false), + AllOr::Some(runs) => { + for &(start, end) in runs { + if sum_run(acc, &values[offset + start..offset + end]) { + return (true, true); + } + } + (false, !runs.is_empty()) + } + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::cast_possible_truncation)] + + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::standard_sum::StandardSum; + use crate::aggregate_fn::fns::standard_sum::standard_sum; + use crate::array_session; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListViewArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::builders::builder_with_capacity; + use crate::dtype::DType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::dtype::PType; + use crate::validity::Validity; + + /// Run a grouped sum through the accumulator. + fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { + let mut acc = GroupedAccumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish() + } + + /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] + /// for valid groups, a null sum for invalid groups. + fn grouped_sum_reference( + elements: &ArrayRef, + ranges: &[(usize, usize)], + group_valid: &[bool], + elem_dtype: &DType, + ) -> VortexResult { + use crate::aggregate_fn::AggregateFnVTable; + + let mut ctx = array_session().create_execution_ctx(); + let sum_dtype = StandardSum + .return_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .expect("sum return dtype"); + let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); + for (i, &(offset, size)) in ranges.iter().enumerate() { + if group_valid[i] { + let slice = elements.slice(offset..offset + size)?; + builder.append_scalar(&standard_sum(&slice, &mut ctx)?)?; + } else { + builder.append_null(); + } + } + Ok(builder.finish()) + } + + fn offsets_sizes(ranges: &[(usize, usize)]) -> (ArrayRef, ArrayRef) { + let offsets = PrimitiveArray::from_iter(ranges.iter().map(|&(o, _)| o as i32)); + let sizes = PrimitiveArray::from_iter(ranges.iter().map(|&(_, s)| s as i32)); + (offsets.into_array(), sizes.into_array()) + } + + fn listview( + elements: ArrayRef, + ranges: &[(usize, usize)], + group_valid: &[bool], + ) -> VortexResult { + let (offsets, sizes) = offsets_sizes(ranges); + let validity = if group_valid.iter().all(|&v| v) { + Validity::NonNullable + } else { + Validity::from_iter(group_valid.iter().copied()) + }; + Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) + } + + #[test] + fn listview_matches_reference_unsigned() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); + let elem_dtype = DType::Primitive(PType::U32, NonNullable); + let ranges = [(0, 2), (2, 1), (3, 3)]; + let valid = [true, true, true]; + + let groups = listview(elements.clone(), &ranges, &valid)?; + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + + // Unsigned input sums to U64. + let direct = PrimitiveArray::from_option_iter([Some(3u64), Some(3u64), Some(15u64)]); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Offsets are not in group order and a group is null: the group validity must be indexed by + // group index, not by element offset. + let elements = + PrimitiveArray::new(buffer![10i32, 20, 30, 40, 50, 60], Validity::NonNullable) + .into_array(); + let elem_dtype = DType::Primitive(PType::I32, NonNullable); + let ranges = [(4, 2), (0, 2), (2, 2)]; + let valid = [true, false, true]; + + let groups = listview(elements.clone(), &ranges, &valid)?; + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + + let direct = PrimitiveArray::from_option_iter([Some(110i64), None, Some(70i64)]); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn listview_interior_and_full_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. + let elements = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) + .into_array(); + let elem_dtype = DType::Primitive(PType::I32, Nullable); + let ranges = [(0, 3), (3, 2), (5, 0), (5, 1)]; + let valid = [true, true, true, true]; + + let groups = listview(elements.clone(), &ranges, &valid)?; + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + + let direct = PrimitiveArray::from_option_iter([Some(4i64), None, None, Some(9i64)]); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn listview_overflow_group_is_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); + let elem_dtype = DType::Primitive(PType::I64, NonNullable); + let ranges = [(0, 2), (2, 2)]; + let valid = [true, true]; + + let groups = listview(elements.clone(), &ranges, &valid)?; + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + + // First group overflows -> null sum; second group sums normally. + let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn listview_float_nan_and_inf() -> VortexResult<()> { + let elements = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, f64::INFINITY, f64::NEG_INFINITY, 4.0], + Validity::NonNullable, + ) + .into_array(); + let elem_dtype = DType::Primitive(PType::F64, NonNullable); + let ranges = [(0, 3), (3, 3)]; + let valid = [true, true]; + + let groups = listview(elements.clone(), &ranges, &valid)?; + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + + // Group 0: NaN skipped -> 3.0. Group 1: INF + -INF = NaN. (Avoid array equality here since + // NaN != NaN; compare element scalars against the reference path instead.) + let mut ctx = array_session().create_execution_ctx(); + let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + let g0 = actual.execute_scalar(0, &mut ctx)?; + assert_eq!(g0.as_primitive().typed_value::(), Some(3.0)); + assert_eq!( + g0.as_primitive().typed_value::(), + expected + .execute_scalar(0, &mut ctx)? + .as_primitive() + .typed_value::() + ); + let g1 = actual.execute_scalar(1, &mut ctx)?; + assert!(g1.as_primitive().typed_value::().unwrap().is_nan()); + assert!( + expected + .execute_scalar(1, &mut ctx)? + .as_primitive() + .typed_value::() + .unwrap() + .is_nan() + ); + Ok(()) + } + + #[test] + fn listview_float_nan_not_skipping() -> VortexResult<()> { + let elements = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, 3.0, 4.0], + Validity::NonNullable, + ) + .into_array(); + let elem_dtype = DType::Primitive(PType::F64, NonNullable); + let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; + + let mut acc = GroupedAccumulator::try_new( + StandardSum, + NumericalAggregateOpts::include_nans(), + elem_dtype, + )?; + let mut ctx2 = array_session().create_execution_ctx(); + acc.accumulate_list(&groups, &mut ctx2)?; + let actual = acc.finish()?; + + let mut ctx = array_session().create_execution_ctx(); + // Group 0 contains a NaN -> NaN sum; group 1 sums normally. + let g0 = actual.execute_scalar(0, &mut ctx)?; + assert!(g0.as_primitive().typed_value::().unwrap().is_nan()); + let g1 = actual.execute_scalar(1, &mut ctx)?; + assert_eq!(g1.as_primitive().typed_value::(), Some(7.0)); + Ok(()) + } + + #[test] + fn fixed_size_overflow_and_nan() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // FixedSize path: first group overflows -> null sum, second sums normally. + let elements = + PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); + let elem_dtype = DType::Primitive(PType::I64, NonNullable); + let groups = FixedSizeListArray::try_new(elements.clone(), 2, Validity::NonNullable, 2)? + .into_array(); + + let actual = grouped_sum_actual(&groups, &elem_dtype)?; + let expected = + grouped_sum_reference(&elements, &[(0, 2), (2, 2)], &[true, true], &elem_dtype)?; + let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs new file mode 100644 index 00000000000..9ab4a8b2abd --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -0,0 +1,465 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod grouped; +pub(crate) use grouped::PrimitiveGroupedStandardSumEncodingKernel; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::SumState; +use crate::aggregate_fn::fns::sum::accumulate_bool; +use crate::aggregate_fn::fns::sum::accumulate_decimal; +use crate::aggregate_fn::fns::sum::accumulate_primitive; +use crate::aggregate_fn::fns::sum::checked_add_i64; +use crate::aggregate_fn::fns::sum::checked_add_u64; +use crate::aggregate_fn::fns::sum::make_zero_state; +use crate::aggregate_fn::fns::sum::multiply_constant; +use crate::aggregate_fn::fns::sum::sum_decimal_dtype; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; +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; +use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::validity::Validity; + +/// Return the sum of an array. The result is null when the array has no valid values or the sum +/// overflows. +/// +/// See [`StandardSum`] for details. +pub fn standard_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Short-circuit using cached array statistics. `Stat::Sum` is zero for an + // array with no valid values instead of null, so we need the null count for a nullable array. + // When it is unknown, fall through and compute rather than trust the cache. + if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { + if !array.dtype().is_nullable() { + return Ok(sum_scalar); + } + match array.statistics().get_as::(Stat::NullCount) { + Precision::Exact(null_count) if null_count == array.len() as u64 => { + return Ok(Scalar::null(sum_scalar.dtype().as_nullable())); + } + Precision::Exact(_) => return Ok(sum_scalar), + _ => {} + } + } + + // Compute using Accumulator + let mut acc = Accumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + acc.accumulate(array, ctx)?; + let result = acc.finish()?; + + // Cache the computed sum as a statistic (only if non-null, i.e. no overflow and at least + // one valid value). + if let Some(val) = result.value().cloned() { + array.statistics().set(Stat::Sum, Precision::Exact(val)); + } + + Ok(result) +} + +/// The same as [`Sum`](crate::aggregate_fn::fns::sum::Sum), except that arrays with no valid elements yield null. +/// +/// Sum aggregates typically follow this behavior. See: +/// - [DuckDB](https://duckdb.org/docs/stable/sql/functions/aggregates.html) +/// - [Arrow](https://docs.rs/arrow/latest/arrow/compute/fn.sum.html) +/// - [DataFusion](https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L370) +#[derive(Clone, Debug)] +pub struct StandardSum; + +impl AggregateFnVTable for StandardSum { + type Options = NumericalAggregateOpts; + type Partial = StandardSumPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.standard_sum"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + // When a sum overflows, we return a sum _value_ of null. Therefore, we all return dtypes + // are nullable. + use Nullability::Nullable; + + Some(match input_dtype { + DType::Bool(_) => DType::Primitive(PType::U64, Nullable), + DType::Primitive(ptype, _) => match ptype { + PType::U8 | PType::U16 | PType::U32 | PType::U64 => { + DType::Primitive(PType::U64, Nullable) + } + PType::I8 | PType::I16 | PType::I32 | PType::I64 => { + DType::Primitive(PType::I64, Nullable) + } + PType::F16 | PType::F32 | PType::F64 => { + // Float sums cannot overflow, but all null floats still end up as null + DType::Primitive(PType::F64, Nullable) + } + }, + DType::Decimal(decimal_dtype, _) => { + DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable) + } + // Unsupported types + _ => return None, + }) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + Some(sum_partial_dtype(self.return_dtype(options, input_dtype)?)) + } + + fn empty_partial( + &self, + options: &Self::Options, + input_dtype: &DType, + ) -> VortexResult { + let return_dtype = self + .return_dtype(options, input_dtype) + .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; + Ok(StandardSumPartial { + return_dtype, + current: None, + overflowed: false, + skip_nans: options.skip_nans, + }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + // Partials are `{sum, overflowed}` structs. A null, non-overflowed sum is the identity. + // A plain scalar is a cached `Stat::Sum`; its null value means that Sum overflowed. + let (sum_value, other_overflowed) = if matches!(other.dtype(), DType::Struct(..)) { + if other.is_null() { + // A null struct partial carries no recoverable state; treat as saturated. + partial.current = None; + partial.overflowed = true; + return Ok(()); + } + let fields = other.as_struct(); + let sum_value = fields + .field("sum") + .ok_or_else(|| vortex_err!("StandardSum partial is missing the `sum` field"))?; + let other_overflowed = fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()) + .unwrap_or(true); + (sum_value, other_overflowed) + } else { + let other_overflowed = other.is_null(); + (other, other_overflowed) + }; + + if partial.overflowed || other_overflowed { + partial.current = None; + partial.overflowed = true; + return Ok(()); + } + if sum_value.is_null() { + // A null, non-overflowed sum is the empty identity. + return Ok(()); + } + + let inner = partial + .current + .get_or_insert_with(|| make_zero_state(&partial.return_dtype)); + let other = sum_value; + let saturated = match inner { + SumState::Unsigned(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + checked_add_u64(acc, val) + } + SumState::Signed(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + checked_add_i64(acc, val) + } + SumState::Float(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + *acc += val; + false + } + SumState::Decimal { value, dtype } => { + let val = other + .as_decimal() + .decimal_value() + .vortex_expect("checked non-null"); + match value.checked_add(&val) { + Some(r) => { + *value = r; + !value.fits_in_precision(*dtype) + } + None => true, + } + } + }; + if saturated { + partial.current = None; + partial.overflowed = true; + } + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::struct_( + sum_partial_dtype(partial.return_dtype.as_nullable()), + vec![ + sum_value_scalar(partial), + Scalar::bool(partial.overflowed, Nullability::NonNullable), + ], + )) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.current = None; + partial.overflowed = false; + } + + #[inline] + fn is_saturated(&self, partial: &Self::Partial) -> bool { + partial.overflowed + || matches!(partial.current.as_ref(), Some(SumState::Float(v)) if v.is_nan()) + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // `Stat::Sum` is produced by `Sum` aggregate, so the default NaN-skipping path can + // consume it directly. + if partial.skip_nans { + return try_accumulate_cached_sum(self, partial, batch); + } + // NaN-including float sums need a NaN-free batch before the cached sum applies; + // everything else takes the default dispatch path. + if !matches!(&partial.return_dtype, DType::Primitive(PType::F64, _)) { + return Ok(false); + } + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => { + // NaN-free batch: the cached NaN-skipping sum (if any) equals the + // NaN-including sum. + try_accumulate_cached_sum(self, partial, batch) + } + Precision::Exact(_) => { + // At least one NaN value (a valid value): the sum is NaN without scanning. + let inner = partial + .current + .get_or_insert_with(|| make_zero_state(&partial.return_dtype)); + let SumState::Float(acc) = inner else { + unreachable!("checked float return dtype") + }; + *acc = f64::NAN; + Ok(true) + } + _ => Ok(false), + } + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if partial.overflowed { + return Ok(()); + } + + // Constants compute scalar * len and combine via combine_partials. + if let Columnar::Constant(c) = batch { + // Valid NaN and `false` constants establish a real zero sum even when they contribute + // nothing to the running value. + if !c.scalar().is_null() && !c.is_empty() && partial.current.is_none() { + partial.current = Some(make_zero_state(&partial.return_dtype)); + } + // NaN constants are treated as missing when skipping NaNs. + if partial.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { + return Ok(()); + } + if let Some(product) = multiply_constant(c.scalar(), c.len(), &partial.return_dtype)? { + self.combine_partials(partial, product)?; + } + return Ok(()); + } + + let skip_nans = partial.skip_nans; + let any_valid = match batch { + Columnar::Canonical(c) => match c { + Canonical::Primitive(p) => { + any_valid(p.as_ref().validity()?, p.as_ref().len(), ctx)? + } + Canonical::Bool(b) => any_valid(b.as_ref().validity()?, b.as_ref().len(), ctx)?, + Canonical::Decimal(d) => any_valid(d.as_ref().validity()?, d.as_ref().len(), ctx)?, + _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), + }, + Columnar::Constant(_) => unreachable!(), + }; + let had_value = partial.current.is_some(); + let mut inner = partial + .current + .take() + .unwrap_or_else(|| make_zero_state(&partial.return_dtype)); + + // Whether the sum is empty is decided by validity alone (NaNs are valid values), while + // the summation reuses [`Sum`]'s accumulation kernels unchanged. + let result = match batch { + Columnar::Canonical(c) => match c { + Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans), + Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx), + _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), + }, + Columnar::Constant(_) => unreachable!(), + }; + + match result { + Ok(false) => { + if had_value || any_valid { + partial.current = Some(inner); + } + } + Ok(true) => partial.overflowed = true, + Err(e) => { + if had_value || any_valid { + partial.current = Some(inner); + } + return Err(e); + } + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + // Empty and overflowed sums have a null `sum` field, while null partial rows (e.g. null + // groups) stay null through the struct field projection's validity intersection. + GetItem.try_new_array(partials.len(), FieldName::from("sum"), [partials]) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(sum_value_scalar(partial)) + } +} + +/// Consume a batch's cached `Stat::Sum` instead of scanning it. The cached sum cannot +/// distinguish its zero identity from a real zero, so the batch's null count decides whether an +/// all-null batch contributes the empty identity; when either statistic is missing the caller +/// falls through to a real scan. +fn try_accumulate_cached_sum( + vtable: &StandardSum, + partial: &mut StandardSumPartial, + batch: &ArrayRef, +) -> VortexResult { + let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) else { + return Ok(false); + }; + match batch.statistics().get_as::(Stat::NullCount) { + Precision::Exact(null_count) if null_count == batch.len() as u64 => { + // No valid values: the batch is the identity. + return Ok(true); + } + Precision::Exact(_) => {} + _ => return Ok(false), + } + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + vtable.combine_partials(partial, sum)?; + Ok(true) +} + +/// Whether a batch contains at least one valid element. +fn any_valid(validity: Validity, len: usize, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(validity.execute_mask(len, ctx)?.true_count() > 0) +} + +/// The group state for a sum aggregate, containing the accumulated value and configuration +/// needed for reset/result without external context. +pub struct StandardSumPartial { + return_dtype: DType, + /// The current accumulated state, or `None` if no valid value has been accumulated or the sum + /// overflowed. [`Self::overflowed`] distinguishes those cases. + current: Option, + /// Whether checked arithmetic overflowed. Once set, this poisons subsequent merges. + overflowed: bool, + /// Whether NaN values in float inputs are skipped. + skip_nans: bool, +} + +/// The partial dtype for a sum whose result is `sum_dtype`: a `{sum, overflowed}` struct. A null, +/// non-overflowed sum is the merge identity, while `overflowed` poisons subsequent merges. The +/// outer struct remains nullable so grouped aggregation can represent null groups separately. +fn sum_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct( + StructFields::new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("overflowed")]), + vec![sum_dtype, DType::Bool(Nullability::NonNullable)], + ), + Nullability::Nullable, + ) +} + +/// The running sum as a nullable scalar of the return dtype (null when empty or overflowed). +fn sum_value_scalar(partial: &StandardSumPartial) -> Scalar { + match &partial.current { + None => Scalar::null(partial.return_dtype.as_nullable()), + Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Decimal { value, .. }) => { + let decimal_dtype = *partial + .return_dtype + .as_decimal_opt() + .vortex_expect("return dtype must be decimal"); + Scalar::decimal(*value, decimal_dtype, Nullability::Nullable) + } + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs new file mode 100644 index 00000000000..62033337321 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs @@ -0,0 +1,794 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for [`StandardSum`]-specific behavior: the `{sum, overflowed}` state algebra, the +//! null-for-zero-valid-values rule across input kinds, NaN and overflow handling, cached-statistic +//! consumption, and grouped aggregation. Plain summation arithmetic is covered by the shared +//! kernels' tests in [`super::super::sum`]. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::standard_sum::StandardSum; +use crate::aggregate_fn::fns::standard_sum::standard_sum; +use crate::aggregate_fn::fns::sum::sum; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +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; +use crate::dtype::Nullability::Nullable; +use crate::dtype::PType; +use crate::dtype::i256; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; +use crate::validity::Validity; + +/// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). +fn sum_with_options(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { + let mut acc = Accumulator::try_new(StandardSum, options, arr.dtype().clone())?; + acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.finish() +} + +// State algebra: the `{sum, overflowed}` monoid. + +#[test] +fn sum_state_empty_is_null() -> VortexResult<()> { + // A state that never saw a valid value finalizes to null, and combining empty states + // stays empty. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let empty = StandardSum.to_scalar(&state)?; + let fields = empty.as_struct(); + assert!(fields.field("sum").is_some_and(|sum| sum.is_null())); + assert_eq!( + fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()), + Some(false) + ); + StandardSum.combine_partials(&mut state, empty)?; + assert!(StandardSum.finalize_scalar(&state)?.is_null()); + Ok(()) +} + +#[test] +fn sum_state_empty_is_identity() -> VortexResult<()> { + // Combining an empty state into a non-empty state changes nothing: `{null, false}` is the + // identity of the `{sum, overflowed}` monoid. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + StandardSum.combine_partials(&mut state, Scalar::primitive(100i64, Nullable))?; + + let empty = StandardSum + .to_scalar(&StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?)?; + StandardSum.combine_partials(&mut state, empty)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert_eq!(result.as_primitive().typed_value::(), Some(100)); + Ok(()) +} + +#[test] +fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { + // Overflow sets the flag and poisons the merge even when combined with later values. + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut overflowed = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + StandardSum.combine_partials(&mut overflowed, Scalar::primitive(i64::MAX, Nullable))?; + StandardSum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; + let overflowed = StandardSum.to_scalar(&overflowed)?; + let fields = overflowed.as_struct(); + assert!(fields.field("sum").is_some_and(|sum| sum.is_null())); + assert_eq!( + fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()), + Some(true) + ); + + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + StandardSum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; + StandardSum.combine_partials(&mut state, overflowed)?; + StandardSum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; + + assert!(StandardSum.finalize_scalar(&state)?.is_null()); + Ok(()) +} + +// The null-for-zero-valid-values rule. + +#[rstest] +#[case::i32(DType::Primitive(PType::I32, Nullability::NonNullable))] +#[case::f64(DType::Primitive(PType::F64, Nullability::NonNullable))] +#[case::bool(DType::Bool(Nullability::NonNullable))] +fn sum_empty_is_null(#[case] dtype: DType) -> VortexResult<()> { + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + assert!(acc.finish()?.is_null()); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::from_option_iter([None::, None, None]).into_array())] +#[case::float(PrimitiveArray::from_option_iter::([None, None, None]).into_array())] +#[case::bool(BoolArray::from_iter([None::, None, None]).into_array())] +#[case::constant_primitive( + ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10).into_array() +)] +#[case::constant_bool(ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array())] +#[case::constant_decimal( + ConstantArray::new(Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullable)), 10) + .into_array() +)] +fn sum_all_null_is_null(#[case] array: ArrayRef) -> VortexResult<()> { + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert!(result.is_null()); + Ok(()) +} + +#[test] +fn sum_all_nan_is_zero_not_null() -> VortexResult<()> { + // NaNs are valid values: with the default `skip_nans` they contribute nothing, but + // the sum is a genuine `0.0`, unlike an all-null array whose sum is null. + let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) +} + +#[test] +fn sum_is_zero_while_standard_sum_is_null() -> VortexResult<()> { + // The persisted statistic keeps the Sum semantics (zero for all-null) that zone + // and chunk merging require, while the StandardSum applies the null-for-empty rule. + let mut ctx = array_session().create_execution_ctx(); + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + assert_eq!( + sum(&arr, &mut ctx)?.as_primitive().typed_value::(), + Some(0) + ); + // The cached Sum statistic must not leak through StandardSum's cache short-circuit. + assert!(standard_sum(&arr, &mut ctx)?.is_null()); + Ok(()) +} + +// Return dtype widening (mirrors `Sum`'s rules; the result is always nullable). + +#[rstest] +#[case::bool( + DType::Bool(Nullability::NonNullable), + DType::Primitive(PType::U64, Nullable) +)] +#[case::i32( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullable) +)] +#[case::u8( + DType::Primitive(PType::U8, Nullability::NonNullable), + DType::Primitive(PType::U64, Nullable) +)] +#[case::f32( + DType::Primitive(PType::F32, Nullability::NonNullable), + DType::Primitive(PType::F64, Nullable) +)] +#[case::decimal( + DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable), + DType::Decimal(DecimalDType::new(20, 2), Nullable) +)] +fn sum_return_dtype_widens(#[case] input: DType, #[case] expected: DType) { + let dtype = StandardSum + .return_dtype(&NumericalAggregateOpts::default(), &input) + .unwrap(); + assert_eq!(dtype, expected); +} + +// One value smoke test per accumulate branch; summation arithmetic is pinned by the +// shared kernels' tests in the `sum` module. + +#[test] +fn sum_primitive_with_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6)); + Ok(()) +} + +#[test] +fn sum_bool_with_nulls() -> VortexResult<()> { + let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(2)); + Ok(()) +} + +#[test] +fn sum_decimal_with_nulls() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32, 400i32], + DecimalDType::new(4, 2), + Validity::from_iter([true, false, true, true]), + ); + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullable), + Some(ScalarValue::from(DecimalValue::from(800i32))), + )?; + assert_eq!(result, expected); + Ok(()) +} + +#[test] +fn sum_constant() -> VortexResult<()> { + let array = ConstantArray::new(5u64, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 50u64.into()); + Ok(()) +} + +#[test] +fn sum_constant_false_is_zero_not_null() -> VortexResult<()> { + let array = ConstantArray::new(false, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) +} + +#[test] +fn sum_multi_batch_and_finish_resets() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + + let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + acc.accumulate(&batch1, &mut ctx)?; + let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(48)); + + // finish resets the state: an untouched accumulator is empty again. + assert!(acc.finish()?.is_null()); + let batch3 = PrimitiveArray::new(buffer![1i32], Validity::NonNullable).into_array(); + acc.accumulate(&batch3, &mut ctx)?; + assert_eq!(acc.finish()?.as_primitive().typed_value::(), Some(1)); + Ok(()) +} + +// Chunked accumulation: the nullable sum must merge across chunks. + +#[test] +fn sum_chunked_floats_with_nulls() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter(vec![Some(1.5f64), None, Some(3.2), Some(4.8)]); + let chunk2 = PrimitiveArray::from_option_iter(vec![Some(2.1f64), Some(5.7), None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = standard_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(17.3)); + Ok(()) +} + +#[test] +fn sum_chunked_all_nulls_is_null() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter::(vec![None, None, None]); + let chunk2 = PrimitiveArray::from_option_iter::(vec![None, None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + let result = standard_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert!(result.is_null()); + Ok(()) +} + +#[test] +fn sum_chunked_empty_chunks() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter(vec![Some(10.5f64), Some(20.3)]); + let chunk2 = ConstantArray::new(Scalar::primitive(0f64, Nullable), 0); + let chunk3 = PrimitiveArray::from_option_iter(vec![Some(5.2f64)]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new( + vec![ + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), + ], + dtype, + )?; + + let result = standard_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(36.0)); + Ok(()) +} + +#[test] +fn sum_chunked_value_survives_empty_chunk() -> VortexResult<()> { + // One valid value in one chunk, followed by an all-null chunk: the value must survive merging + // with the second chunk's null identity. + let chunk1 = PrimitiveArray::from_option_iter::(vec![Some(1)]); + let chunk2 = PrimitiveArray::from_option_iter::(vec![None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = standard_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(1)); + Ok(()) +} + +// NaN handling and its interplay with the nullable sum and cached statistics. + +#[test] +fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(f64::NAN), Some(3.0)]) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); + Ok(()) +} + +#[test] +fn sum_f64_with_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, f64::NAN, 2.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + // The array has no NaNs; a planted exact NaNCount stat proves the NaN poisoning came + // from the stat rather than a scan. + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_uses_cached_stat_sum() -> VortexResult<()> { + // A planted exact `Stat::Sum` with a known null count is consumed instead of a scan + // (the planted value differs from the actual data to prove it). + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); + arr.statistics() + .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) +} + +#[test] +fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact Sum stat is usable as-is. + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); + arr.statistics() + .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) +} + +#[test] +fn sum_constant_nan() -> VortexResult<()> { + let arr = ConstantArray::new(f64::NAN, 4).into_array(); + // NaN constants are skipped by default (a non-empty zero sum) and poison the sum otherwise. + let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_f64_with_infinity() -> VortexResult<()> { + let batch = PrimitiveArray::new( + buffer![1.0f64, f64::INFINITY, f64::NEG_INFINITY, 2.0], + Validity::NonNullable, + ) + .into_array(); + let acc = standard_sum(&batch, &mut array_session().create_execution_ctx())?; + // INFINITY + NEG_INFINITY = NaN, which is treated as saturated + assert!(acc.as_primitive().typed_value::().unwrap().is_nan()); + + let mut acc = Accumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + DType::Primitive(PType::F64, Nullability::NonNullable), + )?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + Ok(()) +} + +// Overflow: a null sum value plus an explicit saturation flag. + +#[test] +fn sum_checked_overflow_is_null_and_saturates() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + assert!(!acc.is_saturated()); + + let batch = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + let result = acc.finish()?; + assert!(result.is_null()); + + // finish resets state, clearing saturation + assert!(!acc.is_saturated()); + Ok(()) +} + +#[test] +fn sum_decimal_i256_overflow() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(76, 0); + let decimal = DecimalArray::new( + buffer![i256::MAX, i256::MAX, i256::MAX], + decimal_dtype, + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!( + result, + Scalar::null(DType::Decimal(decimal_dtype, Nullable)) + ); + Ok(()) +} + +#[test] +fn sum_decimal_near_precision_boundary() -> VortexResult<()> { + // Input precision 4 → return precision min(76, 4+10) = 14. + // Native type for precision 14 is I64 (max precision 18), so 14 < 18. + // Use combine_partials to push state near (but under) 10^14. + let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(99_999_999_999_990i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + // Add a small value that keeps us just under 10^14. + let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); + StandardSum.combine_partials(&mut state, small)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert!(!result.is_null()); + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(99_999_999_999_999))) + ); + Ok(()) +} + +#[rstest] +#[case::positive(99_999_999_999_999i64, 1i64)] +#[case::negative(-99_999_999_999_999i64, -1i64)] +fn sum_decimal_precision_overflow_within_i256( + #[case] near_limit: i64, + #[case] one_more: i64, +) -> VortexResult<()> { + // Input precision 4 → return precision 14. Native I64 (max 18). + // The max representable magnitude for precision 14 is 10^14 - 1: pushing the sum to + // exactly ±10^14 fails fits_in_precision even though i256 arithmetic does not + // overflow. This tests the precision-based saturation path in combine_partials. + let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(near_limit), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + let one_more = Scalar::decimal( + DecimalValue::from(one_more), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, one_more)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert!(result.is_null()); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(14, 0), Nullable) + ); + Ok(()) +} + +#[test] +fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { + // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Input precision 27 → return precision 37. Native for 37 is I128 (max 38), so 37 < 38. + // Use combine_partials to get the state close to 10^37, then accumulate a real array + // that pushes it over. + let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); + let return_dtype = DecimalDType::new(37, 0); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + // Set state to 10^37 - 1 via combine_partials. + let near_limit_val: i128 = 10i128.pow(37) - 1; + let near_limit = Scalar::decimal(DecimalValue::from(near_limit_val), return_dtype, Nullable); + StandardSum.combine_partials(&mut state, near_limit)?; + + // Now accumulate a real i128 array with a single element = 1 to overflow precision. + let decimal = DecimalArray::new(buffer![1i128], DecimalDType::new(27, 0), Validity::AllValid); + let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); + let mut ctx = array_session().create_execution_ctx(); + StandardSum.accumulate(&mut state, &columnar, &mut ctx)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert!(result.is_null()); + Ok(()) +} + +// Grouped aggregation: empty, all-null, and null groups through the lazy finalize. + +fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { + let mut acc = GroupedAccumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish() +} + +#[test] +fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(5i64), None, Some(i64::MAX), Some(1)]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 0, 1, 2, 0].into_array(), + buffer![1i32, 0, 1, 2, 1].into_array(), + Validity::from_iter([true, true, true, true, false]), + )? + .into_array(); + let mut acc = GroupedAccumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + DType::Primitive(PType::I64, Nullable), + )?; + acc.accumulate_list(&groups, &mut ctx)?; + let partials = acc.flush()?; + + let value = partials.execute_scalar(0, &mut ctx)?; + let fields = value.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(5) + ); + assert_eq!( + fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()), + Some(false) + ); + + for index in [1, 2] { + let empty = partials.execute_scalar(index, &mut ctx)?; + let fields = empty.as_struct(); + assert!(fields.field("sum").is_some_and(|sum| sum.is_null())); + assert_eq!( + fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()), + Some(false) + ); + } + + let overflow = partials.execute_scalar(3, &mut ctx)?; + let fields = overflow.as_struct(); + assert!(fields.field("sum").is_some_and(|sum| sum.is_null())); + assert_eq!( + fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()), + Some(true) + ); + assert!(partials.execute_scalar(4, &mut ctx)?.is_null()); + Ok(()) +} + +#[test] +fn grouped_sum_fallback_empty_and_all_null_groups() -> VortexResult<()> { + // Bool elements are rejected by the primitive grouped kernel, forcing the generic + // per-group fallback: empty and all-null groups have null sums there too. + let mut ctx = array_session().create_execution_ctx(); + let elements = BoolArray::from_iter([Some(true), Some(true), None, None]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 2, 2].into_array(), + buffer![2i32, 0, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let result = run_grouped_sum(&groups, &DType::Bool(Nullable))?; + let expected = PrimitiveArray::from_option_iter([Some(2u64), None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_fixed_size_list() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); + let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_with_null_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) + .into_array(); + let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_with_null_group() -> VortexResult<()> { + // A null group must become a null row through the lazy finalize's validity handling. + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) + .into_array(); + let validity = Validity::from_iter([true, false, true]); + let groups = FixedSizeListArray::try_new(elements, 3, validity, 3)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + // The all-null group has a null sum + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_all_nan_is_zero_not_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![f64::NAN, f64::NAN, 3.0, 4.0], Validity::NonNullable) + .into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(0.0f64), Some(7.0)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_finish_resets() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = + GroupedAccumulator::try_new(StandardSum, NumericalAggregateOpts::default(), elem_dtype)?; + + let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; + acc.accumulate_list(&groups1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + + let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); + assert_arrays_eq!(&result1, &expected1, &mut ctx); + + let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; + acc.accumulate_list(&groups2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + + let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); + assert_arrays_eq!(&result2, &expected2, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); + let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new(buffer![1i32, 1, 1], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false, true]); + let groups = ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array(); + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups, &elem_dtype)?; + + // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. + let expected = + PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3993840586..f15356caaa1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -13,7 +13,7 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; -pub(super) fn accumulate_bool( +pub(crate) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..ced92af1d68 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -15,7 +15,7 @@ use crate::scalar::Scalar; /// /// Returns `Ok(None)` if the scalar is null (no contribution to the sum). /// Returns a null scalar on overflow (saturation). -pub(super) fn multiply_constant( +pub(crate) fn multiply_constant( scalar: &Scalar, len: usize, return_dtype: &DType, diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 872e5769a01..197b8f10b04 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -23,7 +23,7 @@ use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. /// Returns Ok(true) if saturated (overflow), Ok(false) if not. -pub(super) fn accumulate_decimal( +pub(crate) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 010371fee4f..2a6b0303d96 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -15,10 +15,15 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use self::bool::accumulate_bool; -use self::constant::multiply_constant; -use self::decimal::accumulate_decimal; -use self::primitive::accumulate_primitive; +// The accumulation kernels are also used by `StandardSum`, which tracks its `seen` bit +// separately from validity; the kernels themselves are unchanged. +pub(crate) use self::bool::accumulate_bool; +pub(crate) use self::constant::multiply_constant; +pub(crate) use self::decimal::accumulate_decimal; +pub(crate) use self::primitive::accumulate_primitive; +pub(crate) use self::primitive::sum_float_all; +pub(crate) use self::primitive::sum_signed_all; +pub(crate) use self::primitive::sum_unsigned_all; use crate::ArrayRef; use crate::Canonical; use crate::Columnar; @@ -69,11 +74,9 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// Sum an array, starting from zero. /// -/// If the sum overflows, a null scalar will be returned. -/// If the array is all-invalid, the sum will be zero. -/// -/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the -/// default) NaN values contribute nothing, otherwise any NaN value poisons the sum to NaN. +/// If the sum overflows, a null scalar will be returned. If the array is all-invalid or empty, the sum will be zero. +/// Note that sum aggregates typically produce null for arrays without at least one valid element. Use +/// [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum) for this behavior and more detail. #[derive(Clone, Debug)] pub struct Sum; @@ -351,7 +354,7 @@ pub enum SumState { }, } -fn make_zero_state(return_dtype: &DType) -> SumState { +pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0), @@ -368,7 +371,7 @@ fn make_zero_state(return_dtype: &DType) -> SumState { /// Checked add for u64, returning true if overflow occurred. #[inline(always)] -fn checked_add_u64(acc: &mut u64, val: u64) -> bool { +pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; @@ -380,7 +383,7 @@ fn checked_add_u64(acc: &mut u64, val: u64) -> bool { /// Checked add for i64, returning true if overflow occurred. #[inline(always)] -fn checked_add_i64(acc: &mut i64, val: i64) -> bool { +pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..930a353f6e1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -21,7 +21,7 @@ use crate::match_each_native_ptype; /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. const SUM_CHUNK: usize = 1 << 16; -pub(super) fn accumulate_primitive( +pub(crate) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -66,7 +66,7 @@ fn accumulate_primitive_all( /// Sum the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values /// are skipped to match the scalar `sum` semantics; otherwise any NaN poisons the accumulator to /// NaN. Floats cannot overflow the accumulator, so this never reports saturation. -pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { +pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { if skip_nans { for &v in slice { if !v.is_nan() { @@ -84,7 +84,7 @@ pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan /// chunks of [`SUM_CHUNK`] with a single checked add per chunk, which lets the inner loop vectorize /// to packed widening adds. `u64` input keeps a per-element checked add since a chunk of `u64`s /// could itself overflow. Returns `true` on overflow. -pub(super) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool +pub(crate) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { @@ -106,7 +106,7 @@ where } /// Signed counterpart of [`sum_unsigned_all`]. -pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool +pub(crate) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c496f7a9bf8..25b270d5b7b 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -29,6 +29,8 @@ use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; +use crate::aggregate_fn::fns::standard_sum::PrimitiveGroupedStandardSumEncodingKernel; +use crate::aggregate_fn::fns::standard_sum::StandardSum; use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; @@ -99,6 +101,7 @@ impl Default for AggregateFnSession { this.register(MinMax); this.register(NanCount); this.register(NullCount); + this.register(StandardSum); this.register(Sum); this.register(UncompressedSizeInBytes); @@ -115,6 +118,11 @@ impl Default for AggregateFnSession { Sum.id(), &PrimitiveGroupedSumEncodingKernel, ); + this.register_grouped_encoding_kernel( + Primitive.id(), + StandardSum.id(), + &PrimitiveGroupedStandardSumEncodingKernel, + ); this } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..31aaee7c39c 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -103,8 +103,14 @@ impl ArrayParentReduceRule for StructGetItemRule { match child.validity()? { Validity::NonNullable | Validity::AllValid => { - // If the struct is non-nullable or all valid, the field's validity is unchanged - Ok(Some(field.clone())) + // The field's values are unchanged, but `get_item` on a *nullable* struct has a + // nullable result dtype even when every struct row is valid, so a non-nullable + // field needs a nullability cast to match. + if child.as_ref().dtype().is_nullable() && !field.dtype().is_nullable() { + field.clone().cast(field.dtype().as_nullable()).map(Some) + } else { + Ok(Some(field.clone())) + } } Validity::AllInvalid => { // If everything is invalid, the field is also all invalid diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index a5f8d219db6..69565597deb 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -1,11 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -16,30 +14,23 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; -use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedAccumulator; -use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::sum::Sum; -use crate::arrays::BoolArray; +use crate::aggregate_fn::fns::standard_sum::StandardSum; use crate::arrays::ConstantArray; -use crate::arrays::FixedSizeList; -use crate::arrays::ListView; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::validity::Validity; /// Sum of the elements in each list of a `List` or `FixedSizeList` typed array. /// /// Follows SQL `SUM` semantics per list, matching DuckDB's `list_sum`: null lists, empty /// lists, and lists whose elements are all null yield a null sum; null elements are skipped. -/// Integer and decimal overflow yields a null sum value, matching [`Sum`]. The result dtype -/// follows [`Sum`]'s widening rules and is always nullable. +/// Integer and decimal overflow yields a null sum value, matching [`StandardSum`]. The result +/// dtype follows [`StandardSum`]'s widening rules and is always nullable. /// /// NaN handling for float elements is controlled by [`NumericalAggregateOpts`]: with /// `skip_nans` (the default) NaN values contribute nothing, otherwise any NaN poisons the @@ -83,7 +74,8 @@ impl ScalarFnVTable for ListSum { DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref(), other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), }; - Sum.return_dtype(options, elem_dtype) + StandardSum + .return_dtype(options, elem_dtype) .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) } @@ -100,7 +92,6 @@ impl ScalarFnVTable for ListSum { other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), }; - // `mask_empty_lists` needs access to list elements validity and sizes let columnar = input.execute::(ctx)?; match columnar { @@ -131,60 +122,17 @@ impl ScalarFnVTable for ListSum { /// Sum each list of a canonical `array` into one value per list. /// -/// Note that we need to nullify sums produced by empty or all-null lists, -/// since grouped sum kernels default to 0 for these. +/// The grouped [`StandardSum`] finalize already yields null for null, empty, and all-null lists +/// (SQL `SUM` semantics), so no post-processing is needed. fn list_sum_impl( canonical: ArrayRef, elem_dtype: DType, options: &NumericalAggregateOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(StandardSum, *options, elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; - let sums = acc.finish()?; - - let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::() { - fsl.into_owned().into() - } else if let Some(lv) = canonical.as_opt::() { - lv.into_owned().into() - } else { - let dtype = canonical.dtype(); - vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}") - }; - - mask_empty_lists(grouped, sums, ctx) -} - -/// Applies a mask to `sums` that nullifies entries produced by lists without at least -/// one valid element. This is necessary because the grouped `Sum` aggregate only produces -/// nulls for null lists and sums that overflow. -fn mask_empty_lists( - grouped: GroupedArray, - sums: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let elements = grouped.elements(); - let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; - let ranges = grouped.group_ranges(ctx)?; - - let has_valid_element: BitBuffer = match elem_mask.bit_buffer() { - AllOr::All => match &ranges { - // fixed-size lists of non-zero width cannot have empty lists. - GroupRanges::FixedSizeList { size, .. } if *size > 0 => return Ok(sums), - GroupRanges::FixedSizeList { len, .. } => BitBuffer::full(false, *len), - GroupRanges::ListView { ranges } => ranges.iter().map(|&(_, size)| size > 0).collect(), - }, - AllOr::None => BitBuffer::full(false, ranges.len()), - AllOr::Some(bits) => ranges - .iter() - .map(|(offset, size)| size > 0 && bits.count_range(offset, offset + size) > 0) - .collect(), - }; - if has_valid_element.true_count() == has_valid_element.len() { - return Ok(sums); - } - - sums.mask(BoolArray::new(has_valid_element, Validity::NonNullable).into_array()) + acc.finish() } #[cfg(test)]