From 33f352dea4c0f12267224f4b59bbabe1039b584d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 9 Jul 2026 15:43:35 -0700 Subject: [PATCH 01/12] Sum aggregate follows SQL semantics: sum of zero valid values is null Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 9 +- vortex-array/src/aggregate_fn/accumulator.rs | 98 ++++-- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 6 +- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 9 +- .../src/aggregate_fn/fns/sum/constant.rs | 11 +- .../src/aggregate_fn/fns/sum/decimal.rs | 10 +- .../src/aggregate_fn/fns/sum/grouped.rs | 69 +++- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 327 ++++++++++++++++-- .../src/aggregate_fn/fns/sum/primitive.rs | 27 +- .../src/arrays/chunked/compute/aggregate.rs | 5 +- .../src/arrays/struct_/compute/rules.rs | 10 +- vortex-array/src/scalar_fn/fns/stat.rs | 35 +- vortex-array/src/stats/array.rs | 7 +- vortex-array/src/stats/expr.rs | 44 ++- vortex-datafusion/src/convert/stats.rs | 26 +- vortex-datafusion/src/v2/source.rs | 6 +- 16 files changed, 561 insertions(+), 138 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..0521b09ede3 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; @@ -164,7 +165,13 @@ where .unwrap(); acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); - divan::black_box(acc.finish().unwrap()) + let result = acc + .finish() + .unwrap() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap() + .into_array(); + divan::black_box(result) } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 69bae4e1053..2447ce0cc57 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -18,6 +18,7 @@ use crate::executor::max_iterations; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; /// Reference-counted type-erased accumulator. @@ -119,26 +120,40 @@ impl DynAccumulator for Accumulator { batch.dtype() ); - // 0. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot, - // consume that exact stat before kernel dispatch or decode. + // 0. Stats bridge: if this aggregate is cached under a Stat slot, consume that exact + // stat before kernel dispatch or decode. The stat may be stored in partial form or, + // for aggregates whose partial differs from their result (e.g. `Sum`'s `{sum, seen}` + // partial vs its monoid sum statistic), in result form — `combine_partials` accepts + // both. Result-form stats cannot express "no valid values", so they are only + // consumed when the batch's null count proves at least one valid value; otherwise + // fall through to a real accumulation. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - let partial = if partial.dtype() == &self.partial_dtype { - partial - } else { - vortex_ensure!( - partial.dtype().eq_ignore_nullability(&self.partial_dtype), - "Aggregate {} read legacy stat {} with dtype {}, expected {}", - self.aggregate_fn, - stat, - partial.dtype(), - self.partial_dtype, - ); - partial.cast(&self.partial_dtype)? - }; - self.vtable.combine_partials(&mut self.partial, partial)?; - return Ok(()); + if partial.dtype() == &self.partial_dtype { + self.vtable.combine_partials(&mut self.partial, partial)?; + return Ok(()); + } + if partial.dtype().eq_ignore_nullability(&self.partial_dtype) { + self.vtable + .combine_partials(&mut self.partial, partial.cast(&self.partial_dtype)?)?; + return Ok(()); + } + if stat.dtype(&self.dtype).as_ref() == Some(partial.dtype()) { + match batch.statistics().get_as::(Stat::NullCount) { + Precision::Exact(null_count) if null_count == batch.len() as u64 => { + // No valid values: the batch contributes the identity; skip the stat. + return Ok(()); + } + Precision::Exact(_) => { + self.vtable.combine_partials(&mut self.partial, partial)?; + return Ok(()); + } + _ => { + // Unknown validity: fall through to a real accumulation. + } + } + } } let session = ctx.session().clone(); @@ -320,8 +335,8 @@ mod tests { } } - /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of - /// `dict_of_seven()` which is `7.0`. + /// Sum partial sentinel `{sum: 42.0, seen: true}` — distinguishable from the natural Sum + /// of `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; impl DynAggregateKernel for SentinelSumPartialKernel { @@ -331,7 +346,7 @@ mod tests { _batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) + Ok(Some(sum_partial(42.0))) } } @@ -359,9 +374,33 @@ mod tests { fn sentinel_partial() -> Scalar { let acc = mean_f64_accumulator().expect("build accumulator"); - let sum = Scalar::primitive(42.0f64, Nullability::Nullable); let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.partial_dtype, vec![sum, count]) + Scalar::struct_(acc.partial_dtype, vec![sum_partial(42.0), count]) + } + + /// A `{sum, seen: true}` Sum partial over f64 input. + fn sum_partial(value: f64) -> Scalar { + let sum_partial_dtype = Sum + .partial_dtype( + &NumericalAggregateOpts::default(), + &DType::Primitive(PType::F64, Nullability::NonNullable), + ) + .expect("sum partial dtype"); + Scalar::struct_( + sum_partial_dtype, + vec![ + Scalar::primitive(value, Nullability::Nullable), + Scalar::bool(true, Nullability::NonNullable), + ], + ) + } + + /// The running sum value of a `{sum, seen}` Sum partial. + fn sum_of(partial: &Scalar) -> Option { + partial + .as_struct() + .field("sum") + .and_then(|s| s.as_primitive().as_::()) } /// Kernel registered for `(Dict, Combined)` fires in preference to @@ -380,10 +419,7 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(42.0) - ); + assert_eq!(sum_of(&s.field("sum").unwrap()), Some(42.0)); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) @@ -407,10 +443,7 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(7.0) - ); + assert_eq!(sum_of(&s.field("sum").unwrap()), Some(7.0)); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) @@ -438,10 +471,7 @@ mod tests { // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the // batch's valid_count, so count is the real 1. - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(42.0) - ); + assert_eq!(sum_of(&s.field("sum").unwrap()), Some(42.0)); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 20b7f15834c..f2ef08b0c5e 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -315,11 +315,13 @@ mod tests { } #[test] - fn mean_all_null_returns_nan() -> VortexResult<()> { + fn mean_all_null_returns_null() -> VortexResult<()> { + // The sum of zero valid values is null (SQL `SUM`), so the mean is null rather than + // the former 0/0 = NaN — matching SQL `AVG` of an all-null column. let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; - assert!(result.as_primitive().as_::().is_some_and(f64::is_nan)); + assert!(result.is_null()); 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..d1f3ce1b6d9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -17,12 +17,14 @@ pub(super) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, + seen: &mut bool, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, @@ -101,16 +103,17 @@ mod tests { &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); Ok(()) } #[test] - fn sum_bool_empty_produces_zero() -> VortexResult<()> { + fn sum_bool_empty_is_null() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..07325e62015 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -127,7 +127,8 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + // SQL `SUM`: an all-null constant has no valid values. + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -151,7 +152,7 @@ mod tests { fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -187,11 +188,7 @@ mod tests { let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, - Scalar::decimal( - DecimalValue::I256(i256::ZERO), - DecimalDType::new(20, 2), - Nullable - ) + Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 872e5769a01..cab120eba11 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -27,8 +27,10 @@ pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, + seen: &mut bool, ) -> VortexResult { let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), @@ -370,7 +372,7 @@ mod tests { let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, small)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -401,7 +403,7 @@ mod tests { Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -430,7 +432,7 @@ mod tests { ); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } @@ -463,7 +465,7 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index efe0825d4d6..a4462108b91 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,6 +1,8 @@ // 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; @@ -16,10 +18,16 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedArray; 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 [`Sum`] kernel for primitive element arrays. #[derive(Debug)] @@ -45,6 +53,10 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { /// [`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, seen}` partial rows (see `Sum`'s partial dtype): `seen` records whether the +/// group contained at least one valid element, so that `finalize` yields null for empty and +/// all-null groups (SQL `SUM`), and null groups are null struct rows. pub(super) fn try_grouped_sum( groups: &GroupedArray, ctx: &mut ExecutionCtx, @@ -80,7 +92,7 @@ fn grouped_sum( .execute_mask(elements.as_ref().len(), ctx)?; let all_valid = matches!(elem_mask.slices(), AllOr::All); - let result = match_each_native_ptype!(elements.ptype(), + let (sums, seen) = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, @@ -98,11 +110,22 @@ fn grouped_sum( } ); - Ok(result.into_array()) + Ok(StructArray::try_new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), + vec![ + sums.into_array(), + BoolArray::new(seen, 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. A group is null when the group -/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). +/// Reduce each group's element slice into a nullable sum plus a per-group `seen` bit. A sum is +/// null when the group itself is invalid or when summing it overflows (`sum_run` returns +/// `true`); `seen` records whether the group contained at least one valid element, decided by +/// validity alone (with `skip_nans`, a valid all-NaN group is still seen and sums to zero). fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -110,24 +133,31 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> PrimitiveArray { +) -> (PrimitiveArray, BitBuffer) { + let mut seen = BitBufferMut::with_capacity(group_ranges.len()); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { + seen.append(false); return None; } let mut acc = A::default(); - let overflow = if all_valid { - sum_run(&mut acc, &values[offset..offset + size]) + 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) }; + seen.append(any_valid); (!overflow).then_some(acc) }); - PrimitiveArray::from_option_iter(sums) + let sums = PrimitiveArray::from_option_iter(sums); + (sums, seen.freeze()) } /// Sum 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], @@ -135,17 +165,17 @@ fn sum_masked_group( size: usize, elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> bool { +) -> (bool, bool) { match elem_mask.slice(offset..offset + size).slices() { - AllOr::All => sum_run(acc, &values[offset..offset + size]), - AllOr::None => false, + 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; + return (true, true); } } - false + (false, !runs.is_empty()) } } } @@ -188,8 +218,9 @@ mod tests { acc.finish() } - /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] for - /// valid groups, a null sum for invalid groups. + /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] + /// (SQL semantics: null for zero valid elements) for valid groups, a null sum for invalid + /// groups. fn grouped_sum_reference( elements: &ArrayRef, ranges: &[(usize, usize)], @@ -200,8 +231,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) - .expect("sum partial dtype"); + .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] { @@ -291,8 +322,8 @@ mod tests { 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), Some(0i64), Some(0i64), Some(9i64)]); + // The all-null and empty groups have null sums (SQL `SUM` semantics). + 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(()) diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 010371fee4f..3cd26c8412f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -23,30 +23,82 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::arrays::ConstantArray; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; use crate::dtype::MAX_PRECISION; 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::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::fns::fill_null::FillNull; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::mask::Mask; -/// Return the sum of an array. +/// The monoid sum of an array: zero when there are no valid values, null on overflow. +/// +/// This is the value stored as `Stat::Sum`. Unlike SQL [`sum`], the empty sum is zero so that +/// per-chunk and per-zone statistics merge as a monoid (`combine(empty, x) = x`); SQL +/// semantics are recovered at read boundaries by combining this stat with the null count. +pub fn stat_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { + return Ok(sum_scalar); + } + + let mut acc = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + acc.accumulate(array, ctx)?; + let partial = acc.partial_scalar()?; + let result = partial + .as_struct() + .field("sum") + .vortex_expect("Sum partial is missing the `sum` field"); + + // Cache the computed sum as a statistic (only if non-null, i.e. no overflow). + if let Some(val) = result.value().cloned() { + array.statistics().set(Stat::Sum, Precision::Exact(val)); + } + + Ok(result) +} + +/// Return the SQL sum of an array: null when the array has no valid values or the sum +/// overflows. /// /// See [`Sum`] for details. pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // Short-circuit using cached array statistics. + // Short-circuit using cached array statistics. `Stat::Sum` is the monoid sum (zero for an + // array with no valid values), so the SQL empty-sum rule needs the null count: when it is + // unknown for a nullable array, fall through and compute rather than trust the cache. if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { - return Ok(sum_scalar); + 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. @@ -59,7 +111,8 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { acc.accumulate(array, ctx)?; let result = acc.finish()?; - // Cache the computed sum as a statistic (only if non-null, i.e. no overflow). + // 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)); } @@ -67,10 +120,16 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// Sum an array, starting from zero. +/// Sum an array following SQL `SUM` semantics. /// -/// If the sum overflows, a null scalar will be returned. -/// If the array is all-invalid, the sum will be zero. +/// A sum over zero valid values yields null (the SQL rule: nulls are eliminated, and the sum +/// of an empty set is null), and a sum that overflows yields null. +/// +/// The partial state is a `{sum, seen}` struct: `sum` is the running monoid value (null once +/// saturated by overflow, which poisons merges), and `seen` records whether any valid value +/// contributed (merged with OR). `finalize` maps `seen == false` to null. `seen` is decided by +/// validity alone: NaN values are valid, so with `skip_nans` a sum over only NaNs is `0`, not +/// null. /// /// 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. @@ -136,7 +195,7 @@ impl AggregateFnVTable for Sum { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + Some(sum_partial_dtype(self.return_dtype(options, input_dtype)?)) } fn empty_partial( @@ -152,19 +211,45 @@ impl AggregateFnVTable for Sum { Ok(SumPartial { return_dtype, current: Some(initial), + seen: false, skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - // A null partial means the sub-accumulator saturated (overflow). + // Partials are `{sum, seen}` structs. A plain (non-struct) scalar is the legacy or + // statistic form: a monoid sum value from an older writer or a cached `Stat::Sum`. It + // cannot distinguish an empty sum from a zero sum, so it is treated as seen. + let (sum_value, other_seen) = if matches!(other.dtype(), DType::Struct(..)) { + if other.is_null() { + // A null struct partial carries no recoverable state; treat as saturated. + partial.seen = true; + partial.current = None; + return Ok(()); + } + let fields = other.as_struct(); + let sum_value = fields + .field("sum") + .ok_or_else(|| vortex_err!("Sum partial is missing the `sum` field"))?; + let other_seen = fields + .field("seen") + .and_then(|seen| seen.as_bool().value()) + .unwrap_or(true); + (sum_value, other_seen) + } else { + (other, true) + }; + + partial.seen |= other_seen; + if sum_value.is_null() { + // A null sum value means the sub-accumulator saturated (overflow). partial.current = None; return Ok(()); } let Some(ref mut inner) = partial.current else { return Ok(()); }; + let other = sum_value; let saturated = match inner { SumState::Unsigned(acc) => { let val = other @@ -209,23 +294,18 @@ impl AggregateFnVTable for Sum { } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(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) - } - }) + Ok(Scalar::struct_( + sum_partial_dtype(partial.return_dtype.as_nullable()), + vec![ + sum_value_scalar(partial), + Scalar::bool(partial.seen, Nullability::NonNullable), + ], + )) } fn reset(&self, partial: &mut Self::Partial) { partial.current = Some(make_zero_state(&partial.return_dtype)); + partial.seen = false; } #[inline] @@ -253,6 +333,16 @@ impl AggregateFnVTable for Sum { // NaN-free batch: the cached NaN-skipping sum (if any) equals the // NaN-including sum. if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { + // The cached sum is a monoid value that cannot carry `seen`; derive it + // from the null count, or fall through to a real scan when unknown. + 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(_) => partial.seen = true, + _ => return Ok(false), + } let sum = if sum.dtype() == &partial.return_dtype { sum } else { @@ -264,7 +354,8 @@ impl AggregateFnVTable for Sum { Ok(false) } Precision::Exact(_) => { - // At least one NaN value: the sum is NaN without scanning the batch. + // At least one NaN value (a valid value): the sum is NaN without scanning. + partial.seen = true; if let Some(SumState::Float(acc)) = partial.current.as_mut() { *acc = f64::NAN; } @@ -282,6 +373,9 @@ impl AggregateFnVTable for Sum { ) -> VortexResult<()> { // Constants compute scalar * len and combine via combine_partials. if let Columnar::Constant(c) = batch { + // Any valid value counts as seen, including NaN and `false` constants that + // contribute nothing to the running sum. + partial.seen |= !c.scalar().is_null() && !c.is_empty(); // 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(()); @@ -300,9 +394,11 @@ impl AggregateFnVTable for Sum { 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), + Canonical::Primitive(p) => { + accumulate_primitive(&mut inner, p, ctx, skip_nans, &mut partial.seen) + } + Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut partial.seen), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut partial.seen), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -320,11 +416,32 @@ impl AggregateFnVTable for Sum { } fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) + // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` + // field (overflow) and null partial rows (e.g. null groups) stay null via the mask's + // validity intersection. + // + // The expressions are built unoptimized: `optimize` costs multiples of the whole + // aggregation at small group counts, and the caller's execution evaluates the lazy + // expression as-is. + let len = partials.len(); + let sum = GetItem.try_new_array(len, FieldName::from("sum"), [partials.clone()])?; + let seen = GetItem.try_new_array(len, FieldName::from("seen"), [partials])?; + let seen = FillNull.try_new_array( + len, + EmptyOptions, + [ + seen, + ConstantArray::new(Scalar::bool(false, Nullability::NonNullable), len).into_array(), + ], + )?; + Mask.try_new_array(len, EmptyOptions, [sum, seen]) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + if !partial.seen { + return Ok(Scalar::null(partial.return_dtype.as_nullable())); + } + Ok(sum_value_scalar(partial)) } } @@ -334,10 +451,55 @@ pub struct SumPartial { return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, + /// Whether at least one valid value has been accumulated. A sum over zero valid values + /// finalizes to null (SQL `SUM` semantics) rather than the monoid zero. + seen: 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, seen}` struct, where +/// `sum` is the running monoid value (null once saturated by overflow, poisoning merges) and +/// `seen` records whether any valid value contributed (merged with OR). Keeping the flag +/// separate from the sum lets merges stay a monoid — the identity is `{0, false}` — while +/// `finalize` maps unseen sums to null. +fn sum_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct( + StructFields::new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), + vec![sum_dtype, DType::Bool(Nullability::NonNullable)], + ), + Nullability::Nullable, + ) +} + +/// The running sum as a nullable scalar of the return dtype (null when saturated by overflow). +fn sum_value_scalar(partial: &SumPartial) -> 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) + } + } +} + +/// Whether the array contains at least one valid element. NaNs are valid values: with +/// `skip_nans` they contribute nothing to the sum, but a sum over only NaNs is `0`, not null. +fn has_valid_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array + .validity()? + .execute_mask(array.len(), ctx)? + .true_count() + > 0) +} + /// The accumulated sum value. // TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the // input type every time? Perhaps? @@ -407,6 +569,7 @@ mod tests { use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::stat_sum; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::BoolArray; @@ -524,6 +687,102 @@ mod tests { // State merge tests (vtable-level) + #[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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let empty = Sum.to_scalar(&state)?; + Sum.combine_partials(&mut state, empty)?; + assert!(Sum.finalize_scalar(&state)?.is_null()); + Ok(()) + } + + #[test] + fn sum_state_empty_is_identity() -> VortexResult<()> { + // Combining an empty state into a seen state changes nothing: `{0, false}` is the + // identity of the `{sum, seen}` monoid. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut state, Scalar::primitive(100i64, Nullable))?; + + let empty = + Sum.to_scalar(&Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?)?; + Sum.combine_partials(&mut state, empty)?; + + let result = Sum.finalize_scalar(&state)?; + assert_eq!(result.as_primitive().typed_value::(), Some(100)); + Ok(()) + } + + #[test] + fn sum_state_overflow_poisons_but_stays_seen() -> VortexResult<()> { + // Overflow (a null `sum` field) poisons the merge even when combined with later + // values: the result is null via the sum value, not via `seen`. + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut overflowed = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut overflowed, Scalar::primitive(i64::MAX, Nullable))?; + Sum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; + let overflowed = Sum.to_scalar(&overflowed)?; + + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; + Sum.combine_partials(&mut state, overflowed)?; + Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; + + assert!(Sum.finalize_scalar(&state)?.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 = sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) + } + + #[test] + fn stat_sum_is_monoid_while_sum_is_sql() -> VortexResult<()> { + // The persisted statistic keeps the monoid semantics (zero for all-null) that zone + // and chunk merging require, while `sum` applies the SQL rule. + let mut ctx = array_session().create_execution_ctx(); + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + assert_eq!( + stat_sum(&arr, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(0) + ); + // The cached monoid statistic must not leak through `sum`'s cache short-circuit. + assert!(sum(&arr, &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 sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); @@ -535,7 +794,7 @@ mod tests { let scalar2 = Scalar::primitive(50i64, Nullable); Sum.combine_partials(&mut state, scalar2)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; Sum.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) @@ -648,7 +907,8 @@ mod tests { 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(0i64), Some(7i64)]).into_array(); + // The all-null group has a null sum (SQL `SUM` semantics). + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -741,7 +1001,7 @@ mod tests { } #[test] - fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> { + fn sum_chunked_floats_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(); @@ -750,7 +1010,8 @@ mod tests { &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + // SQL `SUM`: no valid values across any chunk yields null. + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..36c337e6cb5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -26,12 +26,19 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, skip_nans: bool, + seen: &mut bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => accumulate_primitive_all(inner, p, skip_nans), - AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), + AllOr::All => { + *seen |= !p.as_ref().is_empty(); + accumulate_primitive_all(inner, p, skip_nans) + } + AllOr::Some(slices) => { + *seen |= !slices.is_empty(); + accumulate_primitive_valid(inner, p, slices, skip_nans) + } } } @@ -260,7 +267,8 @@ mod tests { fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); Ok(()) } @@ -268,7 +276,7 @@ mod tests { fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); Ok(()) } @@ -289,20 +297,21 @@ mod tests { } #[test] - fn sum_empty_produces_zero() -> VortexResult<()> { + fn sum_empty_is_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: the sum over no values is null. + assert!(result.is_null()); Ok(()) } #[test] - fn sum_empty_f64_produces_zero() -> VortexResult<()> { + fn sum_empty_f64_is_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + assert!(result.is_null()); Ok(()) } @@ -395,6 +404,8 @@ mod tests { .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(()) diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 149c72f499e..99b6bba44de 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -120,7 +120,8 @@ mod tests { DType::Primitive(PType::I32, Nullability::Nullable), )?; let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); Ok(()) } @@ -158,7 +159,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let chunked = ChunkedArray::try_new(vec![], dtype)?; let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } 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/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 84fc5760495..85ad1ae339e 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -173,12 +173,16 @@ fn stat_array( } .map(ScalarValue::Bool) } else if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) { - array + let stored = array .statistics() .with_typed_stats_set(|stats| stats.get(stat)) // We don't mind whether the stat is approxed or not, since these are row-wise bounds. - .into_inner() - .and_then(Scalar::into_value) + .into_inner(); + match stored { + Some(scalar) if scalar.dtype().eq_ignore_nullability(&dtype) => scalar.into_value(), + Some(scalar) => result_stat_to_state(array, aggregate_fn, scalar, len)?, + None => None, + } } else { tracing::trace!( "No legacy Stat slot for aggregate {}; stat expression will resolve to null", @@ -190,3 +194,28 @@ fn stat_array( let scalar = Scalar::try_new(dtype, value)?; Ok(ConstantArray::new(scalar, len).into_array()) } + +/// Rebuild an aggregate's partial state from a result-form statistic (e.g. `Sum`'s monoid +/// statistic vs its `{sum, seen}` state). A result-form stat cannot express "no valid values", +/// so the array's null count decides between the empty state and a seen state; when the null +/// count is unknown for a nullable array, the stat is unusable and resolves to missing. +fn result_stat_to_state( + array: &ArrayRef, + aggregate_fn: &AggregateFnRef, + stat_scalar: Scalar, + len: usize, +) -> VortexResult> { + let all_null = if array.dtype().is_nullable() { + match array.statistics().get_as::(Stat::NullCount) { + Precision::Exact(null_count) => null_count == len as u64, + _ => return Ok(None), + } + } else { + false + }; + let mut acc = aggregate_fn.accumulator(array.dtype())?; + if !all_null { + acc.combine_partials(stat_scalar)?; + } + Ok(acc.partial_scalar()?.into_value()) +} diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index d3fb7fd11e4..1ceb4bad3f8 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -23,7 +23,7 @@ use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::sum::sum; +use crate::aggregate_fn::fns::sum::stat_sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -172,8 +172,9 @@ impl StatsSetRef<'_> { .dtype(self.dyn_array_ref.dtype()) .is_some() .then(|| { - // Sum is supported for this dtype. - sum(self.dyn_array_ref, ctx) + // Sum is supported for this dtype. The statistic is the monoid sum + // (zero when no valid values), not the SQL sum. + stat_sum(self.dyn_array_ref, ctx) }) .transpose()? } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 2038df0b613..67c06734836 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -89,6 +89,8 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::sum::Sum; use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; @@ -108,6 +110,27 @@ mod tests { static SESSION: LazyLock = LazyLock::new(array_session); + /// The `{sum, seen}` state dtype for a sum over non-nullable i32 input. + fn sum_state_dtype() -> DType { + use crate::aggregate_fn::AggregateFnVTable; + Sum.partial_dtype( + &NumericalAggregateOpts::skip_nans(), + &DType::Primitive(PType::I32, Nullability::NonNullable), + ) + .vortex_expect("sum partial dtype") + } + + /// A seen `{sum, seen}` state scalar with the given running sum. + fn sum_state(value: i64) -> Scalar { + Scalar::struct_( + sum_state_dtype(), + vec![ + Scalar::primitive(value, Nullability::Nullable), + Scalar::bool(true, Nullability::NonNullable), + ], + ) + } + #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { let array = buffer![1i32, 2, 3].into_array(); @@ -122,8 +145,8 @@ mod tests { .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - let expected = - ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); + // Stat expressions resolve to the aggregate's `{sum, seen}` state form. + let expected = ConstantArray::new(sum_state(6), 3).into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) @@ -138,11 +161,7 @@ mod tests { .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - let expected = ConstantArray::new( - Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), - 3, - ) - .into_array(); + let expected = ConstantArray::new(Scalar::null(sum_state_dtype()), 3).into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) @@ -173,10 +192,13 @@ mod tests { let result = result .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - let expected = PrimitiveArray::new( - buffer![3i64, 3, 0, 0, 0], - Validity::from_iter([true, true, false, false, false]), - ) + let expected = ChunkedArray::try_new( + vec![ + ConstantArray::new(sum_state(3), 2).into_array(), + ConstantArray::new(Scalar::null(sum_state_dtype()), 3).into_array(), + ], + sum_state_dtype(), + )? .into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); diff --git a/vortex-datafusion/src/convert/stats.rs b/vortex-datafusion/src/convert/stats.rs index 33a33a78ccf..ccad75ddae6 100644 --- a/vortex-datafusion/src/convert/stats.rs +++ b/vortex-datafusion/src/convert/stats.rs @@ -20,6 +20,7 @@ use crate::convert::TryToDataFusion; pub(crate) fn stats_set_to_df( stats_set: &StatsSet, dtype: &DType, + row_count: Option, ) -> VortexResult { // Update the total size in bytes. let column_size = stats_set.get_as::(Stat::UncompressedSizeInBytes, &PType::U64.into()); @@ -52,8 +53,17 @@ pub(crate) fn stats_set_to_df( .ok() }); - // Find the sum statistic + let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); + + // Find the sum statistic. `Stat::Sum` is the monoid sum (zero when the column has no + // valid values), while DataFusion's `sum_value` is the SQL sum (null in that case), so an + // all-null column must not export its zero; without a known null count the sum cannot be + // interpreted and is withheld. let sum = stats_set.get(Stat::Sum).and_then(|stat_val| { + match (&null_count, row_count) { + (VortexPrecision::Exact(null_count), Some(row_count)) if *null_count < row_count => {} + _ => return None, + } Scalar::try_new( Stat::Sum .dtype(dtype) @@ -65,8 +75,6 @@ pub(crate) fn stats_set_to_df( .ok() }); - let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); - Ok(ColumnStatistics { null_count: null_count.to_df(), min_value: min.to_df(), @@ -97,12 +105,20 @@ mod tests { #[test] fn is_constant_false_does_not_imply_one_distinct_value() -> VortexResult<()> { let false_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(false)); - let false_stats = stats_set_to_df(&false_constant, &DType::Bool(Nullability::NonNullable))?; + let false_stats = stats_set_to_df( + &false_constant, + &DType::Bool(Nullability::NonNullable), + Some(5), + )?; assert_eq!(false_stats.distinct_count, Precision::Absent); let true_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(true)); - let true_stats = stats_set_to_df(&true_constant, &DType::Bool(Nullability::NonNullable))?; + let true_stats = stats_set_to_df( + &true_constant, + &DType::Bool(Nullability::NonNullable), + Some(5), + )?; assert_eq!(true_stats.distinct_count, Precision::Exact(1)); diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 82fae9bf56d..e4e6a870159 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -242,6 +242,10 @@ impl VortexDataSourceBuilder { }; // We now compute initial statistics. + let row_count = match self.data_source.row_count() { + Precision::Exact(n) => usize::try_from(n).ok(), + _ => None, + }; let field_paths: Vec<_> = fields .names() .iter() @@ -256,7 +260,7 @@ impl VortexDataSourceBuilder { .await? .iter() .zip(fields.fields()) - .map(|(stats, dtype)| stats_set_to_df(stats, &dtype)) + .map(|(stats, dtype)| stats_set_to_df(stats, &dtype, row_count)) .collect::>>()?; Ok(VortexDataSource { From a9bf18cfe24d7f89d8ac398a7562f154a9ce057c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 9 Jul 2026 17:02:08 -0700 Subject: [PATCH 02/12] Structural Sum finalize; thread ExecutionCtx through aggregate finalize Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 8 +-- .../src/aggregate_fn/accumulator_grouped.rs | 14 ++-- vortex-array/src/aggregate_fn/combined.rs | 6 +- .../src/aggregate_fn/fns/all_nan/mod.rs | 2 +- .../aggregate_fn/fns/all_non_distinct/mod.rs | 2 +- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 2 +- .../src/aggregate_fn/fns/all_non_null/mod.rs | 2 +- .../src/aggregate_fn/fns/all_null/mod.rs | 2 +- .../src/aggregate_fn/fns/bounded_max/mod.rs | 2 +- .../src/aggregate_fn/fns/bounded_min/mod.rs | 2 +- .../src/aggregate_fn/fns/count/grouped.rs | 4 +- .../src/aggregate_fn/fns/count/mod.rs | 2 +- .../src/aggregate_fn/fns/first/mod.rs | 2 +- .../src/aggregate_fn/fns/is_constant/mod.rs | 2 +- .../src/aggregate_fn/fns/is_sorted/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/last/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/max/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/min/mod.rs | 2 +- .../src/aggregate_fn/fns/min_max/mod.rs | 2 +- .../src/aggregate_fn/fns/nan_count/mod.rs | 2 +- .../src/aggregate_fn/fns/null_count/mod.rs | 2 +- .../src/aggregate_fn/fns/sum/grouped.rs | 10 +-- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 69 +++++++++++-------- .../fns/uncompressed_size_in_bytes/mod.rs | 2 +- vortex-array/src/aggregate_fn/foreign.rs | 2 +- vortex-array/src/aggregate_fn/proto.rs | 2 +- vortex-array/src/aggregate_fn/vtable.rs | 2 +- 27 files changed, 87 insertions(+), 66 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 0521b09ede3..f7827b0620b 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -163,12 +163,12 @@ where list_element_dtype(list_view), ) .unwrap(); - acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) - .unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + acc.accumulate_list(list_view, &mut ctx).unwrap(); let result = acc - .finish() + .finish(&mut ctx) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() .into_array(); divan::black_box(result) diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index b87c04ee204..e6d0fdfaabf 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -238,7 +238,7 @@ pub trait DynGroupedAccumulator: 'static + Send { /// Finish the accumulation and return the final aggregate results for all groups. /// Resets the accumulator state for the next round of accumulation. - fn finish(&mut self) -> VortexResult; + fn finish(&mut self, ctx: &mut ExecutionCtx) -> VortexResult; } impl DynGroupedAccumulator for GroupedAccumulator { @@ -276,9 +276,15 @@ impl DynGroupedAccumulator for GroupedAccumulator { Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) } - fn finish(&mut self) -> VortexResult { - let states = self.flush()?; - let results = self.vtable.finalize(states)?; + fn finish(&mut self, ctx: &mut ExecutionCtx) -> VortexResult { + // The single-batch case (one accumulate_list call) skips the chunked wrapper so that + // finalize implementations can operate on the batch's states directly. + let states = if self.partials.len() == 1 { + self.partials.pop().vortex_expect("checked length") + } else { + self.flush()? + }; + let results = self.vtable.finalize(states, ctx)?; vortex_ensure!( results.dtype() == &self.return_dtype, diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 76ad9877314..86ddb678093 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -245,11 +245,11 @@ impl AggregateFnVTable for Combined { unreachable!("Combined::try_accumulate handles all batches") } - fn finalize(&self, states: ArrayRef) -> VortexResult { + fn finalize(&self, states: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let l_field = states.get_item(FieldName::from(self.0.left_name()))?; let r_field = states.get_item(FieldName::from(self.0.right_name()))?; - let l_finalized = self.0.left().finalize(l_field)?; - let r_finalized = self.0.right().finalize(r_field)?; + let l_finalized = self.0.left().finalize(l_field, ctx)?; + let r_finalized = self.0.right().finalize(r_field, ctx)?; BinaryCombined::finalize(&self.0, l_finalized, r_finalized) } diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index 68a58908018..e60b5f27696 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -121,7 +121,7 @@ impl AggregateFnVTable for AllNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index f030cc810e2..6dede02675c 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -233,7 +233,7 @@ impl AggregateFnVTable for AllNonDistinct { } } - fn finalize(&self, _partials: ArrayRef) -> VortexResult { + fn finalize(&self, _partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { vortex_bail!("AllNonDistinct does not support array finalization"); } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index fe8527da966..f31018f820a 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -111,7 +111,7 @@ impl AggregateFnVTable for AllNonNan { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index c07dbb907c9..b01e1f74e1e 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -101,7 +101,7 @@ impl AggregateFnVTable for AllNonNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index ec64e3d5c43..0c010776c52 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -104,7 +104,7 @@ impl AggregateFnVTable for AllNull { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index bb27d9f9ab1..a7c1a51633e 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -279,7 +279,7 @@ impl AggregateFnVTable for BoundedMax { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { partials.get_item(BOUNDED_MAX_BOUND) } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 7a72442550b..d8295772688 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -197,7 +197,7 @@ impl AggregateFnVTable for BoundedMin { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 39a957530bf..4c70e3e19c2 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -122,7 +122,7 @@ mod tests { elem_dtype.clone(), )?; acc.accumulate_list(groups, ctx)?; - acc.finish() + acc.finish(ctx) } /// Reference valid-counts (non-nullable `U64`), one per group. @@ -235,7 +235,7 @@ mod tests { let mut acc = GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?; acc.accumulate_list(&groups, &mut ctx)?; - let actual = acc.finish()?; + let actual = acc.finish(&mut ctx)?; let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 8f7d68027bc..e05dae91cd4 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -115,7 +115,7 @@ impl AggregateFnVTable for Count { unreachable!("Count::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index c41e2575057..3c2c30539fc 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -117,7 +117,7 @@ impl AggregateFnVTable for First { unreachable!("First::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 490ede7f640..19fb1dbb5e9 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -420,7 +420,7 @@ impl AggregateFnVTable for IsConstant { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index c3ee7f5d4e7..52c3378085a 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -514,7 +514,7 @@ impl AggregateFnVTable for IsSorted { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 63ee54a6efa..7c1def89082 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -115,7 +115,7 @@ impl AggregateFnVTable for Last { unreachable!("Last::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index ccf9f4d899d..8896de2e301 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -205,7 +205,7 @@ impl AggregateFnVTable for Max { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index 488405e8f14..e306897dd48 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -205,7 +205,7 @@ impl AggregateFnVTable for Min { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 51601ce76fd..828fe4bcfa5 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -427,7 +427,7 @@ impl AggregateFnVTable for MinMax { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index 723847160e1..56f9f813edb 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -172,7 +172,7 @@ impl AggregateFnVTable for NanCount { } } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index 031e8f6a09b..9c0427812e5 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -143,7 +143,7 @@ impl AggregateFnVTable for NullCount { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index a4462108b91..4a017bd8559 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -214,8 +214,9 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; - acc.finish() + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish(&mut ctx) } /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] @@ -401,8 +402,9 @@ mod tests { let mut acc = GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; - acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; - let actual = acc.finish()?; + let mut ctx2 = array_session().create_execution_ctx(); + acc.accumulate_list(&groups, &mut ctx2)?; + let actual = acc.finish(&mut ctx2)?; let mut ctx = array_session().create_execution_ctx(); // Group 0 contains a NaN -> NaN sum; group 1 sums normally. diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 3cd26c8412f..14b23df9266 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -29,8 +30,13 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::arrays::ConstantArray; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; +use crate::arrays::Bool; +use crate::arrays::BoolArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::masked::mask_validity_canonical; +use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldName; @@ -45,10 +51,7 @@ use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::mask::Mask; +use crate::validity::Validity; /// The monoid sum of an array: zero when there are no valid values, null on overflow. /// @@ -415,26 +418,35 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` - // field (overflow) and null partial rows (e.g. null groups) stay null via the mask's - // validity intersection. - // - // The expressions are built unoptimized: `optimize` costs multiples of the whole - // aggregation at small group counts, and the caller's execution evaluates the lazy - // expression as-is. + // field (overflow) and null partial rows (e.g. null groups) stay null via the + // validity intersection. Implemented structurally rather than as a `get_item`/`mask` + // expression: expression construction and dispatch cost multiples of the whole + // aggregation at small group counts. let len = partials.len(); - let sum = GetItem.try_new_array(len, FieldName::from("sum"), [partials.clone()])?; - let seen = GetItem.try_new_array(len, FieldName::from("seen"), [partials])?; - let seen = FillNull.try_new_array( - len, - EmptyOptions, - [ - seen, - ConstantArray::new(Scalar::bool(false, Nullability::NonNullable), len).into_array(), - ], - )?; - Mask.try_new_array(len, EmptyOptions, [sum, seen]) + let states = match partials.as_opt::() { + Some(states) => states.into_owned(), + None => partials.execute::(ctx)?, + }; + let struct_mask = states.as_ref().validity()?.execute_mask(len, ctx)?; + let seen = states.unmasked_field_by_name("seen")?.clone(); + let seen = match seen.as_opt::() { + Some(seen) => seen.into_owned(), + None => seen.execute::(ctx)?, + }; + let valid = &struct_mask & &Mask::from_buffer(seen.to_bit_buffer()); + + let sum = states.unmasked_field_by_name("sum")?.clone(); + if valid.all_true() { + // Every partial row is a seen, valid group: the sums are already the result. + return Ok(sum); + } + let sum = sum.execute::(ctx)?; + Ok( + mask_validity_canonical(sum, Validity::from_mask(valid, Nullability::Nullable), ctx)? + .into_array(), + ) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -844,8 +856,9 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; - acc.finish() + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish(&mut ctx) } #[test] @@ -939,7 +952,7 @@ mod tests { 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 result1 = acc.finish(&mut ctx)?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result1, &expected1, &mut ctx); @@ -947,7 +960,7 @@ mod tests { 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 result2 = acc.finish(&mut ctx)?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index eac5b3c33ce..1b22205d87e 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -177,7 +177,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs index feb07e47175..8ef8da01824 100644 --- a/vortex-array/src/aggregate_fn/foreign.rs +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -108,7 +108,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn finalize(&self, _states: ArrayRef) -> VortexResult { + fn finalize(&self, _states: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 92fac87892a..26e206bcae9 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -149,7 +149,7 @@ mod tests { Ok(()) } - fn finalize(&self, partials: ArrayRef) -> VortexResult { + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 49b28dd26d7..9af95717d35 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -152,7 +152,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// /// The provides `states` array has dtype as specified by `state_dtype`, the result array /// must have dtype as specified by `return_dtype`. - fn finalize(&self, states: ArrayRef) -> VortexResult; + fn finalize(&self, states: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; /// Finalize a scalar accumulator state into an aggregate result. /// From 8e8fdc648fa94494d7118f505489104a48cce2f5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 9 Jul 2026 22:25:02 -0700 Subject: [PATCH 03/12] Adapt list_sum to ctx-threaded finish; drop redundant mask_empty_lists and dead code Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 10 ---- vortex-array/src/scalar_fn/fns/list_sum.rs | 59 +------------------- vortex-array/src/stats/expr.rs | 1 - 3 files changed, 3 insertions(+), 67 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 14b23df9266..e0b3ce50070 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -502,16 +502,6 @@ fn sum_value_scalar(partial: &SumPartial) -> Scalar { } } -/// Whether the array contains at least one valid element. NaNs are valid values: with -/// `skip_nans` they contribute nothing to the sum, but a sum over only NaNs is `0`, not null. -fn has_valid_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(array - .validity()? - .execute_mask(array.len(), ctx)? - .true_count() - > 0) -} - /// The accumulated sum value. // TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the // input type every time? Perhaps? diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index a5f8d219db6..b49c60b10ed 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,23 +14,16 @@ 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::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. /// @@ -100,7 +91,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,8 +121,8 @@ 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 [`Sum`] 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, @@ -141,50 +131,7 @@ fn list_sum_impl( ) -> VortexResult { let mut acc = GroupedAccumulator::try_new(Sum, *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(ctx) } #[cfg(test)] diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 67c06734836..0cca359d11f 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -106,7 +106,6 @@ mod tests { use crate::expr::stats::Stat; use crate::scalar::Scalar; use crate::scalar::ScalarValue; - use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(array_session); From c58410af67a2f6383986f0c34905460f575149b5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 14:04:16 -0700 Subject: [PATCH 04/12] Split SQL-semantics Sum from monoid StatSum; only list_sum uses Sum Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 39 +- vortex-array/src/aggregate_fn/accumulator.rs | 106 +- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 24 +- vortex-array/src/aggregate_fn/fns/mod.rs | 1 + .../src/aggregate_fn/fns/stat_sum/bool.rs | 154 +++ .../src/aggregate_fn/fns/stat_sum/constant.rs | 219 +++++ .../src/aggregate_fn/fns/stat_sum/decimal.rs | 470 +++++++++ .../src/aggregate_fn/fns/stat_sum/grouped.rs | 408 ++++++++ .../src/aggregate_fn/fns/stat_sum/mod.rs | 909 ++++++++++++++++++ .../aggregate_fn/fns/stat_sum/primitive.rs | 460 +++++++++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 171 +--- vortex-array/src/aggregate_fn/session.rs | 8 + vortex-array/src/array/erased.rs | 4 +- .../src/arrays/chunked/compute/aggregate.rs | 9 +- .../src/arrays/struct_/compute/rules.rs | 10 +- .../src/compute/conformance/consistency.rs | 6 +- vortex-array/src/expr/stats/mod.rs | 8 +- vortex-array/src/scalar_fn/fns/list_sum.rs | 4 +- vortex-array/src/scalar_fn/fns/stat.rs | 35 +- vortex-array/src/stats/array.rs | 2 +- vortex-array/src/stats/expr.rs | 49 +- vortex-datafusion/src/convert/stats.rs | 26 +- vortex-datafusion/src/v2/source.rs | 6 +- vortex-duckdb/src/convert/expr.rs | 10 +- vortex-layout/src/layouts/file_stats.rs | 4 +- vortex-layout/src/layouts/zoned/schema.rs | 10 +- vortex-layout/src/layouts/zoned/writer.rs | 12 +- 27 files changed, 2825 insertions(+), 339 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs create mode 100644 vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index f7827b0620b..2c6cda34441 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -19,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::stat_sum::StatSum; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; @@ -179,7 +180,7 @@ fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StatSum)); } #[divan::bench] @@ -187,7 +188,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StatSum)); } #[divan::bench] @@ -195,11 +196,43 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StatSum)); } #[divan::bench] fn sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, StatSum)); +} + +#[divan::bench] +fn sql_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, Sum)); +} + +#[divan::bench] +fn sql_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, Sum)); +} + +#[divan::bench] +fn sql_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, Sum)); +} + +#[divan::bench] +fn sql_sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 2447ce0cc57..03a761063da 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -18,7 +18,6 @@ use crate::executor::max_iterations; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; -use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; /// Reference-counted type-erased accumulator. @@ -120,40 +119,26 @@ impl DynAccumulator for Accumulator { batch.dtype() ); - // 0. Stats bridge: if this aggregate is cached under a Stat slot, consume that exact - // stat before kernel dispatch or decode. The stat may be stored in partial form or, - // for aggregates whose partial differs from their result (e.g. `Sum`'s `{sum, seen}` - // partial vs its monoid sum statistic), in result form — `combine_partials` accepts - // both. Result-form stats cannot express "no valid values", so they are only - // consumed when the batch's null count proves at least one valid value; otherwise - // fall through to a real accumulation. + // 0. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot, + // consume that exact stat before kernel dispatch or decode. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - if partial.dtype() == &self.partial_dtype { - self.vtable.combine_partials(&mut self.partial, partial)?; - return Ok(()); - } - if partial.dtype().eq_ignore_nullability(&self.partial_dtype) { - self.vtable - .combine_partials(&mut self.partial, partial.cast(&self.partial_dtype)?)?; - return Ok(()); - } - if stat.dtype(&self.dtype).as_ref() == Some(partial.dtype()) { - match batch.statistics().get_as::(Stat::NullCount) { - Precision::Exact(null_count) if null_count == batch.len() as u64 => { - // No valid values: the batch contributes the identity; skip the stat. - return Ok(()); - } - Precision::Exact(_) => { - self.vtable.combine_partials(&mut self.partial, partial)?; - return Ok(()); - } - _ => { - // Unknown validity: fall through to a real accumulation. - } - } - } + let partial = if partial.dtype() == &self.partial_dtype { + partial + } else { + vortex_ensure!( + partial.dtype().eq_ignore_nullability(&self.partial_dtype), + "Aggregate {} read legacy stat {} with dtype {}, expected {}", + self.aggregate_fn, + stat, + partial.dtype(), + self.partial_dtype, + ); + partial.cast(&self.partial_dtype)? + }; + self.vtable.combine_partials(&mut self.partial, partial)?; + return Ok(()); } let session = ctx.session().clone(); @@ -294,7 +279,7 @@ mod tests { use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; - use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::stat_sum::StatSum; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; use crate::array::VTable; @@ -335,8 +320,8 @@ mod tests { } } - /// Sum partial sentinel `{sum: 42.0, seen: true}` — distinguishable from the natural Sum - /// of `dict_of_seven()` which is `7.0`. + /// StatSum partial sentinel `42.0` — distinguishable from the natural StatSum of + /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; impl DynAggregateKernel for SentinelSumPartialKernel { @@ -346,7 +331,7 @@ mod tests { _batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - Ok(Some(sum_partial(42.0))) + Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) } } @@ -374,33 +359,9 @@ mod tests { fn sentinel_partial() -> Scalar { let acc = mean_f64_accumulator().expect("build accumulator"); + let sum = Scalar::primitive(42.0f64, Nullability::Nullable); let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.partial_dtype, vec![sum_partial(42.0), count]) - } - - /// A `{sum, seen: true}` Sum partial over f64 input. - fn sum_partial(value: f64) -> Scalar { - let sum_partial_dtype = Sum - .partial_dtype( - &NumericalAggregateOpts::default(), - &DType::Primitive(PType::F64, Nullability::NonNullable), - ) - .expect("sum partial dtype"); - Scalar::struct_( - sum_partial_dtype, - vec![ - Scalar::primitive(value, Nullability::Nullable), - Scalar::bool(true, Nullability::NonNullable), - ], - ) - } - - /// The running sum value of a `{sum, seen}` Sum partial. - fn sum_of(partial: &Scalar) -> Option { - partial - .as_struct() - .field("sum") - .and_then(|s| s.as_primitive().as_::()) + Scalar::struct_(acc.partial_dtype, vec![sum, count]) } /// Kernel registered for `(Dict, Combined)` fires in preference to @@ -419,7 +380,10 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - assert_eq!(sum_of(&s.field("sum").unwrap()), Some(42.0)); + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(42.0) + ); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) @@ -443,7 +407,10 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - assert_eq!(sum_of(&s.field("sum").unwrap()), Some(7.0)); + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(7.0) + ); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) @@ -451,7 +418,7 @@ mod tests { Ok(()) } - /// A kernel registered for the inner `(Dict, Sum)` child fires when accumulating a + /// A kernel registered for the inner `(Dict, StatSum)` child fires when accumulating a /// Dict batch through `Combined`. This is the reusable-primitive case the /// refactor enables: no `(Dict, Combined)` kernel is needed. #[test] @@ -460,7 +427,7 @@ mod tests { let session = fresh_session(); session .get::() - .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + .register_aggregate_kernel(Dict.id(), Some(StatSum.id()), &KERNEL); let mut ctx = session.create_execution_ctx(); let mut acc = mean_f64_accumulator()?; @@ -468,10 +435,13 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired + // `StatSum` child returned the sentinel 42.0 — proves the (Dict, StatSum) kernel fired // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the // batch's valid_count, so count is the real 1. - assert_eq!(sum_of(&s.field("sum").unwrap()), Some(42.0)); + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(42.0) + ); assert_eq!( s.field("count").unwrap().as_primitive().as_::(), Some(1) diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index f2ef08b0c5e..6873fbe0d6e 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -17,8 +17,8 @@ use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::CombinedOptions; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::sum::Sum; -use crate::aggregate_fn::fns::sum::sum_decimal_dtype; +use crate::aggregate_fn::fns::stat_sum::StatSum; +use crate::aggregate_fn::fns::stat_sum::sum_decimal_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -49,7 +49,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// Compute the arithmetic mean of an array. /// -/// Implemented as `Sum / Count` via [`BinaryCombined`]. +/// Implemented as `StatSum / Count` via [`BinaryCombined`]. /// /// Booleans and primitive numeric types are cast to f64. Decimals stay decimals. #[derive(Clone, Debug)] @@ -62,7 +62,7 @@ impl Mean { } impl BinaryCombined for Mean { - type Left = Sum; + type Left = StatSum; type Right = Count; fn id(&self) -> AggregateFnId { @@ -70,8 +70,8 @@ impl BinaryCombined for Mean { *ID } - fn left(&self) -> Sum { - Sum + fn left(&self) -> StatSum { + StatSum } fn right(&self) -> Count { @@ -112,7 +112,7 @@ impl BinaryCombined for Mean { let sum = sum_cast.as_primitive().typed_value::(); let count = count_cast.as_primitive().typed_value::(); let value = match (sum, count) { - (None, _) | (_, None) => return Ok(Scalar::null(target)), // Sum overflowed + (None, _) | (_, None) => return Ok(Scalar::null(target)), // StatSum overflowed // A count of zero yields 0/0 = NaN, matching the array `finalize` path: nulls are // skipped during accumulation, so an all-null input is an empty mean, not null. (Some(s), Some(c)) => s / c, @@ -127,7 +127,7 @@ impl BinaryCombined for Mean { fn coerce_args( &self, _options: &PairOptions< - ::Options, + ::Options, ::Options, >, input_dtype: &DType, @@ -140,7 +140,7 @@ impl BinaryCombined for Mean { /// Hint for callers: what to cast the input to before accumulation. /// -/// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't +/// - Bool stays as bool — `StatSum` has a native bool path and bool → f64 isn't /// currently a direct cast in vortex. /// - Primitive numerics → `f64` so the sum and finalize work without overflow. /// - Decimals stay as decimals @@ -315,13 +315,11 @@ mod tests { } #[test] - fn mean_all_null_returns_null() -> VortexResult<()> { - // The sum of zero valid values is null (SQL `SUM`), so the mean is null rather than - // the former 0/0 = NaN — matching SQL `AVG` of an all-null column. + fn mean_all_null_returns_nan() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; - assert!(result.is_null()); + assert!(result.as_primitive().as_::().is_some_and(f64::is_nan)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index da81e67ecb9..a841fbeffed 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 stat_sum; pub mod sum; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs new file mode 100644 index 00000000000..02b9b99f017 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::BitAnd; + +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_mask::AllOr; + +use super::SumState; +use super::checked_add_u64; +use crate::ExecutionCtx; +use crate::arrays::BoolArray; +use crate::arrays::bool::BoolArrayExt; + +pub(super) fn accumulate_bool( + inner: &mut SumState, + b: &BoolArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let SumState::Unsigned(acc) = inner else { + vortex_panic!("expected unsigned sum state for bool input"); + }; + + let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; + let true_count = match mask.bit_buffer() { + AllOr::None => return Ok(false), + AllOr::All => b.bit_buffer_view().true_count() as u64, + AllOr::Some(validity) => b.to_bit_buffer().bitand(validity).true_count() as u64, + }; + + Ok(checked_add_u64(acc, true_count)) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::stat_sum::StatSum; + use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::executor::VortexSessionExecute; + + #[test] + fn sum_bool_all_true() -> VortexResult<()> { + let arr: BoolArray = [true, true, true].into_iter().collect(); + let result = stat_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_mixed() -> VortexResult<()> { + let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); + let result = stat_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_all_false() -> VortexResult<()> { + let arr: BoolArray = [false, false, false].into_iter().collect(); + let result = stat_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) + } + + #[test] + fn sum_bool_with_nulls() -> VortexResult<()> { + let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); + let result = stat_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(2)); + Ok(()) + } + + #[test] + fn sum_bool_all_null() -> VortexResult<()> { + let arr = BoolArray::from_iter([None::, None, None]); + let result = stat_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) + } + + #[test] + fn sum_bool_empty_produces_zero() -> VortexResult<()> { + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) + } + + #[test] + fn sum_bool_finish_resets_state() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + + let batch1: BoolArray = [true, true, false].into_iter().collect(); + acc.accumulate(&batch1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + assert_eq!(result1.as_primitive().typed_value::(), Some(2)); + + let batch2: BoolArray = [false, true].into_iter().collect(); + acc.accumulate(&batch2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + assert_eq!(result2.as_primitive().typed_value::(), Some(1)); + Ok(()) + } + + #[test] + fn sum_bool_return_dtype() -> VortexResult<()> { + let dtype = StatSum + .return_dtype( + &NumericalAggregateOpts::default(), + &DType::Bool(Nullability::NonNullable), + ) + .unwrap(); + assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); + Ok(()) + } + + #[test] + fn sum_boolean_from_iter() -> VortexResult<()> { + let arr = BoolArray::from_iter([true, false, false, true]).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(2)); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs new file mode 100644 index 00000000000..a743696dedb --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs @@ -0,0 +1,219 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; + +/// Compute `scalar * len` for a constant array, returning the product as a sum-typed 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( + scalar: &Scalar, + len: usize, + return_dtype: &DType, +) -> VortexResult> { + if scalar.is_null() || len == 0 { + return Ok(None); + } + + let product = match scalar.dtype() { + DType::Bool(_) => { + let val = scalar + .as_bool() + .value() + .ok_or_else(|| vortex_err!("Expected non-null bool scalar for sum"))?; + if !val { + return Ok(None); + } + Scalar::primitive(len as u64, Nullability::Nullable) + } + DType::Primitive(..) => { + let pvalue = scalar + .as_primitive() + .pvalue() + .ok_or_else(|| vortex_err!("Expected non-null primitive scalar for sum"))?; + match return_dtype { + DType::Primitive(PType::U64, _) => { + let val = pvalue.cast::()?; + match val.checked_mul(len as u64) { + Some(product) => Scalar::primitive(product, Nullability::Nullable), + None => Scalar::null(return_dtype.as_nullable()), + } + } + DType::Primitive(PType::I64, _) => { + let val = pvalue.cast::()?; + match i64::try_from(len).ok().and_then(|l| val.checked_mul(l)) { + Some(product) => Scalar::primitive(product, Nullability::Nullable), + None => Scalar::null(return_dtype.as_nullable()), + } + } + DType::Primitive(PType::F64, _) => { + let val = pvalue.cast::()?; + Scalar::primitive(val * len as f64, Nullability::Nullable) + } + _ => vortex_bail!( + "Unexpected return dtype for primitive sum: {}", + return_dtype + ), + } + } + DType::Decimal(..) => { + let val = scalar + .as_decimal() + .decimal_value() + .ok_or_else(|| vortex_err!("Expected non-null decimal scalar for sum"))?; + let len_decimal = DecimalValue::from(len as i128); + match val.checked_mul(&len_decimal) { + Some(product) => { + let ret_decimal = *return_dtype + .as_decimal_opt() + .ok_or_else(|| vortex_err!("Expected decimal return dtype"))?; + Scalar::decimal(product, ret_decimal, Nullability::Nullable) + } + None => Scalar::null(return_dtype.as_nullable()), + } + } + _ => vortex_bail!("Unsupported constant type for sum: {}", scalar.dtype()), + }; + + Ok(Some(product)) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::array_session; + use crate::arrays::ConstantArray; + 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::Stat; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + + #[test] + fn sum_constant_unsigned() -> VortexResult<()> { + let array = ConstantArray::new(5u64, 10).into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 50u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_signed() -> VortexResult<()> { + let array = ConstantArray::new(-5i64, 10).into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, (-50i64).into()); + Ok(()) + } + + #[test] + fn sum_constant_nullable_value() -> VortexResult<()> { + let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) + .into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::primitive(0u64, Nullable)); + Ok(()) + } + + #[test] + fn sum_constant_bool_false() -> VortexResult<()> { + let array = ConstantArray::new(false, 10).into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 0u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_true() -> VortexResult<()> { + let array = ConstantArray::new(true, 10).into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 10u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_null() -> VortexResult<()> { + let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::primitive(0u64, Nullable)); + Ok(()) + } + + #[test] + fn sum_constant_decimal() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(100), + decimal_dtype, + Nullability::NonNullable, + ), + 5, + ) + .into_array(); + + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(500))) + ); + assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); + Ok(()) + } + + #[test] + fn sum_constant_decimal_null() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) + .into_array(); + + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result, + Scalar::decimal( + DecimalValue::I256(i256::ZERO), + DecimalDType::new(20, 2), + Nullable + ) + ); + Ok(()) + } + + #[test] + fn sum_constant_decimal_large_value() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(999_999_999), + decimal_dtype, + Nullability::NonNullable, + ), + 100, + ) + .into_array(); + + let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) + ); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs new file mode 100644 index 00000000000..0e07ad37c2e --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs @@ -0,0 +1,470 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use num_traits::AsPrimitive; +use num_traits::CheckedAdd; +use num_traits::NumOps; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_mask::Mask; + +use super::SumState; +use crate::ExecutionCtx; +use crate::arrays::DecimalArray; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::match_each_decimal_value_type; +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( + inner: &mut SumState, + d: &DecimalArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + let validity = match &mask { + Mask::AllTrue(_) => None, + Mask::Values(mask_values) => Some(mask_values.bit_buffer()), + Mask::AllFalse(_) => { + return Ok(false); + } + }; + + let SumState::Decimal { value, dtype } = inner else { + vortex_panic!("expected decimal sum state for decimal input"); + }; + + let values_type = DecimalType::smallest_decimal_value_type(dtype); + match_each_decimal_value_type!(d.values_type(), |T| { + match_each_decimal_value_type!(values_type, |I| { + let initial: I = value + .cast() + .vortex_expect("cannot fail to cast initial value"); + match sum_decimal_value(initial, d.buffer::(), validity, *dtype) { + Some(v) => *value = v, + None => return Ok(true), + } + Ok(false) + }) + }) +} + +fn sum_decimal_value( + initial: I, + values: Buffer, + validity: Option<&BitBuffer>, + output_dtype: DecimalDType, +) -> Option +where + T: AsPrimitive, + I: NumOps + CheckedAdd + Copy + NativeDecimalType + 'static, + bool: AsPrimitive, + DecimalValue: From, +{ + let sum = match validity { + Some(v) => sum_decimal_with_validity(values, v, initial), + None => sum_decimal(values, initial), + }; + + sum.map(DecimalValue::from) + // We have to make sure that the decimal value fits the precision of the decimal dtype. + .filter(|v| v.fits_in_precision(output_dtype)) +} + +fn sum_decimal, I: Copy + CheckedAdd + 'static>( + values: Buffer, + initial: I, +) -> Option { + let mut sum = initial; + for v in values.iter() { + let v: I = v.as_(); + sum = CheckedAdd::checked_add(&sum, &v)?; + } + Some(sum) +} + +fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> Option +where + T: AsPrimitive, + I: NumOps + CheckedAdd + Copy + 'static, + bool: AsPrimitive, +{ + let mut sum = initial; + for (v, valid) in values.iter().zip_eq(validity) { + let v: I = v.as_() * valid.as_(); + + sum = CheckedAdd::checked_add(&sum, &v)?; + } + Some(sum) +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::stat_sum::StatSum; + use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::array_session; + use crate::arrays::DecimalArray; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::dtype::Nullability::Nullable; + use crate::dtype::i256; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + use crate::validity::Validity; + + #[test] + fn sum_decimal_basic() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(600i32))), + )?; + + assert_eq!(result, expected); + 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 = stat_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_decimal_negative_values() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, -200i32, 300i32, -50i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(150i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_near_i32_max() -> VortexResult<()> { + let near_max = i32::MAX - 1000; + let decimal = DecimalArray::new( + buffer![near_max, 500i32, 400i32], + DecimalDType::new(10, 2), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = near_max as i64 + 500 + 400; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_large_i64_values() -> VortexResult<()> { + let large_val = i64::MAX / 4; + let decimal = DecimalArray::new( + buffer![large_val, large_val, large_val, large_val + 1], + DecimalDType::new(19, 0), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = (large_val as i128) * 4 + 1; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_preserves_scale() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![12345i32, 67890i32, 11111i32], + DecimalDType::new(6, 4), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(91346i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_single_value() -> VortexResult<()> { + let decimal = + DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(42i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32, 400i32], + DecimalDType::new(4, 2), + Validity::from_iter([false, false, true, false]), + ); + + let result = stat_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(300i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_overflow_detection() -> VortexResult<()> { + let max_val = i128::MAX / 2; + let decimal = DecimalArray::new( + buffer![max_val, max_val, max_val], + DecimalDType::new(38, 0), + Validity::AllValid, + ); + + let result = stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = + i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + 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, + ); + + assert_eq!( + stat_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx() + ) + .vortex_expect("operation should succeed in test"), + 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 = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(99_999_999_999_990i64), + DecimalDType::new(14, 0), + Nullable, + ); + StatSum.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); + StatSum.combine_partials(&mut state, small)?; + + let result = StatSum.to_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(()) + } + + #[test] + fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { + // Input precision 4 → return precision 14. Native I64 (max 18). + // The max representable value for precision 14 is 10^14 - 1. + // When the sum reaches exactly 10^14, fits_in_precision fails 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 = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(99_999_999_999_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StatSum.combine_partials(&mut state, near_limit)?; + + // Push the sum to exactly 10^14, exceeding precision 14. + let one_more = + Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); + StatSum.combine_partials(&mut state, one_more)?; + + let result = StatSum.to_scalar(&state)?; + assert!(result.is_null()); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(14, 0), Nullable) + ); + Ok(()) + } + + #[test] + fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { + // Same setup but with negative values: sum reaches -10^14. + let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); + let mut state = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(-99_999_999_999_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StatSum.combine_partials(&mut state, near_limit)?; + + let one_more = Scalar::decimal( + DecimalValue::from(-1i64), + DecimalDType::new(14, 0), + Nullable, + ); + StatSum.combine_partials(&mut state, one_more)?; + + let result = StatSum.to_scalar(&state)?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { + // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. + // Native for precision 38 is I128 (max 38), so 38 = 38. + // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. + // + // We 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 = StatSum.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); + StatSum.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); + + // Drive accumulate through the vtable directly. + let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); + let mut ctx = array_session().create_execution_ctx(); + StatSum.accumulate(&mut state, &columnar, &mut ctx)?; + + let result = StatSum.to_scalar(&state)?; + assert!(result.is_null()); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs new file mode 100644 index 00000000000..e1affb72229 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs @@ -0,0 +1,408 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::StatSum; +use super::primitive::sum_float_all; +use super::primitive::sum_signed_all; +use super::primitive::sum_unsigned_all; +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::kernels::DynGroupedAggregateKernel; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::dtype::NativePType; +use crate::match_each_native_ptype; + +/// Encoding-specific grouped [`StatSum`] kernel for primitive element arrays. +#[derive(Debug)] +pub(crate) struct PrimitiveGroupedStatSumEncodingKernel; + +impl DynGroupedAggregateKernel for PrimitiveGroupedStatSumEncodingKernel { + 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 [`StatSum`] 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. +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, + )?)) +} + +/// StatSum 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 result = 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(result.into_array()) +} + +/// Reduce each group's element slice into a nullable sum. A group is null when the group +/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). +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 { + let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { + if !group_validity.value(i) { + return None; + } + let mut acc = A::default(); + let overflow = if all_valid { + sum_run(&mut acc, &values[offset..offset + size]) + } else { + sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) + }; + (!overflow).then_some(acc) + }); + PrimitiveArray::from_option_iter(sums) +} + +/// StatSum 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. +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 { + match elem_mask.slice(offset..offset + size).slices() { + AllOr::All => sum_run(acc, &values[offset..offset + size]), + AllOr::None => false, + AllOr::Some(runs) => { + for &(start, end) in runs { + if sum_run(acc, &values[offset + start..offset + end]) { + return true; + } + } + false + } + } +} + +#[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::stat_sum::StatSum; + use crate::aggregate_fn::fns::stat_sum::stat_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( + StatSum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish(&mut ctx) + } + + /// 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 = StatSum + .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .expect("sum partial 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(&stat_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), Some(0i64), Some(0i64), 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( + StatSum, + NumericalAggregateOpts::include_nans(), + elem_dtype, + )?; + let mut ctx2 = array_session().create_execution_ctx(); + acc.accumulate_list(&groups, &mut ctx2)?; + let actual = acc.finish(&mut ctx2)?; + + 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/stat_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs new file mode 100644 index 00000000000..1ca99339cd9 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs @@ -0,0 +1,909 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod bool; +mod constant; +mod decimal; +mod grouped; +mod primitive; +pub(crate) use grouped::PrimitiveGroupedStatSumEncodingKernel; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +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; +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::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::MAX_PRECISION; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; + +/// The monoid sum of an array: zero when there are no valid values, null on overflow. +/// +/// This is the value stored as `Stat::Sum` and merged across chunks, zones, and files: the +/// empty sum is zero so that partials merge as a monoid (`combine(empty, x) = x`). For the +/// SQL rule (a sum over zero valid values is null) use +/// [`sum`](crate::aggregate_fn::fns::sql_sum::sum) / [`SqlSum`](crate::aggregate_fn::fns::sql_sum::SqlSum). +pub fn stat_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Short-circuit using cached array statistics. + if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { + return Ok(sum_scalar); + } + + // Compute using Accumulator. + // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe. + let mut acc = Accumulator::try_new( + StatSum, + 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). + if let Some(val) = result.value().cloned() { + array.statistics().set(Stat::Sum, Precision::Exact(val)); + } + + Ok(result) +} + +/// The monoid sum aggregate: an all-invalid or empty sum is zero, an overflowing sum is null. +/// +/// This aggregate's plain partial is the persisted wire form of sum statistics (zone maps, +/// file stats), so its partial dtype and semantics must stay stable. The SQL `SUM` rule is +/// provided by the separate [`SqlSum`](crate::aggregate_fn::fns::sql_sum::SqlSum) aggregate. +/// +/// 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. +#[derive(Clone, Debug)] +pub struct StatSum; + +// Both Spark and DataFusion use this heuristic. +// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/StatSum.scala#L66 +// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 +pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { + DecimalDType::new( + u8::min(MAX_PRECISION, input.precision() + 10), + input.scale(), + ) +} + +impl AggregateFnVTable for StatSum { + type Options = NumericalAggregateOpts; + type Partial = StatSumPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.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 { + 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))?; + let initial = make_zero_state(&return_dtype); + + Ok(StatSumPartial { + return_dtype, + current: Some(initial), + skip_nans: options.skip_nans, + }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if other.is_null() { + // A null partial means the sub-accumulator saturated (overflow). + partial.current = None; + return Ok(()); + } + let Some(ref mut inner) = partial.current else { + return Ok(()); + }; + 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; + } + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(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) + } + }) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.current = Some(make_zero_state(&partial.return_dtype)); + } + + #[inline] + fn is_saturated(&self, partial: &Self::Partial) -> bool { + match partial.current.as_ref() { + None => true, + Some(SumState::Float(v)) => v.is_nan(), + Some(_) => false, + } + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes + // the default dispatch path. + if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + 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. + if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + self.combine_partials(partial, sum)?; + return Ok(true); + } + Ok(false) + } + Precision::Exact(_) => { + // At least one NaN value: the sum is NaN without scanning the batch. + if let Some(SumState::Float(acc)) = partial.current.as_mut() { + *acc = f64::NAN; + } + Ok(true) + } + _ => Ok(false), + } + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + // Constants compute scalar * len and combine via combine_partials. + if let Columnar::Constant(c) = batch { + // 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 mut inner = match partial.current.take() { + Some(inner) => inner, + None => return Ok(()), + }; + + 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) => partial.current = Some(inner), + Ok(true) => {} // saturated: current stays None + Err(e) => { + partial.current = Some(inner); + return Err(e); + } + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +/// The group state for a sum aggregate, containing the accumulated value and configuration +/// needed for reset/result without external context. +pub struct StatSumPartial { + return_dtype: DType, + /// The current accumulated state, or `None` if saturated (checked overflow). + current: Option, + /// Whether NaN values in float inputs are skipped. + skip_nans: bool, +} + +/// The accumulated sum value. +// TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the +// input type every time? Perhaps? +pub enum SumState { + Unsigned(u64), + Signed(i64), + Float(f64), + Decimal { + value: DecimalValue, + dtype: DecimalDType, + }, +} + +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), + PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0), + PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0), + }, + DType::Decimal(decimal, _) => SumState::Decimal { + value: DecimalValue::zero(decimal), + dtype: *decimal, + }, + _ => vortex_panic!("Unsupported sum type"), + } +} + +/// Checked add for u64, returning true if overflow occurred. +#[inline(always)] +pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { + match acc.checked_add(val) { + Some(r) => { + *acc = r; + false + } + None => true, + } +} + +/// Checked add for i64, returning true if overflow occurred. +#[inline(always)] +pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { + match acc.checked_add(val) { + Some(r) => { + *acc = r; + false + } + None => true, + } +} + +#[cfg(test)] +mod tests { + use num_traits::CheckedAdd; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + 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::stat_sum::StatSum; + use crate::aggregate_fn::fns::stat_sum::stat_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::expr::stats::StatsProvider; + use crate::scalar::DecimalValue; + use crate::scalar::NumericOperator; + use crate::scalar::Scalar; + use crate::validity::Validity; + + /// StatSum an array with an initial value (test-only helper). + fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + if accumulator.is_null() { + return Ok(accumulator.clone()); + } + if accumulator.is_zero() == Some(true) { + return stat_sum(array, &mut ctx); + } + + let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { + vortex_error::vortex_err!("StatSum not supported for dtype: {}", array.dtype()) + })?; + + // For non-float types, try statistics short-circuit with accumulator. + if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float()) + && let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) + { + return add_scalars(&sum_dtype, &sum_scalar, accumulator); + } + + // Compute array sum from zero (also caches stats). + let array_sum = stat_sum(array, &mut ctx)?; + + // Combine with the accumulator. + add_scalars(&sum_dtype, &array_sum, accumulator) + } + + /// Add two sum scalars with overflow checking. + fn add_scalars(sum_dtype: &DType, lhs: &Scalar, rhs: &Scalar) -> VortexResult { + if lhs.is_null() || rhs.is_null() { + return Ok(Scalar::null(sum_dtype.as_nullable())); + } + + Ok(match sum_dtype { + DType::Primitive(ptype, _) if ptype.is_float() => { + let lhs_val = f64::try_from(lhs)?; + let rhs_val = f64::try_from(rhs)?; + Scalar::primitive(lhs_val + rhs_val, Nullable) + } + DType::Primitive(..) => lhs + .as_primitive() + .checked_add(&rhs.as_primitive()) + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), + DType::Decimal(..) => lhs + .as_decimal() + .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), + _ => unreachable!("StatSum will always be a decimal or a primitive dtype"), + }) + } + + // Multi-batch and reset tests + + #[test] + fn sum_multi_batch() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, 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)); + Ok(()) + } + + #[test] + fn sum_finish_resets_state() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + + let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + acc.accumulate(&batch1, &mut ctx)?; + let result1 = acc.finish()?; + assert_eq!(result1.as_primitive().typed_value::(), Some(30)); + + let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + let result2 = acc.finish()?; + assert_eq!(result2.as_primitive().typed_value::(), Some(18)); + Ok(()) + } + + // State merge tests (vtable-level) + + #[test] + fn sum_state_merge() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = StatSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + + let scalar1 = Scalar::primitive(100i64, Nullable); + StatSum.combine_partials(&mut state, scalar1)?; + + let scalar2 = Scalar::primitive(50i64, Nullable); + StatSum.combine_partials(&mut state, scalar2)?; + + let result = StatSum.to_scalar(&state)?; + StatSum.reset(&mut state); + assert_eq!(result.as_primitive().typed_value::(), Some(150)); + Ok(()) + } + + // Stats caching test + + #[test] + fn sum_stats() -> VortexResult<()> { + let array = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([1, 1, 1]).into_array(), + PrimitiveArray::from_iter([2, 2, 2]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + ) + .vortex_expect("operation should succeed in test"); + let array = array.into_array(); + // compute sum with accumulator to populate stats + sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; + + let sum_without_acc = stat_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); + Ok(()) + } + + // Constant float non-multiply test + + #[test] + fn sum_constant_float_non_multiply() -> VortexResult<()> { + let acc = -2048669276050936500000000000f64; + let array = ConstantArray::new(6.1811675e16f64, 25); + let result = sum_with_accumulator(&array.into_array(), &Scalar::primitive(acc, Nullable)) + .vortex_expect("operation should succeed in test"); + assert_eq!( + f64::try_from(&result).vortex_expect("operation should succeed in test"), + -2048669274505644600000000000f64 + ); + Ok(()) + } + + // Grouped sum tests + + fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { + let mut acc = GroupedAccumulator::try_new( + StatSum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish(&mut ctx) + } + + #[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<()> { + 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)?; + + let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn grouped_sum_bool() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); + let groups = + FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Bool(Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).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(StatSum, 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(&mut ctx)?; + + 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(&mut ctx)?; + + 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(()) + } + + // Chunked array tests + + #[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 chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new( + vec![ + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), + ], + dtype, + )?; + + let result = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(20.8)); + Ok(()) + } + + #[test] + fn sum_chunked_floats_all_nulls_is_zero() -> 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 = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result, Scalar::primitive(0f64, Nullable)); + Ok(()) + } + + #[test] + fn sum_chunked_floats_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 = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(36.0)); + Ok(()) + } + + #[test] + fn sum_chunked_int_almost_all_null() -> VortexResult<()> { + 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 = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(1)); + Ok(()) + } + + #[test] + fn sum_chunked_decimals() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let chunk1 = DecimalArray::new( + buffer![100i32, 100i32, 100i32, 100i32, 100i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk2 = DecimalArray::new( + buffer![200i32, 200i32, 200i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new( + vec![ + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), + ], + dtype, + )?; + + let result = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(1700))) + ); + Ok(()) + } + + #[test] + fn sum_chunked_decimals_with_nulls() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let chunk1 = DecimalArray::new( + buffer![100i32, 100i32, 100i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk2 = DecimalArray::new( + buffer![0i32, 0i32], + decimal_dtype, + Validity::from_iter([false, false]), + ); + let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new( + vec![ + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), + ], + dtype, + )?; + + let result = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(700))) + ); + Ok(()) + } + + #[test] + fn sum_chunked_decimals_large() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(3, 0); + let chunk1 = ConstantArray::new( + Scalar::decimal( + DecimalValue::I16(500), + decimal_dtype, + Nullability::NonNullable, + ), + 1, + ); + let chunk2 = ConstantArray::new( + Scalar::decimal( + DecimalValue::I16(600), + decimal_dtype, + Nullability::NonNullable, + ), + 1, + ); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = stat_sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(1100))) + ); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(13, 0), Nullable) + ); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs new file mode 100644 index 00000000000..da34a97e2bd --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs @@ -0,0 +1,460 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; +use num_traits::ToPrimitive; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_mask::AllOr; + +use super::SumState; +use super::checked_add_i64; +use super::checked_add_u64; +use crate::ExecutionCtx; +use crate::arrays::PrimitiveArray; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; + +/// Number of elements summed without an overflow check. Chosen so that a chunk of values narrower +/// 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( + inner: &mut SumState, + p: &PrimitiveArray, + ctx: &mut ExecutionCtx, + skip_nans: bool, +) -> VortexResult { + let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; + match mask.slices() { + AllOr::None => Ok(false), + AllOr::All => accumulate_primitive_all(inner, p, skip_nans), + AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), + } +} + +fn accumulate_primitive_all( + inner: &mut SumState, + p: &PrimitiveArray, + skip_nans: bool, +) -> VortexResult { + match inner { + SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |T| { Ok(sum_unsigned_all(acc, p.as_slice::())) }, + signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, + floating: |_T| { vortex_panic!("unsigned sum state with float input") } + ), + SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, + signed: |T| { Ok(sum_signed_all(acc, p.as_slice::())) }, + floating: |_T| { vortex_panic!("signed sum state with float input") } + ), + SumState::Float(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, + signed: |_T| { vortex_panic!("float sum state with signed input") }, + floating: |T| { + sum_float_all(acc, p.as_slice::(), skip_nans); + Ok(false) + } + ), + SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), + } +} + +/// StatSum 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) { + if skip_nans { + for &v in slice { + if !v.is_nan() { + *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); + } + } + } else { + for &v in slice { + *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); + } + } +} + +/// StatSum all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in +/// 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 +where + T: NativePType + AsPrimitive, +{ + if T::PTYPE == PType::U64 { + for &v in slice { + if checked_add_u64(acc, v.as_()) { + return true; + } + } + return false; + } + for chunk in slice.chunks(SUM_CHUNK) { + let chunk_sum: u64 = chunk.iter().map(|&v| v.as_()).sum(); + if checked_add_u64(acc, chunk_sum) { + return true; + } + } + false +} + +/// Signed counterpart of [`sum_unsigned_all`]. +pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool +where + T: NativePType + AsPrimitive, +{ + if T::PTYPE == PType::I64 { + for &v in slice { + if checked_add_i64(acc, v.as_()) { + return true; + } + } + return false; + } + for chunk in slice.chunks(SUM_CHUNK) { + let chunk_sum: i64 = chunk.iter().map(|&v| v.as_()).sum(); + if checked_add_i64(acc, chunk_sum) { + return true; + } + } + false +} + +/// StatSum the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each +/// run is a slice of fully-valid values, so it reuses the same vectorized reduction as the +/// all-valid path instead of a per-element validity branch. +fn accumulate_primitive_valid( + inner: &mut SumState, + p: &PrimitiveArray, + slices: &[(usize, usize)], + skip_nans: bool, +) -> VortexResult { + match inner { + SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |T| { + let values = p.as_slice::(); + for &(start, end) in slices { + if sum_unsigned_all(acc, &values[start..end]) { + return Ok(true); + } + } + Ok(false) + }, + signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, + floating: |_T| { vortex_panic!("unsigned sum state with float input") } + ), + SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, + signed: |T| { + let values = p.as_slice::(); + for &(start, end) in slices { + if sum_signed_all(acc, &values[start..end]) { + return Ok(true); + } + } + Ok(false) + }, + floating: |_T| { vortex_panic!("signed sum state with float input") } + ), + SumState::Float(acc) => match_each_native_ptype!(p.ptype(), + unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, + signed: |_T| { vortex_panic!("float sum state with signed input") }, + floating: |T| { + let values = p.as_slice::(); + for &(start, end) in slices { + sum_float_all(acc, &values[start..end], skip_nans); + } + Ok(false) + } + ), + SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::stat_sum::StatSum; + use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::array_session; + use crate::arrays::ConstantArray; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + 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; + + #[test] + fn sum_i32() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(10)); + Ok(()) + } + + #[test] + fn sum_u8() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(60)); + Ok(()) + } + + #[test] + fn sum_f64() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); + Ok(()) + } + + #[test] + fn sum_with_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6)); + Ok(()) + } + + #[test] + fn sum_multiple_null_runs() -> VortexResult<()> { + // Several disjoint valid runs separated by nulls exercise the per-run fold. + let arr = PrimitiveArray::from_option_iter([ + Some(1i32), + Some(2), + None, + None, + Some(3), + None, + Some(4), + Some(5), + Some(6), + ]) + .into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(21)); + Ok(()) + } + + #[test] + fn sum_all_null() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) + } + + #[test] + fn sum_all_invalid_float() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::primitive(0f64, Nullable)); + Ok(()) + } + + #[test] + fn sum_buffer_i32() -> VortexResult<()> { + let arr = buffer![1, 1, 1, 1].into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4)); + Ok(()) + } + + #[test] + fn sum_buffer_f64() -> VortexResult<()> { + let arr = buffer![1., 1., 1., 1.].into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4.)); + Ok(()) + } + + #[test] + fn sum_empty_produces_zero() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) + } + + #[test] + fn sum_empty_f64_produces_zero() -> VortexResult<()> { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) + } + + #[test] + fn sum_f64_with_nan() -> VortexResult<()> { + let arr = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], + Validity::NonNullable, + ) + .into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + Ok(()) + } + + #[test] + fn sum_f32_with_nan() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); + Ok(()) + } + + #[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 = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); + Ok(()) + } + + #[test] + fn sum_all_nan() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) + } + + /// StatSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + fn sum_with_options( + arr: &crate::ArrayRef, + options: NumericalAggregateOpts, + ) -> VortexResult { + let mut acc = Accumulator::try_new(StatSum, options, arr.dtype().clone())?; + acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.finish() + } + + #[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_f64_without_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact StatSum 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))); + 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 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 = stat_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( + StatSum, + NumericalAggregateOpts::default(), + DType::Primitive(PType::F64, Nullability::NonNullable), + )?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + Ok(()) + } + + #[test] + fn sum_checked_overflow() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); + let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_checked_overflow_is_saturated() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StatSum, 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()); + + // finish resets state, clearing saturation + drop(acc.finish()?); + assert!(!acc.is_saturated()); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index e0b3ce50070..37a719e9ed4 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -11,7 +11,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -30,6 +29,11 @@ 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::stat_sum::SumState; +use crate::aggregate_fn::fns::stat_sum::checked_add_i64; +use crate::aggregate_fn::fns::stat_sum::checked_add_u64; +use crate::aggregate_fn::fns::stat_sum::make_zero_state; +use crate::aggregate_fn::fns::stat_sum::sum_decimal_dtype; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::Struct; @@ -38,10 +42,8 @@ use crate::arrays::bool::BoolArrayExt; use crate::arrays::masked::mask_validity_canonical; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; -use crate::dtype::DecimalDType; use crate::dtype::FieldName; use crate::dtype::FieldNames; -use crate::dtype::MAX_PRECISION; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; @@ -49,40 +51,9 @@ use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; -use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::validity::Validity; -/// The monoid sum of an array: zero when there are no valid values, null on overflow. -/// -/// This is the value stored as `Stat::Sum`. Unlike SQL [`sum`], the empty sum is zero so that -/// per-chunk and per-zone statistics merge as a monoid (`combine(empty, x) = x`); SQL -/// semantics are recovered at read boundaries by combining this stat with the null count. -pub fn stat_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { - return Ok(sum_scalar); - } - - let mut acc = Accumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - array.dtype().clone(), - )?; - acc.accumulate(array, ctx)?; - let partial = acc.partial_scalar()?; - let result = partial - .as_struct() - .field("sum") - .vortex_expect("Sum partial is missing the `sum` field"); - - // Cache the computed sum as a statistic (only if non-null, i.e. no overflow). - if let Some(val) = result.value().cloned() { - array.statistics().set(Stat::Sum, Precision::Exact(val)); - } - - Ok(result) -} - /// Return the SQL sum of an array: null when the array has no valid values or the sum /// overflows. /// @@ -126,7 +97,10 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// Sum an array following SQL `SUM` semantics. /// /// A sum over zero valid values yields null (the SQL rule: nulls are eliminated, and the sum -/// of an empty set is null), and a sum that overflows yields null. +/// of an empty set is null), and a sum that overflows yields null. This is a distinct +/// aggregate from the monoid [`Sum`](crate::aggregate_fn::fns::sum::Sum), whose empty sum is +/// zero and whose plain partial is the persisted form of sum statistics; `Sum`'s +/// `{sum, seen}` partial is never written to statistics. /// /// The partial state is a `{sum, seen}` struct: `sum` is the running monoid value (null once /// saturated by overflow, which poisons merges), and `seen` records whether any valid value @@ -139,22 +113,12 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { #[derive(Clone, Debug)] pub struct Sum; -// Both Spark and DataFusion use this heuristic. -// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 -// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 -pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { - DecimalDType::new( - u8::min(MAX_PRECISION, input.precision() + 10), - input.scale(), - ) -} - impl AggregateFnVTable for Sum { type Options = NumericalAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.sum"); + static ID: CachedId = CachedId::new("vortex.sql_sum"); *ID } @@ -326,35 +290,22 @@ impl AggregateFnVTable for Sum { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes - // the default dispatch path. - if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + // `Stat::Sum` is the monoid (NaN-skipping) sum, so the default NaN-skipping path can + // consume it directly. `Sum` has no Stat slot of its own — the shortcut lives here + // rather than in the accumulator's stats bridge. + 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.current, Some(SumState::Float(_))) { 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. - if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { - // The cached sum is a monoid value that cannot carry `seen`; derive it - // from the null count, or fall through to a real scan when unknown. - 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(_) => partial.seen = true, - _ => return Ok(false), - } - let sum = if sum.dtype() == &partial.return_dtype { - sum - } else { - sum.cast(&partial.return_dtype)? - }; - self.combine_partials(partial, sum)?; - return Ok(true); - } - Ok(false) + try_accumulate_cached_sum(self, partial, batch) } Precision::Exact(_) => { // At least one NaN value (a valid value): the sum is NaN without scanning. @@ -457,6 +408,34 @@ impl AggregateFnVTable for Sum { } } +/// Consume a batch's cached monoid `Stat::Sum` instead of scanning it. The cached sum cannot +/// carry `seen`, so the batch's null count decides between the identity (all null) and a seen +/// contribution; when either statistic is missing the caller falls through to a real scan. +fn try_accumulate_cached_sum( + vtable: &Sum, + partial: &mut SumPartial, + 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(_) => partial.seen = true, + _ => 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) +} + /// The group state for a sum aggregate, containing the accumulated value and configuration /// needed for reset/result without external context. pub struct SumPartial { @@ -502,58 +481,6 @@ fn sum_value_scalar(partial: &SumPartial) -> Scalar { } } -/// The accumulated sum value. -// TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the -// input type every time? Perhaps? -pub enum SumState { - Unsigned(u64), - Signed(i64), - Float(f64), - Decimal { - value: DecimalValue, - dtype: DecimalDType, - }, -} - -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), - PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0), - PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0), - }, - DType::Decimal(decimal, _) => SumState::Decimal { - value: DecimalValue::zero(decimal), - dtype: *decimal, - }, - _ => vortex_panic!("Unsupported sum type"), - } -} - -/// Checked add for u64, returning true if overflow occurred. -#[inline(always)] -fn checked_add_u64(acc: &mut u64, val: u64) -> bool { - match acc.checked_add(val) { - Some(r) => { - *acc = r; - false - } - None => true, - } -} - -/// Checked add for i64, returning true if overflow occurred. -#[inline(always)] -fn checked_add_i64(acc: &mut i64, val: i64) -> bool { - match acc.checked_add(val) { - Some(r) => { - *acc = r; - false - } - None => true, - } -} - #[cfg(test)] mod tests { use num_traits::CheckedAdd; @@ -570,8 +497,8 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::stat_sum::stat_sum; use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::stat_sum; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::BoolArray; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c496f7a9bf8..693e444b21a 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::stat_sum::PrimitiveGroupedStatSumEncodingKernel; +use crate::aggregate_fn::fns::stat_sum::StatSum; 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(StatSum); this.register(Sum); this.register(UncompressedSizeInBytes); @@ -115,6 +118,11 @@ impl Default for AggregateFnSession { Sum.id(), &PrimitiveGroupedSumEncodingKernel, ); + this.register_grouped_encoding_kernel( + Primitive.id(), + StatSum.id(), + &PrimitiveGroupedStatSumEncodingKernel, + ); this } diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index c3dc2e0eed7..52e9e5224e8 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -28,7 +28,7 @@ use crate::ExecutionResult; use crate::IntoArray; use crate::VTable; use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::sum::sum; +use crate::aggregate_fn::fns::stat_sum::stat_sum; use crate::array::ArrayData; use crate::array::ArrayId; use crate::array::ArrayInner; @@ -334,7 +334,7 @@ impl ArrayRef { Validity::NonNullable | Validity::AllValid => len, Validity::AllInvalid => 0, Validity::Array(a) => { - let array_sum = sum(&a, ctx)?; + let array_sum = stat_sum(&a, ctx)?; array_sum .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 99b6bba44de..9eb24019e69 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -47,7 +47,7 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::stat_sum::StatSum; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -60,7 +60,7 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); let mut acc = Accumulator::try_new( - Sum, + StatSum, NumericalAggregateOpts::default(), batch.dtype().clone(), )?; @@ -120,8 +120,7 @@ mod tests { DType::Primitive(PType::I32, Nullability::Nullable), )?; let result = run_sum(&chunked.into_array())?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -159,7 +158,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let chunked = ChunkedArray::try_new(vec![], dtype)?; let result = run_sum(&chunked.into_array())?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 31aaee7c39c..01981803239 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -103,14 +103,8 @@ impl ArrayParentReduceRule for StructGetItemRule { match child.validity()? { Validity::NonNullable | Validity::AllValid => { - // 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())) - } + // If the struct is non-nullable or all valid, the field's validity is unchanged + Ok(Some(field.clone())) } Validity::AllInvalid => { // If everything is invalid, the field is also all invalid diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 73f98cce4a7..f143e637888 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1002,7 +1002,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::aggregate_fn::fns::sum::sum; + use crate::aggregate_fn::fns::stat_sum::stat_sum; use crate::dtype::DType; let len = array.len(); @@ -1045,7 +1045,9 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { return; } - if let (Ok(slice_sum), Ok(canonical_sum)) = (sum(&sliced, ctx), sum(&canonical_sliced, ctx)) { + if let (Ok(slice_sum), Ok(canonical_sum)) = + (stat_sum(&sliced, ctx), stat_sum(&canonical_sliced, ctx)) + { // Compare sum scalars assert_eq!( slice_sum, canonical_sum, diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 57e11135ef3..68f24284d29 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -188,7 +188,7 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - return aggregate_fn::fns::sum::Sum + return aggregate_fn::fns::stat_sum::StatSum .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); } }) @@ -200,7 +200,9 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), + Self::Sum => { + aggregate_fn::fns::stat_sum::StatSum.bind(NumericalAggregateOpts::skip_nans()) + } Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { @@ -216,7 +218,7 @@ impl Stat { /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those /// aggregates have no stat slot. pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option { - if let Some(options) = aggregate_fn.as_opt::() { + if let Some(options) = aggregate_fn.as_opt::() { return options.skip_nans.then_some(Self::Sum); } if aggregate_fn.is::() { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index b49c60b10ed..b98aa96a38e 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -29,8 +29,8 @@ use crate::scalar_fn::ScalarFnVTable; /// /// 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 [`Sum`]. The result +/// dtype follows [`Sum`]'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 diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 85ad1ae339e..84fc5760495 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -173,16 +173,12 @@ fn stat_array( } .map(ScalarValue::Bool) } else if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) { - let stored = array + array .statistics() .with_typed_stats_set(|stats| stats.get(stat)) // We don't mind whether the stat is approxed or not, since these are row-wise bounds. - .into_inner(); - match stored { - Some(scalar) if scalar.dtype().eq_ignore_nullability(&dtype) => scalar.into_value(), - Some(scalar) => result_stat_to_state(array, aggregate_fn, scalar, len)?, - None => None, - } + .into_inner() + .and_then(Scalar::into_value) } else { tracing::trace!( "No legacy Stat slot for aggregate {}; stat expression will resolve to null", @@ -194,28 +190,3 @@ fn stat_array( let scalar = Scalar::try_new(dtype, value)?; Ok(ConstantArray::new(scalar, len).into_array()) } - -/// Rebuild an aggregate's partial state from a result-form statistic (e.g. `Sum`'s monoid -/// statistic vs its `{sum, seen}` state). A result-form stat cannot express "no valid values", -/// so the array's null count decides between the empty state and a seen state; when the null -/// count is unknown for a nullable array, the stat is unusable and resolves to missing. -fn result_stat_to_state( - array: &ArrayRef, - aggregate_fn: &AggregateFnRef, - stat_scalar: Scalar, - len: usize, -) -> VortexResult> { - let all_null = if array.dtype().is_nullable() { - match array.statistics().get_as::(Stat::NullCount) { - Precision::Exact(null_count) => null_count == len as u64, - _ => return Ok(None), - } - } else { - false - }; - let mut acc = aggregate_fn.accumulator(array.dtype())?; - if !all_null { - acc.combine_partials(stat_scalar)?; - } - Ok(acc.partial_scalar()?.into_value()) -} diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index 1ceb4bad3f8..8c754ecedc3 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -23,7 +23,7 @@ use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::sum::stat_sum; +use crate::aggregate_fn::fns::stat_sum::stat_sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::expr::stats::Precision; use crate::expr::stats::Stat; diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 0cca359d11f..81dafb20c0e 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -14,7 +14,7 @@ use crate::aggregate_fn::fns::all_null::AllNull; 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::sum::Sum; +use crate::aggregate_fn::fns::stat_sum::StatSum; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; pub use crate::scalar_fn::fns::stat::StatFn; @@ -37,7 +37,7 @@ pub fn min_max(expr: Expression) -> Expression { /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. - stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) + stat(expr, StatSum.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. @@ -89,8 +89,6 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; @@ -106,30 +104,10 @@ mod tests { use crate::expr::stats::Stat; use crate::scalar::Scalar; use crate::scalar::ScalarValue; + use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(array_session); - /// The `{sum, seen}` state dtype for a sum over non-nullable i32 input. - fn sum_state_dtype() -> DType { - use crate::aggregate_fn::AggregateFnVTable; - Sum.partial_dtype( - &NumericalAggregateOpts::skip_nans(), - &DType::Primitive(PType::I32, Nullability::NonNullable), - ) - .vortex_expect("sum partial dtype") - } - - /// A seen `{sum, seen}` state scalar with the given running sum. - fn sum_state(value: i64) -> Scalar { - Scalar::struct_( - sum_state_dtype(), - vec![ - Scalar::primitive(value, Nullability::Nullable), - Scalar::bool(true, Nullability::NonNullable), - ], - ) - } - #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { let array = buffer![1i32, 2, 3].into_array(); @@ -144,8 +122,8 @@ mod tests { .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - // Stat expressions resolve to the aggregate's `{sum, seen}` state form. - let expected = ConstantArray::new(sum_state(6), 3).into_array(); + let expected = + ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) @@ -160,7 +138,11 @@ mod tests { .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - let expected = ConstantArray::new(Scalar::null(sum_state_dtype()), 3).into_array(); + let expected = ConstantArray::new( + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + 3, + ) + .into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) @@ -191,13 +173,10 @@ mod tests { let result = result .execute::(&mut SESSION.create_execution_ctx())? .into_array(); - let expected = ChunkedArray::try_new( - vec![ - ConstantArray::new(sum_state(3), 2).into_array(), - ConstantArray::new(Scalar::null(sum_state_dtype()), 3).into_array(), - ], - sum_state_dtype(), - )? + let expected = PrimitiveArray::new( + buffer![3i64, 3, 0, 0, 0], + Validity::from_iter([true, true, false, false, false]), + ) .into_array(); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); diff --git a/vortex-datafusion/src/convert/stats.rs b/vortex-datafusion/src/convert/stats.rs index ccad75ddae6..33a33a78ccf 100644 --- a/vortex-datafusion/src/convert/stats.rs +++ b/vortex-datafusion/src/convert/stats.rs @@ -20,7 +20,6 @@ use crate::convert::TryToDataFusion; pub(crate) fn stats_set_to_df( stats_set: &StatsSet, dtype: &DType, - row_count: Option, ) -> VortexResult { // Update the total size in bytes. let column_size = stats_set.get_as::(Stat::UncompressedSizeInBytes, &PType::U64.into()); @@ -53,17 +52,8 @@ pub(crate) fn stats_set_to_df( .ok() }); - let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); - - // Find the sum statistic. `Stat::Sum` is the monoid sum (zero when the column has no - // valid values), while DataFusion's `sum_value` is the SQL sum (null in that case), so an - // all-null column must not export its zero; without a known null count the sum cannot be - // interpreted and is withheld. + // Find the sum statistic let sum = stats_set.get(Stat::Sum).and_then(|stat_val| { - match (&null_count, row_count) { - (VortexPrecision::Exact(null_count), Some(row_count)) if *null_count < row_count => {} - _ => return None, - } Scalar::try_new( Stat::Sum .dtype(dtype) @@ -75,6 +65,8 @@ pub(crate) fn stats_set_to_df( .ok() }); + let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); + Ok(ColumnStatistics { null_count: null_count.to_df(), min_value: min.to_df(), @@ -105,20 +97,12 @@ mod tests { #[test] fn is_constant_false_does_not_imply_one_distinct_value() -> VortexResult<()> { let false_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(false)); - let false_stats = stats_set_to_df( - &false_constant, - &DType::Bool(Nullability::NonNullable), - Some(5), - )?; + let false_stats = stats_set_to_df(&false_constant, &DType::Bool(Nullability::NonNullable))?; assert_eq!(false_stats.distinct_count, Precision::Absent); let true_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(true)); - let true_stats = stats_set_to_df( - &true_constant, - &DType::Bool(Nullability::NonNullable), - Some(5), - )?; + let true_stats = stats_set_to_df(&true_constant, &DType::Bool(Nullability::NonNullable))?; assert_eq!(true_stats.distinct_count, Precision::Exact(1)); diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index e4e6a870159..82fae9bf56d 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -242,10 +242,6 @@ impl VortexDataSourceBuilder { }; // We now compute initial statistics. - let row_count = match self.data_source.row_count() { - Precision::Exact(n) => usize::try_from(n).ok(), - _ => None, - }; let field_paths: Vec<_> = fields .names() .iter() @@ -260,7 +256,7 @@ impl VortexDataSourceBuilder { .await? .iter() .zip(fields.fields()) - .map(|(stats, dtype)| stats_set_to_df(stats, &dtype, row_count)) + .map(|(stats, dtype)| stats_set_to_df(stats, &dtype)) .collect::>>()?; Ok(VortexDataSource { diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 9acd763bf30..0f1dd90deb8 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::stat_sum::StatSum; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -483,7 +483,7 @@ pub fn try_from_projection_expression( pub enum PushedAggregate { Min, Max, - Sum, + StatSum, Mean, // Also used for ANY_VALUE() which is allowed by definition First, @@ -496,7 +496,7 @@ impl Display for PushedAggregate { match self { PushedAggregate::Min => f.write_str("min"), PushedAggregate::Max => f.write_str("max"), - PushedAggregate::Sum => f.write_str("sum"), + PushedAggregate::StatSum => f.write_str("sum"), PushedAggregate::Mean => f.write_str("mean"), PushedAggregate::First => f.write_str("first"), PushedAggregate::Count => f.write_str("count"), @@ -510,7 +510,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::StatSum => Box::new(Accumulator::try_new(StatSum, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), PairOptions(opts, opts), @@ -535,7 +535,7 @@ pub fn try_from_projection_aggregate( Ok(Some(match agg.aggregate_function.name() { "min" => PushedAggregate::Min, "max" => PushedAggregate::Max, - "sum" | "sum_no_overflow" => PushedAggregate::Sum, + "sum" | "sum_no_overflow" => PushedAggregate::StatSum, "avg" | "mean" => PushedAggregate::Mean, "first" | "any_value" => PushedAggregate::First, "count" => PushedAggregate::Count, diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index c4cd032c686..e2c9ff1f916 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -12,7 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::aggregate_fn::fns::stat_sum::stat_sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; @@ -168,7 +168,7 @@ impl StatsAccumulator { } } Stat::NullCount | Stat::NaNCount | Stat::UncompressedSizeInBytes => { - if let Some(sum_value) = sum(array, ctx)? + if let Some(sum_value) = stat_sum(array, ctx)? .cast(&DType::Primitive(PType::U64, Nullability::Nullable))? .into_value() { diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index d05f5672640..0c1805d8689 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -86,7 +86,7 @@ pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[St .filter_map(|stat| { stat.dtype(column_dtype) .or_else(|| { - // Backward compat: older files may have stored stats (e.g. Sum) + // Backward compat: older files may have stored stats (e.g. StatSum) // for extension types by resolving through the storage dtype. if let DType::Extension(ext) = column_dtype { stat.dtype(ext.storage_dtype()) @@ -157,7 +157,7 @@ mod tests { use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::sum::Sum; + use vortex_array::aggregate_fn::fns::stat_sum::StatSum; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -208,7 +208,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - Sum.bind(NumericalAggregateOpts::skip_nans()), + StatSum.bind(NumericalAggregateOpts::skip_nans()), ], ); @@ -217,7 +217,9 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - Sum.bind(NumericalAggregateOpts::skip_nans()).to_string(), + StatSum + .bind(NumericalAggregateOpts::skip_nans()) + .to_string(), ] ); } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index c1c12d7ddea..f124a0ac56a 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,7 +25,7 @@ use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::stat_sum::StatSum; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -212,11 +212,11 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { }; let mut aggregate_fns = vec![max, min]; - if Sum + if StatSum .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) .is_some() { - aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); + aggregate_fns.push(StatSum.bind(NumericalAggregateOpts::skip_nans())); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); @@ -230,7 +230,7 @@ mod tests { use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::sum::Sum; + use vortex_array::aggregate_fn::fns::stat_sum::StatSum; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::extension::datetime::TimeUnit; @@ -258,7 +258,7 @@ mod tests { assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); - assert!(aggregate_fns[2].is::()); + assert!(aggregate_fns[2].is::()); } #[test] @@ -271,7 +271,7 @@ mod tests { assert!( aggregate_fns .iter() - .all(|aggregate_fn| !aggregate_fn.is::()) + .all(|aggregate_fn| !aggregate_fn.is::()) ); } } From 73819a1e4339ada8beacd63d6d670b1fe592e876 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 15:31:17 -0700 Subject: [PATCH 05/12] Rename monoid sum aggregate to Total (SQLite precedent); wire id stays vortex.sum Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 10 +-- vortex-array/src/aggregate_fn/accumulator.rs | 10 +-- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 18 ++--- vortex-array/src/aggregate_fn/fns/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 20 +++--- .../fns/{stat_sum => total}/bool.rs | 22 +++--- .../fns/{stat_sum => total}/constant.rs | 20 +++--- .../fns/{stat_sum => total}/decimal.rs | 56 +++++++-------- .../fns/{stat_sum => total}/grouped.rs | 33 ++++----- .../fns/{stat_sum => total}/mod.rs | 72 +++++++++---------- .../fns/{stat_sum => total}/primitive.rs | 54 +++++++------- vortex-array/src/aggregate_fn/session.rs | 10 +-- vortex-array/src/array/erased.rs | 4 +- .../src/arrays/chunked/compute/aggregate.rs | 4 +- .../src/compute/conformance/consistency.rs | 5 +- vortex-array/src/expr/stats/mod.rs | 8 +-- vortex-array/src/stats/array.rs | 4 +- vortex-array/src/stats/expr.rs | 4 +- vortex-duckdb/src/convert/expr.rs | 10 +-- vortex-layout/src/layouts/file_stats.rs | 4 +- vortex-layout/src/layouts/zoned/schema.rs | 10 ++- vortex-layout/src/layouts/zoned/writer.rs | 12 ++-- 22 files changed, 191 insertions(+), 201 deletions(-) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/bool.rs (89%) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/constant.rs (89%) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/decimal.rs (90%) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/grouped.rs (93%) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/mod.rs (94%) rename vortex-array/src/aggregate_fn/fns/{stat_sum => total}/primitive.rs (87%) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 2c6cda34441..d931c8579cf 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -19,8 +19,8 @@ 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::stat_sum::StatSum; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::total::Total; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -180,7 +180,7 @@ fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, StatSum)); + .bench_refs(|input| grouped_accumulator(input, Total)); } #[divan::bench] @@ -188,7 +188,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, StatSum)); + .bench_refs(|input| grouped_accumulator(input, Total)); } #[divan::bench] @@ -196,7 +196,7 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, StatSum)); + .bench_refs(|input| grouped_accumulator(input, Total)); } #[divan::bench] @@ -204,7 +204,7 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, StatSum)); + .bench_refs(|input| grouped_accumulator(input, Total)); } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 03a761063da..a714e5ff6a1 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -279,7 +279,7 @@ mod tests { use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; - use crate::aggregate_fn::fns::stat_sum::StatSum; + use crate::aggregate_fn::fns::total::Total; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; use crate::array::VTable; @@ -320,7 +320,7 @@ mod tests { } } - /// StatSum partial sentinel `42.0` — distinguishable from the natural StatSum of + /// Total partial sentinel `42.0` — distinguishable from the natural Total of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -418,7 +418,7 @@ mod tests { Ok(()) } - /// A kernel registered for the inner `(Dict, StatSum)` child fires when accumulating a + /// A kernel registered for the inner `(Dict, Total)` child fires when accumulating a /// Dict batch through `Combined`. This is the reusable-primitive case the /// refactor enables: no `(Dict, Combined)` kernel is needed. #[test] @@ -427,7 +427,7 @@ mod tests { let session = fresh_session(); session .get::() - .register_aggregate_kernel(Dict.id(), Some(StatSum.id()), &KERNEL); + .register_aggregate_kernel(Dict.id(), Some(Total.id()), &KERNEL); let mut ctx = session.create_execution_ctx(); let mut acc = mean_f64_accumulator()?; @@ -435,7 +435,7 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - // `StatSum` child returned the sentinel 42.0 — proves the (Dict, StatSum) kernel fired + // `Total` child returned the sentinel 42.0 — proves the (Dict, Total) kernel fired // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the // batch's valid_count, so count is the real 1. assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 6873fbe0d6e..74e35cb1df0 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -17,8 +17,8 @@ use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::CombinedOptions; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::stat_sum::StatSum; -use crate::aggregate_fn::fns::stat_sum::sum_decimal_dtype; +use crate::aggregate_fn::fns::total::Total; +use crate::aggregate_fn::fns::total::sum_decimal_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -49,7 +49,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// Compute the arithmetic mean of an array. /// -/// Implemented as `StatSum / Count` via [`BinaryCombined`]. +/// Implemented as `Total / Count` via [`BinaryCombined`]. /// /// Booleans and primitive numeric types are cast to f64. Decimals stay decimals. #[derive(Clone, Debug)] @@ -62,7 +62,7 @@ impl Mean { } impl BinaryCombined for Mean { - type Left = StatSum; + type Left = Total; type Right = Count; fn id(&self) -> AggregateFnId { @@ -70,8 +70,8 @@ impl BinaryCombined for Mean { *ID } - fn left(&self) -> StatSum { - StatSum + fn left(&self) -> Total { + Total } fn right(&self) -> Count { @@ -112,7 +112,7 @@ impl BinaryCombined for Mean { let sum = sum_cast.as_primitive().typed_value::(); let count = count_cast.as_primitive().typed_value::(); let value = match (sum, count) { - (None, _) | (_, None) => return Ok(Scalar::null(target)), // StatSum overflowed + (None, _) | (_, None) => return Ok(Scalar::null(target)), // Total overflowed // A count of zero yields 0/0 = NaN, matching the array `finalize` path: nulls are // skipped during accumulation, so an all-null input is an empty mean, not null. (Some(s), Some(c)) => s / c, @@ -127,7 +127,7 @@ impl BinaryCombined for Mean { fn coerce_args( &self, _options: &PairOptions< - ::Options, + ::Options, ::Options, >, input_dtype: &DType, @@ -140,7 +140,7 @@ impl BinaryCombined for Mean { /// Hint for callers: what to cast the input to before accumulation. /// -/// - Bool stays as bool — `StatSum` has a native bool path and bool → f64 isn't +/// - Bool stays as bool — `Total` has a native bool path and bool → f64 isn't /// currently a direct cast in vortex. /// - Primitive numerics → `f64` so the sum and finalize work without overflow. /// - Decimals stay as decimals diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index a841fbeffed..2bb845c9d11 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -19,6 +19,6 @@ pub mod min; pub mod min_max; pub mod nan_count; pub mod null_count; -pub mod stat_sum; pub mod sum; +pub mod total; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 37a719e9ed4..5539113fc77 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -29,11 +29,11 @@ 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::stat_sum::SumState; -use crate::aggregate_fn::fns::stat_sum::checked_add_i64; -use crate::aggregate_fn::fns::stat_sum::checked_add_u64; -use crate::aggregate_fn::fns::stat_sum::make_zero_state; -use crate::aggregate_fn::fns::stat_sum::sum_decimal_dtype; +use crate::aggregate_fn::fns::total::SumState; +use crate::aggregate_fn::fns::total::checked_add_i64; +use crate::aggregate_fn::fns::total::checked_add_u64; +use crate::aggregate_fn::fns::total::make_zero_state; +use crate::aggregate_fn::fns::total::sum_decimal_dtype; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::Struct; @@ -98,7 +98,7 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// /// A sum over zero valid values yields null (the SQL rule: nulls are eliminated, and the sum /// of an empty set is null), and a sum that overflows yields null. This is a distinct -/// aggregate from the monoid [`Sum`](crate::aggregate_fn::fns::sum::Sum), whose empty sum is +/// aggregate from the monoid [`Total`](crate::aggregate_fn::fns::total::Total), whose empty sum is /// zero and whose plain partial is the persisted form of sum statistics; `Sum`'s /// `{sum, seen}` partial is never written to statistics. /// @@ -497,9 +497,9 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::stat_sum; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -676,15 +676,13 @@ mod tests { } #[test] - fn stat_sum_is_monoid_while_sum_is_sql() -> VortexResult<()> { + fn total_is_monoid_while_sum_is_sql() -> VortexResult<()> { // The persisted statistic keeps the monoid semantics (zero for all-null) that zone // and chunk merging require, while `sum` applies the SQL rule. let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); assert_eq!( - stat_sum(&arr, &mut ctx)? - .as_primitive() - .typed_value::(), + total(&arr, &mut ctx)?.as_primitive().typed_value::(), Some(0) ); // The cached monoid statistic must not leak through `sum`'s cache short-circuit. diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs b/vortex-array/src/aggregate_fn/fns/total/bool.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs rename to vortex-array/src/aggregate_fn/fns/total/bool.rs index 02b9b99f017..9fc24f14b78 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/total/bool.rs @@ -41,8 +41,8 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::BoolArray; use crate::dtype::DType; @@ -53,7 +53,7 @@ mod tests { #[test] fn sum_bool_all_true() -> VortexResult<()> { let arr: BoolArray = [true, true, true].into_iter().collect(); - let result = stat_sum( + let result = total( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -64,7 +64,7 @@ mod tests { #[test] fn sum_bool_mixed() -> VortexResult<()> { let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); - let result = stat_sum( + let result = total( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -75,7 +75,7 @@ mod tests { #[test] fn sum_bool_all_false() -> VortexResult<()> { let arr: BoolArray = [false, false, false].into_iter().collect(); - let result = stat_sum( + let result = total( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -86,7 +86,7 @@ mod tests { #[test] fn sum_bool_with_nulls() -> VortexResult<()> { let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); - let result = stat_sum( + let result = total( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -97,7 +97,7 @@ mod tests { #[test] fn sum_bool_all_null() -> VortexResult<()> { let arr = BoolArray::from_iter([None::, None, None]); - let result = stat_sum( + let result = total( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -108,7 +108,7 @@ mod tests { #[test] fn sum_bool_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -118,7 +118,7 @@ mod tests { fn sum_bool_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -134,7 +134,7 @@ mod tests { #[test] fn sum_bool_return_dtype() -> VortexResult<()> { - let dtype = StatSum + let dtype = Total .return_dtype( &NumericalAggregateOpts::default(), &DType::Bool(Nullability::NonNullable), @@ -147,7 +147,7 @@ mod tests { #[test] fn sum_boolean_from_iter() -> VortexResult<()> { let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(2)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs b/vortex-array/src/aggregate_fn/fns/total/constant.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs rename to vortex-array/src/aggregate_fn/fns/total/constant.rs index a743696dedb..758e48c8521 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/total/constant.rs @@ -93,7 +93,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::ConstantArray; use crate::dtype::DType; @@ -109,7 +109,7 @@ mod tests { #[test] fn sum_constant_unsigned() -> VortexResult<()> { let array = ConstantArray::new(5u64, 10).into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 50u64.into()); Ok(()) } @@ -117,7 +117,7 @@ mod tests { #[test] fn sum_constant_signed() -> VortexResult<()> { let array = ConstantArray::new(-5i64, 10).into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, (-50i64).into()); Ok(()) } @@ -126,7 +126,7 @@ mod tests { fn sum_constant_nullable_value() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -134,7 +134,7 @@ mod tests { #[test] fn sum_constant_bool_false() -> VortexResult<()> { let array = ConstantArray::new(false, 10).into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 0u64.into()); Ok(()) } @@ -142,7 +142,7 @@ mod tests { #[test] fn sum_constant_bool_true() -> VortexResult<()> { let array = ConstantArray::new(true, 10).into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 10u64.into()); Ok(()) } @@ -150,7 +150,7 @@ mod tests { #[test] fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -168,7 +168,7 @@ mod tests { ) .into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), @@ -184,7 +184,7 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) .into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, Scalar::decimal( @@ -209,7 +209,7 @@ mod tests { ) .into_array(); - let result = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let result = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/total/decimal.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs rename to vortex-array/src/aggregate_fn/fns/total/decimal.rs index 0e07ad37c2e..df0d3d7f4de 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/total/decimal.rs @@ -115,8 +115,8 @@ mod tests { use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::DecimalArray; use crate::dtype::DType; @@ -137,7 +137,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -159,7 +159,7 @@ mod tests { Validity::from_iter([true, false, true, true]), ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -181,7 +181,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -204,7 +204,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -228,7 +228,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -251,7 +251,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -270,7 +270,7 @@ mod tests { let decimal = DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -292,7 +292,7 @@ mod tests { Validity::from_iter([false, false, true, false]), ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -315,7 +315,7 @@ mod tests { Validity::AllValid, ); - let result = stat_sum( + let result = total( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -341,7 +341,7 @@ mod tests { ); assert_eq!( - stat_sum( + total( &decimal.into_array(), &mut array_session().create_execution_ctx() ) @@ -357,20 +357,20 @@ mod tests { // 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 = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), DecimalDType::new(14, 0), Nullable, ); - StatSum.combine_partials(&mut state, near_limit)?; + Total.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); - StatSum.combine_partials(&mut state, small)?; + Total.combine_partials(&mut state, small)?; - let result = StatSum.to_scalar(&state)?; + let result = Total.to_scalar(&state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -387,21 +387,21 @@ mod tests { // 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 = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - StatSum.combine_partials(&mut state, near_limit)?; + Total.combine_partials(&mut state, near_limit)?; // Push the sum to exactly 10^14, exceeding precision 14. let one_more = Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); - StatSum.combine_partials(&mut state, one_more)?; + Total.combine_partials(&mut state, one_more)?; - let result = StatSum.to_scalar(&state)?; + let result = Total.to_scalar(&state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -414,23 +414,23 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - StatSum.combine_partials(&mut state, near_limit)?; + Total.combine_partials(&mut state, near_limit)?; let one_more = Scalar::decimal( DecimalValue::from(-1i64), DecimalDType::new(14, 0), Nullable, ); - StatSum.combine_partials(&mut state, one_more)?; + Total.combine_partials(&mut state, one_more)?; - let result = StatSum.to_scalar(&state)?; + let result = Total.to_scalar(&state)?; assert!(result.is_null()); Ok(()) } @@ -446,13 +446,13 @@ mod tests { // 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 = StatSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Total.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); - StatSum.combine_partials(&mut state, near_limit)?; + Total.combine_partials(&mut state, near_limit)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -461,9 +461,9 @@ mod tests { // Drive accumulate through the vtable directly. let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); let mut ctx = array_session().create_execution_ctx(); - StatSum.accumulate(&mut state, &columnar, &mut ctx)?; + Total.accumulate(&mut state, &columnar, &mut ctx)?; - let result = StatSum.to_scalar(&state)?; + let result = Total.to_scalar(&state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/total/grouped.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs rename to vortex-array/src/aggregate_fn/fns/total/grouped.rs index e1affb72229..04dade28147 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/total/grouped.rs @@ -5,7 +5,7 @@ use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; -use super::StatSum; +use super::Total; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; @@ -21,25 +21,25 @@ use crate::arrays::PrimitiveArray; use crate::dtype::NativePType; use crate::match_each_native_ptype; -/// Encoding-specific grouped [`StatSum`] kernel for primitive element arrays. +/// Encoding-specific grouped [`Total`] kernel for primitive element arrays. #[derive(Debug)] -pub(crate) struct PrimitiveGroupedStatSumEncodingKernel; +pub(crate) struct PrimitiveGroupedTotalEncodingKernel; -impl DynGroupedAggregateKernel for PrimitiveGroupedStatSumEncodingKernel { +impl DynGroupedAggregateKernel for PrimitiveGroupedTotalEncodingKernel { fn grouped_aggregate( &self, aggregate_fn: &AggregateFnRef, groups: &GroupedArray, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some(options) = aggregate_fn.as_opt::() else { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); }; try_grouped_sum(groups, ctx, options.skip_nans) } } -/// Grouped [`StatSum`] implementation for canonical primitive elements. +/// Grouped [`Total`] 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 @@ -66,7 +66,7 @@ pub(super) fn try_grouped_sum( )?)) } -/// StatSum each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. +/// Total each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. fn grouped_sum( elements: &PrimitiveArray, group_ranges: &GroupRanges, @@ -126,7 +126,7 @@ fn collect_sums( PrimitiveArray::from_option_iter(sums) } -/// StatSum the valid elements of a single group, using the contiguous valid runs of the element mask +/// Total 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. fn sum_masked_group( acc: &mut A, @@ -163,8 +163,8 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; @@ -180,7 +180,7 @@ mod tests { /// Run a grouped sum through the accumulator. fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = GroupedAccumulator::try_new( - StatSum, + Total, NumericalAggregateOpts::default(), elem_dtype.clone(), )?; @@ -200,14 +200,14 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; let mut ctx = array_session().create_execution_ctx(); - let sum_dtype = StatSum + let sum_dtype = Total .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) .expect("sum partial 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(&stat_sum(&slice, &mut ctx)?)?; + builder.append_scalar(&total(&slice, &mut ctx)?)?; } else { builder.append_null(); } @@ -369,11 +369,8 @@ mod tests { let elem_dtype = DType::Primitive(PType::F64, NonNullable); let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; - let mut acc = GroupedAccumulator::try_new( - StatSum, - NumericalAggregateOpts::include_nans(), - elem_dtype, - )?; + let mut acc = + GroupedAccumulator::try_new(Total, NumericalAggregateOpts::include_nans(), elem_dtype)?; let mut ctx2 = array_session().create_execution_ctx(); acc.accumulate_list(&groups, &mut ctx2)?; let actual = acc.finish(&mut ctx2)?; diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/total/mod.rs similarity index 94% rename from vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs rename to vortex-array/src/aggregate_fn/fns/total/mod.rs index 1ca99339cd9..937dfc02522 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/total/mod.rs @@ -6,7 +6,7 @@ mod constant; mod decimal; mod grouped; mod primitive; -pub(crate) use grouped::PrimitiveGroupedStatSumEncodingKernel; +pub(crate) use grouped::PrimitiveGroupedTotalEncodingKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -45,17 +45,17 @@ use crate::scalar::Scalar; /// This is the value stored as `Stat::Sum` and merged across chunks, zones, and files: the /// empty sum is zero so that partials merge as a monoid (`combine(empty, x) = x`). For the /// SQL rule (a sum over zero valid values is null) use -/// [`sum`](crate::aggregate_fn::fns::sql_sum::sum) / [`SqlSum`](crate::aggregate_fn::fns::sql_sum::SqlSum). -pub fn stat_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { +/// [`sum`](crate::aggregate_fn::fns::sum::sum) / [`Sum`](crate::aggregate_fn::fns::sum::Sum). +pub fn total(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // Short-circuit using cached array statistics. if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { return Ok(sum_scalar); } - // Compute using Accumulator. + // Compute using Accumulator. // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe. let mut acc = Accumulator::try_new( - StatSum, + Total, NumericalAggregateOpts::default(), array.dtype().clone(), )?; @@ -74,15 +74,15 @@ pub fn stat_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult DecimalDType { DecimalDType::new( @@ -91,9 +91,9 @@ pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { ) } -impl AggregateFnVTable for StatSum { +impl AggregateFnVTable for Total { type Options = NumericalAggregateOpts; - type Partial = StatSumPartial; + type Partial = TotalPartial; fn id(&self) -> AggregateFnId { static ID: CachedId = CachedId::new("vortex.sum"); @@ -153,7 +153,7 @@ impl AggregateFnVTable for StatSum { .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; let initial = make_zero_state(&return_dtype); - Ok(StatSumPartial { + Ok(TotalPartial { return_dtype, current: Some(initial), skip_nans: options.skip_nans, @@ -334,7 +334,7 @@ impl AggregateFnVTable for StatSum { /// The group state for a sum aggregate, containing the accumulated value and configuration /// needed for reset/result without external context. -pub struct StatSumPartial { +pub struct TotalPartial { return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, @@ -410,8 +410,8 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -435,18 +435,18 @@ mod tests { use crate::scalar::Scalar; use crate::validity::Validity; - /// StatSum an array with an initial value (test-only helper). + /// Total an array with an initial value (test-only helper). fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); if accumulator.is_null() { return Ok(accumulator.clone()); } if accumulator.is_zero() == Some(true) { - return stat_sum(array, &mut ctx); + return total(array, &mut ctx); } let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("StatSum not supported for dtype: {}", array.dtype()) + vortex_error::vortex_err!("Total not supported for dtype: {}", array.dtype()) })?; // For non-float types, try statistics short-circuit with accumulator. @@ -457,7 +457,7 @@ mod tests { } // Compute array sum from zero (also caches stats). - let array_sum = stat_sum(array, &mut ctx)?; + let array_sum = total(array, &mut ctx)?; // Combine with the accumulator. add_scalars(&sum_dtype, &array_sum, accumulator) @@ -485,7 +485,7 @@ mod tests { .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) .map(Scalar::from) .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - _ => unreachable!("StatSum will always be a decimal or a primitive dtype"), + _ => unreachable!("Total will always be a decimal or a primitive dtype"), }) } @@ -495,7 +495,7 @@ mod tests { fn sum_multi_batch() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -512,7 +512,7 @@ mod tests { fn sum_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -531,16 +531,16 @@ mod tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = StatSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); - StatSum.combine_partials(&mut state, scalar1)?; + Total.combine_partials(&mut state, scalar1)?; let scalar2 = Scalar::primitive(50i64, Nullable); - StatSum.combine_partials(&mut state, scalar2)?; + Total.combine_partials(&mut state, scalar2)?; - let result = StatSum.to_scalar(&state)?; - StatSum.reset(&mut state); + let result = Total.to_scalar(&state)?; + Total.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) } @@ -561,7 +561,7 @@ mod tests { // compute sum with accumulator to populate stats sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; - let sum_without_acc = stat_sum(&array, &mut array_session().create_execution_ctx())?; + let sum_without_acc = total(&array, &mut array_session().create_execution_ctx())?; assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); Ok(()) } @@ -585,7 +585,7 @@ mod tests { fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = GroupedAccumulator::try_new( - StatSum, + Total, NumericalAggregateOpts::default(), elem_dtype.clone(), )?; @@ -678,7 +678,7 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = - GroupedAccumulator::try_new(StatSum, NumericalAggregateOpts::default(), elem_dtype)?; + GroupedAccumulator::try_new(Total, NumericalAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); @@ -737,7 +737,7 @@ mod tests { dtype, )?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -751,7 +751,7 @@ mod tests { 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 = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -774,7 +774,7 @@ mod tests { dtype, )?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -789,7 +789,7 @@ mod tests { let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -821,7 +821,7 @@ mod tests { dtype, )?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -857,7 +857,7 @@ mod tests { dtype, )?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -891,7 +891,7 @@ mod tests { let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let result = stat_sum( + let result = total( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; diff --git a/vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/total/primitive.rs similarity index 87% rename from vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs rename to vortex-array/src/aggregate_fn/fns/total/primitive.rs index da34a97e2bd..36fe136e8c2 100644 --- a/vortex-array/src/aggregate_fn/fns/stat_sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/total/primitive.rs @@ -63,7 +63,7 @@ fn accumulate_primitive_all( } } -/// StatSum the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values +/// Total 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) { @@ -80,7 +80,7 @@ pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan } } -/// StatSum all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in +/// Total all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in /// 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. @@ -127,7 +127,7 @@ where false } -/// StatSum the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each +/// Total the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each /// run is a slice of fully-valid values, so it reuses the same vectorized reduction as the /// all-valid path instead of a per-element validity branch. fn accumulate_primitive_valid( @@ -188,8 +188,8 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -206,7 +206,7 @@ mod tests { #[test] fn sum_i32() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(10)); Ok(()) } @@ -214,7 +214,7 @@ mod tests { #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(60)); Ok(()) } @@ -223,7 +223,7 @@ mod tests { fn sum_f64() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); Ok(()) } @@ -231,7 +231,7 @@ mod tests { #[test] fn sum_with_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6)); Ok(()) } @@ -251,7 +251,7 @@ mod tests { Some(6), ]) .into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(21)); Ok(()) } @@ -259,7 +259,7 @@ mod tests { #[test] fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -267,7 +267,7 @@ mod tests { #[test] fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } @@ -275,7 +275,7 @@ mod tests { #[test] fn sum_buffer_i32() -> VortexResult<()> { let arr = buffer![1, 1, 1, 1].into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4)); Ok(()) } @@ -283,7 +283,7 @@ mod tests { #[test] fn sum_buffer_f64() -> VortexResult<()> { let arr = buffer![1., 1., 1., 1.].into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4.)); Ok(()) } @@ -291,7 +291,7 @@ mod tests { #[test] fn sum_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -300,7 +300,7 @@ mod tests { #[test] fn sum_empty_f64_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) @@ -313,7 +313,7 @@ mod tests { Validity::NonNullable, ) .into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -322,7 +322,7 @@ mod tests { fn sum_f32_with_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); Ok(()) } @@ -331,7 +331,7 @@ mod tests { 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 = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); Ok(()) } @@ -340,17 +340,17 @@ mod tests { fn sum_all_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) } - /// StatSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + /// Total an array with explicit [`NumericalAggregateOpts`] (test-only helper). fn sum_with_options( arr: &crate::ArrayRef, options: NumericalAggregateOpts, ) -> VortexResult { - let mut acc = Accumulator::try_new(StatSum, options, arr.dtype().clone())?; + let mut acc = Accumulator::try_new(Total, options, arr.dtype().clone())?; acc.accumulate(arr, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -388,7 +388,7 @@ mod tests { #[test] fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { - // With an exact NaNCount of zero, the planted exact StatSum stat is usable as-is. + // With an exact NaNCount of zero, the planted exact Total stat is usable as-is. let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); arr.statistics() @@ -419,12 +419,12 @@ mod tests { Validity::NonNullable, ) .into_array(); - let acc = stat_sum(&batch, &mut array_session().create_execution_ctx())?; + let acc = total(&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( - StatSum, + Total, NumericalAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; @@ -436,7 +436,7 @@ mod tests { #[test] fn sum_checked_overflow() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = stat_sum(&arr, &mut array_session().create_execution_ctx())?; + let result = total(&arr, &mut array_session().create_execution_ctx())?; assert!(result.is_null()); Ok(()) } @@ -444,7 +444,7 @@ mod tests { #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StatSum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; assert!(!acc.is_saturated()); let batch = diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 693e444b21a..0e1aa992b02 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -29,10 +29,10 @@ 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::stat_sum::PrimitiveGroupedStatSumEncodingKernel; -use crate::aggregate_fn::fns::stat_sum::StatSum; use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::total::PrimitiveGroupedTotalEncodingKernel; +use crate::aggregate_fn::fns::total::Total; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; @@ -101,7 +101,7 @@ impl Default for AggregateFnSession { this.register(MinMax); this.register(NanCount); this.register(NullCount); - this.register(StatSum); + this.register(Total); this.register(Sum); this.register(UncompressedSizeInBytes); @@ -120,8 +120,8 @@ impl Default for AggregateFnSession { ); this.register_grouped_encoding_kernel( Primitive.id(), - StatSum.id(), - &PrimitiveGroupedStatSumEncodingKernel, + Total.id(), + &PrimitiveGroupedTotalEncodingKernel, ); this diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 52e9e5224e8..c9028ade519 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -28,7 +28,7 @@ use crate::ExecutionResult; use crate::IntoArray; use crate::VTable; use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::stat_sum::stat_sum; +use crate::aggregate_fn::fns::total::total; use crate::array::ArrayData; use crate::array::ArrayId; use crate::array::ArrayInner; @@ -334,7 +334,7 @@ impl ArrayRef { Validity::NonNullable | Validity::AllValid => len, Validity::AllInvalid => 0, Validity::Array(a) => { - let array_sum = stat_sum(&a, ctx)?; + let array_sum = total(&a, ctx)?; array_sum .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 9eb24019e69..664176cb680 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -47,7 +47,7 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::stat_sum::StatSum; + use crate::aggregate_fn::fns::total::Total; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -60,7 +60,7 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); let mut acc = Accumulator::try_new( - StatSum, + Total, NumericalAggregateOpts::default(), batch.dtype().clone(), )?; diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index f143e637888..d2e5535fbc6 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1002,7 +1002,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::aggregate_fn::fns::stat_sum::stat_sum; + use crate::aggregate_fn::fns::total::total; use crate::dtype::DType; let len = array.len(); @@ -1045,8 +1045,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { return; } - if let (Ok(slice_sum), Ok(canonical_sum)) = - (stat_sum(&sliced, ctx), stat_sum(&canonical_sliced, ctx)) + if let (Ok(slice_sum), Ok(canonical_sum)) = (total(&sliced, ctx), total(&canonical_sliced, ctx)) { // Compare sum scalars assert_eq!( diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 68f24284d29..20054213813 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -188,7 +188,7 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - return aggregate_fn::fns::stat_sum::StatSum + return aggregate_fn::fns::total::Total .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); } }) @@ -200,9 +200,7 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => { - aggregate_fn::fns::stat_sum::StatSum.bind(NumericalAggregateOpts::skip_nans()) - } + Self::Sum => aggregate_fn::fns::total::Total.bind(NumericalAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { @@ -218,7 +216,7 @@ impl Stat { /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those /// aggregates have no stat slot. pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option { - if let Some(options) = aggregate_fn.as_opt::() { + if let Some(options) = aggregate_fn.as_opt::() { return options.skip_nans.then_some(Self::Sum); } if aggregate_fn.is::() { diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index 8c754ecedc3..ad4776745a6 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -23,7 +23,7 @@ use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::stat_sum::stat_sum; +use crate::aggregate_fn::fns::total::total; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -174,7 +174,7 @@ impl StatsSetRef<'_> { .then(|| { // Sum is supported for this dtype. The statistic is the monoid sum // (zero when no valid values), not the SQL sum. - stat_sum(self.dyn_array_ref, ctx) + total(self.dyn_array_ref, ctx) }) .transpose()? } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 81dafb20c0e..ed24e54c06b 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -14,7 +14,7 @@ use crate::aggregate_fn::fns::all_null::AllNull; 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::stat_sum::StatSum; +use crate::aggregate_fn::fns::total::Total; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; pub use crate::scalar_fn::fns::stat::StatFn; @@ -37,7 +37,7 @@ pub fn min_max(expr: Expression) -> Expression { /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. - stat(expr, StatSum.bind(NumericalAggregateOpts::skip_nans())) + stat(expr, Total.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 0f1dd90deb8..d2620f41d9d 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::stat_sum::StatSum; +use vortex::aggregate_fn::fns::total::Total; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -483,7 +483,7 @@ pub fn try_from_projection_expression( pub enum PushedAggregate { Min, Max, - StatSum, + Total, Mean, // Also used for ANY_VALUE() which is allowed by definition First, @@ -496,7 +496,7 @@ impl Display for PushedAggregate { match self { PushedAggregate::Min => f.write_str("min"), PushedAggregate::Max => f.write_str("max"), - PushedAggregate::StatSum => f.write_str("sum"), + PushedAggregate::Total => f.write_str("sum"), PushedAggregate::Mean => f.write_str("mean"), PushedAggregate::First => f.write_str("first"), PushedAggregate::Count => f.write_str("count"), @@ -510,7 +510,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::StatSum => Box::new(Accumulator::try_new(StatSum, opts, dtype)?), + Self::Total => Box::new(Accumulator::try_new(Total, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), PairOptions(opts, opts), @@ -535,7 +535,7 @@ pub fn try_from_projection_aggregate( Ok(Some(match agg.aggregate_function.name() { "min" => PushedAggregate::Min, "max" => PushedAggregate::Max, - "sum" | "sum_no_overflow" => PushedAggregate::StatSum, + "sum" | "sum_no_overflow" => PushedAggregate::Total, "avg" | "mean" => PushedAggregate::Mean, "first" | "any_value" => PushedAggregate::First, "count" => PushedAggregate::Count, diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index e2c9ff1f916..4702008f01a 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -12,7 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::stat_sum::stat_sum; +use vortex_array::aggregate_fn::fns::total::total; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; @@ -168,7 +168,7 @@ impl StatsAccumulator { } } Stat::NullCount | Stat::NaNCount | Stat::UncompressedSizeInBytes => { - if let Some(sum_value) = stat_sum(array, ctx)? + if let Some(sum_value) = total(array, ctx)? .cast(&DType::Primitive(PType::U64, Nullability::Nullable))? .into_value() { diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 0c1805d8689..f02a8edc86d 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -86,7 +86,7 @@ pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[St .filter_map(|stat| { stat.dtype(column_dtype) .or_else(|| { - // Backward compat: older files may have stored stats (e.g. StatSum) + // Backward compat: older files may have stored stats (e.g. Total) // for extension types by resolving through the storage dtype. if let DType::Extension(ext) = column_dtype { stat.dtype(ext.storage_dtype()) @@ -157,7 +157,7 @@ mod tests { use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::stat_sum::StatSum; + use vortex_array::aggregate_fn::fns::total::Total; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -208,7 +208,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - StatSum.bind(NumericalAggregateOpts::skip_nans()), + Total.bind(NumericalAggregateOpts::skip_nans()), ], ); @@ -217,9 +217,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - StatSum - .bind(NumericalAggregateOpts::skip_nans()) - .to_string(), + Total.bind(NumericalAggregateOpts::skip_nans()).to_string(), ] ); } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index f124a0ac56a..5a20fd9b486 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,7 +25,7 @@ use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::stat_sum::StatSum; +use vortex_array::aggregate_fn::fns::total::Total; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -212,11 +212,11 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { }; let mut aggregate_fns = vec![max, min]; - if StatSum + if Total .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) .is_some() { - aggregate_fns.push(StatSum.bind(NumericalAggregateOpts::skip_nans())); + aggregate_fns.push(Total.bind(NumericalAggregateOpts::skip_nans())); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); @@ -230,7 +230,7 @@ mod tests { use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::stat_sum::StatSum; + use vortex_array::aggregate_fn::fns::total::Total; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::extension::datetime::TimeUnit; @@ -258,7 +258,7 @@ mod tests { assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); - assert!(aggregate_fns[2].is::()); + assert!(aggregate_fns[2].is::()); } #[test] @@ -271,7 +271,7 @@ mod tests { assert!( aggregate_fns .iter() - .all(|aggregate_fn| !aggregate_fn.is::()) + .all(|aggregate_fn| !aggregate_fn.is::()) ); } } From 11b5d68e88b95c62fa39a14fb8919a0e705fded6 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 17:44:54 -0700 Subject: [PATCH 06/12] Monoid Sum keeps its name and wire id; SQL-semantics aggregate is StandardSum, used only by list_sum Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 26 +- vortex-array/src/aggregate_fn/accumulator.rs | 10 +- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 18 +- vortex-array/src/aggregate_fn/fns/mod.rs | 2 +- .../fns/{total => standard_sum}/bool.rs | 31 +- .../fns/{total => standard_sum}/constant.rs | 31 +- .../fns/{total => standard_sum}/decimal.rs | 62 ++- .../fns/{total => standard_sum}/grouped.rs | 102 ++-- .../fns/{total => standard_sum}/mod.rs | 482 ++++++++++++------ .../fns/{total => standard_sum}/primitive.rs | 81 +-- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 9 +- .../src/aggregate_fn/fns/sum/constant.rs | 11 +- .../src/aggregate_fn/fns/sum/decimal.rs | 10 +- .../src/aggregate_fn/fns/sum/grouped.rs | 69 +-- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 412 +++++---------- .../src/aggregate_fn/fns/sum/primitive.rs | 27 +- vortex-array/src/aggregate_fn/session.rs | 10 +- vortex-array/src/array/erased.rs | 4 +- .../src/arrays/chunked/compute/aggregate.rs | 4 +- .../src/compute/conformance/consistency.rs | 5 +- vortex-array/src/expr/stats/mod.rs | 6 +- vortex-array/src/scalar_fn/fns/list_sum.rs | 13 +- vortex-array/src/stats/array.rs | 7 +- vortex-array/src/stats/expr.rs | 4 +- vortex-duckdb/src/convert/expr.rs | 10 +- vortex-layout/src/layouts/file_stats.rs | 4 +- vortex-layout/src/layouts/zoned/schema.rs | 8 +- vortex-layout/src/layouts/zoned/writer.rs | 12 +- 28 files changed, 742 insertions(+), 728 deletions(-) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/bool.rs (83%) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/constant.rs (84%) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/decimal.rs (89%) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/grouped.rs (80%) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/mod.rs (62%) rename vortex-array/src/aggregate_fn/fns/{total => standard_sum}/primitive.rs (82%) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index d931c8579cf..1a037c49978 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -19,8 +19,8 @@ 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::aggregate_fn::fns::total::Total; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -180,7 +180,7 @@ fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Total)); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -188,7 +188,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Total)); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -196,7 +196,7 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Total)); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -204,39 +204,39 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Total)); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] -fn sql_sum_i32_nullable_all_valid(bencher: Bencher) { +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, Sum)); + .bench_refs(|input| grouped_accumulator(input, StandardSum)); } #[divan::bench] -fn sql_sum_i32_clustered_nulls(bencher: Bencher) { +fn standard_sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StandardSum)); } #[divan::bench] -fn sql_sum_f64_all_valid(bencher: Bencher) { +fn standard_sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StandardSum)); } #[divan::bench] -fn sql_sum_f64_clustered_nulls(bencher: Bencher) { +fn standard_sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, StandardSum)); } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index a714e5ff6a1..69bae4e1053 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -279,7 +279,7 @@ mod tests { use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; - use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; use crate::array::VTable; @@ -320,7 +320,7 @@ mod tests { } } - /// Total partial sentinel `42.0` — distinguishable from the natural Total of + /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -418,7 +418,7 @@ mod tests { Ok(()) } - /// A kernel registered for the inner `(Dict, Total)` child fires when accumulating a + /// A kernel registered for the inner `(Dict, Sum)` child fires when accumulating a /// Dict batch through `Combined`. This is the reusable-primitive case the /// refactor enables: no `(Dict, Combined)` kernel is needed. #[test] @@ -427,7 +427,7 @@ mod tests { let session = fresh_session(); session .get::() - .register_aggregate_kernel(Dict.id(), Some(Total.id()), &KERNEL); + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); let mut ctx = session.create_execution_ctx(); let mut acc = mean_f64_accumulator()?; @@ -435,7 +435,7 @@ mod tests { let partial = acc.flush()?; let s = partial.as_struct(); - // `Total` child returned the sentinel 42.0 — proves the (Dict, Total) kernel fired + // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the // batch's valid_count, so count is the real 1. assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 74e35cb1df0..20b7f15834c 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -17,8 +17,8 @@ use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::CombinedOptions; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::total::Total; -use crate::aggregate_fn::fns::total::sum_decimal_dtype; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::sum_decimal_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -49,7 +49,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// Compute the arithmetic mean of an array. /// -/// Implemented as `Total / Count` via [`BinaryCombined`]. +/// Implemented as `Sum / Count` via [`BinaryCombined`]. /// /// Booleans and primitive numeric types are cast to f64. Decimals stay decimals. #[derive(Clone, Debug)] @@ -62,7 +62,7 @@ impl Mean { } impl BinaryCombined for Mean { - type Left = Total; + type Left = Sum; type Right = Count; fn id(&self) -> AggregateFnId { @@ -70,8 +70,8 @@ impl BinaryCombined for Mean { *ID } - fn left(&self) -> Total { - Total + fn left(&self) -> Sum { + Sum } fn right(&self) -> Count { @@ -112,7 +112,7 @@ impl BinaryCombined for Mean { let sum = sum_cast.as_primitive().typed_value::(); let count = count_cast.as_primitive().typed_value::(); let value = match (sum, count) { - (None, _) | (_, None) => return Ok(Scalar::null(target)), // Total overflowed + (None, _) | (_, None) => return Ok(Scalar::null(target)), // Sum overflowed // A count of zero yields 0/0 = NaN, matching the array `finalize` path: nulls are // skipped during accumulation, so an all-null input is an empty mean, not null. (Some(s), Some(c)) => s / c, @@ -127,7 +127,7 @@ impl BinaryCombined for Mean { fn coerce_args( &self, _options: &PairOptions< - ::Options, + ::Options, ::Options, >, input_dtype: &DType, @@ -140,7 +140,7 @@ impl BinaryCombined for Mean { /// Hint for callers: what to cast the input to before accumulation. /// -/// - Bool stays as bool — `Total` has a native bool path and bool → f64 isn't +/// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't /// currently a direct cast in vortex. /// - Primitive numerics → `f64` so the sum and finalize work without overflow. /// - Decimals stay as decimals diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index 2bb845c9d11..f7a6b29953a 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -19,6 +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 total; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/total/bool.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs similarity index 83% rename from vortex-array/src/aggregate_fn/fns/total/bool.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs index 9fc24f14b78..4bafb6f31fb 100644 --- a/vortex-array/src/aggregate_fn/fns/total/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs @@ -17,12 +17,14 @@ pub(super) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, + seen: &mut bool, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, @@ -41,8 +43,8 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::standard_sum::StandardSum; + use crate::aggregate_fn::fns::standard_sum::standard_sum; use crate::array_session; use crate::arrays::BoolArray; use crate::dtype::DType; @@ -53,7 +55,7 @@ mod tests { #[test] fn sum_bool_all_true() -> VortexResult<()> { let arr: BoolArray = [true, true, true].into_iter().collect(); - let result = total( + let result = standard_sum( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -64,7 +66,7 @@ mod tests { #[test] fn sum_bool_mixed() -> VortexResult<()> { let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); - let result = total( + let result = standard_sum( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -75,7 +77,7 @@ mod tests { #[test] fn sum_bool_all_false() -> VortexResult<()> { let arr: BoolArray = [false, false, false].into_iter().collect(); - let result = total( + let result = standard_sum( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -86,7 +88,7 @@ mod tests { #[test] fn sum_bool_with_nulls() -> VortexResult<()> { let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); - let result = total( + let result = standard_sum( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -97,20 +99,21 @@ mod tests { #[test] fn sum_bool_all_null() -> VortexResult<()> { let arr = BoolArray::from_iter([None::, None, None]); - let result = total( + let result = standard_sum( &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); Ok(()) } #[test] - fn sum_bool_empty_produces_zero() -> VortexResult<()> { + fn sum_bool_empty_is_null() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -118,7 +121,7 @@ mod tests { fn sum_bool_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -134,7 +137,7 @@ mod tests { #[test] fn sum_bool_return_dtype() -> VortexResult<()> { - let dtype = Total + let dtype = StandardSum .return_dtype( &NumericalAggregateOpts::default(), &DType::Bool(Nullability::NonNullable), @@ -147,7 +150,7 @@ mod tests { #[test] fn sum_boolean_from_iter() -> VortexResult<()> { let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(2)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/total/constant.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs similarity index 84% rename from vortex-array/src/aggregate_fn/fns/total/constant.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs index 758e48c8521..9155b6a5c4c 100644 --- a/vortex-array/src/aggregate_fn/fns/total/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs @@ -93,7 +93,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::standard_sum::standard_sum; use crate::array_session; use crate::arrays::ConstantArray; use crate::dtype::DType; @@ -109,7 +109,7 @@ mod tests { #[test] fn sum_constant_unsigned() -> VortexResult<()> { let array = ConstantArray::new(5u64, 10).into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 50u64.into()); Ok(()) } @@ -117,7 +117,7 @@ mod tests { #[test] fn sum_constant_signed() -> VortexResult<()> { let array = ConstantArray::new(-5i64, 10).into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, (-50i64).into()); Ok(()) } @@ -126,15 +126,16 @@ mod tests { fn sum_constant_nullable_value() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + // SQL `SUM`: an all-null constant has no valid values. + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } #[test] fn sum_constant_bool_false() -> VortexResult<()> { let array = ConstantArray::new(false, 10).into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 0u64.into()); Ok(()) } @@ -142,7 +143,7 @@ mod tests { #[test] fn sum_constant_bool_true() -> VortexResult<()> { let array = ConstantArray::new(true, 10).into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 10u64.into()); Ok(()) } @@ -150,8 +151,8 @@ mod tests { #[test] fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -168,7 +169,7 @@ mod tests { ) .into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), @@ -184,14 +185,10 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) .into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, - Scalar::decimal( - DecimalValue::I256(i256::ZERO), - DecimalDType::new(20, 2), - Nullable - ) + Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) ); Ok(()) } @@ -209,7 +206,7 @@ mod tests { ) .into_array(); - let result = total(&array, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) diff --git a/vortex-array/src/aggregate_fn/fns/total/decimal.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/total/decimal.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs index df0d3d7f4de..40b8f38236b 100644 --- a/vortex-array/src/aggregate_fn/fns/total/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs @@ -27,8 +27,10 @@ pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, + seen: &mut bool, ) -> VortexResult { let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), @@ -115,8 +117,8 @@ mod tests { use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::standard_sum::StandardSum; + use crate::aggregate_fn::fns::standard_sum::standard_sum; use crate::array_session; use crate::arrays::DecimalArray; use crate::dtype::DType; @@ -137,7 +139,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -159,7 +161,7 @@ mod tests { Validity::from_iter([true, false, true, true]), ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -181,7 +183,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -204,7 +206,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -228,7 +230,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -251,7 +253,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -270,7 +272,7 @@ mod tests { let decimal = DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -292,7 +294,7 @@ mod tests { Validity::from_iter([false, false, true, false]), ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -315,7 +317,7 @@ mod tests { Validity::AllValid, ); - let result = total( + let result = standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -341,7 +343,7 @@ mod tests { ); assert_eq!( - total( + standard_sum( &decimal.into_array(), &mut array_session().create_execution_ctx() ) @@ -357,20 +359,21 @@ mod tests { // 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 = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + 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, ); - Total.combine_partials(&mut state, near_limit)?; + 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); - Total.combine_partials(&mut state, small)?; + StandardSum.combine_partials(&mut state, small)?; - let result = Total.to_scalar(&state)?; + let result = StandardSum.finalize_scalar(&state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -387,21 +390,22 @@ mod tests { // 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 = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = + StandardSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - Total.combine_partials(&mut state, near_limit)?; + StandardSum.combine_partials(&mut state, near_limit)?; // Push the sum to exactly 10^14, exceeding precision 14. let one_more = Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); - Total.combine_partials(&mut state, one_more)?; + StandardSum.combine_partials(&mut state, one_more)?; - let result = Total.to_scalar(&state)?; + let result = StandardSum.finalize_scalar(&state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -414,23 +418,24 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = + StandardSum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - Total.combine_partials(&mut state, near_limit)?; + StandardSum.combine_partials(&mut state, near_limit)?; let one_more = Scalar::decimal( DecimalValue::from(-1i64), DecimalDType::new(14, 0), Nullable, ); - Total.combine_partials(&mut state, one_more)?; + StandardSum.combine_partials(&mut state, one_more)?; - let result = Total.to_scalar(&state)?; + let result = StandardSum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } @@ -446,13 +451,14 @@ mod tests { // 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 = Total.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + 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); - Total.combine_partials(&mut state, near_limit)?; + StandardSum.combine_partials(&mut state, near_limit)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -461,9 +467,9 @@ mod tests { // Drive accumulate through the vtable directly. let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); let mut ctx = array_session().create_execution_ctx(); - Total.accumulate(&mut state, &columnar, &mut ctx)?; + StandardSum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Total.to_scalar(&state)?; + let result = StandardSum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/total/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs similarity index 80% rename from vortex-array/src/aggregate_fn/fns/total/grouped.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs index 04dade28147..04f5f857588 100644 --- a/vortex-array/src/aggregate_fn/fns/total/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -1,11 +1,13 @@ // 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::Total; +use super::StandardSum; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; @@ -16,35 +18,45 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedArray; 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 [`Total`] kernel for primitive element arrays. +/// Encoding-specific grouped [`StandardSum`] kernel for primitive element arrays. #[derive(Debug)] -pub(crate) struct PrimitiveGroupedTotalEncodingKernel; +pub(crate) struct PrimitiveGroupedStandardSumEncodingKernel; -impl DynGroupedAggregateKernel for PrimitiveGroupedTotalEncodingKernel { +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 { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); }; try_grouped_sum(groups, ctx, options.skip_nans) } } -/// Grouped [`Total`] implementation for canonical primitive elements. +/// 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, seen}` partial rows (see `StandardSum`'s partial dtype): `seen` records whether the +/// group contained at least one valid element, so that `finalize` yields null for empty and +/// all-null groups (SQL `SUM`), and null groups are null struct rows. pub(super) fn try_grouped_sum( groups: &GroupedArray, ctx: &mut ExecutionCtx, @@ -66,7 +78,7 @@ pub(super) fn try_grouped_sum( )?)) } -/// Total each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. +/// StandardSum each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. fn grouped_sum( elements: &PrimitiveArray, group_ranges: &GroupRanges, @@ -80,7 +92,7 @@ fn grouped_sum( .execute_mask(elements.as_ref().len(), ctx)?; let all_valid = matches!(elem_mask.slices(), AllOr::All); - let result = match_each_native_ptype!(elements.ptype(), + let (sums, seen) = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, @@ -98,11 +110,22 @@ fn grouped_sum( } ); - Ok(result.into_array()) + Ok(StructArray::try_new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), + vec![ + sums.into_array(), + BoolArray::new(seen, 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. A group is null when the group -/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). +/// Reduce each group's element slice into a nullable sum plus a per-group `seen` bit. A sum is +/// null when the group itself is invalid or when summing it overflows (`sum_run` returns +/// `true`); `seen` records whether the group contained at least one valid element, decided by +/// validity alone (with `skip_nans`, a valid all-NaN group is still seen and sums to zero). fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -110,24 +133,31 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> PrimitiveArray { +) -> (PrimitiveArray, BitBuffer) { + let mut seen = BitBufferMut::with_capacity(group_ranges.len()); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { + seen.append(false); return None; } let mut acc = A::default(); - let overflow = if all_valid { - sum_run(&mut acc, &values[offset..offset + size]) + 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) }; + seen.append(any_valid); (!overflow).then_some(acc) }); - PrimitiveArray::from_option_iter(sums) + let sums = PrimitiveArray::from_option_iter(sums); + (sums, seen.freeze()) } -/// Total the valid elements of a single group, using the contiguous valid runs of the element mask +/// 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], @@ -135,17 +165,17 @@ fn sum_masked_group( size: usize, elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> bool { +) -> (bool, bool) { match elem_mask.slice(offset..offset + size).slices() { - AllOr::All => sum_run(acc, &values[offset..offset + size]), - AllOr::None => false, + 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; + return (true, true); } } - false + (false, !runs.is_empty()) } } } @@ -163,8 +193,8 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; - use crate::aggregate_fn::fns::total::total; + 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; @@ -180,7 +210,7 @@ mod tests { /// Run a grouped sum through the accumulator. fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = GroupedAccumulator::try_new( - Total, + StandardSum, NumericalAggregateOpts::default(), elem_dtype.clone(), )?; @@ -189,8 +219,9 @@ mod tests { acc.finish(&mut ctx) } - /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] for - /// valid groups, a null sum for invalid groups. + /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] + /// (SQL semantics: null for zero valid elements) for valid groups, a null sum for invalid + /// groups. fn grouped_sum_reference( elements: &ArrayRef, ranges: &[(usize, usize)], @@ -200,14 +231,14 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; let mut ctx = array_session().create_execution_ctx(); - let sum_dtype = Total - .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) - .expect("sum partial dtype"); + 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(&total(&slice, &mut ctx)?)?; + builder.append_scalar(&standard_sum(&slice, &mut ctx)?)?; } else { builder.append_null(); } @@ -292,8 +323,8 @@ mod tests { 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), Some(0i64), Some(0i64), Some(9i64)]); + // The all-null and empty groups have null sums (SQL `SUM` semantics). + 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(()) @@ -369,8 +400,11 @@ mod tests { let elem_dtype = DType::Primitive(PType::F64, NonNullable); let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; - let mut acc = - GroupedAccumulator::try_new(Total, NumericalAggregateOpts::include_nans(), elem_dtype)?; + 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(&mut ctx2)?; diff --git a/vortex-array/src/aggregate_fn/fns/total/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs similarity index 62% rename from vortex-array/src/aggregate_fn/fns/total/mod.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 937dfc02522..74b58fee298 100644 --- a/vortex-array/src/aggregate_fn/fns/total/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -6,12 +6,12 @@ mod constant; mod decimal; mod grouped; mod primitive; -pub(crate) use grouped::PrimitiveGroupedTotalEncodingKernel; +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_error::vortex_panic; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -23,46 +23,70 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::SumState; +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::sum_decimal_dtype; +use crate::arrays::Bool; +use crate::arrays::BoolArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::masked::mask_validity_canonical; +use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::MAX_PRECISION; +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::DecimalValue; use crate::scalar::Scalar; +use crate::validity::Validity; -/// The monoid sum of an array: zero when there are no valid values, null on overflow. +/// Return the SQL sum of an array: null when the array has no valid values or the sum +/// overflows. /// -/// This is the value stored as `Stat::Sum` and merged across chunks, zones, and files: the -/// empty sum is zero so that partials merge as a monoid (`combine(empty, x) = x`). For the -/// SQL rule (a sum over zero valid values is null) use -/// [`sum`](crate::aggregate_fn::fns::sum::sum) / [`Sum`](crate::aggregate_fn::fns::sum::Sum). -pub fn total(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // Short-circuit using cached array statistics. +/// See [`StandardSum`] for details. +pub fn standard_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Short-circuit using cached array statistics. `Stat::Sum` is the monoid sum (zero for an + // array with no valid values), so the SQL empty-sum rule needs the null count: when it is + // unknown for a nullable array, fall through and compute rather than trust the cache. if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { - return Ok(sum_scalar); + 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. + // Compute using Accumulator. // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe. let mut acc = Accumulator::try_new( - Total, + 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). + // 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)); } @@ -70,33 +94,31 @@ pub fn total(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// The monoid sum aggregate: an all-invalid or empty sum is zero, an overflowing sum is null. +/// StandardSum an array following SQL `SUM` semantics. /// -/// This aggregate's plain partial is the persisted wire form of sum statistics (zone maps, -/// file stats), so its partial dtype and semantics must stay stable. The SQL `SUM` rule is -/// provided by the separate [`Sum`](crate::aggregate_fn::fns::sum::Sum) aggregate. +/// A sum over zero valid values yields null (the SQL rule: nulls are eliminated, and the sum +/// of an empty set is null), and a sum that overflows yields null. This is a distinct +/// aggregate from the monoid [`Sum`](crate::aggregate_fn::fns::sum::Sum), whose empty sum is +/// zero and whose plain partial is the persisted form of sum statistics; `StandardSum`'s +/// `{sum, seen}` partial is never written to statistics. +/// +/// The partial state is a `{sum, seen}` struct: `sum` is the running monoid value (null once +/// saturated by overflow, which poisons merges), and `seen` records whether any valid value +/// contributed (merged with OR). `finalize` maps `seen == false` to null. `seen` is decided by +/// validity alone: NaN values are valid, so with `skip_nans` a sum over only NaNs is `0`, not +/// null. /// /// 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. #[derive(Clone, Debug)] -pub struct Total; - -// Both Spark and DataFusion use this heuristic. -// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Total.scala#L66 -// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 -pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { - DecimalDType::new( - u8::min(MAX_PRECISION, input.precision() + 10), - input.scale(), - ) -} +pub struct StandardSum; -impl AggregateFnVTable for Total { +impl AggregateFnVTable for StandardSum { type Options = NumericalAggregateOpts; - type Partial = TotalPartial; + type Partial = StandardSumPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.sum"); + static ID: CachedId = CachedId::new("vortex.standard_sum"); *ID } @@ -140,7 +162,7 @@ impl AggregateFnVTable for Total { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + Some(sum_partial_dtype(self.return_dtype(options, input_dtype)?)) } fn empty_partial( @@ -153,22 +175,48 @@ impl AggregateFnVTable for Total { .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; let initial = make_zero_state(&return_dtype); - Ok(TotalPartial { + Ok(StandardSumPartial { return_dtype, current: Some(initial), + seen: false, skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - // A null partial means the sub-accumulator saturated (overflow). + // Partials are `{sum, seen}` structs. A plain (non-struct) scalar is the legacy or + // statistic form: a monoid sum value from an older writer or a cached `Stat::Sum`. It + // cannot distinguish an empty sum from a zero sum, so it is treated as seen. + let (sum_value, other_seen) = if matches!(other.dtype(), DType::Struct(..)) { + if other.is_null() { + // A null struct partial carries no recoverable state; treat as saturated. + partial.seen = true; + partial.current = None; + 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_seen = fields + .field("seen") + .and_then(|seen| seen.as_bool().value()) + .unwrap_or(true); + (sum_value, other_seen) + } else { + (other, true) + }; + + partial.seen |= other_seen; + if sum_value.is_null() { + // A null sum value means the sub-accumulator saturated (overflow). partial.current = None; return Ok(()); } let Some(ref mut inner) = partial.current else { return Ok(()); }; + let other = sum_value; let saturated = match inner { SumState::Unsigned(acc) => { let val = other @@ -213,23 +261,18 @@ impl AggregateFnVTable for Total { } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(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) - } - }) + Ok(Scalar::struct_( + sum_partial_dtype(partial.return_dtype.as_nullable()), + vec![ + sum_value_scalar(partial), + Scalar::bool(partial.seen, Nullability::NonNullable), + ], + )) } fn reset(&self, partial: &mut Self::Partial) { partial.current = Some(make_zero_state(&partial.return_dtype)); + partial.seen = false; } #[inline] @@ -247,28 +290,26 @@ impl AggregateFnVTable for Total { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes - // the default dispatch path. - if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + // `Stat::Sum` is the monoid (NaN-skipping) sum, so the default NaN-skipping path can + // consume it directly. `StandardSum` has no Stat slot of its own — the shortcut lives here + // rather than in the accumulator's stats bridge. + 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.current, Some(SumState::Float(_))) { 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. - if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { - let sum = if sum.dtype() == &partial.return_dtype { - sum - } else { - sum.cast(&partial.return_dtype)? - }; - self.combine_partials(partial, sum)?; - return Ok(true); - } - Ok(false) + try_accumulate_cached_sum(self, partial, batch) } Precision::Exact(_) => { - // At least one NaN value: the sum is NaN without scanning the batch. + // At least one NaN value (a valid value): the sum is NaN without scanning. + partial.seen = true; if let Some(SumState::Float(acc)) = partial.current.as_mut() { *acc = f64::NAN; } @@ -286,6 +327,9 @@ impl AggregateFnVTable for Total { ) -> VortexResult<()> { // Constants compute scalar * len and combine via combine_partials. if let Columnar::Constant(c) = batch { + // Any valid value counts as seen, including NaN and `false` constants that + // contribute nothing to the running sum. + partial.seen |= !c.scalar().is_null() && !c.is_empty(); // 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(()); @@ -304,9 +348,11 @@ impl AggregateFnVTable for Total { 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), + Canonical::Primitive(p) => { + accumulate_primitive(&mut inner, p, ctx, skip_nans, &mut partial.seen) + } + Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut partial.seen), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut partial.seen), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -323,74 +369,115 @@ impl AggregateFnVTable for Total { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(partials) + fn finalize(&self, partials: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` + // field (overflow) and null partial rows (e.g. null groups) stay null via the + // validity intersection. Implemented structurally rather than as a `get_item`/`mask` + // expression: expression construction and dispatch cost multiples of the whole + // aggregation at small group counts. + let len = partials.len(); + let states = match partials.as_opt::() { + Some(states) => states.into_owned(), + None => partials.execute::(ctx)?, + }; + let struct_mask = states.as_ref().validity()?.execute_mask(len, ctx)?; + let seen = states.unmasked_field_by_name("seen")?.clone(); + let seen = match seen.as_opt::() { + Some(seen) => seen.into_owned(), + None => seen.execute::(ctx)?, + }; + let valid = &struct_mask & &Mask::from_buffer(seen.to_bit_buffer()); + + let sum = states.unmasked_field_by_name("sum")?.clone(); + if valid.all_true() { + // Every partial row is a seen, valid group: the sums are already the result. + return Ok(sum); + } + let sum = sum.execute::(ctx)?; + Ok( + mask_validity_canonical(sum, Validity::from_mask(valid, Nullability::Nullable), ctx)? + .into_array(), + ) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + if !partial.seen { + return Ok(Scalar::null(partial.return_dtype.as_nullable())); + } + Ok(sum_value_scalar(partial)) + } +} + +/// Consume a batch's cached monoid `Stat::Sum` instead of scanning it. The cached sum cannot +/// carry `seen`, so the batch's null count decides between the identity (all null) and a seen +/// contribution; 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(_) => partial.seen = true, + _ => 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) } /// The group state for a sum aggregate, containing the accumulated value and configuration /// needed for reset/result without external context. -pub struct TotalPartial { +pub struct StandardSumPartial { return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, + /// Whether at least one valid value has been accumulated. A sum over zero valid values + /// finalizes to null (SQL `SUM` semantics) rather than the monoid zero. + seen: bool, /// Whether NaN values in float inputs are skipped. skip_nans: bool, } -/// The accumulated sum value. -// TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the -// input type every time? Perhaps? -pub enum SumState { - Unsigned(u64), - Signed(i64), - Float(f64), - Decimal { - value: DecimalValue, - dtype: DecimalDType, - }, -} - -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), - PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0), - PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0), - }, - DType::Decimal(decimal, _) => SumState::Decimal { - value: DecimalValue::zero(decimal), - dtype: *decimal, - }, - _ => vortex_panic!("Unsupported sum type"), - } -} - -/// Checked add for u64, returning true if overflow occurred. -#[inline(always)] -pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { - match acc.checked_add(val) { - Some(r) => { - *acc = r; - false - } - None => true, - } +/// The partial dtype for a sum whose result is `sum_dtype`: a `{sum, seen}` struct, where +/// `sum` is the running monoid value (null once saturated by overflow, poisoning merges) and +/// `seen` records whether any valid value contributed (merged with OR). Keeping the flag +/// separate from the sum lets merges stay a monoid — the identity is `{0, false}` — while +/// `finalize` maps unseen sums to null. +fn sum_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct( + StructFields::new( + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), + vec![sum_dtype, DType::Bool(Nullability::NonNullable)], + ), + Nullability::Nullable, + ) } -/// Checked add for i64, returning true if overflow occurred. -#[inline(always)] -pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { - match acc.checked_add(val) { - Some(r) => { - *acc = r; - false +/// The running sum as a nullable scalar of the return dtype (null when saturated by overflow). +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) } - None => true, } } @@ -410,8 +497,9 @@ mod tests { use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::standard_sum::StandardSum; + use crate::aggregate_fn::fns::standard_sum::standard_sum; + use crate::aggregate_fn::fns::sum::sum as monoid_sum; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -435,18 +523,18 @@ mod tests { use crate::scalar::Scalar; use crate::validity::Validity; - /// Total an array with an initial value (test-only helper). + /// StandardSum an array with an initial value (test-only helper). fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); if accumulator.is_null() { return Ok(accumulator.clone()); } if accumulator.is_zero() == Some(true) { - return total(array, &mut ctx); + return standard_sum(array, &mut ctx); } let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("Total not supported for dtype: {}", array.dtype()) + vortex_error::vortex_err!("StandardSum not supported for dtype: {}", array.dtype()) })?; // For non-float types, try statistics short-circuit with accumulator. @@ -457,7 +545,7 @@ mod tests { } // Compute array sum from zero (also caches stats). - let array_sum = total(array, &mut ctx)?; + let array_sum = standard_sum(array, &mut ctx)?; // Combine with the accumulator. add_scalars(&sum_dtype, &array_sum, accumulator) @@ -485,7 +573,7 @@ mod tests { .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) .map(Scalar::from) .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - _ => unreachable!("Total will always be a decimal or a primitive dtype"), + _ => unreachable!("StandardSum will always be a decimal or a primitive dtype"), }) } @@ -495,7 +583,7 @@ mod tests { fn sum_multi_batch() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + 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)?; @@ -512,7 +600,7 @@ mod tests { fn sum_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + 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)?; @@ -528,19 +616,116 @@ mod tests { // State merge tests (vtable-level) + #[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)?; + 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 seen state changes nothing: `{0, false}` is the + // identity of the `{sum, seen}` 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_poisons_but_stays_seen() -> VortexResult<()> { + // Overflow (a null `sum` field) poisons the merge even when combined with later + // values: the result is null via the sum value, not via `seen`. + 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 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(()) + } + + #[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_monoid_while_standard_sum_is_sql() -> VortexResult<()> { + // The persisted statistic keeps the monoid semantics (zero for all-null) that zone + // and chunk merging require, while the SQL `sum` 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!( + monoid_sum(&arr, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(0) + ); + // The cached monoid statistic must not leak through `sum`'s cache short-circuit. + assert!(standard_sum(&arr, &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 sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Total.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); - Total.combine_partials(&mut state, scalar1)?; + StandardSum.combine_partials(&mut state, scalar1)?; let scalar2 = Scalar::primitive(50i64, Nullable); - Total.combine_partials(&mut state, scalar2)?; + StandardSum.combine_partials(&mut state, scalar2)?; - let result = Total.to_scalar(&state)?; - Total.reset(&mut state); + let result = StandardSum.finalize_scalar(&state)?; + StandardSum.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) } @@ -561,7 +746,7 @@ mod tests { // compute sum with accumulator to populate stats sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; - let sum_without_acc = total(&array, &mut array_session().create_execution_ctx())?; + let sum_without_acc = standard_sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); Ok(()) } @@ -585,7 +770,7 @@ mod tests { fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = GroupedAccumulator::try_new( - Total, + StandardSum, NumericalAggregateOpts::default(), elem_dtype.clone(), )?; @@ -653,7 +838,8 @@ mod tests { 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(0i64), Some(7i64)]).into_array(); + // The all-null group has a null sum (SQL `SUM` semantics). + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -677,8 +863,11 @@ mod tests { 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(Total, NumericalAggregateOpts::default(), elem_dtype)?; + let mut acc = GroupedAccumulator::try_new( + StandardSum, + NumericalAggregateOpts::default(), + elem_dtype, + )?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); @@ -737,7 +926,7 @@ mod tests { dtype, )?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -746,16 +935,17 @@ mod tests { } #[test] - fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> { + fn sum_chunked_floats_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 = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + // SQL `SUM`: no valid values across any chunk yields null. + assert!(result.is_null()); Ok(()) } @@ -774,7 +964,7 @@ mod tests { dtype, )?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -789,7 +979,7 @@ mod tests { let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -821,7 +1011,7 @@ mod tests { dtype, )?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -857,7 +1047,7 @@ mod tests { dtype, )?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; @@ -891,7 +1081,7 @@ mod tests { let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let result = total( + let result = standard_sum( &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; diff --git a/vortex-array/src/aggregate_fn/fns/total/primitive.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs similarity index 82% rename from vortex-array/src/aggregate_fn/fns/total/primitive.rs rename to vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs index 36fe136e8c2..6935cd10b23 100644 --- a/vortex-array/src/aggregate_fn/fns/total/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs @@ -26,12 +26,19 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, skip_nans: bool, + seen: &mut bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => accumulate_primitive_all(inner, p, skip_nans), - AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), + AllOr::All => { + *seen |= !p.as_ref().is_empty(); + accumulate_primitive_all(inner, p, skip_nans) + } + AllOr::Some(slices) => { + *seen |= !slices.is_empty(); + accumulate_primitive_valid(inner, p, slices, skip_nans) + } } } @@ -63,7 +70,7 @@ fn accumulate_primitive_all( } } -/// Total the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values +/// StandardSum 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) { @@ -80,7 +87,7 @@ pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan } } -/// Total all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in +/// StandardSum all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in /// 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. @@ -127,7 +134,7 @@ where false } -/// Total the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each +/// StandardSum the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each /// run is a slice of fully-valid values, so it reuses the same vectorized reduction as the /// all-valid path instead of a per-element validity branch. fn accumulate_primitive_valid( @@ -188,8 +195,8 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::standard_sum::StandardSum; + use crate::aggregate_fn::fns::standard_sum::standard_sum; use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -206,7 +213,7 @@ mod tests { #[test] fn sum_i32() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(10)); Ok(()) } @@ -214,7 +221,7 @@ mod tests { #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(60)); Ok(()) } @@ -223,7 +230,7 @@ mod tests { fn sum_f64() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); Ok(()) } @@ -231,7 +238,7 @@ mod tests { #[test] fn sum_with_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6)); Ok(()) } @@ -251,7 +258,7 @@ mod tests { Some(6), ]) .into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(21)); Ok(()) } @@ -259,23 +266,24 @@ mod tests { #[test] fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); Ok(()) } #[test] fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); Ok(()) } #[test] fn sum_buffer_i32() -> VortexResult<()> { let arr = buffer![1, 1, 1, 1].into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4)); Ok(()) } @@ -283,26 +291,27 @@ mod tests { #[test] fn sum_buffer_f64() -> VortexResult<()> { let arr = buffer![1., 1., 1., 1.].into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4.)); Ok(()) } #[test] - fn sum_empty_produces_zero() -> VortexResult<()> { + fn sum_empty_is_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + // SQL `SUM`: the sum over no values is null. + assert!(result.is_null()); Ok(()) } #[test] - fn sum_empty_f64_produces_zero() -> VortexResult<()> { + fn sum_empty_f64_is_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + assert!(result.is_null()); Ok(()) } @@ -313,7 +322,7 @@ mod tests { Validity::NonNullable, ) .into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -322,7 +331,7 @@ mod tests { fn sum_f32_with_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); Ok(()) } @@ -331,7 +340,7 @@ mod tests { 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 = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); Ok(()) } @@ -340,17 +349,17 @@ mod tests { fn sum_all_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) } - /// Total an array with explicit [`NumericalAggregateOpts`] (test-only helper). + /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). fn sum_with_options( arr: &crate::ArrayRef, options: NumericalAggregateOpts, ) -> VortexResult { - let mut acc = Accumulator::try_new(Total, options, arr.dtype().clone())?; + let mut acc = Accumulator::try_new(StandardSum, options, arr.dtype().clone())?; acc.accumulate(arr, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -388,13 +397,15 @@ mod tests { #[test] fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { - // With an exact NaNCount of zero, the planted exact Total stat is usable as-is. + // With an exact NaNCount of zero, the planted exact StandardSum 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(()) @@ -419,12 +430,12 @@ mod tests { Validity::NonNullable, ) .into_array(); - let acc = total(&batch, &mut array_session().create_execution_ctx())?; + 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( - Total, + StandardSum, NumericalAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; @@ -436,7 +447,7 @@ mod tests { #[test] fn sum_checked_overflow() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = total(&arr, &mut array_session().create_execution_ctx())?; + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; assert!(result.is_null()); Ok(()) } @@ -444,7 +455,7 @@ mod tests { #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Total, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; assert!(!acc.is_saturated()); let batch = diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index d1f3ce1b6d9..b3993840586 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -17,14 +17,12 @@ pub(super) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, - seen: &mut bool, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, @@ -103,17 +101,16 @@ mod tests { &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } #[test] - fn sum_bool_empty_is_null() -> VortexResult<()> { + fn sum_bool_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 07325e62015..0f366620e5c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -127,8 +127,7 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: an all-null constant has no valid values. - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -152,7 +151,7 @@ mod tests { fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -188,7 +187,11 @@ mod tests { let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, - Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) + Scalar::decimal( + DecimalValue::I256(i256::ZERO), + DecimalDType::new(20, 2), + Nullable + ) ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index cab120eba11..872e5769a01 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -27,10 +27,8 @@ pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, - seen: &mut bool, ) -> VortexResult { let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), @@ -372,7 +370,7 @@ mod tests { let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, small)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -403,7 +401,7 @@ mod tests { Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -432,7 +430,7 @@ mod tests { ); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; assert!(result.is_null()); Ok(()) } @@ -465,7 +463,7 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 4a017bd8559..fb9c566631d 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,8 +1,6 @@ // 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; @@ -18,16 +16,10 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedArray; 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 [`Sum`] kernel for primitive element arrays. #[derive(Debug)] @@ -53,10 +45,6 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { /// [`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, seen}` partial rows (see `Sum`'s partial dtype): `seen` records whether the -/// group contained at least one valid element, so that `finalize` yields null for empty and -/// all-null groups (SQL `SUM`), and null groups are null struct rows. pub(super) fn try_grouped_sum( groups: &GroupedArray, ctx: &mut ExecutionCtx, @@ -92,7 +80,7 @@ fn grouped_sum( .execute_mask(elements.as_ref().len(), ctx)?; let all_valid = matches!(elem_mask.slices(), AllOr::All); - let (sums, seen) = match_each_native_ptype!(elements.ptype(), + let result = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, @@ -110,22 +98,11 @@ fn grouped_sum( } ); - Ok(StructArray::try_new( - FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), - vec![ - sums.into_array(), - BoolArray::new(seen, Validity::NonNullable).into_array(), - ], - group_validity.len(), - Validity::from_mask(group_validity.clone(), Nullability::Nullable), - )? - .into_array()) + Ok(result.into_array()) } -/// Reduce each group's element slice into a nullable sum plus a per-group `seen` bit. A sum is -/// null when the group itself is invalid or when summing it overflows (`sum_run` returns -/// `true`); `seen` records whether the group contained at least one valid element, decided by -/// validity alone (with `skip_nans`, a valid all-NaN group is still seen and sums to zero). +/// Reduce each group's element slice into a nullable sum. A group is null when the group +/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -133,31 +110,24 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> (PrimitiveArray, BitBuffer) { - let mut seen = BitBufferMut::with_capacity(group_ranges.len()); +) -> PrimitiveArray { let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - seen.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) + let overflow = if all_valid { + sum_run(&mut acc, &values[offset..offset + size]) } else { sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) }; - seen.append(any_valid); (!overflow).then_some(acc) }); - let sums = PrimitiveArray::from_option_iter(sums); - (sums, seen.freeze()) + PrimitiveArray::from_option_iter(sums) } /// Sum 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], @@ -165,17 +135,17 @@ fn sum_masked_group( size: usize, elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> (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::All => sum_run(acc, &values[offset..offset + size]), + AllOr::None => false, AllOr::Some(runs) => { for &(start, end) in runs { if sum_run(acc, &values[offset + start..offset + end]) { - return (true, true); + return true; } } - (false, !runs.is_empty()) + false } } } @@ -219,9 +189,8 @@ mod tests { acc.finish(&mut ctx) } - /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] - /// (SQL semantics: null for zero valid elements) for valid groups, a null sum for invalid - /// groups. + /// 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)], @@ -232,8 +201,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .return_dtype(&NumericalAggregateOpts::default(), elem_dtype) - .expect("sum return dtype"); + .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .expect("sum partial dtype"); let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); for (i, &(offset, size)) in ranges.iter().enumerate() { if group_valid[i] { @@ -323,8 +292,8 @@ mod tests { let actual = grouped_sum_actual(&groups, &elem_dtype)?; let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - // The all-null and empty groups have null sums (SQL `SUM` semantics). - let direct = PrimitiveArray::from_option_iter([Some(4i64), None, None, Some(9i64)]); + let direct = + PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); 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/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 5539113fc77..54a123ef765 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -11,7 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::Mask; +use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -23,56 +23,34 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; -use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::total::SumState; -use crate::aggregate_fn::fns::total::checked_add_i64; -use crate::aggregate_fn::fns::total::checked_add_u64; -use crate::aggregate_fn::fns::total::make_zero_state; -use crate::aggregate_fn::fns::total::sum_decimal_dtype; -use crate::arrays::Bool; -use crate::arrays::BoolArray; -use crate::arrays::Struct; -use crate::arrays::StructArray; -use crate::arrays::bool::BoolArrayExt; -use crate::arrays::masked::mask_validity_canonical; -use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; +use crate::dtype::DecimalDType; +use crate::dtype::MAX_PRECISION; 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::DecimalValue; use crate::scalar::Scalar; -use crate::validity::Validity; -/// Return the SQL sum of an array: null when the array has no valid values or the sum -/// overflows. +/// The monoid sum of an array: zero when there are no valid values, null on overflow. /// -/// See [`Sum`] for details. +/// This is the value stored as `Stat::Sum` and merged across chunks, zones, and files: the +/// empty sum is zero so that partials merge as a monoid (`combine(empty, x) = x`). For the +/// SQL rule (a sum over zero valid values is null) use +/// [`standard_sum`](crate::aggregate_fn::fns::standard_sum::standard_sum) / +/// [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum). pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // Short-circuit using cached array statistics. `Stat::Sum` is the monoid sum (zero for an - // array with no valid values), so the SQL empty-sum rule needs the null count: when it is - // unknown for a nullable array, fall through and compute rather than trust the cache. + // Short-circuit using cached array statistics. 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), - _ => {} - } + return Ok(sum_scalar); } // Compute using Accumulator. @@ -85,8 +63,7 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { 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). + // Cache the computed sum as a statistic (only if non-null, i.e. no overflow). if let Some(val) = result.value().cloned() { array.statistics().set(Stat::Sum, Precision::Exact(val)); } @@ -94,31 +71,34 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// Sum an array following SQL `SUM` semantics. +/// The monoid sum aggregate: an all-invalid or empty sum is zero, an overflowing sum is null. /// -/// A sum over zero valid values yields null (the SQL rule: nulls are eliminated, and the sum -/// of an empty set is null), and a sum that overflows yields null. This is a distinct -/// aggregate from the monoid [`Total`](crate::aggregate_fn::fns::total::Total), whose empty sum is -/// zero and whose plain partial is the persisted form of sum statistics; `Sum`'s -/// `{sum, seen}` partial is never written to statistics. -/// -/// The partial state is a `{sum, seen}` struct: `sum` is the running monoid value (null once -/// saturated by overflow, which poisons merges), and `seen` records whether any valid value -/// contributed (merged with OR). `finalize` maps `seen == false` to null. `seen` is decided by -/// validity alone: NaN values are valid, so with `skip_nans` a sum over only NaNs is `0`, not -/// null. +/// This aggregate's plain partial is the persisted wire form of sum statistics (zone maps, +/// file stats), so its partial dtype and semantics must stay stable. The SQL `SUM` rule is +/// provided by the separate [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum) +/// aggregate. /// /// 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. #[derive(Clone, Debug)] pub struct Sum; +// Both Spark and DataFusion use this heuristic. +// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 +// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 +pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { + DecimalDType::new( + u8::min(MAX_PRECISION, input.precision() + 10), + input.scale(), + ) +} + impl AggregateFnVTable for Sum { type Options = NumericalAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.sql_sum"); + static ID: CachedId = CachedId::new("vortex.sum"); *ID } @@ -162,7 +142,7 @@ impl AggregateFnVTable for Sum { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - Some(sum_partial_dtype(self.return_dtype(options, input_dtype)?)) + self.return_dtype(options, input_dtype) } fn empty_partial( @@ -178,45 +158,19 @@ impl AggregateFnVTable for Sum { Ok(SumPartial { return_dtype, current: Some(initial), - seen: false, skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - // Partials are `{sum, seen}` structs. A plain (non-struct) scalar is the legacy or - // statistic form: a monoid sum value from an older writer or a cached `Stat::Sum`. It - // cannot distinguish an empty sum from a zero sum, so it is treated as seen. - let (sum_value, other_seen) = if matches!(other.dtype(), DType::Struct(..)) { - if other.is_null() { - // A null struct partial carries no recoverable state; treat as saturated. - partial.seen = true; - partial.current = None; - return Ok(()); - } - let fields = other.as_struct(); - let sum_value = fields - .field("sum") - .ok_or_else(|| vortex_err!("Sum partial is missing the `sum` field"))?; - let other_seen = fields - .field("seen") - .and_then(|seen| seen.as_bool().value()) - .unwrap_or(true); - (sum_value, other_seen) - } else { - (other, true) - }; - - partial.seen |= other_seen; - if sum_value.is_null() { - // A null sum value means the sub-accumulator saturated (overflow). + if other.is_null() { + // A null partial means the sub-accumulator saturated (overflow). partial.current = None; return Ok(()); } let Some(ref mut inner) = partial.current else { return Ok(()); }; - let other = sum_value; let saturated = match inner { SumState::Unsigned(acc) => { let val = other @@ -261,18 +215,23 @@ impl AggregateFnVTable for Sum { } 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.seen, Nullability::NonNullable), - ], - )) + Ok(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) + } + }) } fn reset(&self, partial: &mut Self::Partial) { partial.current = Some(make_zero_state(&partial.return_dtype)); - partial.seen = false; } #[inline] @@ -290,26 +249,28 @@ impl AggregateFnVTable for Sum { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - // `Stat::Sum` is the monoid (NaN-skipping) sum, so the default NaN-skipping path can - // consume it directly. `Sum` has no Stat slot of its own — the shortcut lives here - // rather than in the accumulator's stats bridge. - 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.current, Some(SumState::Float(_))) { + // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes + // the default dispatch path. + if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { 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) + if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + self.combine_partials(partial, sum)?; + return Ok(true); + } + Ok(false) } Precision::Exact(_) => { - // At least one NaN value (a valid value): the sum is NaN without scanning. - partial.seen = true; + // At least one NaN value: the sum is NaN without scanning the batch. if let Some(SumState::Float(acc)) = partial.current.as_mut() { *acc = f64::NAN; } @@ -327,9 +288,6 @@ impl AggregateFnVTable for Sum { ) -> VortexResult<()> { // Constants compute scalar * len and combine via combine_partials. if let Columnar::Constant(c) = batch { - // Any valid value counts as seen, including NaN and `false` constants that - // contribute nothing to the running sum. - partial.seen |= !c.scalar().is_null() && !c.is_empty(); // 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(()); @@ -348,11 +306,9 @@ impl AggregateFnVTable for Sum { let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => { - accumulate_primitive(&mut inner, p, ctx, skip_nans, &mut partial.seen) - } - Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut partial.seen), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut partial.seen), + 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!(), @@ -369,71 +325,13 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn finalize(&self, partials: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` - // field (overflow) and null partial rows (e.g. null groups) stay null via the - // validity intersection. Implemented structurally rather than as a `get_item`/`mask` - // expression: expression construction and dispatch cost multiples of the whole - // aggregation at small group counts. - let len = partials.len(); - let states = match partials.as_opt::() { - Some(states) => states.into_owned(), - None => partials.execute::(ctx)?, - }; - let struct_mask = states.as_ref().validity()?.execute_mask(len, ctx)?; - let seen = states.unmasked_field_by_name("seen")?.clone(); - let seen = match seen.as_opt::() { - Some(seen) => seen.into_owned(), - None => seen.execute::(ctx)?, - }; - let valid = &struct_mask & &Mask::from_buffer(seen.to_bit_buffer()); - - let sum = states.unmasked_field_by_name("sum")?.clone(); - if valid.all_true() { - // Every partial row is a seen, valid group: the sums are already the result. - return Ok(sum); - } - let sum = sum.execute::(ctx)?; - Ok( - mask_validity_canonical(sum, Validity::from_mask(valid, Nullability::Nullable), ctx)? - .into_array(), - ) + fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(partials) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - if !partial.seen { - return Ok(Scalar::null(partial.return_dtype.as_nullable())); - } - Ok(sum_value_scalar(partial)) - } -} - -/// Consume a batch's cached monoid `Stat::Sum` instead of scanning it. The cached sum cannot -/// carry `seen`, so the batch's null count decides between the identity (all null) and a seen -/// contribution; when either statistic is missing the caller falls through to a real scan. -fn try_accumulate_cached_sum( - vtable: &Sum, - partial: &mut SumPartial, - 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(_) => partial.seen = true, - _ => return Ok(false), + self.to_scalar(partial) } - let sum = if sum.dtype() == &partial.return_dtype { - sum - } else { - sum.cast(&partial.return_dtype)? - }; - vtable.combine_partials(partial, sum)?; - Ok(true) } /// The group state for a sum aggregate, containing the accumulated value and configuration @@ -442,42 +340,59 @@ pub struct SumPartial { return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, - /// Whether at least one valid value has been accumulated. A sum over zero valid values - /// finalizes to null (SQL `SUM` semantics) rather than the monoid zero. - seen: 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, seen}` struct, where -/// `sum` is the running monoid value (null once saturated by overflow, poisoning merges) and -/// `seen` records whether any valid value contributed (merged with OR). Keeping the flag -/// separate from the sum lets merges stay a monoid — the identity is `{0, false}` — while -/// `finalize` maps unseen sums to null. -fn sum_partial_dtype(sum_dtype: DType) -> DType { - DType::Struct( - StructFields::new( - FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), - vec![sum_dtype, DType::Bool(Nullability::NonNullable)], - ), - Nullability::Nullable, - ) +/// The accumulated sum value. +// TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the +// input type every time? Perhaps? +pub enum SumState { + Unsigned(u64), + Signed(i64), + Float(f64), + Decimal { + value: DecimalValue, + dtype: DecimalDType, + }, } -/// The running sum as a nullable scalar of the return dtype (null when saturated by overflow). -fn sum_value_scalar(partial: &SumPartial) -> 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) +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), + PType::I8 | PType::I16 | PType::I32 | PType::I64 => SumState::Signed(0), + PType::F16 | PType::F32 | PType::F64 => SumState::Float(0.0), + }, + DType::Decimal(decimal, _) => SumState::Decimal { + value: DecimalValue::zero(decimal), + dtype: *decimal, + }, + _ => vortex_panic!("Unsupported sum type"), + } +} + +/// Checked add for u64, returning true if overflow occurred. +#[inline(always)] +pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { + match acc.checked_add(val) { + Some(r) => { + *acc = r; + false } + None => true, + } +} + +/// Checked add for i64, returning true if overflow occurred. +#[inline(always)] +pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { + match acc.checked_add(val) { + Some(r) => { + *acc = r; + false + } + None => true, } } @@ -499,7 +414,6 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::aggregate_fn::fns::total::total; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -616,100 +530,6 @@ mod tests { // State merge tests (vtable-level) - #[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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - let empty = Sum.to_scalar(&state)?; - Sum.combine_partials(&mut state, empty)?; - assert!(Sum.finalize_scalar(&state)?.is_null()); - Ok(()) - } - - #[test] - fn sum_state_empty_is_identity() -> VortexResult<()> { - // Combining an empty state into a seen state changes nothing: `{0, false}` is the - // identity of the `{sum, seen}` monoid. - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - Sum.combine_partials(&mut state, Scalar::primitive(100i64, Nullable))?; - - let empty = - Sum.to_scalar(&Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?)?; - Sum.combine_partials(&mut state, empty)?; - - let result = Sum.finalize_scalar(&state)?; - assert_eq!(result.as_primitive().typed_value::(), Some(100)); - Ok(()) - } - - #[test] - fn sum_state_overflow_poisons_but_stays_seen() -> VortexResult<()> { - // Overflow (a null `sum` field) poisons the merge even when combined with later - // values: the result is null via the sum value, not via `seen`. - let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut overflowed = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - Sum.combine_partials(&mut overflowed, Scalar::primitive(i64::MAX, Nullable))?; - Sum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; - let overflowed = Sum.to_scalar(&overflowed)?; - - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - Sum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; - Sum.combine_partials(&mut state, overflowed)?; - Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; - - assert!(Sum.finalize_scalar(&state)?.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 = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); - Ok(()) - } - - #[test] - fn total_is_monoid_while_sum_is_sql() -> VortexResult<()> { - // The persisted statistic keeps the monoid semantics (zero for all-null) that zone - // and chunk merging require, while `sum` applies the SQL rule. - let mut ctx = array_session().create_execution_ctx(); - let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - assert_eq!( - total(&arr, &mut ctx)?.as_primitive().typed_value::(), - Some(0) - ); - // The cached monoid statistic must not leak through `sum`'s cache short-circuit. - assert!(sum(&arr, &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 sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); @@ -721,7 +541,7 @@ mod tests { let scalar2 = Scalar::primitive(50i64, Nullable); Sum.combine_partials(&mut state, scalar2)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; Sum.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) @@ -835,8 +655,7 @@ mod tests { 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 (SQL `SUM` semantics). - let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); + let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -929,7 +748,7 @@ mod tests { } #[test] - fn sum_chunked_floats_all_nulls_is_null() -> VortexResult<()> { + fn sum_chunked_floats_all_nulls_is_zero() -> 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(); @@ -938,8 +757,7 @@ mod tests { &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - // SQL `SUM`: no valid values across any chunk yields null. - assert!(result.is_null()); + assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 36c337e6cb5..87d8da4b143 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -26,19 +26,12 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, skip_nans: bool, - seen: &mut bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => { - *seen |= !p.as_ref().is_empty(); - accumulate_primitive_all(inner, p, skip_nans) - } - AllOr::Some(slices) => { - *seen |= !slices.is_empty(); - accumulate_primitive_valid(inner, p, slices, skip_nans) - } + AllOr::All => accumulate_primitive_all(inner, p, skip_nans), + AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), } } @@ -267,8 +260,7 @@ mod tests { fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -276,7 +268,7 @@ mod tests { fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); + assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } @@ -297,21 +289,20 @@ mod tests { } #[test] - fn sum_empty_is_null() -> VortexResult<()> { + fn sum_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - // SQL `SUM`: the sum over no values is null. - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } #[test] - fn sum_empty_f64_is_null() -> VortexResult<()> { + fn sum_empty_f64_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) } @@ -404,8 +395,6 @@ mod tests { .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(()) diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 0e1aa992b02..25b270d5b7b 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -29,10 +29,10 @@ 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::total::PrimitiveGroupedTotalEncodingKernel; -use crate::aggregate_fn::fns::total::Total; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; @@ -101,7 +101,7 @@ impl Default for AggregateFnSession { this.register(MinMax); this.register(NanCount); this.register(NullCount); - this.register(Total); + this.register(StandardSum); this.register(Sum); this.register(UncompressedSizeInBytes); @@ -120,8 +120,8 @@ impl Default for AggregateFnSession { ); this.register_grouped_encoding_kernel( Primitive.id(), - Total.id(), - &PrimitiveGroupedTotalEncodingKernel, + StandardSum.id(), + &PrimitiveGroupedStandardSumEncodingKernel, ); this diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index c9028ade519..c3dc2e0eed7 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -28,7 +28,7 @@ use crate::ExecutionResult; use crate::IntoArray; use crate::VTable; use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::total::total; +use crate::aggregate_fn::fns::sum::sum; use crate::array::ArrayData; use crate::array::ArrayId; use crate::array::ArrayInner; @@ -334,7 +334,7 @@ impl ArrayRef { Validity::NonNullable | Validity::AllValid => len, Validity::AllInvalid => 0, Validity::Array(a) => { - let array_sum = total(&a, ctx)?; + let array_sum = sum(&a, ctx)?; array_sum .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 664176cb680..149c72f499e 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -47,7 +47,7 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::total::Total; + use crate::aggregate_fn::fns::sum::Sum; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -60,7 +60,7 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); let mut acc = Accumulator::try_new( - Total, + Sum, NumericalAggregateOpts::default(), batch.dtype().clone(), )?; diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index d2e5535fbc6..73f98cce4a7 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1002,7 +1002,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::aggregate_fn::fns::total::total; + use crate::aggregate_fn::fns::sum::sum; use crate::dtype::DType; let len = array.len(); @@ -1045,8 +1045,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { return; } - if let (Ok(slice_sum), Ok(canonical_sum)) = (total(&sliced, ctx), total(&canonical_sliced, ctx)) - { + if let (Ok(slice_sum), Ok(canonical_sum)) = (sum(&sliced, ctx), sum(&canonical_sliced, ctx)) { // Compare sum scalars assert_eq!( slice_sum, canonical_sum, diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 20054213813..57e11135ef3 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -188,7 +188,7 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - return aggregate_fn::fns::total::Total + return aggregate_fn::fns::sum::Sum .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); } }) @@ -200,7 +200,7 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::total::Total.bind(NumericalAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { @@ -216,7 +216,7 @@ impl Stat { /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those /// aggregates have no stat slot. pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option { - if let Some(options) = aggregate_fn.as_opt::() { + if let Some(options) = aggregate_fn.as_opt::() { return options.skip_nans.then_some(Self::Sum); } if aggregate_fn.is::() { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index b98aa96a38e..02e4a700cc6 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -16,7 +16,7 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::standard_sum::StandardSum; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::scalar_fn::Arity; @@ -29,8 +29,8 @@ use crate::scalar_fn::ScalarFnVTable; /// /// 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 @@ -74,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}")) } @@ -121,7 +122,7 @@ impl ScalarFnVTable for ListSum { /// Sum each list of a canonical `array` into one value per list. /// -/// The grouped [`Sum`] finalize already yields null for null, empty, and all-null lists +/// 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, @@ -129,7 +130,7 @@ fn list_sum_impl( 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)?; acc.finish(ctx) } diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index ad4776745a6..d3fb7fd11e4 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -23,7 +23,7 @@ use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::total::total; +use crate::aggregate_fn::fns::sum::sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -172,9 +172,8 @@ impl StatsSetRef<'_> { .dtype(self.dyn_array_ref.dtype()) .is_some() .then(|| { - // Sum is supported for this dtype. The statistic is the monoid sum - // (zero when no valid values), not the SQL sum. - total(self.dyn_array_ref, ctx) + // Sum is supported for this dtype. + sum(self.dyn_array_ref, ctx) }) .transpose()? } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index ed24e54c06b..2038df0b613 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -14,7 +14,7 @@ use crate::aggregate_fn::fns::all_null::AllNull; 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::total::Total; +use crate::aggregate_fn::fns::sum::Sum; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; pub use crate::scalar_fn::fns::stat::StatFn; @@ -37,7 +37,7 @@ pub fn min_max(expr: Expression) -> Expression { /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. - stat(expr, Total.bind(NumericalAggregateOpts::skip_nans())) + stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index d2620f41d9d..9acd763bf30 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::total::Total; +use vortex::aggregate_fn::fns::sum::Sum; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -483,7 +483,7 @@ pub fn try_from_projection_expression( pub enum PushedAggregate { Min, Max, - Total, + Sum, Mean, // Also used for ANY_VALUE() which is allowed by definition First, @@ -496,7 +496,7 @@ impl Display for PushedAggregate { match self { PushedAggregate::Min => f.write_str("min"), PushedAggregate::Max => f.write_str("max"), - PushedAggregate::Total => f.write_str("sum"), + PushedAggregate::Sum => f.write_str("sum"), PushedAggregate::Mean => f.write_str("mean"), PushedAggregate::First => f.write_str("first"), PushedAggregate::Count => f.write_str("count"), @@ -510,7 +510,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::Total => Box::new(Accumulator::try_new(Total, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), PairOptions(opts, opts), @@ -535,7 +535,7 @@ pub fn try_from_projection_aggregate( Ok(Some(match agg.aggregate_function.name() { "min" => PushedAggregate::Min, "max" => PushedAggregate::Max, - "sum" | "sum_no_overflow" => PushedAggregate::Total, + "sum" | "sum_no_overflow" => PushedAggregate::Sum, "avg" | "mean" => PushedAggregate::Mean, "first" | "any_value" => PushedAggregate::First, "count" => PushedAggregate::Count, diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 4702008f01a..c4cd032c686 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -12,7 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::total::total; +use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; @@ -168,7 +168,7 @@ impl StatsAccumulator { } } Stat::NullCount | Stat::NaNCount | Stat::UncompressedSizeInBytes => { - if let Some(sum_value) = total(array, ctx)? + if let Some(sum_value) = sum(array, ctx)? .cast(&DType::Primitive(PType::U64, Nullability::Nullable))? .into_value() { diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index f02a8edc86d..d05f5672640 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -86,7 +86,7 @@ pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[St .filter_map(|stat| { stat.dtype(column_dtype) .or_else(|| { - // Backward compat: older files may have stored stats (e.g. Total) + // Backward compat: older files may have stored stats (e.g. Sum) // for extension types by resolving through the storage dtype. if let DType::Extension(ext) = column_dtype { stat.dtype(ext.storage_dtype()) @@ -157,7 +157,7 @@ mod tests { use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::total::Total; + use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -208,7 +208,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - Total.bind(NumericalAggregateOpts::skip_nans()), + Sum.bind(NumericalAggregateOpts::skip_nans()), ], ); @@ -217,7 +217,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - Total.bind(NumericalAggregateOpts::skip_nans()).to_string(), + Sum.bind(NumericalAggregateOpts::skip_nans()).to_string(), ] ); } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 5a20fd9b486..c1c12d7ddea 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,7 +25,7 @@ use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::total::Total; +use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -212,11 +212,11 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { }; let mut aggregate_fns = vec![max, min]; - if Total + if Sum .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) .is_some() { - aggregate_fns.push(Total.bind(NumericalAggregateOpts::skip_nans())); + aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); @@ -230,7 +230,7 @@ mod tests { use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::total::Total; + use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::extension::datetime::TimeUnit; @@ -258,7 +258,7 @@ mod tests { assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); - assert!(aggregate_fns[2].is::()); + assert!(aggregate_fns[2].is::()); } #[test] @@ -271,7 +271,7 @@ mod tests { assert!( aggregate_fns .iter() - .all(|aggregate_fn| !aggregate_fn.is::()) + .all(|aggregate_fn| !aggregate_fn.is::()) ); } } From ae3dad2839c2bd6fa478fc87f4021b587eb714c1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 17:53:40 -0700 Subject: [PATCH 07/12] Drop ExecutionCtx from aggregate finalize; StandardSum finalize returns a lazy masked array Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 43 ++++++++---- .../src/aggregate_fn/accumulator_grouped.rs | 14 ++-- vortex-array/src/aggregate_fn/combined.rs | 6 +- .../src/aggregate_fn/fns/all_nan/mod.rs | 2 +- .../aggregate_fn/fns/all_non_distinct/mod.rs | 2 +- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 2 +- .../src/aggregate_fn/fns/all_non_null/mod.rs | 2 +- .../src/aggregate_fn/fns/all_null/mod.rs | 2 +- .../src/aggregate_fn/fns/bounded_max/mod.rs | 2 +- .../src/aggregate_fn/fns/bounded_min/mod.rs | 2 +- .../src/aggregate_fn/fns/count/grouped.rs | 4 +- .../src/aggregate_fn/fns/count/mod.rs | 2 +- .../src/aggregate_fn/fns/first/mod.rs | 2 +- .../src/aggregate_fn/fns/is_constant/mod.rs | 2 +- .../src/aggregate_fn/fns/is_sorted/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/last/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/max/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/min/mod.rs | 2 +- .../src/aggregate_fn/fns/min_max/mod.rs | 2 +- .../src/aggregate_fn/fns/nan_count/mod.rs | 2 +- .../src/aggregate_fn/fns/null_count/mod.rs | 2 +- .../aggregate_fn/fns/standard_sum/grouped.rs | 4 +- .../src/aggregate_fn/fns/standard_sum/mod.rs | 68 ++++++++----------- .../src/aggregate_fn/fns/sum/grouped.rs | 10 ++- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 11 ++- .../fns/uncompressed_size_in_bytes/mod.rs | 2 +- vortex-array/src/aggregate_fn/foreign.rs | 2 +- vortex-array/src/aggregate_fn/proto.rs | 2 +- vortex-array/src/aggregate_fn/vtable.rs | 2 +- .../src/arrays/struct_/compute/rules.rs | 10 ++- vortex-array/src/scalar_fn/fns/list_sum.rs | 2 +- 31 files changed, 108 insertions(+), 106 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 1a037c49978..d8dc172ac4f 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -164,15 +164,9 @@ where 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(&mut ctx) - .unwrap() - .execute::(&mut ctx) - .unwrap() - .into_array(); - divan::black_box(result) + acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) + .unwrap(); + divan::black_box(acc.finish().unwrap()) } #[divan::bench] @@ -207,12 +201,35 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Sum)); } +/// 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 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)); + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); } #[divan::bench] @@ -220,7 +237,7 @@ 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)); + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); } #[divan::bench] @@ -228,7 +245,7 @@ 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)); + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); } #[divan::bench] @@ -236,7 +253,7 @@ 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)); + .bench_refs(|input| grouped_accumulator_canonical(input, StandardSum)); } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index e6d0fdfaabf..b87c04ee204 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -238,7 +238,7 @@ pub trait DynGroupedAccumulator: 'static + Send { /// Finish the accumulation and return the final aggregate results for all groups. /// Resets the accumulator state for the next round of accumulation. - fn finish(&mut self, ctx: &mut ExecutionCtx) -> VortexResult; + fn finish(&mut self) -> VortexResult; } impl DynGroupedAccumulator for GroupedAccumulator { @@ -276,15 +276,9 @@ impl DynGroupedAccumulator for GroupedAccumulator { Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) } - fn finish(&mut self, ctx: &mut ExecutionCtx) -> VortexResult { - // The single-batch case (one accumulate_list call) skips the chunked wrapper so that - // finalize implementations can operate on the batch's states directly. - let states = if self.partials.len() == 1 { - self.partials.pop().vortex_expect("checked length") - } else { - self.flush()? - }; - let results = self.vtable.finalize(states, ctx)?; + fn finish(&mut self) -> VortexResult { + let states = self.flush()?; + let results = self.vtable.finalize(states)?; vortex_ensure!( results.dtype() == &self.return_dtype, diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 86ddb678093..76ad9877314 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -245,11 +245,11 @@ impl AggregateFnVTable for Combined { unreachable!("Combined::try_accumulate handles all batches") } - fn finalize(&self, states: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, states: ArrayRef) -> VortexResult { let l_field = states.get_item(FieldName::from(self.0.left_name()))?; let r_field = states.get_item(FieldName::from(self.0.right_name()))?; - let l_finalized = self.0.left().finalize(l_field, ctx)?; - let r_finalized = self.0.right().finalize(r_field, ctx)?; + let l_finalized = self.0.left().finalize(l_field)?; + let r_finalized = self.0.right().finalize(r_field)?; BinaryCombined::finalize(&self.0, l_finalized, r_finalized) } diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index e60b5f27696..68a58908018 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -121,7 +121,7 @@ impl AggregateFnVTable for AllNan { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index 6dede02675c..f030cc810e2 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -233,7 +233,7 @@ impl AggregateFnVTable for AllNonDistinct { } } - fn finalize(&self, _partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, _partials: ArrayRef) -> VortexResult { vortex_bail!("AllNonDistinct does not support array finalization"); } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index f31018f820a..fe8527da966 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -111,7 +111,7 @@ impl AggregateFnVTable for AllNonNan { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index b01e1f74e1e..c07dbb907c9 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -101,7 +101,7 @@ impl AggregateFnVTable for AllNonNull { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index 0c010776c52..ec64e3d5c43 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -104,7 +104,7 @@ impl AggregateFnVTable for AllNull { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index a7c1a51633e..bb27d9f9ab1 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -279,7 +279,7 @@ impl AggregateFnVTable for BoundedMax { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { partials.get_item(BOUNDED_MAX_BOUND) } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index d8295772688..7a72442550b 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -197,7 +197,7 @@ impl AggregateFnVTable for BoundedMin { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 4c70e3e19c2..39a957530bf 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -122,7 +122,7 @@ mod tests { elem_dtype.clone(), )?; acc.accumulate_list(groups, ctx)?; - acc.finish(ctx) + acc.finish() } /// Reference valid-counts (non-nullable `U64`), one per group. @@ -235,7 +235,7 @@ mod tests { let mut acc = GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?; acc.accumulate_list(&groups, &mut ctx)?; - let actual = acc.finish(&mut ctx)?; + let actual = acc.finish()?; let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index e05dae91cd4..8f7d68027bc 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -115,7 +115,7 @@ impl AggregateFnVTable for Count { unreachable!("Count::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 3c2c30539fc..c41e2575057 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -117,7 +117,7 @@ impl AggregateFnVTable for First { unreachable!("First::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 19fb1dbb5e9..490ede7f640 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -420,7 +420,7 @@ impl AggregateFnVTable for IsConstant { } } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index f5dab74a7d9..6256dc22327 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -514,7 +514,7 @@ impl AggregateFnVTable for IsSorted { } } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { partials.get_item(NAMES.get(0).vortex_expect("out of bounds").clone()) } diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 7c1def89082..63ee54a6efa 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -115,7 +115,7 @@ impl AggregateFnVTable for Last { unreachable!("Last::try_accumulate handles all arrays") } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index 8896de2e301..ccf9f4d899d 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -205,7 +205,7 @@ impl AggregateFnVTable for Max { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index e306897dd48..488405e8f14 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -205,7 +205,7 @@ impl AggregateFnVTable for Min { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 828fe4bcfa5..51601ce76fd 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -427,7 +427,7 @@ impl AggregateFnVTable for MinMax { } } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index 56f9f813edb..723847160e1 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -172,7 +172,7 @@ impl AggregateFnVTable for NanCount { } } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index 9c0427812e5..031e8f6a09b 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -143,7 +143,7 @@ impl AggregateFnVTable for NullCount { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs index 04f5f857588..63e836083ce 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -216,7 +216,7 @@ mod tests { )?; let mut ctx = array_session().create_execution_ctx(); acc.accumulate_list(groups, &mut ctx)?; - acc.finish(&mut ctx) + acc.finish() } /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] @@ -407,7 +407,7 @@ mod tests { )?; let mut ctx2 = array_session().create_execution_ctx(); acc.accumulate_list(&groups, &mut ctx2)?; - let actual = acc.finish(&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. diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 74b58fee298..0d0563f99dd 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -11,7 +11,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -34,13 +33,8 @@ 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::sum_decimal_dtype; -use crate::arrays::Bool; -use crate::arrays::BoolArray; -use crate::arrays::Struct; -use crate::arrays::StructArray; -use crate::arrays::bool::BoolArrayExt; -use crate::arrays::masked::mask_validity_canonical; -use crate::arrays::struct_::StructArrayExt; +use crate::arrays::ConstantArray; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -52,7 +46,10 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; -use crate::validity::Validity; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::fns::fill_null::FillNull; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::mask::Mask; /// Return the SQL sum of an array: null when the array has no valid values or the sum /// overflows. @@ -94,7 +91,7 @@ pub fn standard_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` - // field (overflow) and null partial rows (e.g. null groups) stay null via the - // validity intersection. Implemented structurally rather than as a `get_item`/`mask` - // expression: expression construction and dispatch cost multiples of the whole - // aggregation at small group counts. + // field (overflow) and null partial rows (e.g. null groups) stay null via the mask's + // validity intersection. + // + // The expressions are built unoptimized: `optimize` costs multiples of the whole + // aggregation at small group counts, and the caller's execution evaluates the lazy + // expression as-is. let len = partials.len(); - let states = match partials.as_opt::() { - Some(states) => states.into_owned(), - None => partials.execute::(ctx)?, - }; - let struct_mask = states.as_ref().validity()?.execute_mask(len, ctx)?; - let seen = states.unmasked_field_by_name("seen")?.clone(); - let seen = match seen.as_opt::() { - Some(seen) => seen.into_owned(), - None => seen.execute::(ctx)?, - }; - let valid = &struct_mask & &Mask::from_buffer(seen.to_bit_buffer()); - - let sum = states.unmasked_field_by_name("sum")?.clone(); - if valid.all_true() { - // Every partial row is a seen, valid group: the sums are already the result. - return Ok(sum); - } - let sum = sum.execute::(ctx)?; - Ok( - mask_validity_canonical(sum, Validity::from_mask(valid, Nullability::Nullable), ctx)? - .into_array(), - ) + let sum = GetItem.try_new_array(len, FieldName::from("sum"), [partials.clone()])?; + let seen = GetItem.try_new_array(len, FieldName::from("seen"), [partials])?; + let seen = FillNull.try_new_array( + len, + EmptyOptions, + [ + seen, + ConstantArray::new(Scalar::bool(false, Nullability::NonNullable), len).into_array(), + ], + )?; + Mask.try_new_array(len, EmptyOptions, [sum, seen]) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { @@ -776,7 +764,7 @@ mod tests { )?; let mut ctx = array_session().create_execution_ctx(); acc.accumulate_list(groups, &mut ctx)?; - acc.finish(&mut ctx) + acc.finish() } #[test] @@ -873,7 +861,7 @@ mod tests { 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(&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); @@ -881,7 +869,7 @@ mod tests { 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(&mut ctx)?; + let result2 = acc.finish()?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index fb9c566631d..efe0825d4d6 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -184,9 +184,8 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - let mut ctx = array_session().create_execution_ctx(); - acc.accumulate_list(groups, &mut ctx)?; - acc.finish(&mut ctx) + acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; + acc.finish() } /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] for @@ -371,9 +370,8 @@ mod tests { let mut acc = GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; - let mut ctx2 = array_session().create_execution_ctx(); - acc.accumulate_list(&groups, &mut ctx2)?; - let actual = acc.finish(&mut ctx2)?; + acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; + let actual = acc.finish()?; let mut ctx = array_session().create_execution_ctx(); // Group 0 contains a NaN -> NaN sum; group 1 sums normally. diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 54a123ef765..9e71bc8edf2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -325,7 +325,7 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } @@ -591,9 +591,8 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - let mut ctx = array_session().create_execution_ctx(); - acc.accumulate_list(groups, &mut ctx)?; - acc.finish(&mut ctx) + acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; + acc.finish() } #[test] @@ -686,7 +685,7 @@ mod tests { 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(&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); @@ -694,7 +693,7 @@ mod tests { 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(&mut ctx)?; + let result2 = acc.finish()?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 3e78f0298ec..6b2feb380a7 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -177,7 +177,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs index 8ef8da01824..feb07e47175 100644 --- a/vortex-array/src/aggregate_fn/foreign.rs +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -108,7 +108,7 @@ impl AggregateFnVTable for ForeignAggregateFnVTable { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } - fn finalize(&self, _states: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, _states: ArrayRef) -> VortexResult { vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) } diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 26e206bcae9..92fac87892a 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -149,7 +149,7 @@ mod tests { Ok(()) } - fn finalize(&self, partials: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 9af95717d35..49b28dd26d7 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -152,7 +152,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// /// The provides `states` array has dtype as specified by `state_dtype`, the result array /// must have dtype as specified by `return_dtype`. - fn finalize(&self, states: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + fn finalize(&self, states: ArrayRef) -> VortexResult; /// Finalize a scalar accumulator state into an aggregate result. /// 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 02e4a700cc6..69565597deb 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -132,7 +132,7 @@ fn list_sum_impl( ) -> VortexResult { let mut acc = GroupedAccumulator::try_new(StandardSum, *options, elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; - acc.finish(ctx) + acc.finish() } #[cfg(test)] From 8d5cd560c9c9a440d2b39a693adfbba85f5a9f29 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 18:25:58 -0700 Subject: [PATCH 08/12] Share sum accumulation kernels with StandardSum via a seen out-param Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 72 +- .../src/aggregate_fn/fns/standard_sum/bool.rs | 157 --- .../aggregate_fn/fns/standard_sum/constant.rs | 216 ---- .../aggregate_fn/fns/standard_sum/decimal.rs | 476 --------- .../aggregate_fn/fns/standard_sum/grouped.rs | 6 +- .../src/aggregate_fn/fns/standard_sum/mod.rs | 923 +++++++++++++++++- .../fns/standard_sum/primitive.rs | 471 --------- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 4 +- .../src/aggregate_fn/fns/sum/constant.rs | 2 +- .../src/aggregate_fn/fns/sum/decimal.rs | 4 +- .../src/aggregate_fn/fns/sum/design.md | 140 +++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 45 +- .../src/aggregate_fn/fns/sum/primitive.rs | 19 +- 13 files changed, 1171 insertions(+), 1364 deletions(-) delete mode 100644 vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs delete mode 100644 vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs delete mode 100644 vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs delete mode 100644 vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs create mode 100644 vortex-array/src/aggregate_fn/fns/sum/design.md diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index d8dc172ac4f..1dc57c431c2 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -201,6 +201,38 @@ 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 @@ -225,7 +257,39 @@ where } #[divan::bench] -fn standard_sum_i32_nullable_all_valid(bencher: Bencher) { +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) @@ -233,7 +297,7 @@ fn standard_sum_i32_nullable_all_valid(bencher: Bencher) { } #[divan::bench] -fn standard_sum_i32_clustered_nulls(bencher: Bencher) { +fn canonical_standard_sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) @@ -241,7 +305,7 @@ fn standard_sum_i32_clustered_nulls(bencher: Bencher) { } #[divan::bench] -fn standard_sum_f64_all_valid(bencher: Bencher) { +fn canonical_standard_sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) @@ -249,7 +313,7 @@ fn standard_sum_f64_all_valid(bencher: Bencher) { } #[divan::bench] -fn standard_sum_f64_clustered_nulls(bencher: Bencher) { +fn canonical_standard_sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs deleted file mode 100644 index 4bafb6f31fb..00000000000 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/bool.rs +++ /dev/null @@ -1,157 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::ops::BitAnd; - -use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::AllOr; - -use super::SumState; -use super::checked_add_u64; -use crate::ExecutionCtx; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; - -pub(super) fn accumulate_bool( - inner: &mut SumState, - b: &BoolArray, - ctx: &mut ExecutionCtx, - seen: &mut bool, -) -> VortexResult { - let SumState::Unsigned(acc) = inner else { - vortex_panic!("expected unsigned sum state for bool input"); - }; - - let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; - let true_count = match mask.bit_buffer() { - AllOr::None => return Ok(false), - AllOr::All => b.bit_buffer_view().true_count() as u64, - AllOr::Some(validity) => b.to_bit_buffer().bitand(validity).true_count() as u64, - }; - - Ok(checked_add_u64(acc, true_count)) -} - -#[cfg(test)] -mod tests { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - 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::BoolArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::executor::VortexSessionExecute; - - #[test] - fn sum_bool_all_true() -> VortexResult<()> { - let arr: BoolArray = [true, true, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_mixed() -> VortexResult<()> { - let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_all_false() -> VortexResult<()> { - let arr: BoolArray = [false, false, false].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); - 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_bool_all_null() -> VortexResult<()> { - let arr = BoolArray::from_iter([None::, None, None]); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_empty_is_null() -> VortexResult<()> { - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - - let batch1: BoolArray = [true, true, false].into_iter().collect(); - acc.accumulate(&batch1.into_array(), &mut ctx)?; - let result1 = acc.finish()?; - assert_eq!(result1.as_primitive().typed_value::(), Some(2)); - - let batch2: BoolArray = [false, true].into_iter().collect(); - acc.accumulate(&batch2.into_array(), &mut ctx)?; - let result2 = acc.finish()?; - assert_eq!(result2.as_primitive().typed_value::(), Some(1)); - Ok(()) - } - - #[test] - fn sum_bool_return_dtype() -> VortexResult<()> { - let dtype = StandardSum - .return_dtype( - &NumericalAggregateOpts::default(), - &DType::Bool(Nullability::NonNullable), - ) - .unwrap(); - assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); - Ok(()) - } - - #[test] - fn sum_boolean_from_iter() -> VortexResult<()> { - let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(2)); - Ok(()) - } -} diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs deleted file mode 100644 index 9155b6a5c4c..00000000000 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/constant.rs +++ /dev/null @@ -1,216 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; - -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; - -/// Compute `scalar * len` for a constant array, returning the product as a sum-typed 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( - scalar: &Scalar, - len: usize, - return_dtype: &DType, -) -> VortexResult> { - if scalar.is_null() || len == 0 { - return Ok(None); - } - - let product = match scalar.dtype() { - DType::Bool(_) => { - let val = scalar - .as_bool() - .value() - .ok_or_else(|| vortex_err!("Expected non-null bool scalar for sum"))?; - if !val { - return Ok(None); - } - Scalar::primitive(len as u64, Nullability::Nullable) - } - DType::Primitive(..) => { - let pvalue = scalar - .as_primitive() - .pvalue() - .ok_or_else(|| vortex_err!("Expected non-null primitive scalar for sum"))?; - match return_dtype { - DType::Primitive(PType::U64, _) => { - let val = pvalue.cast::()?; - match val.checked_mul(len as u64) { - Some(product) => Scalar::primitive(product, Nullability::Nullable), - None => Scalar::null(return_dtype.as_nullable()), - } - } - DType::Primitive(PType::I64, _) => { - let val = pvalue.cast::()?; - match i64::try_from(len).ok().and_then(|l| val.checked_mul(l)) { - Some(product) => Scalar::primitive(product, Nullability::Nullable), - None => Scalar::null(return_dtype.as_nullable()), - } - } - DType::Primitive(PType::F64, _) => { - let val = pvalue.cast::()?; - Scalar::primitive(val * len as f64, Nullability::Nullable) - } - _ => vortex_bail!( - "Unexpected return dtype for primitive sum: {}", - return_dtype - ), - } - } - DType::Decimal(..) => { - let val = scalar - .as_decimal() - .decimal_value() - .ok_or_else(|| vortex_err!("Expected non-null decimal scalar for sum"))?; - let len_decimal = DecimalValue::from(len as i128); - match val.checked_mul(&len_decimal) { - Some(product) => { - let ret_decimal = *return_dtype - .as_decimal_opt() - .ok_or_else(|| vortex_err!("Expected decimal return dtype"))?; - Scalar::decimal(product, ret_decimal, Nullability::Nullable) - } - None => Scalar::null(return_dtype.as_nullable()), - } - } - _ => vortex_bail!("Unsupported constant type for sum: {}", scalar.dtype()), - }; - - Ok(Some(product)) -} - -#[cfg(test)] -mod tests { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::standard_sum::standard_sum; - use crate::array_session; - use crate::arrays::ConstantArray; - 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::Stat; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - - #[test] - fn sum_constant_unsigned() -> 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_signed() -> VortexResult<()> { - let array = ConstantArray::new(-5i64, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, (-50i64).into()); - Ok(()) - } - - #[test] - fn sum_constant_nullable_value() -> VortexResult<()> { - let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) - .into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: an all-null constant has no valid values. - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_bool_false() -> VortexResult<()> { - let array = ConstantArray::new(false, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 0u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_true() -> VortexResult<()> { - let array = ConstantArray::new(true, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 10u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_null() -> VortexResult<()> { - let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_decimal() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(100), - decimal_dtype, - Nullability::NonNullable, - ), - 5, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(500))) - ); - assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); - Ok(()) - } - - #[test] - fn sum_constant_decimal_null() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result, - Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) - ); - Ok(()) - } - - #[test] - fn sum_constant_decimal_large_value() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(999_999_999), - decimal_dtype, - Nullability::NonNullable, - ), - 100, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) - ); - Ok(()) - } -} diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs deleted file mode 100644 index 40b8f38236b..00000000000 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/decimal.rs +++ /dev/null @@ -1,476 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use itertools::Itertools; -use num_traits::AsPrimitive; -use num_traits::CheckedAdd; -use num_traits::NumOps; -use vortex_buffer::BitBuffer; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::Mask; - -use super::SumState; -use crate::ExecutionCtx; -use crate::arrays::DecimalArray; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; -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( - inner: &mut SumState, - d: &DecimalArray, - ctx: &mut ExecutionCtx, - seen: &mut bool, -) -> VortexResult { - let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; - let validity = match &mask { - Mask::AllTrue(_) => None, - Mask::Values(mask_values) => Some(mask_values.bit_buffer()), - Mask::AllFalse(_) => { - return Ok(false); - } - }; - - let SumState::Decimal { value, dtype } = inner else { - vortex_panic!("expected decimal sum state for decimal input"); - }; - - let values_type = DecimalType::smallest_decimal_value_type(dtype); - match_each_decimal_value_type!(d.values_type(), |T| { - match_each_decimal_value_type!(values_type, |I| { - let initial: I = value - .cast() - .vortex_expect("cannot fail to cast initial value"); - match sum_decimal_value(initial, d.buffer::(), validity, *dtype) { - Some(v) => *value = v, - None => return Ok(true), - } - Ok(false) - }) - }) -} - -fn sum_decimal_value( - initial: I, - values: Buffer, - validity: Option<&BitBuffer>, - output_dtype: DecimalDType, -) -> Option -where - T: AsPrimitive, - I: NumOps + CheckedAdd + Copy + NativeDecimalType + 'static, - bool: AsPrimitive, - DecimalValue: From, -{ - let sum = match validity { - Some(v) => sum_decimal_with_validity(values, v, initial), - None => sum_decimal(values, initial), - }; - - sum.map(DecimalValue::from) - // We have to make sure that the decimal value fits the precision of the decimal dtype. - .filter(|v| v.fits_in_precision(output_dtype)) -} - -fn sum_decimal, I: Copy + CheckedAdd + 'static>( - values: Buffer, - initial: I, -) -> Option { - let mut sum = initial; - for v in values.iter() { - let v: I = v.as_(); - sum = CheckedAdd::checked_add(&sum, &v)?; - } - Some(sum) -} - -fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> Option -where - T: AsPrimitive, - I: NumOps + CheckedAdd + Copy + 'static, - bool: AsPrimitive, -{ - let mut sum = initial; - for (v, valid) in values.iter().zip_eq(validity) { - let v: I = v.as_() * valid.as_(); - - sum = CheckedAdd::checked_add(&sum, &v)?; - } - Some(sum) -} - -#[cfg(test)] -mod tests { - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::AggregateFnVTable; - 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::DecimalArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::i256; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - #[test] - fn sum_decimal_basic() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(600i32))), - )?; - - assert_eq!(result, expected); - 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_decimal_negative_values() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, -200i32, 300i32, -50i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(150i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_near_i32_max() -> VortexResult<()> { - let near_max = i32::MAX - 1000; - let decimal = DecimalArray::new( - buffer![near_max, 500i32, 400i32], - DecimalDType::new(10, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = near_max as i64 + 500 + 400; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_large_i64_values() -> VortexResult<()> { - let large_val = i64::MAX / 4; - let decimal = DecimalArray::new( - buffer![large_val, large_val, large_val, large_val + 1], - DecimalDType::new(19, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = (large_val as i128) * 4 + 1; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_preserves_scale() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![12345i32, 67890i32, 11111i32], - DecimalDType::new(6, 4), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(91346i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_single_value() -> VortexResult<()> { - let decimal = - DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(42i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32, 400i32], - DecimalDType::new(4, 2), - Validity::from_iter([false, false, true, false]), - ); - - 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(300i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_overflow_detection() -> VortexResult<()> { - let max_val = i128::MAX / 2; - let decimal = DecimalArray::new( - buffer![max_val, max_val, max_val], - DecimalDType::new(38, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = - i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - 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, - ); - - assert_eq!( - standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx() - ) - .vortex_expect("operation should succeed in test"), - 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(()) - } - - #[test] - fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { - // Input precision 4 → return precision 14. Native I64 (max 18). - // The max representable value for precision 14 is 10^14 - 1. - // When the sum reaches exactly 10^14, fits_in_precision fails 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(99_999_999_999_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - // Push the sum to exactly 10^14, exceeding precision 14. - let one_more = - Scalar::decimal(DecimalValue::from(1i64), 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_precision_overflow_negative() -> VortexResult<()> { - // Same setup but with negative values: sum reaches -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_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - let one_more = Scalar::decimal( - DecimalValue::from(-1i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, one_more)?; - - let result = StandardSum.finalize_scalar(&state)?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { - // Test precision overflow via the accumulate_decimal path (not combine_partials). - // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. - // Native for precision 38 is I128 (max 38), so 38 = 38. - // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. - // - // We 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); - - // Drive accumulate through the vtable directly. - 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(()) - } -} diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs index 63e836083ce..450b962f5e5 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -8,15 +8,15 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use super::StandardSum; -use super::primitive::sum_float_all; -use super::primitive::sum_signed_all; -use super::primitive::sum_unsigned_all; 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; diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 0d0563f99dd..749ad224d27 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -1,11 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod bool; -mod constant; -mod decimal; mod grouped; -mod primitive; pub(crate) use grouped::PrimitiveGroupedStandardSumEncodingKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -14,10 +10,6 @@ use vortex_error::vortex_err; 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; use crate::ArrayRef; use crate::Canonical; use crate::Columnar; @@ -29,9 +21,13 @@ 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::ConstantArray; use crate::arrays::scalar_fn::ScalarFnFactoryExt; @@ -1084,4 +1080,915 @@ mod tests { ); Ok(()) } + + mod bool_inputs { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + 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::BoolArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::executor::VortexSessionExecute; + + #[test] + fn sum_bool_all_true() -> VortexResult<()> { + let arr: BoolArray = [true, true, true].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_mixed() -> VortexResult<()> { + let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_all_false() -> VortexResult<()> { + let arr: BoolArray = [false, false, false].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + 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_bool_all_null() -> VortexResult<()> { + let arr = BoolArray::from_iter([None::, None, None]); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_bool_empty_is_null() -> VortexResult<()> { + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = + Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_bool_finish_resets_state() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = + Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + + let batch1: BoolArray = [true, true, false].into_iter().collect(); + acc.accumulate(&batch1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + assert_eq!(result1.as_primitive().typed_value::(), Some(2)); + + let batch2: BoolArray = [false, true].into_iter().collect(); + acc.accumulate(&batch2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + assert_eq!(result2.as_primitive().typed_value::(), Some(1)); + Ok(()) + } + + #[test] + fn sum_bool_return_dtype() -> VortexResult<()> { + let dtype = StandardSum + .return_dtype( + &NumericalAggregateOpts::default(), + &DType::Bool(Nullability::NonNullable), + ) + .unwrap(); + assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); + Ok(()) + } + + #[test] + fn sum_boolean_from_iter() -> VortexResult<()> { + let arr = BoolArray::from_iter([true, false, false, true]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(2)); + Ok(()) + } + } + + mod constant_inputs { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::fns::standard_sum::standard_sum; + use crate::array_session; + use crate::arrays::ConstantArray; + 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::Stat; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + + #[test] + fn sum_constant_unsigned() -> 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_signed() -> VortexResult<()> { + let array = ConstantArray::new(-5i64, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, (-50i64).into()); + Ok(()) + } + + #[test] + fn sum_constant_nullable_value() -> VortexResult<()> { + let array = + ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) + .into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + // SQL `SUM`: an all-null constant has no valid values. + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + Ok(()) + } + + #[test] + fn sum_constant_bool_false() -> VortexResult<()> { + let array = ConstantArray::new(false, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 0u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_true() -> VortexResult<()> { + let array = ConstantArray::new(true, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 10u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_null() -> VortexResult<()> { + let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + Ok(()) + } + + #[test] + fn sum_constant_decimal() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(100), + decimal_dtype, + Nullability::NonNullable, + ), + 5, + ) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(500))) + ); + assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); + Ok(()) + } + + #[test] + fn sum_constant_decimal_null() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = + ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result, + Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) + ); + Ok(()) + } + + #[test] + fn sum_constant_decimal_large_value() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(999_999_999), + decimal_dtype, + Nullability::NonNullable, + ), + 100, + ) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) + ); + Ok(()) + } + } + + mod decimal_inputs { + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateFnVTable; + 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::DecimalArray; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::dtype::Nullability::Nullable; + use crate::dtype::i256; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + use crate::validity::Validity; + + #[test] + fn sum_decimal_basic() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(600i32))), + )?; + + assert_eq!(result, expected); + 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_decimal_negative_values() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, -200i32, 300i32, -50i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(150i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_near_i32_max() -> VortexResult<()> { + let near_max = i32::MAX - 1000; + let decimal = DecimalArray::new( + buffer![near_max, 500i32, 400i32], + DecimalDType::new(10, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = near_max as i64 + 500 + 400; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_large_i64_values() -> VortexResult<()> { + let large_val = i64::MAX / 4; + let decimal = DecimalArray::new( + buffer![large_val, large_val, large_val, large_val + 1], + DecimalDType::new(19, 0), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = (large_val as i128) * 4 + 1; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_preserves_scale() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![12345i32, 67890i32, 11111i32], + DecimalDType::new(6, 4), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(91346i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_single_value() -> VortexResult<()> { + let decimal = + DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(42i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32, 400i32], + DecimalDType::new(4, 2), + Validity::from_iter([false, false, true, false]), + ); + + 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(300i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_overflow_detection() -> VortexResult<()> { + let max_val = i128::MAX / 2; + let decimal = DecimalArray::new( + buffer![max_val, max_val, max_val], + DecimalDType::new(38, 0), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = + i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + 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, + ); + + assert_eq!( + standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx() + ) + .vortex_expect("operation should succeed in test"), + 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(()) + } + + #[test] + fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { + // Input precision 4 → return precision 14. Native I64 (max 18). + // The max representable value for precision 14 is 10^14 - 1. + // When the sum reaches exactly 10^14, fits_in_precision fails 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(99_999_999_999_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + // Push the sum to exactly 10^14, exceeding precision 14. + let one_more = + Scalar::decimal(DecimalValue::from(1i64), 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_precision_overflow_negative() -> VortexResult<()> { + // Same setup but with negative values: sum reaches -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_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + let one_more = Scalar::decimal( + DecimalValue::from(-1i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, one_more)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { + // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. + // Native for precision 38 is I128 (max 38), so 38 = 38. + // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. + // + // We 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); + + // Drive accumulate through the vtable directly. + 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(()) + } + } + + mod primitive_inputs { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::DynAccumulator; + 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::ConstantArray; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + 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; + + #[test] + fn sum_i32() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(10)); + Ok(()) + } + + #[test] + fn sum_u8() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(60)); + Ok(()) + } + + #[test] + fn sum_f64() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); + Ok(()) + } + + #[test] + fn sum_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_multiple_null_runs() -> VortexResult<()> { + // Several disjoint valid runs separated by nulls exercise the per-run fold. + let arr = PrimitiveArray::from_option_iter([ + Some(1i32), + Some(2), + None, + None, + Some(3), + None, + Some(4), + Some(5), + Some(6), + ]) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(21)); + Ok(()) + } + + #[test] + fn sum_all_null() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + // SQL `SUM`: no valid values yields null. + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_all_invalid_float() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); + Ok(()) + } + + #[test] + fn sum_buffer_i32() -> VortexResult<()> { + let arr = buffer![1, 1, 1, 1].into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4)); + Ok(()) + } + + #[test] + fn sum_buffer_f64() -> VortexResult<()> { + let arr = buffer![1., 1., 1., 1.].into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4.)); + Ok(()) + } + + #[test] + fn sum_empty_is_null() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = + Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + // SQL `SUM`: the sum over no values is null. + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_empty_f64_is_null() -> VortexResult<()> { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = + Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_f64_with_nan() -> VortexResult<()> { + let arr = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], + Validity::NonNullable, + ) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + Ok(()) + } + + #[test] + fn sum_f32_with_nan() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); + Ok(()) + } + + #[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_all_nan() -> VortexResult<()> { + 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(()) + } + + /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + fn sum_with_options( + arr: &crate::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() + } + + #[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_f64_without_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact StandardSum 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 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(()) + } + + #[test] + fn sum_checked_overflow() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_checked_overflow_is_saturated() -> 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()); + + // finish resets state, clearing saturation + drop(acc.finish()?); + assert!(!acc.is_saturated()); + Ok(()) + } + } } diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs deleted file mode 100644 index 6935cd10b23..00000000000 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/primitive.rs +++ /dev/null @@ -1,471 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use num_traits::AsPrimitive; -use num_traits::ToPrimitive; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::AllOr; - -use super::SumState; -use super::checked_add_i64; -use super::checked_add_u64; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::match_each_native_ptype; - -/// Number of elements summed without an overflow check. Chosen so that a chunk of values narrower -/// 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( - inner: &mut SumState, - p: &PrimitiveArray, - ctx: &mut ExecutionCtx, - skip_nans: bool, - seen: &mut bool, -) -> VortexResult { - let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; - match mask.slices() { - AllOr::None => Ok(false), - AllOr::All => { - *seen |= !p.as_ref().is_empty(); - accumulate_primitive_all(inner, p, skip_nans) - } - AllOr::Some(slices) => { - *seen |= !slices.is_empty(); - accumulate_primitive_valid(inner, p, slices, skip_nans) - } - } -} - -fn accumulate_primitive_all( - inner: &mut SumState, - p: &PrimitiveArray, - skip_nans: bool, -) -> VortexResult { - match inner { - SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |T| { Ok(sum_unsigned_all(acc, p.as_slice::())) }, - signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, - floating: |_T| { vortex_panic!("unsigned sum state with float input") } - ), - SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, - signed: |T| { Ok(sum_signed_all(acc, p.as_slice::())) }, - floating: |_T| { vortex_panic!("signed sum state with float input") } - ), - SumState::Float(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, - signed: |_T| { vortex_panic!("float sum state with signed input") }, - floating: |T| { - sum_float_all(acc, p.as_slice::(), skip_nans); - Ok(false) - } - ), - SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), - } -} - -/// StandardSum 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) { - if skip_nans { - for &v in slice { - if !v.is_nan() { - *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); - } - } - } else { - for &v in slice { - *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); - } - } -} - -/// StandardSum all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in -/// 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 -where - T: NativePType + AsPrimitive, -{ - if T::PTYPE == PType::U64 { - for &v in slice { - if checked_add_u64(acc, v.as_()) { - return true; - } - } - return false; - } - for chunk in slice.chunks(SUM_CHUNK) { - let chunk_sum: u64 = chunk.iter().map(|&v| v.as_()).sum(); - if checked_add_u64(acc, chunk_sum) { - return true; - } - } - false -} - -/// Signed counterpart of [`sum_unsigned_all`]. -pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool -where - T: NativePType + AsPrimitive, -{ - if T::PTYPE == PType::I64 { - for &v in slice { - if checked_add_i64(acc, v.as_()) { - return true; - } - } - return false; - } - for chunk in slice.chunks(SUM_CHUNK) { - let chunk_sum: i64 = chunk.iter().map(|&v| v.as_()).sum(); - if checked_add_i64(acc, chunk_sum) { - return true; - } - } - false -} - -/// StandardSum the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each -/// run is a slice of fully-valid values, so it reuses the same vectorized reduction as the -/// all-valid path instead of a per-element validity branch. -fn accumulate_primitive_valid( - inner: &mut SumState, - p: &PrimitiveArray, - slices: &[(usize, usize)], - skip_nans: bool, -) -> VortexResult { - match inner { - SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |T| { - let values = p.as_slice::(); - for &(start, end) in slices { - if sum_unsigned_all(acc, &values[start..end]) { - return Ok(true); - } - } - Ok(false) - }, - signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, - floating: |_T| { vortex_panic!("unsigned sum state with float input") } - ), - SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, - signed: |T| { - let values = p.as_slice::(); - for &(start, end) in slices { - if sum_signed_all(acc, &values[start..end]) { - return Ok(true); - } - } - Ok(false) - }, - floating: |_T| { vortex_panic!("signed sum state with float input") } - ), - SumState::Float(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, - signed: |_T| { vortex_panic!("float sum state with signed input") }, - floating: |T| { - let values = p.as_slice::(); - for &(start, end) in slices { - sum_float_all(acc, &values[start..end], skip_nans); - } - Ok(false) - } - ), - SumState::Decimal { .. } => vortex_panic!("decimal sum state with primitive input"), - } -} - -#[cfg(test)] -mod tests { - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - 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::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - 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; - - #[test] - fn sum_i32() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(10)); - Ok(()) - } - - #[test] - fn sum_u8() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(60)); - Ok(()) - } - - #[test] - fn sum_f64() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); - Ok(()) - } - - #[test] - fn sum_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_multiple_null_runs() -> VortexResult<()> { - // Several disjoint valid runs separated by nulls exercise the per-run fold. - let arr = PrimitiveArray::from_option_iter([ - Some(1i32), - Some(2), - None, - None, - Some(3), - None, - Some(4), - Some(5), - Some(6), - ]) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(21)); - Ok(()) - } - - #[test] - fn sum_all_null() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_all_invalid_float() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); - Ok(()) - } - - #[test] - fn sum_buffer_i32() -> VortexResult<()> { - let arr = buffer![1, 1, 1, 1].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4)); - Ok(()) - } - - #[test] - fn sum_buffer_f64() -> VortexResult<()> { - let arr = buffer![1., 1., 1., 1.].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4.)); - Ok(()) - } - - #[test] - fn sum_empty_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - // SQL `SUM`: the sum over no values is null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_empty_f64_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_f64_with_nan() -> VortexResult<()> { - let arr = PrimitiveArray::new( - buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], - Validity::NonNullable, - ) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - Ok(()) - } - - #[test] - fn sum_f32_with_nan() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); - Ok(()) - } - - #[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_all_nan() -> VortexResult<()> { - 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(()) - } - - /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). - fn sum_with_options( - arr: &crate::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() - } - - #[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_f64_without_nan_not_skipping() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { - // With an exact NaNCount of zero, the planted exact StandardSum 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 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(()) - } - - #[test] - fn sum_checked_overflow() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_checked_overflow_is_saturated() -> 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()); - - // finish resets state, clearing saturation - drop(acc.finish()?); - assert!(!acc.is_saturated()); - 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..3008763ebc9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -13,16 +13,18 @@ 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, + seen: &mut bool, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, 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..62193f43984 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -23,12 +23,14 @@ 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, + seen: &mut bool, ) -> VortexResult { let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + *seen |= mask.true_count() > 0; let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), diff --git a/vortex-array/src/aggregate_fn/fns/sum/design.md b/vortex-array/src/aggregate_fn/fns/sum/design.md new file mode 100644 index 00000000000..00b75228689 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum/design.md @@ -0,0 +1,140 @@ +# SQL semantics for the `Sum` aggregate + +## Goal + +`Sum` of zero valid values yields **null** (SQL `SUM`, ISO 9075-2 §10.9), everywhere: +scalar `sum()`, grouped aggregation, and `list_sum` — replacing today's contract of +"all-invalid sums to zero". Overflow keeps yielding null (unchanged). This matches +arrow-rs (`sum` returns `None` for empty/all-null), DuckDB (`SumState.isset` → +`ReturnNull`), and DataFusion (`SumAccumulator { sum: Option }`). + +## Why this is not a one-line change + +The current partial state is a single nullable sum value, and **null is already taken**: +a null partial means overflow and must *poison* merges (`combine(null, x) = null`), while +"empty" must be the *identity* (`combine(empty, x) = x`). One symbol cannot be both, and +the partial value crosses two persistence boundaries where merging happens: + +1. **Array/layout statistics**: `Stat::Sum` is serialized as a dedicated flatbuffer field + (`vortex-array/src/stats/flatbuffers.rs`) and merged as final scalars via + `StatsSet::merge_sum` (`stats/stats_set.rs`, `checked_add`). +2. **Zoned layouts**: `vortex-layout/src/layouts/zoned/builder.rs` writes + `accumulator.partial_scalar()` per zone; readers fold these back through + `combine_partials` (see the "read legacy stat" path in + `vortex-array/src/aggregate_fn/accumulator.rs`). + +Additionally, `Accumulator::accumulate` short-circuits through cached `Stat::Sum` values +(`sum/mod.rs`, `try_accumulate`), so a batch's cached monoid sum can substitute for +actually reading the batch. Any design must keep these paths correct and mutually +consistent. + +## Design options + +### A. Struct partials (recommended) + +Change `SumPartial`'s wire form to `Struct { sum: , seen: bool }`, +where `seen` records whether at least one valid value contributed (DuckDB's `isset`; +`sum = null` still means overflow). A valid-*count* would be redundant — `NullCount` is +already persisted alongside the sum in every stats pipeline, so any boundary that needs +a count derives it from existing stats — and the bool keeps the grouped kernel's partial +column bit-packed (1 bit per group instead of 64). + +The algebra becomes unambiguous and total: + +- identity: `{sum: 0, seen: false}` +- combine: `{sum: a.sum ⊕ b.sum, seen: a.seen ∨ b.seen}` where `⊕` is + null-poisoning checked add +- finalize: `!seen → null`, else `sum` (null if overflow) +- checks: `combine(overflow, empty) = {null, true} → null`; + `combine(empty, empty) = {0, false} → null`; `combine(v, empty) = {v, true} → v` + +(If the field is named `all_null` instead, the combine flips to AND; `seen`/OR avoids +the double negative.) + +Consequences: + +- **Grouped machinery needs no special-casing at all.** The per-group fallback + accumulator naturally produces `seen = false` for empty/all-null groups; `finalize` + nulls them. The grouped sum kernel emits a struct column whose `seen` field is + exactly the has-valid-element bitmap it already computes (via + `BitBuffer::count_range` over the materialized element mask). + `null_for_empty_groups`, `mask_empty_lists`, and every fix-up variant dissolve. +- **`list_sum_impl` becomes accumulate + finish**, no post-pass. +- **Mean** (`Combined`): its empty result changes from `NaN` to `null` + (SQL-correct; matches DuckDB and Postgres). Its `finalize_scalar` comment and tests + must be updated deliberately. +- **Persistence formats change**, which is the cost: + - zoned stats: new files write struct partials; the existing legacy-stat shim in + `accumulator.rs` (which already casts old stat dtypes) grows one rule: a legacy + plain-primitive partial is read as `{sum: v, seen: true}` (preserves old behavior + for old files, including treating legacy null as overflow-poison). + - array stats flatbuffer: either keep the `sum` field as the *finalized* value and add + a sibling field/valid-count consultation, or version the field. **This needs a + format-owner decision** — the one open question to settle before implementing. + Note `NullCount` is already persisted alongside, so "empty" is derivable at read + time as `null_count == len` without any new field, which may make the flatbuffer + change unnecessary: keep `Stat::Sum` as the monoid scalar on disk and apply + `finalize` semantics at the read boundary using `NullCount`. +- Forward-compat: old readers encountering new zoned struct partials will fail the + partial-dtype ensure and should degrade to "stat unavailable" rather than error — + verify and test. + +### B. In-memory `seen` bit only (fallback option, no format change) + +`SumPartial` gains a non-serialized `seen: bool`; `finish`/`finalize_scalar` null when +unseen; persisted forms stay plain monoid sums; every stat-fed path (`try_accumulate` +short-circuit, zoned reads, `merge_sum`) marks `seen` conservatively and boundaries that +present SQL SUM consult `NullCount` (`null_count == row_count → null`). + +Cheaper, but the grouped fallback still can't express empty-group nulls through +`to_scalar` (partials must stay monoid for the zoned writer), so the grouped machinery +still needs an explicit empty-group mechanism — i.e. this degenerates into the +`null_for_empty_groups` design we already have stashed, plus boundary guards. Choose B +only if the format-owner conversation for A stalls. + +## Plan (assuming A) + +1. **Core algebra** (`aggregate_fn/fns/sum/mod.rs`) + - `SumPartial { sum: Option, seen: bool }`; update `empty_partial`, + `combine_partials`, `to_scalar`, `reset`, `is_saturated`, `finalize`, + `finalize_scalar`; `partial_dtype` → struct; `return_dtype` unchanged. + - `try_accumulate` stat short-circuit: derive `seen` from the batch's + `NullCount`/len when consuming a cached `Stat::Sum`; if unavailable, fall through + to real accumulation instead of guessing. + - Rustdoc: sum of zero valid values is null; overflow is null; NaN handling + unchanged (NaNs are valid values — an all-NaN sum under `skip_nans` is `0`, and + `list_sum`'s `test_all_nan_list_sums_to_zero` pins this). +2. **Grouped paths** (`accumulator_grouped.rs`, `fns/sum/grouped.rs`, `fns/count/grouped.rs`) + - Kernel emits `{sum, seen}` struct rows (`seen` via `BitBuffer::count_range > 0` + over the materialized element mask — never `Mask::slice` per group). + - Fallback: no changes beyond the struct builder working generically. + - `finalize(states)`: struct → nullable sums in one pass (validity intersection). +3. **`list_sum`** — delete `mask_empty_lists` and the grouped-view plumbing; the + existing 21 tests must pass unchanged (they pin the public semantics). +4. **Stats boundaries** + - Legacy shim in `accumulator.rs` for old zoned partials. + - `stats_set.rs::merge_sum`: decide monoid-on-disk vs struct (see A) and align. + - `vortex-datafusion/src/convert/stats.rs`: export `sum_value` as null/absent when + `null_count == row_count` — fixes the pre-existing wrong-results edge where + DataFusion could answer `SELECT SUM(x)` as `0` for an all-null column. + - `sum()`'s `Stat::Sum` cache read: apply the same `NullCount` guard. +5. **Mean** — accept null-for-empty (update `finalize_scalar` comment + tests). +6. **Tests** + - Algebra: combine across empty/value/overflow in all orders; struct round-trip. + - Grouped: empty/all-null groups null on kernel and fallback (bool elements) paths. + - Legacy: old-format zoned partial read shim; forward-compat degradation. + - End-to-end: file with an all-null zone → correct file-level sum; all-null file → + null; DataFusion stats export edge. + - Existing `list_sum` suite passes unmodified. +7. **Benches** — `aggregate_grouped` + `list_sum` A/B against develop (expect: nullable + grouped cases improve from deleting the fix-up; small kernel cost for the count + column; quantify). The stashed zero-fill kernel rewrite (stash@{0} on the list-sum + worktree) composes with this and can follow as its own PR. + +## Open questions (resolve before/while implementing) + +1. Flatbuffer `Stat::Sum` field: keep monoid-on-disk + `NullCount` at boundaries, or + version the field to carry count? (Recommend monoid-on-disk — no format change to + array stats; only zoned partials change shape, where a legacy shim already exists.) +2. Forward-compat policy for old readers on new zoned partials: skip-stat vs error. +3. Mean: confirm null-for-empty is wanted (SQL says yes). diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 9e71bc8edf2..a8cf52f165b 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 shared with `StandardSum`, which layers `seen` tracking on +// the same summation; [`Sum`] itself ignores the seen bit. +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; @@ -40,13 +45,9 @@ use crate::expr::stats::StatsProviderExt; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -/// The monoid sum of an array: zero when there are no valid values, null on overflow. +/// Return the sum of an array. /// -/// This is the value stored as `Stat::Sum` and merged across chunks, zones, and files: the -/// empty sum is zero so that partials merge as a monoid (`combine(empty, x) = x`). For the -/// SQL rule (a sum over zero valid values is null) use -/// [`standard_sum`](crate::aggregate_fn::fns::standard_sum::standard_sum) / -/// [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum). +/// See [`Sum`] for details. pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // Short-circuit using cached array statistics. if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { @@ -71,15 +72,15 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// The monoid sum aggregate: an all-invalid or empty sum is zero, an overflowing sum is null. +/// Sum an array, starting from zero. /// -/// This aggregate's plain partial is the persisted wire form of sum statistics (zone maps, -/// file stats), so its partial dtype and semantics must stay stable. The SQL `SUM` rule is -/// provided by the separate [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum) -/// aggregate. +/// 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. 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) /// -/// 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. +/// For a sum aggregate with more standard behavior, see [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum). #[derive(Clone, Debug)] pub struct Sum; @@ -304,11 +305,15 @@ impl AggregateFnVTable for Sum { None => return Ok(()), }; + // The shared kernels track a `seen` bit for `StandardSum`; the monoid sum ignores it. + let mut seen = false; 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), + Canonical::Primitive(p) => { + accumulate_primitive(&mut inner, p, ctx, skip_nans, &mut seen) + } + Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut seen), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut seen), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..c44b0b7f566 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -21,17 +21,24 @@ 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, skip_nans: bool, + seen: &mut bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => accumulate_primitive_all(inner, p, skip_nans), - AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), + AllOr::All => { + *seen |= !p.as_ref().is_empty(); + accumulate_primitive_all(inner, p, skip_nans) + } + AllOr::Some(slices) => { + *seen |= !slices.is_empty(); + accumulate_primitive_valid(inner, p, slices, skip_nans) + } } } @@ -66,7 +73,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 +91,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 +113,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, { From 5fdc4314f46aef7e15c29f945c7eeda2750c2152 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 18:34:30 -0700 Subject: [PATCH 09/12] Keep Sum kernels unchanged; StandardSum tracks seen from validity in accumulate Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/standard_sum/mod.rs | 21 ++- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 2 - .../src/aggregate_fn/fns/sum/decimal.rs | 2 - .../src/aggregate_fn/fns/sum/design.md | 140 ------------------ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 14 +- .../src/aggregate_fn/fns/sum/primitive.rs | 11 +- 6 files changed, 25 insertions(+), 165 deletions(-) delete mode 100644 vortex-array/src/aggregate_fn/fns/sum/design.md diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 749ad224d27..367bc074560 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -46,6 +46,7 @@ use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::mask::Mask; +use crate::validity::Validity; /// Return the SQL sum of an array: null when the array has no valid values or the sum /// overflows. @@ -339,13 +340,22 @@ impl AggregateFnVTable for StandardSum { None => return Ok(()), }; + // `seen` is decided by validity alone (NaNs are valid values), so it is tracked here + // and 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, &mut partial.seen) + partial.seen |= any_valid(p.as_ref().validity()?, p.as_ref().len(), ctx)?; + accumulate_primitive(&mut inner, p, ctx, skip_nans) + } + Canonical::Bool(b) => { + partial.seen |= any_valid(b.as_ref().validity()?, b.as_ref().len(), ctx)?; + accumulate_bool(&mut inner, b, ctx) + } + Canonical::Decimal(d) => { + partial.seen |= any_valid(d.as_ref().validity()?, d.as_ref().len(), ctx)?; + accumulate_decimal(&mut inner, d, ctx) } - Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut partial.seen), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut partial.seen), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), @@ -420,6 +430,11 @@ fn try_accumulate_cached_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 { diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index 3008763ebc9..f15356caaa1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -17,14 +17,12 @@ pub(crate) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, - seen: &mut bool, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 62193f43984..197b8f10b04 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -27,10 +27,8 @@ pub(crate) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, - seen: &mut bool, ) -> VortexResult { let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - *seen |= mask.true_count() > 0; let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), diff --git a/vortex-array/src/aggregate_fn/fns/sum/design.md b/vortex-array/src/aggregate_fn/fns/sum/design.md deleted file mode 100644 index 00b75228689..00000000000 --- a/vortex-array/src/aggregate_fn/fns/sum/design.md +++ /dev/null @@ -1,140 +0,0 @@ -# SQL semantics for the `Sum` aggregate - -## Goal - -`Sum` of zero valid values yields **null** (SQL `SUM`, ISO 9075-2 §10.9), everywhere: -scalar `sum()`, grouped aggregation, and `list_sum` — replacing today's contract of -"all-invalid sums to zero". Overflow keeps yielding null (unchanged). This matches -arrow-rs (`sum` returns `None` for empty/all-null), DuckDB (`SumState.isset` → -`ReturnNull`), and DataFusion (`SumAccumulator { sum: Option }`). - -## Why this is not a one-line change - -The current partial state is a single nullable sum value, and **null is already taken**: -a null partial means overflow and must *poison* merges (`combine(null, x) = null`), while -"empty" must be the *identity* (`combine(empty, x) = x`). One symbol cannot be both, and -the partial value crosses two persistence boundaries where merging happens: - -1. **Array/layout statistics**: `Stat::Sum` is serialized as a dedicated flatbuffer field - (`vortex-array/src/stats/flatbuffers.rs`) and merged as final scalars via - `StatsSet::merge_sum` (`stats/stats_set.rs`, `checked_add`). -2. **Zoned layouts**: `vortex-layout/src/layouts/zoned/builder.rs` writes - `accumulator.partial_scalar()` per zone; readers fold these back through - `combine_partials` (see the "read legacy stat" path in - `vortex-array/src/aggregate_fn/accumulator.rs`). - -Additionally, `Accumulator::accumulate` short-circuits through cached `Stat::Sum` values -(`sum/mod.rs`, `try_accumulate`), so a batch's cached monoid sum can substitute for -actually reading the batch. Any design must keep these paths correct and mutually -consistent. - -## Design options - -### A. Struct partials (recommended) - -Change `SumPartial`'s wire form to `Struct { sum: , seen: bool }`, -where `seen` records whether at least one valid value contributed (DuckDB's `isset`; -`sum = null` still means overflow). A valid-*count* would be redundant — `NullCount` is -already persisted alongside the sum in every stats pipeline, so any boundary that needs -a count derives it from existing stats — and the bool keeps the grouped kernel's partial -column bit-packed (1 bit per group instead of 64). - -The algebra becomes unambiguous and total: - -- identity: `{sum: 0, seen: false}` -- combine: `{sum: a.sum ⊕ b.sum, seen: a.seen ∨ b.seen}` where `⊕` is - null-poisoning checked add -- finalize: `!seen → null`, else `sum` (null if overflow) -- checks: `combine(overflow, empty) = {null, true} → null`; - `combine(empty, empty) = {0, false} → null`; `combine(v, empty) = {v, true} → v` - -(If the field is named `all_null` instead, the combine flips to AND; `seen`/OR avoids -the double negative.) - -Consequences: - -- **Grouped machinery needs no special-casing at all.** The per-group fallback - accumulator naturally produces `seen = false` for empty/all-null groups; `finalize` - nulls them. The grouped sum kernel emits a struct column whose `seen` field is - exactly the has-valid-element bitmap it already computes (via - `BitBuffer::count_range` over the materialized element mask). - `null_for_empty_groups`, `mask_empty_lists`, and every fix-up variant dissolve. -- **`list_sum_impl` becomes accumulate + finish**, no post-pass. -- **Mean** (`Combined`): its empty result changes from `NaN` to `null` - (SQL-correct; matches DuckDB and Postgres). Its `finalize_scalar` comment and tests - must be updated deliberately. -- **Persistence formats change**, which is the cost: - - zoned stats: new files write struct partials; the existing legacy-stat shim in - `accumulator.rs` (which already casts old stat dtypes) grows one rule: a legacy - plain-primitive partial is read as `{sum: v, seen: true}` (preserves old behavior - for old files, including treating legacy null as overflow-poison). - - array stats flatbuffer: either keep the `sum` field as the *finalized* value and add - a sibling field/valid-count consultation, or version the field. **This needs a - format-owner decision** — the one open question to settle before implementing. - Note `NullCount` is already persisted alongside, so "empty" is derivable at read - time as `null_count == len` without any new field, which may make the flatbuffer - change unnecessary: keep `Stat::Sum` as the monoid scalar on disk and apply - `finalize` semantics at the read boundary using `NullCount`. -- Forward-compat: old readers encountering new zoned struct partials will fail the - partial-dtype ensure and should degrade to "stat unavailable" rather than error — - verify and test. - -### B. In-memory `seen` bit only (fallback option, no format change) - -`SumPartial` gains a non-serialized `seen: bool`; `finish`/`finalize_scalar` null when -unseen; persisted forms stay plain monoid sums; every stat-fed path (`try_accumulate` -short-circuit, zoned reads, `merge_sum`) marks `seen` conservatively and boundaries that -present SQL SUM consult `NullCount` (`null_count == row_count → null`). - -Cheaper, but the grouped fallback still can't express empty-group nulls through -`to_scalar` (partials must stay monoid for the zoned writer), so the grouped machinery -still needs an explicit empty-group mechanism — i.e. this degenerates into the -`null_for_empty_groups` design we already have stashed, plus boundary guards. Choose B -only if the format-owner conversation for A stalls. - -## Plan (assuming A) - -1. **Core algebra** (`aggregate_fn/fns/sum/mod.rs`) - - `SumPartial { sum: Option, seen: bool }`; update `empty_partial`, - `combine_partials`, `to_scalar`, `reset`, `is_saturated`, `finalize`, - `finalize_scalar`; `partial_dtype` → struct; `return_dtype` unchanged. - - `try_accumulate` stat short-circuit: derive `seen` from the batch's - `NullCount`/len when consuming a cached `Stat::Sum`; if unavailable, fall through - to real accumulation instead of guessing. - - Rustdoc: sum of zero valid values is null; overflow is null; NaN handling - unchanged (NaNs are valid values — an all-NaN sum under `skip_nans` is `0`, and - `list_sum`'s `test_all_nan_list_sums_to_zero` pins this). -2. **Grouped paths** (`accumulator_grouped.rs`, `fns/sum/grouped.rs`, `fns/count/grouped.rs`) - - Kernel emits `{sum, seen}` struct rows (`seen` via `BitBuffer::count_range > 0` - over the materialized element mask — never `Mask::slice` per group). - - Fallback: no changes beyond the struct builder working generically. - - `finalize(states)`: struct → nullable sums in one pass (validity intersection). -3. **`list_sum`** — delete `mask_empty_lists` and the grouped-view plumbing; the - existing 21 tests must pass unchanged (they pin the public semantics). -4. **Stats boundaries** - - Legacy shim in `accumulator.rs` for old zoned partials. - - `stats_set.rs::merge_sum`: decide monoid-on-disk vs struct (see A) and align. - - `vortex-datafusion/src/convert/stats.rs`: export `sum_value` as null/absent when - `null_count == row_count` — fixes the pre-existing wrong-results edge where - DataFusion could answer `SELECT SUM(x)` as `0` for an all-null column. - - `sum()`'s `Stat::Sum` cache read: apply the same `NullCount` guard. -5. **Mean** — accept null-for-empty (update `finalize_scalar` comment + tests). -6. **Tests** - - Algebra: combine across empty/value/overflow in all orders; struct round-trip. - - Grouped: empty/all-null groups null on kernel and fallback (bool elements) paths. - - Legacy: old-format zoned partial read shim; forward-compat degradation. - - End-to-end: file with an all-null zone → correct file-level sum; all-null file → - null; DataFusion stats export edge. - - Existing `list_sum` suite passes unmodified. -7. **Benches** — `aggregate_grouped` + `list_sum` A/B against develop (expect: nullable - grouped cases improve from deleting the fix-up; small kernel cost for the count - column; quantify). The stashed zero-fill kernel rewrite (stash@{0} on the list-sum - worktree) composes with this and can follow as its own PR. - -## Open questions (resolve before/while implementing) - -1. Flatbuffer `Stat::Sum` field: keep monoid-on-disk + `NullCount` at boundaries, or - version the field to carry count? (Recommend monoid-on-disk — no format change to - array stats; only zoned partials change shape, where a legacy shim already exists.) -2. Forward-compat policy for old readers on new zoned partials: skip-stat vs error. -3. Mean: confirm null-for-empty is wanted (SQL says yes). diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index a8cf52f165b..360e0600885 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -15,8 +15,8 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -// The accumulation kernels are shared with `StandardSum`, which layers `seen` tracking on -// the same summation; [`Sum`] itself ignores the seen bit. +// 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; @@ -305,15 +305,11 @@ impl AggregateFnVTable for Sum { None => return Ok(()), }; - // The shared kernels track a `seen` bit for `StandardSum`; the monoid sum ignores it. - let mut seen = false; let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => { - accumulate_primitive(&mut inner, p, ctx, skip_nans, &mut seen) - } - Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx, &mut seen), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx, &mut seen), + 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!(), diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index c44b0b7f566..930a353f6e1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -26,19 +26,12 @@ pub(crate) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, skip_nans: bool, - seen: &mut bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => { - *seen |= !p.as_ref().is_empty(); - accumulate_primitive_all(inner, p, skip_nans) - } - AllOr::Some(slices) => { - *seen |= !slices.is_empty(); - accumulate_primitive_valid(inner, p, slices, skip_nans) - } + AllOr::All => accumulate_primitive_all(inner, p, skip_nans), + AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), } } From bce8e7d117a1e37dbe38eb4d0772e41d412217f0 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 19:22:17 -0700 Subject: [PATCH 10/12] Move StandardSum tests to tests.rs; simplify docs Signed-off-by: Matt Katz --- .../aggregate_fn/fns/standard_sum/grouped.rs | 6 +- .../src/aggregate_fn/fns/standard_sum/mod.rs | 1576 +---------------- .../aggregate_fn/fns/standard_sum/tests.rs | 1504 ++++++++++++++++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 8 +- 4 files changed, 1527 insertions(+), 1567 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs index 450b962f5e5..1e18392745d 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -56,7 +56,7 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedStandardSumEncodingKernel { /// /// Produces `{sum, seen}` partial rows (see `StandardSum`'s partial dtype): `seen` records whether the /// group contained at least one valid element, so that `finalize` yields null for empty and -/// all-null groups (SQL `SUM`), and null groups are null struct rows. +/// all-null groups, and null groups are null struct rows. pub(super) fn try_grouped_sum( groups: &GroupedArray, ctx: &mut ExecutionCtx, @@ -220,8 +220,7 @@ mod tests { } /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] - /// (SQL semantics: null for zero valid elements) for valid groups, a null sum for invalid - /// groups. + /// for valid groups, a null sum for invalid groups. fn grouped_sum_reference( elements: &ArrayRef, ranges: &[(usize, usize)], @@ -323,7 +322,6 @@ mod tests { let actual = grouped_sum_actual(&groups, &elem_dtype)?; let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - // The all-null and empty groups have null sums (SQL `SUM` semantics). 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); diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 367bc074560..6648a454a04 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -48,14 +48,14 @@ use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::mask::Mask; use crate::validity::Validity; -/// Return the SQL sum of an array: null when the array has no valid values or the sum +/// 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 the monoid sum (zero for an - // array with no valid values), so the SQL empty-sum rule needs the null count: when it is - // unknown for a nullable array, fall through and compute rather than trust the cache. + // 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); @@ -69,8 +69,7 @@ pub fn standard_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult. - // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe. + // Compute using Accumulator let mut acc = Accumulator::try_new( StandardSum, NumericalAggregateOpts::default(), @@ -88,22 +87,12 @@ pub fn standard_sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult<()> { - // Partials are `{sum, seen}` structs. A plain (non-struct) scalar is the legacy or - // statistic form: a monoid sum value from an older writer or a cached `Stat::Sum`. It - // cannot distinguish an empty sum from a zero sum, so it is treated as seen. + // Partials are `{sum, seen}` structs. A plain scalar is a cached `Stat::Sum` value, + // which cannot distinguish an empty sum from a zero sum, so it is treated as seen. let (sum_value, other_seen) = if matches!(other.dtype(), DType::Struct(..)) { if other.is_null() { // A null struct partial carries no recoverable state; treat as saturated. @@ -284,9 +272,8 @@ impl AggregateFnVTable for StandardSum { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - // `Stat::Sum` is the monoid (NaN-skipping) sum, so the default NaN-skipping path can - // consume it directly. `StandardSum` has no Stat slot of its own — the shortcut lives here - // rather than in the accumulator's stats bridge. + // `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); } @@ -373,7 +360,7 @@ impl AggregateFnVTable for StandardSum { } fn finalize(&self, partials: ArrayRef) -> VortexResult { - // Entries that saw no valid values finalize to null (SQL `SUM`), while a null `sum` + // Entries that saw no valid values finalize to null, while a null `sum` // field (overflow) and null partial rows (e.g. null groups) stay null via the mask's // validity intersection. // @@ -402,7 +389,7 @@ impl AggregateFnVTable for StandardSum { } } -/// Consume a batch's cached monoid `Stat::Sum` instead of scanning it. The cached sum cannot +/// Consume a batch's cached `Stat::Sum` instead of scanning it. The cached sum cannot /// carry `seen`, so the batch's null count decides between the identity (all null) and a seen /// contribution; when either statistic is missing the caller falls through to a real scan. fn try_accumulate_cached_sum( @@ -442,14 +429,14 @@ pub struct StandardSumPartial { /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, /// Whether at least one valid value has been accumulated. A sum over zero valid values - /// finalizes to null (SQL `SUM` semantics) rather than the monoid zero. + /// finalizes to null rather than zero. seen: 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, seen}` struct, where -/// `sum` is the running monoid value (null once saturated by overflow, poisoning merges) and +/// `sum` is the running sum, and is null once saturated by overflow, poisoning merges, and /// `seen` records whether any valid value contributed (merged with OR). Keeping the flag /// separate from the sum lets merges stay a monoid — the identity is `{0, false}` — while /// `finalize` maps unseen sums to null. @@ -481,1529 +468,4 @@ fn sum_value_scalar(partial: &StandardSumPartial) -> Scalar { } #[cfg(test)] -mod tests { - use num_traits::CheckedAdd; - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - 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 as monoid_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::expr::stats::StatsProvider; - use crate::scalar::DecimalValue; - use crate::scalar::NumericOperator; - use crate::scalar::Scalar; - use crate::validity::Validity; - - /// StandardSum an array with an initial value (test-only helper). - fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); - if accumulator.is_null() { - return Ok(accumulator.clone()); - } - if accumulator.is_zero() == Some(true) { - return standard_sum(array, &mut ctx); - } - - let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("StandardSum not supported for dtype: {}", array.dtype()) - })?; - - // For non-float types, try statistics short-circuit with accumulator. - if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float()) - && let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) - { - return add_scalars(&sum_dtype, &sum_scalar, accumulator); - } - - // Compute array sum from zero (also caches stats). - let array_sum = standard_sum(array, &mut ctx)?; - - // Combine with the accumulator. - add_scalars(&sum_dtype, &array_sum, accumulator) - } - - /// Add two sum scalars with overflow checking. - fn add_scalars(sum_dtype: &DType, lhs: &Scalar, rhs: &Scalar) -> VortexResult { - if lhs.is_null() || rhs.is_null() { - return Ok(Scalar::null(sum_dtype.as_nullable())); - } - - Ok(match sum_dtype { - DType::Primitive(ptype, _) if ptype.is_float() => { - let lhs_val = f64::try_from(lhs)?; - let rhs_val = f64::try_from(rhs)?; - Scalar::primitive(lhs_val + rhs_val, Nullable) - } - DType::Primitive(..) => lhs - .as_primitive() - .checked_add(&rhs.as_primitive()) - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - DType::Decimal(..) => lhs - .as_decimal() - .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - _ => unreachable!("StandardSum will always be a decimal or a primitive dtype"), - }) - } - - // Multi-batch and reset tests - - #[test] - fn sum_multi_batch() -> 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)); - Ok(()) - } - - #[test] - fn sum_finish_resets_state() -> 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 result1 = acc.finish()?; - assert_eq!(result1.as_primitive().typed_value::(), Some(30)); - - let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); - acc.accumulate(&batch2, &mut ctx)?; - let result2 = acc.finish()?; - assert_eq!(result2.as_primitive().typed_value::(), Some(18)); - Ok(()) - } - - // State merge tests (vtable-level) - - #[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)?; - 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 seen state changes nothing: `{0, false}` is the - // identity of the `{sum, seen}` 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_poisons_but_stays_seen() -> VortexResult<()> { - // Overflow (a null `sum` field) poisons the merge even when combined with later - // values: the result is null via the sum value, not via `seen`. - 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 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(()) - } - - #[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_monoid_while_standard_sum_is_sql() -> VortexResult<()> { - // The persisted statistic keeps the monoid semantics (zero for all-null) that zone - // and chunk merging require, while the SQL `sum` 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!( - monoid_sum(&arr, &mut ctx)? - .as_primitive() - .typed_value::(), - Some(0) - ); - // The cached monoid statistic must not leak through `sum`'s cache short-circuit. - assert!(standard_sum(&arr, &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 sum_state_merge() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; - - let scalar1 = Scalar::primitive(100i64, Nullable); - StandardSum.combine_partials(&mut state, scalar1)?; - - let scalar2 = Scalar::primitive(50i64, Nullable); - StandardSum.combine_partials(&mut state, scalar2)?; - - let result = StandardSum.finalize_scalar(&state)?; - StandardSum.reset(&mut state); - assert_eq!(result.as_primitive().typed_value::(), Some(150)); - Ok(()) - } - - // Stats caching test - - #[test] - fn sum_stats() -> VortexResult<()> { - let array = ChunkedArray::try_new( - vec![ - PrimitiveArray::from_iter([1, 1, 1]).into_array(), - PrimitiveArray::from_iter([2, 2, 2]).into_array(), - ], - DType::Primitive(PType::I32, Nullability::NonNullable), - ) - .vortex_expect("operation should succeed in test"); - let array = array.into_array(); - // compute sum with accumulator to populate stats - sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; - - let sum_without_acc = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); - Ok(()) - } - - // Constant float non-multiply test - - #[test] - fn sum_constant_float_non_multiply() -> VortexResult<()> { - let acc = -2048669276050936500000000000f64; - let array = ConstantArray::new(6.1811675e16f64, 25); - let result = sum_with_accumulator(&array.into_array(), &Scalar::primitive(acc, Nullable)) - .vortex_expect("operation should succeed in test"); - assert_eq!( - f64::try_from(&result).vortex_expect("operation should succeed in test"), - -2048669274505644600000000000f64 - ); - Ok(()) - } - - // Grouped sum tests - - 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_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<()> { - 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 (SQL `SUM` semantics). - let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut ctx); - Ok(()) - } - - #[test] - fn grouped_sum_bool() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); - let groups = - FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Bool(Nullability::NonNullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; - - let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).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(()) - } - - // Chunked array tests - - #[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 chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]); - 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(20.8)); - Ok(()) - } - - #[test] - fn sum_chunked_floats_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(), - )?; - // SQL `SUM`: no valid values across any chunk yields null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_chunked_floats_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_int_almost_all_null() -> VortexResult<()> { - 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(()) - } - - #[test] - fn sum_chunked_decimals() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let chunk1 = DecimalArray::new( - buffer![100i32, 100i32, 100i32, 100i32, 100i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk2 = DecimalArray::new( - buffer![200i32, 200i32, 200i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(1700))) - ); - Ok(()) - } - - #[test] - fn sum_chunked_decimals_with_nulls() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let chunk1 = DecimalArray::new( - buffer![100i32, 100i32, 100i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk2 = DecimalArray::new( - buffer![0i32, 0i32], - decimal_dtype, - Validity::from_iter([false, false]), - ); - let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(700))) - ); - Ok(()) - } - - #[test] - fn sum_chunked_decimals_large() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(3, 0); - let chunk1 = ConstantArray::new( - Scalar::decimal( - DecimalValue::I16(500), - decimal_dtype, - Nullability::NonNullable, - ), - 1, - ); - let chunk2 = ConstantArray::new( - Scalar::decimal( - DecimalValue::I16(600), - decimal_dtype, - Nullability::NonNullable, - ), - 1, - ); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(1100))) - ); - assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(13, 0), Nullable) - ); - Ok(()) - } - - mod bool_inputs { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - 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::BoolArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::executor::VortexSessionExecute; - - #[test] - fn sum_bool_all_true() -> VortexResult<()> { - let arr: BoolArray = [true, true, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_mixed() -> VortexResult<()> { - let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_all_false() -> VortexResult<()> { - let arr: BoolArray = [false, false, false].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); - 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_bool_all_null() -> VortexResult<()> { - let arr = BoolArray::from_iter([None::, None, None]); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_empty_is_null() -> VortexResult<()> { - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = - Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = - Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - - let batch1: BoolArray = [true, true, false].into_iter().collect(); - acc.accumulate(&batch1.into_array(), &mut ctx)?; - let result1 = acc.finish()?; - assert_eq!(result1.as_primitive().typed_value::(), Some(2)); - - let batch2: BoolArray = [false, true].into_iter().collect(); - acc.accumulate(&batch2.into_array(), &mut ctx)?; - let result2 = acc.finish()?; - assert_eq!(result2.as_primitive().typed_value::(), Some(1)); - Ok(()) - } - - #[test] - fn sum_bool_return_dtype() -> VortexResult<()> { - let dtype = StandardSum - .return_dtype( - &NumericalAggregateOpts::default(), - &DType::Bool(Nullability::NonNullable), - ) - .unwrap(); - assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); - Ok(()) - } - - #[test] - fn sum_boolean_from_iter() -> VortexResult<()> { - let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(2)); - Ok(()) - } - } - - mod constant_inputs { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::standard_sum::standard_sum; - use crate::array_session; - use crate::arrays::ConstantArray; - 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::Stat; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - - #[test] - fn sum_constant_unsigned() -> 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_signed() -> VortexResult<()> { - let array = ConstantArray::new(-5i64, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, (-50i64).into()); - Ok(()) - } - - #[test] - fn sum_constant_nullable_value() -> VortexResult<()> { - let array = - ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) - .into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: an all-null constant has no valid values. - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_bool_false() -> VortexResult<()> { - let array = ConstantArray::new(false, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 0u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_true() -> VortexResult<()> { - let array = ConstantArray::new(true, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 10u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_null() -> VortexResult<()> { - let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_decimal() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(100), - decimal_dtype, - Nullability::NonNullable, - ), - 5, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(500))) - ); - assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); - Ok(()) - } - - #[test] - fn sum_constant_decimal_null() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = - ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result, - Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) - ); - Ok(()) - } - - #[test] - fn sum_constant_decimal_large_value() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(999_999_999), - decimal_dtype, - Nullability::NonNullable, - ), - 100, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) - ); - Ok(()) - } - } - - mod decimal_inputs { - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::AggregateFnVTable; - 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::DecimalArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::i256; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - #[test] - fn sum_decimal_basic() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(600i32))), - )?; - - assert_eq!(result, expected); - 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_decimal_negative_values() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, -200i32, 300i32, -50i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(150i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_near_i32_max() -> VortexResult<()> { - let near_max = i32::MAX - 1000; - let decimal = DecimalArray::new( - buffer![near_max, 500i32, 400i32], - DecimalDType::new(10, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = near_max as i64 + 500 + 400; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_large_i64_values() -> VortexResult<()> { - let large_val = i64::MAX / 4; - let decimal = DecimalArray::new( - buffer![large_val, large_val, large_val, large_val + 1], - DecimalDType::new(19, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = (large_val as i128) * 4 + 1; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_preserves_scale() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![12345i32, 67890i32, 11111i32], - DecimalDType::new(6, 4), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(91346i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_single_value() -> VortexResult<()> { - let decimal = - DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(42i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32, 400i32], - DecimalDType::new(4, 2), - Validity::from_iter([false, false, true, false]), - ); - - 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(300i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_overflow_detection() -> VortexResult<()> { - let max_val = i128::MAX / 2; - let decimal = DecimalArray::new( - buffer![max_val, max_val, max_val], - DecimalDType::new(38, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = - i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - 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, - ); - - assert_eq!( - standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx() - ) - .vortex_expect("operation should succeed in test"), - 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(()) - } - - #[test] - fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { - // Input precision 4 → return precision 14. Native I64 (max 18). - // The max representable value for precision 14 is 10^14 - 1. - // When the sum reaches exactly 10^14, fits_in_precision fails 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(99_999_999_999_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - // Push the sum to exactly 10^14, exceeding precision 14. - let one_more = - Scalar::decimal(DecimalValue::from(1i64), 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_precision_overflow_negative() -> VortexResult<()> { - // Same setup but with negative values: sum reaches -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_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - let one_more = Scalar::decimal( - DecimalValue::from(-1i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, one_more)?; - - let result = StandardSum.finalize_scalar(&state)?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { - // Test precision overflow via the accumulate_decimal path (not combine_partials). - // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. - // Native for precision 38 is I128 (max 38), so 38 = 38. - // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. - // - // We 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); - - // Drive accumulate through the vtable directly. - 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(()) - } - } - - mod primitive_inputs { - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - 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::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - 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; - - #[test] - fn sum_i32() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(10)); - Ok(()) - } - - #[test] - fn sum_u8() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(60)); - Ok(()) - } - - #[test] - fn sum_f64() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); - Ok(()) - } - - #[test] - fn sum_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_multiple_null_runs() -> VortexResult<()> { - // Several disjoint valid runs separated by nulls exercise the per-run fold. - let arr = PrimitiveArray::from_option_iter([ - Some(1i32), - Some(2), - None, - None, - Some(3), - None, - Some(4), - Some(5), - Some(6), - ]) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(21)); - Ok(()) - } - - #[test] - fn sum_all_null() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - // SQL `SUM`: no valid values yields null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_all_invalid_float() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); - Ok(()) - } - - #[test] - fn sum_buffer_i32() -> VortexResult<()> { - let arr = buffer![1, 1, 1, 1].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4)); - Ok(()) - } - - #[test] - fn sum_buffer_f64() -> VortexResult<()> { - let arr = buffer![1., 1., 1., 1.].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4.)); - Ok(()) - } - - #[test] - fn sum_empty_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = - Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - // SQL `SUM`: the sum over no values is null. - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_empty_f64_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = - Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_f64_with_nan() -> VortexResult<()> { - let arr = PrimitiveArray::new( - buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], - Validity::NonNullable, - ) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - Ok(()) - } - - #[test] - fn sum_f32_with_nan() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); - Ok(()) - } - - #[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_all_nan() -> VortexResult<()> { - 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(()) - } - - /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). - fn sum_with_options( - arr: &crate::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() - } - - #[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_f64_without_nan_not_skipping() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { - // With an exact NaNCount of zero, the planted exact StandardSum 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 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(()) - } - - #[test] - fn sum_checked_overflow() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_checked_overflow_is_saturated() -> 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()); - - // finish resets state, clearing saturation - drop(acc.finish()?); - assert!(!acc.is_saturated()); - Ok(()) - } - } -} +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..8f6a1c19365 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs @@ -0,0 +1,1504 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::CheckedAdd; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +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::expr::stats::StatsProvider; +use crate::scalar::DecimalValue; +use crate::scalar::NumericOperator; +use crate::scalar::Scalar; +use crate::validity::Validity; + +/// StandardSum an array with an initial value (test-only helper). +fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + if accumulator.is_null() { + return Ok(accumulator.clone()); + } + if accumulator.is_zero() == Some(true) { + return standard_sum(array, &mut ctx); + } + + let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { + vortex_error::vortex_err!("StandardSum not supported for dtype: {}", array.dtype()) + })?; + + // For non-float types, try statistics short-circuit with accumulator. + if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float()) + && let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) + { + return add_scalars(&sum_dtype, &sum_scalar, accumulator); + } + + // Compute array sum from zero (also caches stats). + let array_sum = standard_sum(array, &mut ctx)?; + + // Combine with the accumulator. + add_scalars(&sum_dtype, &array_sum, accumulator) +} + +/// Add two sum scalars with overflow checking. +fn add_scalars(sum_dtype: &DType, lhs: &Scalar, rhs: &Scalar) -> VortexResult { + if lhs.is_null() || rhs.is_null() { + return Ok(Scalar::null(sum_dtype.as_nullable())); + } + + Ok(match sum_dtype { + DType::Primitive(ptype, _) if ptype.is_float() => { + let lhs_val = f64::try_from(lhs)?; + let rhs_val = f64::try_from(rhs)?; + Scalar::primitive(lhs_val + rhs_val, Nullable) + } + DType::Primitive(..) => lhs + .as_primitive() + .checked_add(&rhs.as_primitive()) + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), + DType::Decimal(..) => lhs + .as_decimal() + .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), + _ => unreachable!("StandardSum will always be a decimal or a primitive dtype"), + }) +} + +// Multi-batch and reset tests + +#[test] +fn sum_multi_batch() -> 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)); + Ok(()) +} + +#[test] +fn sum_finish_resets_state() -> 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 result1 = acc.finish()?; + assert_eq!(result1.as_primitive().typed_value::(), Some(30)); + + let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + let result2 = acc.finish()?; + assert_eq!(result2.as_primitive().typed_value::(), Some(18)); + Ok(()) +} + +// State merge tests (vtable-level) + +#[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)?; + 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 seen state changes nothing: `{0, false}` is the + // identity of the `{sum, seen}` 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_poisons_but_stays_seen() -> VortexResult<()> { + // Overflow (a null `sum` field) poisons the merge even when combined with later + // values: the result is null via the sum value, not via `seen`. + 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 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(()) +} + +#[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(()) +} + +#[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 sum_state_merge() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + + let scalar1 = Scalar::primitive(100i64, Nullable); + StandardSum.combine_partials(&mut state, scalar1)?; + + let scalar2 = Scalar::primitive(50i64, Nullable); + StandardSum.combine_partials(&mut state, scalar2)?; + + let result = StandardSum.finalize_scalar(&state)?; + StandardSum.reset(&mut state); + assert_eq!(result.as_primitive().typed_value::(), Some(150)); + Ok(()) +} + +// Stats caching test + +#[test] +fn sum_stats() -> VortexResult<()> { + let array = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([1, 1, 1]).into_array(), + PrimitiveArray::from_iter([2, 2, 2]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + ) + .vortex_expect("operation should succeed in test"); + let array = array.into_array(); + // compute sum with accumulator to populate stats + sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; + + let sum_without_acc = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); + Ok(()) +} + +// Constant float non-multiply test + +#[test] +fn sum_constant_float_non_multiply() -> VortexResult<()> { + let acc = -2048669276050936500000000000f64; + let array = ConstantArray::new(6.1811675e16f64, 25); + let result = sum_with_accumulator(&array.into_array(), &Scalar::primitive(acc, Nullable)) + .vortex_expect("operation should succeed in test"); + assert_eq!( + f64::try_from(&result).vortex_expect("operation should succeed in test"), + -2048669274505644600000000000f64 + ); + Ok(()) +} + +// Grouped sum tests + +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_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<()> { + 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_bool() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); + let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Bool(Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).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(()) +} + +// Chunked array tests + +#[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 chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]); + 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(20.8)); + Ok(()) +} + +#[test] +fn sum_chunked_floats_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_floats_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_int_almost_all_null() -> VortexResult<()> { + 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(()) +} + +#[test] +fn sum_chunked_decimals() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let chunk1 = DecimalArray::new( + buffer![100i32, 100i32, 100i32, 100i32, 100i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk2 = DecimalArray::new( + buffer![200i32, 200i32, 200i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid); + 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(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(1700))) + ); + Ok(()) +} + +#[test] +fn sum_chunked_decimals_with_nulls() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let chunk1 = DecimalArray::new( + buffer![100i32, 100i32, 100i32], + decimal_dtype, + Validity::AllValid, + ); + let chunk2 = DecimalArray::new( + buffer![0i32, 0i32], + decimal_dtype, + Validity::from_iter([false, false]), + ); + let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid); + 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(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(700))) + ); + Ok(()) +} + +#[test] +fn sum_chunked_decimals_large() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(3, 0); + let chunk1 = ConstantArray::new( + Scalar::decimal( + DecimalValue::I16(500), + decimal_dtype, + Nullability::NonNullable, + ), + 1, + ); + let chunk2 = ConstantArray::new( + Scalar::decimal( + DecimalValue::I16(600), + decimal_dtype, + Nullability::NonNullable, + ), + 1, + ); + 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(), + )?; + let decimal_result = result.as_decimal(); + assert_eq!( + decimal_result.decimal_value(), + Some(DecimalValue::I256(i256::from_i128(1100))) + ); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(13, 0), Nullable) + ); + Ok(()) +} + +mod bool_inputs { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + 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::BoolArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::executor::VortexSessionExecute; + + #[test] + fn sum_bool_all_true() -> VortexResult<()> { + let arr: BoolArray = [true, true, true].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_mixed() -> VortexResult<()> { + let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(3)); + Ok(()) + } + + #[test] + fn sum_bool_all_false() -> VortexResult<()> { + let arr: BoolArray = [false, false, false].into_iter().collect(); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + 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_bool_all_null() -> VortexResult<()> { + let arr = BoolArray::from_iter([None::, None, None]); + let result = standard_sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_bool_empty_is_null() -> VortexResult<()> { + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_bool_finish_resets_state() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Bool(Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + + let batch1: BoolArray = [true, true, false].into_iter().collect(); + acc.accumulate(&batch1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + assert_eq!(result1.as_primitive().typed_value::(), Some(2)); + + let batch2: BoolArray = [false, true].into_iter().collect(); + acc.accumulate(&batch2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + assert_eq!(result2.as_primitive().typed_value::(), Some(1)); + Ok(()) + } + + #[test] + fn sum_bool_return_dtype() -> VortexResult<()> { + let dtype = StandardSum + .return_dtype( + &NumericalAggregateOpts::default(), + &DType::Bool(Nullability::NonNullable), + ) + .unwrap(); + assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); + Ok(()) + } + + #[test] + fn sum_boolean_from_iter() -> VortexResult<()> { + let arr = BoolArray::from_iter([true, false, false, true]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(2)); + Ok(()) + } +} + +mod constant_inputs { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::fns::standard_sum::standard_sum; + use crate::array_session; + use crate::arrays::ConstantArray; + 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::Stat; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + + #[test] + fn sum_constant_unsigned() -> 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_signed() -> VortexResult<()> { + let array = ConstantArray::new(-5i64, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, (-50i64).into()); + Ok(()) + } + + #[test] + fn sum_constant_nullable_value() -> VortexResult<()> { + let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) + .into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + Ok(()) + } + + #[test] + fn sum_constant_bool_false() -> VortexResult<()> { + let array = ConstantArray::new(false, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 0u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_true() -> VortexResult<()> { + let array = ConstantArray::new(true, 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 10u64.into()); + Ok(()) + } + + #[test] + fn sum_constant_bool_null() -> VortexResult<()> { + let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); + Ok(()) + } + + #[test] + fn sum_constant_decimal() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(100), + decimal_dtype, + Nullability::NonNullable, + ), + 5, + ) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(500))) + ); + assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); + Ok(()) + } + + #[test] + fn sum_constant_decimal_null() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result, + Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) + ); + Ok(()) + } + + #[test] + fn sum_constant_decimal_large_value() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(10, 2); + let array = ConstantArray::new( + Scalar::decimal( + DecimalValue::I64(999_999_999), + decimal_dtype, + Nullability::NonNullable, + ), + 100, + ) + .into_array(); + + let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) + ); + Ok(()) + } +} + +mod decimal_inputs { + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateFnVTable; + 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::DecimalArray; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::dtype::Nullability::Nullable; + use crate::dtype::i256; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + use crate::validity::Validity; + + #[test] + fn sum_decimal_basic() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(600i32))), + )?; + + assert_eq!(result, expected); + 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_decimal_negative_values() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, -200i32, 300i32, -50i32], + DecimalDType::new(4, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(150i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_near_i32_max() -> VortexResult<()> { + let near_max = i32::MAX - 1000; + let decimal = DecimalArray::new( + buffer![near_max, 500i32, 400i32], + DecimalDType::new(10, 2), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = near_max as i64 + 500 + 400; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_large_i64_values() -> VortexResult<()> { + let large_val = i64::MAX / 4; + let decimal = DecimalArray::new( + buffer![large_val, large_val, large_val, large_val + 1], + DecimalDType::new(19, 0), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = (large_val as i128) * 4 + 1; + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_preserves_scale() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![12345i32, 67890i32, 11111i32], + DecimalDType::new(6, 4), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(91346i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_single_value() -> VortexResult<()> { + let decimal = + DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(42i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32, 400i32], + DecimalDType::new(4, 2), + Validity::from_iter([false, false, true, false]), + ); + + 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(300i32))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_overflow_detection() -> VortexResult<()> { + let max_val = i128::MAX / 2; + let decimal = DecimalArray::new( + buffer![max_val, max_val, max_val], + DecimalDType::new(38, 0), + Validity::AllValid, + ); + + let result = standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + + let expected_sum = + i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), + Some(ScalarValue::from(DecimalValue::from(expected_sum))), + )?; + + assert_eq!(result, expected); + 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, + ); + + assert_eq!( + standard_sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx() + ) + .vortex_expect("operation should succeed in test"), + 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(()) + } + + #[test] + fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { + // Input precision 4 → return precision 14. Native I64 (max 18). + // The max representable value for precision 14 is 10^14 - 1. + // When the sum reaches exactly 10^14, fits_in_precision fails 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(99_999_999_999_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + // Push the sum to exactly 10^14, exceeding precision 14. + let one_more = + Scalar::decimal(DecimalValue::from(1i64), 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_precision_overflow_negative() -> VortexResult<()> { + // Same setup but with negative values: sum reaches -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_999i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, near_limit)?; + + let one_more = Scalar::decimal( + DecimalValue::from(-1i64), + DecimalDType::new(14, 0), + Nullable, + ); + StandardSum.combine_partials(&mut state, one_more)?; + + let result = StandardSum.finalize_scalar(&state)?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { + // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. + // Native for precision 38 is I128 (max 38), so 38 = 38. + // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. + // + // We 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); + + // Drive accumulate through the vtable directly. + 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(()) + } +} + +mod primitive_inputs { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::DynAccumulator; + 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::ConstantArray; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + 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; + + #[test] + fn sum_i32() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(10)); + Ok(()) + } + + #[test] + fn sum_u8() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(60)); + Ok(()) + } + + #[test] + fn sum_f64() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); + Ok(()) + } + + #[test] + fn sum_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_multiple_null_runs() -> VortexResult<()> { + // Several disjoint valid runs separated by nulls exercise the per-run fold. + let arr = PrimitiveArray::from_option_iter([ + Some(1i32), + Some(2), + None, + None, + Some(3), + None, + Some(4), + Some(5), + Some(6), + ]) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(21)); + Ok(()) + } + + #[test] + fn sum_all_null() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_all_invalid_float() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); + Ok(()) + } + + #[test] + fn sum_buffer_i32() -> VortexResult<()> { + let arr = buffer![1, 1, 1, 1].into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4)); + Ok(()) + } + + #[test] + fn sum_buffer_f64() -> VortexResult<()> { + let arr = buffer![1., 1., 1., 1.].into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().as_::(), Some(4.)); + Ok(()) + } + + #[test] + fn sum_empty_is_null() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_empty_f64_is_null() -> VortexResult<()> { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; + let result = acc.finish()?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_f64_with_nan() -> VortexResult<()> { + let arr = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], + Validity::NonNullable, + ) + .into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + Ok(()) + } + + #[test] + fn sum_f32_with_nan() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); + Ok(()) + } + + #[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_all_nan() -> VortexResult<()> { + 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(()) + } + + /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + fn sum_with_options( + arr: &crate::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() + } + + #[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_f64_without_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact StandardSum 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 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(()) + } + + #[test] + fn sum_checked_overflow() -> VortexResult<()> { + let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); + let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; + assert!(result.is_null()); + Ok(()) + } + + #[test] + fn sum_checked_overflow_is_saturated() -> 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()); + + // finish resets state, clearing saturation + drop(acc.finish()?); + assert!(!acc.is_saturated()); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 360e0600885..2a6b0303d96 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -75,12 +75,8 @@ 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 or empty, the sum will be zero. -/// Note that sum aggregates typically produce null for arrays without at least one valid element. 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) -/// -/// For a sum aggregate with more standard behavior, see [`StandardSum`](crate::aggregate_fn::fns::standard_sum::StandardSum). +/// 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; From 420719a78564498ef0476685acb034a38419e537 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 20:02:35 -0700 Subject: [PATCH 11/12] Consolidate StandardSum tests: parameterize null-for-empty matrix, drop arithmetic duplicated by shared-kernel tests Signed-off-by: Matt Katz --- .../aggregate_fn/fns/standard_sum/tests.rs | 1673 +++++------------ 1 file changed, 432 insertions(+), 1241 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs index 8f6a1c19365..0c6a0856343 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs @@ -1,9 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::CheckedAdd; +//! Tests for [`StandardSum`]-specific behavior: the `{sum, seen}` state algebra, the +//! null-for-zero-valid-values rule across input kinds, NaN and overflow interplay with +//! `seen`, 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::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; @@ -35,104 +39,19 @@ use crate::dtype::PType; use crate::dtype::i256; use crate::expr::stats::Precision; use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; use crate::scalar::DecimalValue; -use crate::scalar::NumericOperator; use crate::scalar::Scalar; +use crate::scalar::ScalarValue; use crate::validity::Validity; -/// StandardSum an array with an initial value (test-only helper). -fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); - if accumulator.is_null() { - return Ok(accumulator.clone()); - } - if accumulator.is_zero() == Some(true) { - return standard_sum(array, &mut ctx); - } - - let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("StandardSum not supported for dtype: {}", array.dtype()) - })?; - - // For non-float types, try statistics short-circuit with accumulator. - if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float()) - && let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) - { - return add_scalars(&sum_dtype, &sum_scalar, accumulator); - } - - // Compute array sum from zero (also caches stats). - let array_sum = standard_sum(array, &mut ctx)?; - - // Combine with the accumulator. - add_scalars(&sum_dtype, &array_sum, accumulator) -} - -/// Add two sum scalars with overflow checking. -fn add_scalars(sum_dtype: &DType, lhs: &Scalar, rhs: &Scalar) -> VortexResult { - if lhs.is_null() || rhs.is_null() { - return Ok(Scalar::null(sum_dtype.as_nullable())); - } - - Ok(match sum_dtype { - DType::Primitive(ptype, _) if ptype.is_float() => { - let lhs_val = f64::try_from(lhs)?; - let rhs_val = f64::try_from(rhs)?; - Scalar::primitive(lhs_val + rhs_val, Nullable) - } - DType::Primitive(..) => lhs - .as_primitive() - .checked_add(&rhs.as_primitive()) - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - DType::Decimal(..) => lhs - .as_decimal() - .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - _ => unreachable!("StandardSum will always be a decimal or a primitive dtype"), - }) -} - -// Multi-batch and reset tests - -#[test] -fn sum_multi_batch() -> 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)); - Ok(()) -} - -#[test] -fn sum_finish_resets_state() -> 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 result1 = acc.finish()?; - assert_eq!(result1.as_primitive().typed_value::(), Some(30)); - - let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); - acc.accumulate(&batch2, &mut ctx)?; - let result2 = acc.finish()?; - assert_eq!(result2.as_primitive().typed_value::(), Some(18)); - Ok(()) +/// 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 merge tests (vtable-level) +// State algebra: the `{sum, seen}` monoid. #[test] fn sum_state_empty_is_null() -> VortexResult<()> { @@ -182,6 +101,36 @@ fn sum_state_overflow_poisons_but_stays_seen() -> VortexResult<()> { 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 @@ -195,7 +144,7 @@ fn sum_all_nan_is_zero_not_null() -> VortexResult<()> { #[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. + // 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!( @@ -207,80 +156,402 @@ fn sum_is_zero_while_standard_sum_is_null() -> VortexResult<()> { 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 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(); +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(()) +} - 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); +#[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_state_merge() -> VortexResult<()> { +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_multi_batch_and_finish_resets() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = StandardSum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let scalar1 = Scalar::primitive(100i64, Nullable); - StandardSum.combine_partials(&mut state, scalar1)?; + 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)); - let scalar2 = Scalar::primitive(50i64, Nullable); - StandardSum.combine_partials(&mut state, scalar2)?; + // finish resets the state, including `seen`: 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(()) +} - let result = StandardSum.finalize_scalar(&state)?; - StandardSum.reset(&mut state); - assert_eq!(result.as_primitive().typed_value::(), Some(150)); +// Chunked accumulation: `seen` 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(()) } -// Stats caching test +#[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_stats() -> VortexResult<()> { - let array = ChunkedArray::try_new( +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![ - PrimitiveArray::from_iter([1, 1, 1]).into_array(), - PrimitiveArray::from_iter([2, 2, 2]).into_array(), + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), ], - DType::Primitive(PType::I32, Nullability::NonNullable), + 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_seen_from_one_chunk() -> VortexResult<()> { + // One valid value in one chunk, an all-null chunk after: seen from the first chunk + // must survive merging with the second's 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 `seen` and the 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 seen, 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, ) - .vortex_expect("operation should succeed in test"); - let array = array.into_array(); - // compute sum with accumulator to populate stats - sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; + .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 sum_without_acc = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); + 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(()) } -// Constant float non-multiply test +// Overflow: a null sum value, distinct from the unseen null. #[test] -fn sum_constant_float_non_multiply() -> VortexResult<()> { - let acc = -2048669276050936500000000000f64; - let array = ConstantArray::new(6.1811675e16f64, 25); - let result = sum_with_accumulator(&array.into_array(), &Scalar::primitive(acc, Nullable)) - .vortex_expect("operation should succeed in 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!( - f64::try_from(&result).vortex_expect("operation should succeed in test"), - -2048669274505644600000000000f64 + result, + Scalar::null(DType::Decimal(decimal_dtype, Nullable)) ); Ok(()) } -// Grouped sum tests +#[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( @@ -293,6 +564,26 @@ fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult 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(); @@ -326,6 +617,7 @@ fn grouped_sum_with_null_elements() -> VortexResult<()> { #[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) @@ -357,20 +649,6 @@ fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { Ok(()) } -#[test] -fn grouped_sum_bool() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); - let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Bool(Nullability::NonNullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; - - let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).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(); @@ -415,1090 +693,3 @@ fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<( assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } - -// Chunked array tests - -#[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 chunk3 = PrimitiveArray::from_option_iter(vec![None, Some(1.0f64), Some(2.5), None]); - 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(20.8)); - Ok(()) -} - -#[test] -fn sum_chunked_floats_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_floats_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_int_almost_all_null() -> VortexResult<()> { - 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(()) -} - -#[test] -fn sum_chunked_decimals() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let chunk1 = DecimalArray::new( - buffer![100i32, 100i32, 100i32, 100i32, 100i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk2 = DecimalArray::new( - buffer![200i32, 200i32, 200i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk3 = DecimalArray::new(buffer![300i32, 300i32], decimal_dtype, Validity::AllValid); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(1700))) - ); - Ok(()) -} - -#[test] -fn sum_chunked_decimals_with_nulls() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let chunk1 = DecimalArray::new( - buffer![100i32, 100i32, 100i32], - decimal_dtype, - Validity::AllValid, - ); - let chunk2 = DecimalArray::new( - buffer![0i32, 0i32], - decimal_dtype, - Validity::from_iter([false, false]), - ); - let chunk3 = DecimalArray::new(buffer![200i32, 200i32], decimal_dtype, Validity::AllValid); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(700))) - ); - Ok(()) -} - -#[test] -fn sum_chunked_decimals_large() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(3, 0); - let chunk1 = ConstantArray::new( - Scalar::decimal( - DecimalValue::I16(500), - decimal_dtype, - Nullability::NonNullable, - ), - 1, - ); - let chunk2 = ConstantArray::new( - Scalar::decimal( - DecimalValue::I16(600), - decimal_dtype, - Nullability::NonNullable, - ), - 1, - ); - 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(), - )?; - let decimal_result = result.as_decimal(); - assert_eq!( - decimal_result.decimal_value(), - Some(DecimalValue::I256(i256::from_i128(1100))) - ); - assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(13, 0), Nullable) - ); - Ok(()) -} - -mod bool_inputs { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - 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::BoolArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::executor::VortexSessionExecute; - - #[test] - fn sum_bool_all_true() -> VortexResult<()> { - let arr: BoolArray = [true, true, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_mixed() -> VortexResult<()> { - let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(3)); - Ok(()) - } - - #[test] - fn sum_bool_all_false() -> VortexResult<()> { - let arr: BoolArray = [false, false, false].into_iter().collect(); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); - 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_bool_all_null() -> VortexResult<()> { - let arr = BoolArray::from_iter([None::, None, None]); - let result = standard_sum( - &arr.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_empty_is_null() -> VortexResult<()> { - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_bool_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - - let batch1: BoolArray = [true, true, false].into_iter().collect(); - acc.accumulate(&batch1.into_array(), &mut ctx)?; - let result1 = acc.finish()?; - assert_eq!(result1.as_primitive().typed_value::(), Some(2)); - - let batch2: BoolArray = [false, true].into_iter().collect(); - acc.accumulate(&batch2.into_array(), &mut ctx)?; - let result2 = acc.finish()?; - assert_eq!(result2.as_primitive().typed_value::(), Some(1)); - Ok(()) - } - - #[test] - fn sum_bool_return_dtype() -> VortexResult<()> { - let dtype = StandardSum - .return_dtype( - &NumericalAggregateOpts::default(), - &DType::Bool(Nullability::NonNullable), - ) - .unwrap(); - assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); - Ok(()) - } - - #[test] - fn sum_boolean_from_iter() -> VortexResult<()> { - let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(2)); - Ok(()) - } -} - -mod constant_inputs { - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::standard_sum::standard_sum; - use crate::array_session; - use crate::arrays::ConstantArray; - 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::Stat; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - - #[test] - fn sum_constant_unsigned() -> 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_signed() -> VortexResult<()> { - let array = ConstantArray::new(-5i64, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, (-50i64).into()); - Ok(()) - } - - #[test] - fn sum_constant_nullable_value() -> VortexResult<()> { - let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) - .into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_bool_false() -> VortexResult<()> { - let array = ConstantArray::new(false, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 0u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_true() -> VortexResult<()> { - let array = ConstantArray::new(true, 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, 10u64.into()); - Ok(()) - } - - #[test] - fn sum_constant_bool_null() -> VortexResult<()> { - let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); - Ok(()) - } - - #[test] - fn sum_constant_decimal() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(100), - decimal_dtype, - Nullability::NonNullable, - ), - 5, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(500))) - ); - assert_eq!(result.dtype(), &Stat::Sum.dtype(array.dtype()).unwrap()); - Ok(()) - } - - #[test] - fn sum_constant_decimal_null() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result, - Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) - ); - Ok(()) - } - - #[test] - fn sum_constant_decimal_large_value() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(10, 2); - let array = ConstantArray::new( - Scalar::decimal( - DecimalValue::I64(999_999_999), - decimal_dtype, - Nullability::NonNullable, - ), - 100, - ) - .into_array(); - - let result = standard_sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) - ); - Ok(()) - } -} - -mod decimal_inputs { - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::AggregateFnVTable; - 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::DecimalArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::i256; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - #[test] - fn sum_decimal_basic() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(600i32))), - )?; - - assert_eq!(result, expected); - 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_decimal_negative_values() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, -200i32, 300i32, -50i32], - DecimalDType::new(4, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(14, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(150i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_near_i32_max() -> VortexResult<()> { - let near_max = i32::MAX - 1000; - let decimal = DecimalArray::new( - buffer![near_max, 500i32, 400i32], - DecimalDType::new(10, 2), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = near_max as i64 + 500 + 400; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(20, 2), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_large_i64_values() -> VortexResult<()> { - let large_val = i64::MAX / 4; - let decimal = DecimalArray::new( - buffer![large_val, large_val, large_val, large_val + 1], - DecimalDType::new(19, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = (large_val as i128) * 4 + 1; - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(29, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_preserves_scale() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![12345i32, 67890i32, 11111i32], - DecimalDType::new(6, 4), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(16, 4), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(91346i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_single_value() -> VortexResult<()> { - let decimal = - DecimalArray::new(buffer![42i32], DecimalDType::new(3, 1), Validity::AllValid); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(13, 1), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(42i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_all_nulls_except_one() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![100i32, 200i32, 300i32, 400i32], - DecimalDType::new(4, 2), - Validity::from_iter([false, false, true, false]), - ); - - 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(300i32))), - )?; - - assert_eq!(result, expected); - Ok(()) - } - - #[test] - fn sum_decimal_overflow_detection() -> VortexResult<()> { - let max_val = i128::MAX / 2; - let decimal = DecimalArray::new( - buffer![max_val, max_val, max_val], - DecimalDType::new(38, 0), - Validity::AllValid, - ); - - let result = standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx(), - )?; - - let expected_sum = - i256::from_i128(max_val) + i256::from_i128(max_val) + i256::from_i128(max_val); - let expected = Scalar::try_new( - DType::Decimal(DecimalDType::new(48, 0), Nullability::NonNullable), - Some(ScalarValue::from(DecimalValue::from(expected_sum))), - )?; - - assert_eq!(result, expected); - 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, - ); - - assert_eq!( - standard_sum( - &decimal.into_array(), - &mut array_session().create_execution_ctx() - ) - .vortex_expect("operation should succeed in test"), - 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(()) - } - - #[test] - fn sum_decimal_precision_overflow_within_i256() -> VortexResult<()> { - // Input precision 4 → return precision 14. Native I64 (max 18). - // The max representable value for precision 14 is 10^14 - 1. - // When the sum reaches exactly 10^14, fits_in_precision fails 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(99_999_999_999_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - // Push the sum to exactly 10^14, exceeding precision 14. - let one_more = - Scalar::decimal(DecimalValue::from(1i64), 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_precision_overflow_negative() -> VortexResult<()> { - // Same setup but with negative values: sum reaches -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_999i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, near_limit)?; - - let one_more = Scalar::decimal( - DecimalValue::from(-1i64), - DecimalDType::new(14, 0), - Nullable, - ); - StandardSum.combine_partials(&mut state, one_more)?; - - let result = StandardSum.finalize_scalar(&state)?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { - // Test precision overflow via the accumulate_decimal path (not combine_partials). - // Input precision 28 (I128 storage) → return precision min(76, 38) = 38. - // Native for precision 38 is I128 (max 38), so 38 = 38. - // Use precision 27 → return 37. Native for 37 is I128 (max 38), so 37 < 38. - // - // We 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); - - // Drive accumulate through the vtable directly. - 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(()) - } -} - -mod primitive_inputs { - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - 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::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - 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; - - #[test] - fn sum_i32() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(10)); - Ok(()) - } - - #[test] - fn sum_u8() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(60)); - Ok(()) - } - - #[test] - fn sum_f64() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); - Ok(()) - } - - #[test] - fn sum_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_multiple_null_runs() -> VortexResult<()> { - // Several disjoint valid runs separated by nulls exercise the per-run fold. - let arr = PrimitiveArray::from_option_iter([ - Some(1i32), - Some(2), - None, - None, - Some(3), - None, - Some(4), - Some(5), - Some(6), - ]) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(21)); - Ok(()) - } - - #[test] - fn sum_all_null() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_all_invalid_float() -> VortexResult<()> { - let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); - Ok(()) - } - - #[test] - fn sum_buffer_i32() -> VortexResult<()> { - let arr = buffer![1, 1, 1, 1].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4)); - Ok(()) - } - - #[test] - fn sum_buffer_f64() -> VortexResult<()> { - let arr = buffer![1., 1., 1., 1.].into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().as_::(), Some(4.)); - Ok(()) - } - - #[test] - fn sum_empty_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_empty_f64_is_null() -> VortexResult<()> { - let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(StandardSum, NumericalAggregateOpts::default(), dtype)?; - let result = acc.finish()?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_f64_with_nan() -> VortexResult<()> { - let arr = PrimitiveArray::new( - buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0], - Validity::NonNullable, - ) - .into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - Ok(()) - } - - #[test] - fn sum_f32_with_nan() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); - Ok(()) - } - - #[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_all_nan() -> VortexResult<()> { - 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(()) - } - - /// StandardSum an array with explicit [`NumericalAggregateOpts`] (test-only helper). - fn sum_with_options( - arr: &crate::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() - } - - #[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_f64_without_nan_not_skipping() -> VortexResult<()> { - let arr = - PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; - assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); - 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_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { - // With an exact NaNCount of zero, the planted exact StandardSum 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 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(()) - } - - #[test] - fn sum_checked_overflow() -> VortexResult<()> { - let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = standard_sum(&arr, &mut array_session().create_execution_ctx())?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_checked_overflow_is_saturated() -> 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()); - - // finish resets state, clearing saturation - drop(acc.finish()?); - assert!(!acc.is_saturated()); - Ok(()) - } -} From 1ca56ffb6d471fe40317a8e22f30569519ab254f Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 13 Jul 2026 22:58:35 -0700 Subject: [PATCH 12/12] Refactor StandardSum partial state Signed-off-by: Matt Katz --- .../aggregate_fn/fns/standard_sum/grouped.rs | 29 ++- .../src/aggregate_fn/fns/standard_sum/mod.rs | 186 +++++++++--------- .../aggregate_fn/fns/standard_sum/tests.rs | 135 +++++++++++-- 3 files changed, 221 insertions(+), 129 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs index 1e18392745d..a4c0fb46da9 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/grouped.rs @@ -54,9 +54,9 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedStandardSumEncodingKernel { /// 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, seen}` partial rows (see `StandardSum`'s partial dtype): `seen` records whether the -/// group contained at least one valid element, so that `finalize` yields null for empty and -/// all-null groups, and null groups are null struct rows. +/// 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, @@ -92,7 +92,7 @@ fn grouped_sum( .execute_mask(elements.as_ref().len(), ctx)?; let all_valid = matches!(elem_mask.slices(), AllOr::All); - let (sums, seen) = match_each_native_ptype!(elements.ptype(), + 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, @@ -111,10 +111,10 @@ fn grouped_sum( ); Ok(StructArray::try_new( - FieldNames::from_iter([FieldName::from("sum"), FieldName::from("seen")]), + FieldNames::from_iter([FieldName::from("sum"), FieldName::from("overflowed")]), vec![ sums.into_array(), - BoolArray::new(seen, Validity::NonNullable).into_array(), + BoolArray::new(overflowed, Validity::NonNullable).into_array(), ], group_validity.len(), Validity::from_mask(group_validity.clone(), Nullability::Nullable), @@ -122,10 +122,9 @@ fn grouped_sum( .into_array()) } -/// Reduce each group's element slice into a nullable sum plus a per-group `seen` bit. A sum is -/// null when the group itself is invalid or when summing it overflows (`sum_run` returns -/// `true`); `seen` records whether the group contained at least one valid element, decided by -/// validity alone (with `skip_nans`, a valid all-NaN group is still seen and sums to zero). +/// 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, @@ -134,10 +133,10 @@ fn collect_sums( all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, ) -> (PrimitiveArray, BitBuffer) { - let mut seen = BitBufferMut::with_capacity(group_ranges.len()); + let mut overflowed = BitBufferMut::with_capacity(group_ranges.len()); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - seen.append(false); + overflowed.append(false); return None; } let mut acc = A::default(); @@ -146,11 +145,11 @@ fn collect_sums( } else { sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) }; - seen.append(any_valid); - (!overflow).then_some(acc) + overflowed.append(overflow); + (!overflow && any_valid).then_some(acc) }); let sums = PrimitiveArray::from_option_iter(sums); - (sums, seen.freeze()) + (sums, overflowed.freeze()) } /// StandardSum the valid elements of a single group, using the contiguous valid runs of the element mask diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs index 6648a454a04..9ab4a8b2abd 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/mod.rs @@ -14,7 +14,6 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; -use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; @@ -29,7 +28,6 @@ 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::ConstantArray; use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::FieldName; @@ -42,10 +40,7 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::mask::Mask; use crate::validity::Validity; /// Return the sum of an array. The result is null when the array has no valid values or the sum @@ -156,48 +151,51 @@ impl AggregateFnVTable for StandardSum { let return_dtype = self .return_dtype(options, input_dtype) .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; - let initial = make_zero_state(&return_dtype); - Ok(StandardSumPartial { return_dtype, - current: Some(initial), - seen: false, + current: None, + overflowed: false, skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - // Partials are `{sum, seen}` structs. A plain scalar is a cached `Stat::Sum` value, - // which cannot distinguish an empty sum from a zero sum, so it is treated as seen. - let (sum_value, other_seen) = if matches!(other.dtype(), DType::Struct(..)) { + // 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.seen = true; 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_seen = fields - .field("seen") - .and_then(|seen| seen.as_bool().value()) + let other_overflowed = fields + .field("overflowed") + .and_then(|overflowed| overflowed.as_bool().value()) .unwrap_or(true); - (sum_value, other_seen) + (sum_value, other_overflowed) } else { - (other, true) + let other_overflowed = other.is_null(); + (other, other_overflowed) }; - partial.seen |= other_seen; - if sum_value.is_null() { - // A null sum value means the sub-accumulator saturated (overflow). + if partial.overflowed || other_overflowed { partial.current = None; + partial.overflowed = true; return Ok(()); } - let Some(ref mut inner) = partial.current else { + 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) => { @@ -238,6 +236,7 @@ impl AggregateFnVTable for StandardSum { }; if saturated { partial.current = None; + partial.overflowed = true; } Ok(()) } @@ -247,23 +246,20 @@ impl AggregateFnVTable for StandardSum { sum_partial_dtype(partial.return_dtype.as_nullable()), vec![ sum_value_scalar(partial), - Scalar::bool(partial.seen, Nullability::NonNullable), + Scalar::bool(partial.overflowed, Nullability::NonNullable), ], )) } fn reset(&self, partial: &mut Self::Partial) { - partial.current = Some(make_zero_state(&partial.return_dtype)); - partial.seen = false; + partial.current = None; + partial.overflowed = false; } #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { - match partial.current.as_ref() { - None => true, - Some(SumState::Float(v)) => v.is_nan(), - Some(_) => false, - } + partial.overflowed + || matches!(partial.current.as_ref(), Some(SumState::Float(v)) if v.is_nan()) } fn try_accumulate( @@ -279,7 +275,7 @@ impl AggregateFnVTable for StandardSum { } // NaN-including float sums need a NaN-free batch before the cached sum applies; // everything else takes the default dispatch path. - if !matches!(partial.current, Some(SumState::Float(_))) { + if !matches!(&partial.return_dtype, DType::Primitive(PType::F64, _)) { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { @@ -290,10 +286,13 @@ impl AggregateFnVTable for StandardSum { } Precision::Exact(_) => { // At least one NaN value (a valid value): the sum is NaN without scanning. - partial.seen = true; - if let Some(SumState::Float(acc)) = partial.current.as_mut() { - *acc = f64::NAN; - } + 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), @@ -306,11 +305,17 @@ impl AggregateFnVTable for StandardSum { 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 { - // Any valid value counts as seen, including NaN and `false` constants that - // contribute nothing to the running sum. - partial.seen |= !c.scalar().is_null() && !c.is_empty(); + // 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(()); @@ -322,37 +327,46 @@ impl AggregateFnVTable for StandardSum { } let skip_nans = partial.skip_nans; - let mut inner = match partial.current.take() { - Some(inner) => inner, - None => return Ok(()), - }; - - // `seen` is decided by validity alone (NaNs are valid values), so it is tracked here - // and the summation reuses [`Sum`]'s accumulation kernels unchanged. - let result = match batch { + let any_valid = match batch { Columnar::Canonical(c) => match c { Canonical::Primitive(p) => { - partial.seen |= any_valid(p.as_ref().validity()?, p.as_ref().len(), ctx)?; - accumulate_primitive(&mut inner, p, ctx, skip_nans) - } - Canonical::Bool(b) => { - partial.seen |= any_valid(b.as_ref().validity()?, b.as_ref().len(), ctx)?; - accumulate_bool(&mut inner, b, ctx) - } - Canonical::Decimal(d) => { - partial.seen |= any_valid(d.as_ref().validity()?, d.as_ref().len(), ctx)?; - accumulate_decimal(&mut inner, d, ctx) + 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) => partial.current = Some(inner), - Ok(true) => {} // saturated: current stays None + Ok(false) => { + if had_value || any_valid { + partial.current = Some(inner); + } + } + Ok(true) => partial.overflowed = true, Err(e) => { - partial.current = Some(inner); + if had_value || any_valid { + partial.current = Some(inner); + } return Err(e); } } @@ -360,38 +374,20 @@ impl AggregateFnVTable for StandardSum { } fn finalize(&self, partials: ArrayRef) -> VortexResult { - // Entries that saw no valid values finalize to null, while a null `sum` - // field (overflow) and null partial rows (e.g. null groups) stay null via the mask's - // validity intersection. - // - // The expressions are built unoptimized: `optimize` costs multiples of the whole - // aggregation at small group counts, and the caller's execution evaluates the lazy - // expression as-is. - let len = partials.len(); - let sum = GetItem.try_new_array(len, FieldName::from("sum"), [partials.clone()])?; - let seen = GetItem.try_new_array(len, FieldName::from("seen"), [partials])?; - let seen = FillNull.try_new_array( - len, - EmptyOptions, - [ - seen, - ConstantArray::new(Scalar::bool(false, Nullability::NonNullable), len).into_array(), - ], - )?; - Mask.try_new_array(len, EmptyOptions, [sum, seen]) + // 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 { - if !partial.seen { - return Ok(Scalar::null(partial.return_dtype.as_nullable())); - } Ok(sum_value_scalar(partial)) } } /// Consume a batch's cached `Stat::Sum` instead of scanning it. The cached sum cannot -/// carry `seen`, so the batch's null count decides between the identity (all null) and a seen -/// contribution; when either statistic is missing the caller falls through to a real scan. +/// 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, @@ -405,7 +401,7 @@ fn try_accumulate_cached_sum( // No valid values: the batch is the identity. return Ok(true); } - Precision::Exact(_) => partial.seen = true, + Precision::Exact(_) => {} _ => return Ok(false), } let sum = if sum.dtype() == &partial.return_dtype { @@ -426,31 +422,29 @@ fn any_valid(validity: Validity, len: usize, ctx: &mut ExecutionCtx) -> VortexRe /// needed for reset/result without external context. pub struct StandardSumPartial { return_dtype: DType, - /// The current accumulated state, or `None` if saturated (checked overflow). + /// 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 at least one valid value has been accumulated. A sum over zero valid values - /// finalizes to null rather than zero. - seen: bool, + /// 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, seen}` struct, where -/// `sum` is the running sum, and is null once saturated by overflow, poisoning merges, and -/// `seen` records whether any valid value contributed (merged with OR). Keeping the flag -/// separate from the sum lets merges stay a monoid — the identity is `{0, false}` — while -/// `finalize` maps unseen sums to null. +/// 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("seen")]), + 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 saturated by overflow). +/// 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()), diff --git a/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs index 0c6a0856343..62033337321 100644 --- a/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/standard_sum/tests.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Tests for [`StandardSum`]-specific behavior: the `{sum, seen}` state algebra, the -//! null-for-zero-valid-values rule across input kinds, NaN and overflow interplay with -//! `seen`, cached-statistic consumption, and grouped aggregation. Plain summation -//! arithmetic is covered by the shared kernels' tests in [`super::super::sum`]. +//! 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; @@ -51,7 +51,7 @@ fn sum_with_options(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexRe acc.finish() } -// State algebra: the `{sum, seen}` monoid. +// State algebra: the `{sum, overflowed}` monoid. #[test] fn sum_state_empty_is_null() -> VortexResult<()> { @@ -60,6 +60,14 @@ fn sum_state_empty_is_null() -> VortexResult<()> { 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(()) @@ -67,8 +75,8 @@ fn sum_state_empty_is_null() -> VortexResult<()> { #[test] fn sum_state_empty_is_identity() -> VortexResult<()> { - // Combining an empty state into a seen state changes nothing: `{0, false}` is the - // identity of the `{sum, seen}` monoid. + // 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))?; @@ -83,14 +91,21 @@ fn sum_state_empty_is_identity() -> VortexResult<()> { } #[test] -fn sum_state_overflow_poisons_but_stays_seen() -> VortexResult<()> { - // Overflow (a null `sum` field) poisons the merge even when combined with later - // values: the result is null via the sum value, not via `seen`. +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))?; @@ -235,6 +250,14 @@ fn sum_constant() -> VortexResult<()> { 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(); @@ -248,7 +271,7 @@ fn sum_multi_batch_and_finish_resets() -> VortexResult<()> { let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(48)); - // finish resets the state, including `seen`: an untouched accumulator is empty again. + // 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)?; @@ -256,7 +279,7 @@ fn sum_multi_batch_and_finish_resets() -> VortexResult<()> { Ok(()) } -// Chunked accumulation: `seen` must merge across chunks. +// Chunked accumulation: the nullable sum must merge across chunks. #[test] fn sum_chunked_floats_with_nulls() -> VortexResult<()> { @@ -311,9 +334,9 @@ fn sum_chunked_empty_chunks() -> VortexResult<()> { } #[test] -fn sum_chunked_seen_from_one_chunk() -> VortexResult<()> { - // One valid value in one chunk, an all-null chunk after: seen from the first chunk - // must survive merging with the second's identity. +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(); @@ -327,7 +350,7 @@ fn sum_chunked_seen_from_one_chunk() -> VortexResult<()> { Ok(()) } -// NaN handling and its interplay with `seen` and the cached statistics. +// NaN handling and its interplay with the nullable sum and cached statistics. #[test] fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { @@ -391,7 +414,7 @@ fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { #[test] fn sum_constant_nan() -> VortexResult<()> { let arr = ConstantArray::new(f64::NAN, 4).into_array(); - // NaN constants are skipped by default (a seen, zero sum) and poison the sum otherwise. + // 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)); @@ -421,7 +444,7 @@ fn sum_f64_with_infinity() -> VortexResult<()> { Ok(()) } -// Overflow: a null sum value, distinct from the unseen null. +// Overflow: a null sum value plus an explicit saturation flag. #[test] fn sum_checked_overflow_is_null_and_saturates() -> VortexResult<()> { @@ -564,6 +587,66 @@ fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult 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 @@ -649,6 +732,22 @@ fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { 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();