From 98705269bada49d7afd663a3d532d8e833f6b62a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 24 Aug 2026 14:39:39 -0400 Subject: [PATCH 1/5] feat: add sum_v2 aggregate Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/mod.rs | 1 + vortex-array/src/aggregate_fn/fns/sum/bool.rs | 2 +- .../src/aggregate_fn/fns/sum/constant.rs | 2 +- .../src/aggregate_fn/fns/sum/decimal.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 13 +- .../src/aggregate_fn/fns/sum/primitive.rs | 8 +- .../src/aggregate_fn/fns/sum_v2/grouped.rs | 192 +++++++++ .../src/aggregate_fn/fns/sum_v2/mod.rs | 398 ++++++++++++++++++ .../src/aggregate_fn/fns/sum_v2/tests.rs | 265 ++++++++++++ vortex-array/src/aggregate_fn/session.rs | 8 + vortex-duckdb/src/convert/expr.rs | 4 +- .../slt/duckdb/aggregates_edge_cases.slt | 3 +- .../slt/duckdb/nan_aggregates.slt | 2 +- 13 files changed, 883 insertions(+), 17 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs create mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs create mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index da81e67ecb9..c4e8d1b72ca 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -20,4 +20,5 @@ pub mod min_max; pub mod nan_count; pub mod null_count; pub mod sum; +pub mod sum_v2; pub mod uncompressed_size_in_bytes; 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 7ad62eacb87..1ceb529bb10 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -15,10 +15,13 @@ 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; +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; @@ -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), 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/fns/sum_v2/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs new file mode 100644 index 00000000000..5c5e1d7fa6a --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs @@ -0,0 +1,192 @@ +// 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::Mask; + +use super::SumV2; +use super::sum_v2_partial_fields; +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::NativePType; +use crate::dtype::Nullability; +use crate::match_each_native_ptype; +use crate::validity::Validity; + +/// Encoding-specific grouped [`SumV2`] kernel for primitive element arrays. +#[derive(Debug)] +pub(crate) struct PrimitiveGroupedSumV2EncodingKernel; + +impl DynGroupedAggregateKernel for PrimitiveGroupedSumV2EncodingKernel { + 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 [`SumV2`] 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_v2` 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. +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, + )?)) +} + +/// Sum 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 = elem_mask.all_true(); + + let (sums, is_overflow, is_empty) = 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 }) + } + ); + + let partial_fields = sum_v2_partial_fields(sums.dtype().clone()); + + // SAFETY: all three children have one value per group and match `partial_fields`; the struct + // validity is derived from the same group count. + Ok(unsafe { + StructArray::new_unchecked( + vec![ + sums.into_array(), + BoolArray::new(is_overflow, Validity::NonNullable).into_array(), + BoolArray::new(is_empty, Validity::NonNullable).into_array(), + ], + partial_fields, + group_validity.len(), + Validity::from_mask(group_validity.clone(), Nullability::Nullable), + ) + } + .into_array()) +} + +/// Reduce each group's element slice into a non-null sum, overflow bitmap, and empty bitmap. +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, BitBuffer) { + let group_count = group_ranges.len(); + let mut is_overflow = BitBufferMut::new_unset(group_count); + let mut is_empty = BitBufferMut::new_unset(group_count); + let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { + if !group_validity.value(i) { + return A::default(); + } + 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) + }; + if overflow { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_overflow.set_unchecked(i) }; + } + if !any_valid { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_empty.set_unchecked(i) }; + } + acc + }); + let sums = PrimitiveArray::from_iter(sums); + (sums, is_overflow.freeze(), is_empty.freeze()) +} + +/// Sum valid runs in one group, returning `(overflow, any_valid)`. +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 { + Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), + Mask::AllFalse(_) => (false, false), + Mask::Values(mask_values) => { + let validity = mask_values + .bit_buffer() + .as_view() + .slice(offset..offset + size); + let mut any_valid = false; + for (start, end) in validity.set_slices() { + any_valid = true; + if sum_run(acc, &values[offset + start..offset + end]) { + return (true, true); + } + } + (false, any_valid) + } + } +} diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs new file mode 100644 index 00000000000..8dea7c77890 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -0,0 +1,398 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod grouped; + +pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +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::Sum; +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::make_zero_state; +use crate::aggregate_fn::fns::sum::multiply_constant; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProviderExt; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; + +const SUM_FIELD: &str = "sum"; +const IS_OVERFLOW_FIELD: &str = "is_overflow"; +const IS_EMPTY_FIELD: &str = "is_empty"; + +/// Return the SQL-style sum of an array. +/// +/// Unlike [`sum`](crate::aggregate_fn::fns::sum::sum), this returns null when the array has no +/// valid values. +pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let mut acc = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + acc.accumulate(array, ctx)?; + acc.finish() +} + +/// Sum an array, returning null when it has no valid values or if the sum overflows. +/// +/// This aggregate intentionally has a distinct ID and partial representation from the legacy +/// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex +/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans` +/// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise, +/// any NaN value poisons the sum to NaN. +#[derive(Clone, Copy, Debug)] +pub struct SumV2; + +impl AggregateFnVTable for SumV2 { + type Options = NumericalAggregateOpts; + type Partial = SumV2Partial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.sum_v2"); + *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 { + Sum.return_dtype(options, input_dtype) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + .map(sum_v2_partial_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_v2 dtype: {}", input_dtype))?; + let sum = make_zero_state(&return_dtype); + Ok(SumV2Partial { + return_dtype, + sum, + is_overflow: false, + is_empty: true, + skip_nans: options.skip_nans, + }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + let (other_sum, other_is_overflow, other_is_empty) = decode_partial_scalar(other)?; + validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; + + if partial.is_overflow { + return Ok(()); + } + if other_is_overflow { + partial.is_overflow = true; + partial.is_empty = false; + return Ok(()); + } + if other_is_empty { + return Ok(()); + } + + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; + partial.is_empty = false; + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::struct_( + sum_v2_partial_dtype(partial.return_dtype.clone()), + vec![ + sum_state_scalar(partial, Nullability::NonNullable), + Scalar::bool(partial.is_overflow, Nullability::NonNullable), + Scalar::bool(partial.is_empty, Nullability::NonNullable), + ], + )) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.sum = make_zero_state(&partial.return_dtype); + partial.is_overflow = false; + partial.is_empty = true; + } + + #[inline] + fn is_saturated(&self, partial: &Self::Partial) -> bool { + partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + if partial.skip_nans || !matches!(&partial.sum, SumState::Float(_)) { + return Ok(false); + } + + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => Ok(false), + Precision::Exact(_) => { + let SumState::Float(sum) = &mut partial.sum else { + unreachable!("checked float sum state") + }; + *sum = f64::NAN; + partial.is_empty = false; + Ok(true) + } + _ => Ok(false), + } + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if partial.is_overflow { + return Ok(()); + } + + if let Columnar::Constant(constant) = batch { + if !constant.scalar().is_null() && !constant.is_empty() { + partial.is_empty = false; + } + if partial.skip_nans + && constant + .scalar() + .as_primitive_opt() + .is_some_and(|primitive| primitive.is_nan()) + { + return Ok(()); + } + if let Some(product) = + multiply_constant(constant.scalar(), constant.len(), &partial.return_dtype)? + { + if product.is_null() { + partial.is_overflow = true; + partial.is_empty = false; + } else { + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &product)?; + partial.is_empty = false; + } + } + return Ok(()); + } + + let any_valid = partial.is_empty && has_valid_value(batch, ctx)?; + let result = match batch { + Columnar::Canonical(canonical) => match canonical { + Canonical::Primitive(array) => { + accumulate_primitive(&mut partial.sum, array, ctx, partial.skip_nans) + } + Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx), + Canonical::Decimal(array) => accumulate_decimal(&mut partial.sum, array, ctx), + _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()), + }, + Columnar::Constant(_) => unreachable!(), + }; + + if any_valid { + partial.is_empty = false; + } + if result? { + partial.is_overflow = true; + partial.is_empty = false; + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + let sum = partials.get_item(SUM_FIELD)?; + let is_invalid = partials + .get_item(IS_OVERFLOW_FIELD)? + .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? + .fill_null(true)?; + sum.mask(is_invalid.not()?) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + if partial.is_overflow || partial.is_empty { + return Ok(Scalar::null(partial.return_dtype.as_nullable())); + } + Ok(sum_state_scalar(partial, Nullability::Nullable)) + } +} + +/// In-memory state for SumV2 accumulation. +pub struct SumV2Partial { + return_dtype: DType, + sum: SumState, + is_overflow: bool, + is_empty: bool, + skip_nans: bool, +} + +fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult { + let (validity, len) = match batch { + Columnar::Canonical(Canonical::Primitive(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(Canonical::Bool(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(Canonical::Decimal(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(_) => return Ok(false), + Columnar::Constant(constant) => { + return Ok(!constant.is_empty() && !constant.scalar().is_null()); + } + }; + Ok(validity.execute_mask(len, ctx)?.true_count() > 0) +} + +fn decode_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { + vortex_ensure!(!scalar.is_null(), "SumV2 partial must not be null"); + + let Some(fields) = scalar.as_struct_opt() else { + vortex_bail!("SumV2 partial must be a struct, got {}", scalar.dtype()); + }; + let sum = fields + .field(SUM_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the sum field"))?; + let is_overflow = bool::try_from( + &fields + .field(IS_OVERFLOW_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_overflow field"))?, + )?; + let is_empty = bool::try_from( + &fields + .field(IS_EMPTY_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_empty field"))?, + )?; + + Ok((sum, is_overflow, is_empty)) +} + +fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + sum.dtype().nullability() == Nullability::NonNullable + && sum.dtype().eq_ignore_nullability(return_dtype), + "SumV2 partial value has dtype {}, expected {}", + sum.dtype(), + return_dtype.as_nonnullable(), + ); + Ok(()) +} + +fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { + Ok(match state { + SumState::Unsigned(sum) => checked_add_u64(sum, u64::try_from(other)?), + SumState::Signed(sum) => checked_add_i64(sum, i64::try_from(other)?), + SumState::Float(sum) => { + *sum += f64::try_from(other)?; + false + } + SumState::Decimal { value, dtype } => { + let other = DecimalValue::try_from(other)?; + match value.checked_add(&other) { + Some(result) if result.fits_in_precision(*dtype) => { + *value = result; + false + } + Some(_) | None => true, + } + } + }) +} + +fn sum_v2_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct(sum_v2_partial_fields(sum_dtype), Nullability::Nullable) +} + +fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields { + StructFields::new( + FieldNames::from_iter([ + FieldName::from(SUM_FIELD), + FieldName::from(IS_OVERFLOW_FIELD), + FieldName::from(IS_EMPTY_FIELD), + ]), + vec![ + sum_dtype.as_nonnullable(), + DType::Bool(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + ) +} + +fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar { + match &partial.sum { + SumState::Unsigned(value) => Scalar::primitive(*value, nullability), + SumState::Signed(value) => Scalar::primitive(*value, nullability), + SumState::Float(value) => Scalar::primitive(*value, nullability), + SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), + } +} + +#[inline(always)] +fn checked_add_u64(sum: &mut u64, value: u64) -> bool { + match sum.checked_add(value) { + Some(result) => { + *sum = result; + false + } + None => true, + } +} + +#[inline(always)] +fn checked_add_i64(sum: &mut i64, value: i64) -> bool { + match sum.checked_add(value) { + Some(result) => { + *sum = result; + false + } + None => true, + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs new file mode 100644 index 00000000000..2f16e533f7c --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -0,0 +1,265 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_proto::expr as pb; + +use super::SumV2; +use super::sum_v2; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::AggregateFnVTableExt; +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::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::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::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; +use crate::validity::Validity; + +fn sum_with_options(array: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { + let mut accumulator = Accumulator::try_new(SumV2, options, array.dtype().clone())?; + accumulator.accumulate(array, &mut array_session().create_execution_ctx())?; + accumulator.finish() +} + +#[test] +fn distinct_id_and_struct_partial() -> VortexResult<()> { + let options = NumericalAggregateOpts::default(); + let aggregate = SumV2.bind(options); + assert_eq!(aggregate.id().as_ref(), "vortex.sum_v2"); + + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let partial_dtype = SumV2 + .partial_dtype(&options, &input_dtype) + .expect("supported sum_v2 dtype"); + assert_eq!(partial_dtype.nullability(), Nullable); + let fields = partial_dtype.as_struct_fields(); + assert_eq!(fields.names().as_ref(), &["sum", "is_overflow", "is_empty"]); + assert_eq!( + fields.field("sum"), + Some(DType::Primitive(PType::I64, Nullability::NonNullable)) + ); + + let proto = aggregate.serialize_proto()?; + let decoded = pb::AggregateFn::decode(proto.encode_to_vec().as_slice())?; + let round_tripped = AggregateFnRef::from_proto(&decoded, &array_session())?; + assert_eq!(round_tripped, aggregate); + Ok(()) +} + +#[test] +fn legacy_sum_is_unchanged() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let options = NumericalAggregateOpts::default(); + assert_eq!(Sum.bind(options).id().as_ref(), "vortex.sum"); + assert!(matches!( + Sum.partial_dtype(&options, &dtype), + Some(DType::Primitive(..)) + )); + + let mut legacy = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone())?; + let mut v2 = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + + assert_eq!( + legacy.finish()?.as_primitive().typed_value::(), + Some(0) + ); + assert!(v2.finish()?.is_null()); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::from_option_iter([None::, None]).into_array())] +#[case::float(PrimitiveArray::from_option_iter::([None, None]).into_array())] +#[case::bool(BoolArray::from_iter([None::, None]).into_array())] +#[case::constant( + ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 4).into_array() +)] +#[case::decimal( + DecimalArray::new( + buffer![0i64, 0], + DecimalDType::new(10, 2), + Validity::AllInvalid, + ) + .into_array() +)] +fn all_null_is_null(#[case] array: ArrayRef) -> VortexResult<()> { + assert!(sum_v2(&array, &mut array_session().create_execution_ctx())?.is_null()); + Ok(()) +} + +#[test] +fn valid_zero_is_not_empty() -> VortexResult<()> { + let zeros = PrimitiveArray::from_iter([0i32, 0]).into_array(); + assert_eq!( + sum_v2(&zeros, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0) + ); + + let false_values = BoolArray::from_iter([false, false]).into_array(); + assert_eq!( + sum_v2(&false_values, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0) + ); + Ok(()) +} + +#[test] +fn all_nan_is_nonempty_zero_when_skipped() -> VortexResult<()> { + let array = + PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); + assert_eq!( + sum_v2(&array, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0.0) + ); + + let included = sum_with_options(&array, NumericalAggregateOpts::include_nans())?; + assert!( + included + .as_primitive() + .typed_value::() + .expect("non-null sum") + .is_nan() + ); + Ok(()) +} + +#[test] +fn ignores_ambiguous_legacy_sum_stat() -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter([None::, None]).into_array(); + array + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42i64))); + + assert!(sum_v2(&array, &mut array_session().create_execution_ctx())?.is_null()); + Ok(()) +} + +#[test] +fn overflow_is_explicit_and_absorbing() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut accumulator = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + let batch = PrimitiveArray::new(buffer![i64::MAX, 1, 7], Validity::NonNullable).into_array(); + accumulator.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(accumulator.is_saturated()); + + let partial = accumulator.partial_scalar()?; + let fields = partial.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|flag| flag.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|flag| flag.as_bool().value()), + Some(false) + ); + assert!(accumulator.finish()?.is_null()); + Ok(()) +} + +#[test] +fn empty_chunk_is_a_merge_identity() -> VortexResult<()> { + let value = PrimitiveArray::from_option_iter([Some(7i32)]).into_array(); + let empty = PrimitiveArray::from_option_iter([None::]).into_array(); + let chunked = + ChunkedArray::try_new(vec![value, empty], DType::Primitive(PType::I32, Nullable))? + .into_array(); + + assert_eq!( + sum_v2(&chunked, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(7) + ); + Ok(()) +} + +#[test] +fn grouped_primitive_tracks_empty_overflow_and_null_groups() -> VortexResult<()> { + 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 accumulator = GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + DType::Primitive(PType::I64, Nullable), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate_list(&groups, &mut ctx)?; + + let result = accumulator.finish()?; + let expected = + PrimitiveArray::from_option_iter([Some(5i64), None, None, None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_bool_fallback_tracks_empty_groups() -> VortexResult<()> { + 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 mut accumulator = GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + DType::Bool(Nullable), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate_list(&groups, &mut ctx)?; + + let result = accumulator.finish()?; + let expected = PrimitiveArray::from_option_iter([Some(2u64), None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 0fd07afc9c4..65b00d81031 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -33,6 +33,8 @@ use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::PrimitiveGroupedSumV2EncodingKernel; +use crate::aggregate_fn::fns::sum_v2::SumV2; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; @@ -111,6 +113,7 @@ impl Default for AggregateFnSession { this.register(NanCount); this.register(NullCount); this.register(Sum); + this.register(SumV2); this.register(UncompressedSizeInBytes); // Register the built-in aggregate kernels. @@ -126,6 +129,11 @@ impl Default for AggregateFnSession { Sum.id(), &PrimitiveGroupedSumEncodingKernel, ); + this.register_grouped_encoding_kernel( + Primitive.id(), + SumV2.id(), + &PrimitiveGroupedSumV2EncodingKernel, + ); this } diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 66bb7aa6df9..5da738ef8eb 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -17,7 +17,7 @@ use vortex::aggregate_fn::fns::first::First; use vortex::aggregate_fn::fns::max::Max; use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::sum::Sum; +use vortex::aggregate_fn::fns::sum_v2::SumV2; use vortex::arrow::ArrowSessionExt; use vortex::dtype::DType; use vortex::dtype::Nullability; @@ -535,7 +535,7 @@ impl PushedAggregate { Ok(match self { Self::Min => Box::new(Accumulator::try_new(Min, opts, dtype)?), Self::Max => Box::new(Accumulator::try_new(Max, opts, dtype)?), - Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new(SumV2, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), PairOptions(opts, opts), diff --git a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt index fbf9d2b74fe..c0c7f9128e9 100644 --- a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt +++ b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt @@ -13,8 +13,7 @@ statement ok COPY (SELECT CAST(x AS INTEGER) AS x FROM (VALUES (CAST(NULL AS INTEGER)),(CAST(NULL AS INTEGER))) t(x)) TO '${WORK_DIR}/i-null.vortex'; -# TODO(https://github.com/vortex-data/vortex/issues/9084): this should be fixed (expected NULL) query I SELECT sum(x) FROM '${WORK_DIR}/i-null.vortex'; ---- -0 +NULL diff --git a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt index 59931886824..b6640cf0f39 100644 --- a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt +++ b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt @@ -123,7 +123,7 @@ COPY ( query IRR SELECT count(x), sum(x), max(x) FROM '${WORK_DIR}/all-null.vortex'; ---- -0 0 NULL +0 NULL NULL query R SELECT min(x) FROM '${WORK_DIR}/all-null.vortex'; From 255659fcc00cf1eb5d7ab3eda5b1926394168522 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 24 Aug 2026 15:07:13 -0400 Subject: [PATCH 2/5] bench: add sum_v2 coverage Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 121 ++++++++++++++++++++++ vortex-array/benches/aggregate_sum.rs | 117 +++++++++++++++++++++ 2 files changed, 238 insertions(+) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..64820997778 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; @@ -19,6 +20,7 @@ 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::sum::Sum; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -199,6 +201,125 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Sum)); } +#[divan::bench] +fn sum_v2_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +/// Execute the lazy finalize result to canonical form so the benchmark includes 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_sum_v2_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + #[divan::bench] fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index 8972dfe3365..31856e22f9f 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -7,6 +7,7 @@ use divan::Bencher; use rand::prelude::*; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::expr::stats::Stat; @@ -36,6 +37,20 @@ fn sum_i32(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(1); + let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_u32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(2); @@ -50,6 +65,20 @@ fn sum_u32(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_u32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(2); + let data: Vec = (0..N).map(|_| rng.random_range(0..2000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_i64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(3); @@ -64,6 +93,20 @@ fn sum_i64(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(3); + let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_f64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(6); @@ -78,6 +121,20 @@ fn sum_f64(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_f64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(6); + let data: Vec = (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_f64_nulls_clustered(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(7); @@ -100,6 +157,28 @@ fn sum_f64_nulls_clustered(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_f64_nulls_clustered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(7); + let data: Vec> = (0..N) + .map(|i| { + if (i / 64) % 10 == 0 { + None + } else { + Some(rng.random_range(-1000.0..1000.0)) + } + }) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + // Clustered nulls: long runs of valid values broken up by occasional null blocks. This is the // case the run-based valid path is expected to accelerate. #[divan::bench] @@ -124,6 +203,28 @@ fn sum_i32_nulls_clustered(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i32_nulls_clustered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(4); + let data: Vec> = (0..N) + .map(|i| { + if (i / 64) % 10 == 0 { + None + } else { + Some(rng.random_range(-1000..1000)) + } + }) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + // Scattered nulls: ~50% nulls placed at random, producing many short runs. This is the worst case // for a run-based valid path, used to guard against regressions versus a per-element loop. #[divan::bench] @@ -141,3 +242,19 @@ fn sum_i32_nulls_scattered(bencher: Bencher) { }) .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } + +#[divan::bench] +fn sum_v2_i32_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(5); + let data: Vec> = (0..N) + .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000..1000))) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} From 72679f0004e97c640c69802440c4045aa1d87cc8 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 24 Aug 2026 15:18:25 -0400 Subject: [PATCH 3/5] refactor: remove sum_v2 inline hints Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs | 3 --- 1 file changed, 3 deletions(-) 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 8dea7c77890..b22e98c09c1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -156,7 +156,6 @@ impl AggregateFnVTable for SumV2 { partial.is_empty = true; } - #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) } @@ -372,7 +371,6 @@ fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar } } -#[inline(always)] fn checked_add_u64(sum: &mut u64, value: u64) -> bool { match sum.checked_add(value) { Some(result) => { @@ -383,7 +381,6 @@ fn checked_add_u64(sum: &mut u64, value: u64) -> bool { } } -#[inline(always)] fn checked_add_i64(sum: &mut i64, value: i64) -> bool { match sum.checked_add(value) { Some(result) => { From 9c95149a1f9bbac9a346b96a08cbde744982e52c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 24 Aug 2026 16:05:38 -0400 Subject: [PATCH 4/5] perf: fast-path sum_v2 struct finalization Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/sum_v2/mod.rs | 27 +++++++++++++ .../src/aggregate_fn/fns/sum_v2/tests.rs | 38 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index b22e98c09c1..e1287448a8f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -12,6 +12,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayRef; +use crate::ArrayView; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; @@ -27,6 +28,8 @@ 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::Struct; +use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; @@ -39,6 +42,7 @@ use crate::expr::stats::StatsProviderExt; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; const SUM_FIELD: &str = "sum"; const IS_OVERFLOW_FIELD: &str = "is_overflow"; @@ -244,6 +248,10 @@ impl AggregateFnVTable for SumV2 { } fn finalize(&self, partials: ArrayRef) -> VortexResult { + if let Some(partials) = partials.as_opt::() { + return finalize_struct(partials); + } + let sum = partials.get_item(SUM_FIELD)?; let is_invalid = partials .get_item(IS_OVERFLOW_FIELD)? @@ -260,6 +268,25 @@ impl AggregateFnVTable for SumV2 { } } +fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult { + let sum = partials.unmasked_field_by_name(SUM_FIELD)?.clone(); + let is_overflow = partials.unmasked_field_by_name(IS_OVERFLOW_FIELD)?.clone(); + let is_empty = partials.unmasked_field_by_name(IS_EMPTY_FIELD)?.clone(); + + let is_invalid = is_overflow + .binary(is_empty, Operator::Or)? + .fill_null(true)?; + let mut is_valid = is_invalid.not()?; + match partials.struct_validity() { + Validity::NonNullable | Validity::AllValid => {} + validity => { + is_valid = is_valid.binary(validity.to_array(partials.len()), Operator::And)?; + } + } + + sum.mask(is_valid) +} + /// In-memory state for SumV2 accumulation. pub struct SumV2Partial { return_dtype: DType, diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 2f16e533f7c..3c3fe6400c0 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -28,9 +28,11 @@ use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::Nullability::Nullable; use crate::dtype::PType; @@ -263,3 +265,39 @@ fn grouped_bool_fallback_tracks_empty_groups() -> VortexResult<()> { assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } + +#[rstest] +#[case::all_valid( + Validity::AllValid, + [Some(10i64), None, None, Some(40)] +)] +#[case::all_invalid(Validity::AllInvalid, [None, None, None, None])] +#[case::some_invalid( + Validity::from_iter([true, true, true, false]), + [Some(10i64), None, None, None] +)] +fn finalize_struct_applies_partial_and_struct_validity( + #[case] validity: Validity, + #[case] expected: [Option; 4], +) -> VortexResult<()> { + let partials = StructArray::try_new( + FieldNames::from(["sum", "is_overflow", "is_empty"]), + vec![ + PrimitiveArray::from_iter([10i64, 20, 30, 40]).into_array(), + BoolArray::from_iter([false, true, false, false]).into_array(), + BoolArray::from_iter([false, false, true, false]).into_array(), + ], + 4, + validity, + )? + .into_array(); + + let result = SumV2.finalize(partials)?; + let expected = PrimitiveArray::from_option_iter(expected).into_array(); + assert_arrays_eq!( + &result, + &expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) +} From 1a9234af09141a1aa539c93ccb36b902a010971c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 25 Aug 2026 13:11:38 -0400 Subject: [PATCH 5/5] add partial edge case tests Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/sum_v2/tests.rs | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 3c3fe6400c0..65e54a5bb6a 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -216,6 +216,87 @@ fn empty_chunk_is_a_merge_identity() -> VortexResult<()> { Ok(()) } +#[test] +fn combine_partials_empty_is_identity() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut empty = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let empty_partial = empty.partial_scalar()?; + + empty.combine_partials(empty_partial.clone())?; + assert!(empty.final_scalar()?.is_null()); + + let mut value = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let batch = PrimitiveArray::from_iter([7i64]).into_array(); + value.accumulate(&batch, &mut array_session().create_execution_ctx())?; + let value_partial = value.partial_scalar()?; + + empty.combine_partials(value_partial.clone())?; + assert_eq!( + empty.final_scalar()?.as_primitive().typed_value::(), + Some(7) + ); + + let mut value_then_empty = + Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + value_then_empty.combine_partials(value_partial)?; + value_then_empty.combine_partials(empty_partial)?; + assert_eq!( + value_then_empty + .final_scalar()? + .as_primitive() + .typed_value::(), + Some(7) + ); + Ok(()) +} + +#[test] +fn combine_partials_overflow_is_absorbing() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut max = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let max_batch = PrimitiveArray::from_iter([i64::MAX]).into_array(); + max.accumulate(&max_batch, &mut array_session().create_execution_ctx())?; + + let mut one = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let one_batch = PrimitiveArray::from_iter([1i64]).into_array(); + one.accumulate(&one_batch, &mut array_session().create_execution_ctx())?; + + let mut combined = + Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + combined.combine_partials(max.partial_scalar()?)?; + combined.combine_partials(one.partial_scalar()?)?; + assert!(combined.is_saturated()); + assert!(combined.final_scalar()?.is_null()); + + combined.combine_partials(max.partial_scalar()?)?; + let overflow_partial = combined.partial_scalar()?; + let fields = overflow_partial.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|flag| flag.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|flag| flag.as_bool().value()), + Some(false) + ); + + let mut propagated = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + propagated.combine_partials(overflow_partial)?; + assert!(propagated.is_saturated()); + assert!(propagated.final_scalar()?.is_null()); + Ok(()) +} + #[test] fn grouped_primitive_tracks_empty_overflow_and_null_groups() -> VortexResult<()> { let elements =