diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..312c5c9c184 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -11,12 +11,12 @@ 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; 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::sum::Sum; use vortex_array::arrays::ListViewArray; @@ -154,14 +154,12 @@ fn list_element_dtype(list_view: &ArrayRef) -> DType { fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef where - V: AggregateFnVTable + Clone, + V: AggregateFnVTable + Clone, + V::Options: Default, { - let mut acc = GroupedAccumulator::try_new( - vtable, - NumericalAggregateOpts::default(), - list_element_dtype(list_view), - ) - .unwrap(); + let mut acc = + GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) + .unwrap(); acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); divan::black_box(acc.finish().unwrap()) @@ -199,6 +197,59 @@ 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, + V::Options: Default, +{ + let mut acc = + GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) + .unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + acc.accumulate_list(list_view, &mut ctx).unwrap(); + let result = acc + .finish() + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + divan::black_box(result) +} + +#[divan::bench] +fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + #[divan::bench] fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 69bae4e1053..52255481a43 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -119,9 +119,18 @@ 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. + // Stat::Sum stores a finalized scalar rather than a partial. Give Sum's vtable the first + // chance to consume that cache before dispatching an encoding kernel. + let checked_cached_sum = Stat::from_aggregate_fn(&self.aggregate_fn) == Some(Stat::Sum) + && batch.statistics().get(Stat::Sum).is_exact(); + if checked_cached_sum && self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { + return Ok(()); + } + + // 0. Cached stats bridge: consume an exact partial from the aggregate's Stat slot before + // kernel dispatch or decode. Sum is handled above because its cache stores a result. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) + && stat != Stat::Sum && let Precision::Exact(partial) = batch.statistics().get(stat) { let partial = if partial.dtype() == &self.partial_dtype { @@ -167,7 +176,7 @@ impl DynAccumulator for Accumulator { } // 2. Allow the vtable to short-circuit on the raw array before decompression. - if self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { + if !checked_cached_sum && self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { return Ok(()); } @@ -280,6 +289,7 @@ mod tests { 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::sum::SumAggregateOpts; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; use crate::array::VTable; @@ -288,7 +298,10 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; use crate::scalar::Scalar; + use crate::scalar::ScalarValue; /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` @@ -320,7 +333,7 @@ mod tests { } } - /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of + /// Sum partial sentinel `{sum: 42.0, is_overflow: false, is_empty: false}` — distinguishable from the natural Sum of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -331,7 +344,7 @@ mod tests { _batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) + Ok(Some(sum_partial(42.0))) } } @@ -350,7 +363,7 @@ mod tests { Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -359,11 +372,24 @@ mod tests { fn sentinel_partial() -> Scalar { let acc = mean_f64_accumulator().expect("build accumulator"); - let sum = Scalar::primitive(42.0f64, Nullability::Nullable); + let sum = sum_partial(42.0); let count = Scalar::primitive(1u64, Nullability::NonNullable); Scalar::struct_(acc.partial_dtype, vec![sum, count]) } + fn sum_partial(value: f64) -> Scalar { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype).expect("sum"); + Scalar::struct_( + acc.partial_dtype, + vec![ + Scalar::primitive(value, Nullability::NonNullable), + Scalar::bool(false, Nullability::NonNullable), + Scalar::bool(false, Nullability::NonNullable), + ], + ) + } + /// Kernel registered for `(Dict, Combined)` fires in preference to /// `Combined::try_accumulate`'s fan-out path — proves the dispatch reorder. #[test] @@ -381,7 +407,13 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(42.0) ); assert_eq!( @@ -408,7 +440,13 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(7.0) ); assert_eq!( @@ -439,7 +477,13 @@ mod tests { // 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_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(42.0) ); assert_eq!( @@ -448,4 +492,26 @@ mod tests { ); Ok(()) } + + #[test] + fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { + static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let batch = dict_of_seven(); + batch + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(11.0f64))); + + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + acc.accumulate(&batch, &mut ctx)?; + + assert_eq!(acc.finish()?.as_primitive().as_::(), Some(11.0)); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 286f4194cd5..8bd26dad961 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -273,7 +273,10 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn flush(&mut self) -> VortexResult { - let states = std::mem::take(&mut self.partials); + let mut states = std::mem::take(&mut self.partials); + if states.len() == 1 { + return Ok(states.pop().vortex_expect("checked one partial")); + } Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) } diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index da93d2352a3..16a1b5431fe 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -19,6 +19,7 @@ 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::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum_decimal_dtype; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; @@ -40,7 +41,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), array.dtype().clone(), @@ -319,7 +320,7 @@ mod tests { let keep_nans = NumericalAggregateOpts::include_nans(); let mut acc = Accumulator::try_new( Mean::combined(), - PairOptions(keep_nans, keep_nans), + PairOptions(keep_nans.into(), keep_nans), array.dtype().clone(), )?; acc.accumulate(&array, &mut ctx)?; @@ -411,7 +412,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -446,7 +447,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), DType::Primitive(PType::F64, Nullability::Nullable), @@ -474,7 +475,7 @@ mod tests { let mut acc = GroupedAccumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), DType::Primitive(PType::F64, Nullability::Nullable), diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3993840586..4e33a567dd2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -6,23 +6,22 @@ use std::ops::BitAnd; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; +use vortex_mask::Mask; 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, + mask: &Mask, ) -> 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, @@ -40,8 +39,8 @@ mod tests { 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::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::BoolArray; @@ -101,16 +100,16 @@ mod tests { &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } #[test] - fn sum_bool_empty_produces_zero() -> VortexResult<()> { + fn sum_bool_empty_produces_null() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -118,7 +117,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(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -136,7 +135,7 @@ mod tests { fn sum_bool_return_dtype() -> VortexResult<()> { let dtype = Sum .return_dtype( - &NumericalAggregateOpts::default(), + &SumAggregateOpts::default(), &DType::Bool(Nullability::NonNullable), ) .unwrap(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..14da9d60799 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -127,7 +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())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -151,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::primitive(0u64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -187,11 +187,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..03c3f116fc5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -13,7 +13,6 @@ 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; @@ -26,10 +25,9 @@ use crate::scalar::DecimalValue; pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, - ctx: &mut ExecutionCtx, + mask: &Mask, ) -> VortexResult { - let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - let validity = match &mask { + let validity = match mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), Mask::AllFalse(_) => { @@ -111,11 +109,12 @@ mod tests { use vortex_error::VortexExpect; use vortex_error::VortexResult; + use super::super::sum_result_partial_scalar; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::DecimalArray; @@ -129,6 +128,11 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; + fn partial_with_value(value: Scalar) -> VortexResult { + let return_dtype = value.dtype().as_nullable(); + sum_result_partial_scalar(value, &return_dtype, false) + } + #[test] fn sum_decimal_basic() -> VortexResult<()> { let decimal = DecimalArray::new( @@ -357,23 +361,31 @@ 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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(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); - Sum.combine_partials(&mut state, small)?; + Sum.combine_partials(&mut state, partial_with_value(small)?)?; let result = Sum.to_scalar(&state)?; - assert!(!result.is_null()); + let fields = result.as_struct(); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(false) + ); assert_eq!( - result.as_decimal().decimal_value(), + fields + .field("sum") + .and_then(|sum| sum.as_decimal().decimal_value()), Some(DecimalValue::I256(i256::from_i128(99_999_999_999_999))) ); Ok(()) @@ -387,25 +399,27 @@ 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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(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); - Sum.combine_partials(&mut state, one_more)?; + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; let result = Sum.to_scalar(&state)?; - assert!(result.is_null()); assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(14, 0), Nullable) + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) ); Ok(()) } @@ -414,24 +428,30 @@ 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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; let one_more = Scalar::decimal( DecimalValue::from(-1i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, one_more)?; + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; let result = Sum.to_scalar(&state)?; - assert!(result.is_null()); + assert_eq!( + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); Ok(()) } @@ -446,13 +466,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 = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::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); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -464,7 +484,13 @@ mod tests { Sum.accumulate(&mut state, &columnar, &mut ctx)?; let result = Sum.to_scalar(&state)?; - assert!(result.is_null()); + assert_eq!( + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); 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..64b0ce81995 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,14 +1,16 @@ // 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::Sum; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; +use super::sum_partial_fields; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -16,10 +18,14 @@ 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::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)] @@ -78,31 +84,47 @@ fn grouped_sum( .as_ref() .validity()? .execute_mask(elements.as_ref().len(), ctx)?; - let all_valid = matches!(elem_mask.slices(), AllOr::All); + let all_valid = elem_mask.all_true(); - let result = match_each_native_ptype!(elements.ptype(), + let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); - collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, - sum_unsigned_all) + 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) + 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, + 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()) + let partial_fields = sum_partial_fields(sums.dtype().clone()); + + // SAFETY: all three children have one value per group and match `partial_fields`; the struct + // validity is derived from the same group count. + Ok(unsafe { + StructArray::new_unchecked( + vec![ + sums.into_array(), + BoolArray::new(is_overflow, Validity::NonNullable).into_array(), + BoolArray::new(is_empty, Validity::NonNullable).into_array(), + ], + partial_fields, + group_validity.len(), + Validity::from_mask(group_validity.clone(), Nullability::Nullable), + ) + } + .into_array()) } -/// Reduce each group's element slice into a 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 non-null sum, overflow bitmap, and empty bitmap. fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -110,24 +132,37 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> PrimitiveArray { +) -> (PrimitiveArray, BitBuffer, BitBuffer) { + let group_count = group_ranges.len(); + let mut is_overflow = BitBufferMut::new_unset(group_count); + let mut is_empty = BitBufferMut::new_unset(group_count); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - return None; + return A::default(); } 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) }; - (!overflow).then_some(acc) + if overflow { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_overflow.set_unchecked(i) }; + } + if !any_valid { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_empty.set_unchecked(i) }; + } + acc }); - PrimitiveArray::from_option_iter(sums) + let sums = PrimitiveArray::from_iter(sums); + (sums, is_overflow.freeze(), is_empty.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. +/// Sum valid runs in one group, returning `(overflow, any_valid)`. fn sum_masked_group( acc: &mut A, values: &[T], @@ -135,17 +170,23 @@ fn sum_masked_group( 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 { +) -> (bool, bool) { + match elem_mask { + Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), + Mask::AllFalse(_) => (false, false), + Mask::Values(mask_values) => { + let validity = mask_values + .bit_buffer() + .as_view() + .slice(offset..offset + size); + let mut any_valid = false; + for (start, end) in validity.set_slices() { + any_valid = true; if sum_run(acc, &values[offset + start..offset + end]) { - return true; + return (true, true); } } - false + (false, any_valid) } } } @@ -162,8 +203,8 @@ mod tests { use crate::VortexSessionExecute; 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::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::FixedSizeListArray; @@ -179,11 +220,8 @@ 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( - Sum, - NumericalAggregateOpts::default(), - elem_dtype.clone(), - )?; + let mut acc = + GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -200,8 +238,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(&SumAggregateOpts::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 +329,7 @@ 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)]); + 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,7 +406,7 @@ mod tests { let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; let mut acc = - GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; + GroupedAccumulator::try_new(Sum, SumAggregateOpts::include_nans(), elem_dtype)?; acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; let actual = acc.finish()?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 7ad62eacb87..89b19b5575f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -6,12 +6,18 @@ mod constant; mod decimal; mod grouped; mod primitive; +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; -use vortex_error::VortexExpect; +use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_proto::expr as pb; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -23,43 +29,48 @@ 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::StructArray; +use crate::builtins::ArrayBuiltins; 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::fns::operators::Operator; +use crate::validity::Validity; + +const SUM_FIELD: &str = "sum"; +const IS_OVERFLOW_FIELD: &str = "is_overflow"; +const IS_EMPTY_FIELD: &str = "is_empty"; /// Return the sum of an array. /// /// 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) { - return Ok(sum_scalar); + if let Precision::Exact(sum) = array.statistics().get(Stat::Sum) { + return Ok(sum); } - // 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( - Sum, - NumericalAggregateOpts::default(), - array.dtype().clone(), - )?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::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)); } @@ -67,16 +78,96 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// Sum an array, starting from zero. -/// -/// If the sum overflows, a null scalar will be returned. -/// If the array is all-invalid, the sum will be zero. -/// -/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the -/// default) NaN values contribute nothing, otherwise any NaN value poisons the sum to NaN. -#[derive(Clone, Debug)] +/// Sum an array, returning null when it has no valid values or if the sum overflows. +#[derive(Clone, Copy, Debug)] pub struct Sum; +/// Options for [`Sum`]. +/// +/// Sums always use a struct partial that can distinguish an empty input from a zero sum. The +/// `struct_partial` field only describes the representation persisted by an existing serialized +/// aggregate; callers should normally construct these options with [`Default`], +/// [`SumAggregateOpts::skip_nans`], or [`SumAggregateOpts::include_nans`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct SumAggregateOpts { + /// Whether NaN values are skipped (treated as missing) during aggregation. + pub skip_nans: bool, + /// Whether persisted partials use the `{ sum, is_overflow, is_empty }` struct representation. + /// + /// This is false only when deserializing metadata for legacy scalar Sum partials. It does not + /// alter the representation or semantics of a live accumulator. + pub struct_partial: bool, +} + +impl SumAggregateOpts { + /// Options that skip NaN values and use the canonical struct partial. + pub const fn skip_nans() -> Self { + Self { + skip_nans: true, + struct_partial: true, + } + } + + /// Options that include NaN values and use the canonical struct partial. + pub const fn include_nans() -> Self { + Self { + skip_nans: false, + struct_partial: true, + } + } + + /// Serialize these options to protobuf-encoded metadata bytes. + /// + /// This preserves the persisted partial representation so a legacy layout can be reserialized + /// without changing the metadata describing its existing zones child. + pub fn serialize(&self) -> Vec { + pb::SumAggregateOpts { + skip_nans: self.skip_nans, + struct_partial: Some(self.struct_partial), + } + .encode_to_vec() + } + + /// Deserialize these options from protobuf-encoded metadata bytes. + /// + /// Historical Sum options were serialized as [`NumericalAggregateOpts`]. They use the same + /// wire representation for `skip_nans` and omit `struct_partial`, which selects the legacy + /// scalar partial representation. + pub fn deserialize(metadata: &[u8]) -> VortexResult { + let options = pb::SumAggregateOpts::decode(metadata)?; + Ok(Self { + skip_nans: options.skip_nans, + struct_partial: options.struct_partial.unwrap_or(false), + }) + } +} + +impl Default for SumAggregateOpts { + fn default() -> Self { + Self::skip_nans() + } +} + +impl From for SumAggregateOpts { + fn from(options: NumericalAggregateOpts) -> Self { + Self { + skip_nans: options.skip_nans, + struct_partial: true, + } + } +} + +impl Display for SumAggregateOpts { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + // The partial representation is a storage-compatibility detail. Keeping it out of the + // display preserves the stats-table field name across old and new Sum partials. + if !self.skip_nans { + write!(f, "skip_nans=false")?; + } + Ok(()) + } +} + // 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 @@ -88,7 +179,7 @@ pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { } impl AggregateFnVTable for Sum { - type Options = NumericalAggregateOpts; + type Options = SumAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { @@ -105,12 +196,12 @@ impl AggregateFnVTable for Sum { metadata: &[u8], _session: &VortexSession, ) -> VortexResult { - NumericalAggregateOpts::deserialize(metadata) + SumAggregateOpts::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. + // When a sum overflows, we return a null sum value. Therefore, all return dtypes are + // nullable. use Nullability::Nullable; Some(match input_dtype { @@ -136,7 +227,8 @@ impl AggregateFnVTable for Sum { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + let return_dtype = self.return_dtype(options, input_dtype)?; + Some(sum_partial_dtype(return_dtype)) } fn empty_partial( @@ -147,94 +239,57 @@ impl AggregateFnVTable for Sum { 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); - + let sum = make_zero_state(&return_dtype); Ok(SumPartial { return_dtype, - current: Some(initial), + sum, + is_overflow: false, + is_empty: true, 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; + let (other_sum, other_is_overflow, other_is_empty) = decode_sum_partial_scalar(other)?; + validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; + + if partial.is_overflow { return Ok(()); } - let Some(ref mut inner) = partial.current else { + if other_is_overflow { + partial.is_overflow = true; + partial.is_empty = false; 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; } + if other_is_empty { + return Ok(()); + } + + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; + partial.is_empty = false; Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(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.clone()), + vec![ + sum_state_scalar(partial, Nullability::NonNullable), + Scalar::bool(partial.is_overflow, Nullability::NonNullable), + Scalar::bool(partial.is_empty, Nullability::NonNullable), + ], + )) } fn reset(&self, partial: &mut Self::Partial) { - partial.current = Some(make_zero_state(&partial.return_dtype)); + partial.sum = make_zero_state(&partial.return_dtype); + partial.is_overflow = false; + partial.is_empty = true; } #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { - match partial.current.as_ref() { - None => true, - Some(SumState::Float(v)) => v.is_nan(), - Some(_) => false, - } + partial.is_overflow || matches!(&partial.sum, SumState::Float(v) if v.is_nan()) } fn try_accumulate( @@ -243,31 +298,27 @@ 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(_))) { + if partial.skip_nans { + return try_accumulate_cached_sum(self, partial, batch); + } + + // NaN-aware short-circuits only apply to NaN-including float sums. + if !matches!(&partial.sum, 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. - if let Some(SumState::Float(acc)) = partial.current.as_mut() { - *acc = f64::NAN; - } + let SumState::Float(acc) = &mut partial.sum else { + unreachable!("checked float sum state") + }; + *acc = f64::NAN; + partial.is_empty = false; Ok(true) } _ => Ok(false), @@ -280,60 +331,85 @@ impl AggregateFnVTable for Sum { batch: &Columnar, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + if partial.is_overflow { + return Ok(()); + } + // Constants compute scalar * len and combine via combine_partials. if let Columnar::Constant(c) = batch { + if !c.scalar().is_null() && !c.is_empty() { + partial.is_empty = false; + } // 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)? { + let product = sum_result_partial_scalar(product, &partial.return_dtype, false)?; 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 { + let (result, any_valid) = 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) => { + let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; + let any_valid = partial.is_empty && mask.true_count() > 0; + ( + accumulate_primitive(&mut partial.sum, p, &mask, skip_nans), + any_valid, + ) + } + Canonical::Bool(b) => { + let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; + let any_valid = partial.is_empty && mask.true_count() > 0; + (accumulate_bool(&mut partial.sum, b, &mask), any_valid) + } + Canonical::Decimal(d) => { + let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + let any_valid = partial.is_empty && mask.true_count() > 0; + (accumulate_decimal(&mut partial.sum, d, &mask), any_valid) + } _ => 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); - } + if any_valid { + partial.is_empty = false; + } + if result? { + partial.is_overflow = true; + partial.is_empty = false; } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) + let sum = partials.get_item(SUM_FIELD)?; + let is_invalid = partials + .get_item(IS_OVERFLOW_FIELD)? + .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? + .fill_null(true)?; + sum.mask(is_invalid.not()?) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + Ok(sum_value_scalar(partial)) } } -/// The group state for a sum aggregate, containing the accumulated value and configuration -/// needed for reset/result without external context. +/// In-memory state for sum accumulation. pub struct SumPartial { return_dtype: DType, - /// The current accumulated state, or `None` if saturated (checked overflow). - current: Option, + /// The non-null running sum, initialized to zero. + sum: SumState, + /// Whether checked arithmetic overflowed. + is_overflow: bool, + /// Whether no valid value has been accumulated. + is_empty: bool, /// Whether NaN values in float inputs are skipped. skip_nans: bool, } @@ -351,6 +427,60 @@ pub enum SumState { }, } +fn decode_sum_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { + vortex_ensure!(!scalar.is_null(), "Sum partial must not be null"); + + let Some(fields) = scalar.as_struct_opt() else { + vortex_bail!("Sum partial must be a struct, got {}", scalar.dtype()); + }; + let sum = fields + .field(SUM_FIELD) + .ok_or_else(|| vortex_err!("Sum partial is missing the `{SUM_FIELD}` field"))?; + let is_overflow = + bool::try_from(&fields.field(IS_OVERFLOW_FIELD).ok_or_else(|| { + vortex_err!("Sum partial is missing the `{IS_OVERFLOW_FIELD}` field") + })?)?; + let is_empty = bool::try_from( + &fields + .field(IS_EMPTY_FIELD) + .ok_or_else(|| vortex_err!("Sum partial is missing the `{IS_EMPTY_FIELD}` field"))?, + )?; + + Ok((sum, is_overflow, is_empty)) +} + +fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + sum.dtype().nullability() == Nullability::NonNullable + && sum.dtype().eq_ignore_nullability(return_dtype), + "Sum partial value has dtype {}, expected {}", + sum.dtype(), + return_dtype.as_nonnullable(), + ); + Ok(()) +} + +fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { + Ok(match state { + SumState::Unsigned(acc) => checked_add_u64(acc, u64::try_from(other)?), + SumState::Signed(acc) => checked_add_i64(acc, i64::try_from(other)?), + SumState::Float(acc) => { + *acc += f64::try_from(other)?; + false + } + SumState::Decimal { value, dtype } => { + let other = DecimalValue::try_from(other)?; + match value.checked_add(&other) { + Some(result) if result.fits_in_precision(*dtype) => { + *value = result; + false + } + Some(_) | None => true, + } + } + }) +} + fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { @@ -366,6 +496,124 @@ fn make_zero_state(return_dtype: &DType) -> SumState { } } +fn sum_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct(sum_partial_fields(sum_dtype), Nullability::Nullable) +} + +fn sum_partial_fields(sum_dtype: DType) -> StructFields { + StructFields::new( + FieldNames::from_iter([ + FieldName::from(SUM_FIELD), + FieldName::from(IS_OVERFLOW_FIELD), + FieldName::from(IS_EMPTY_FIELD), + ]), + vec![ + sum_dtype.as_nonnullable(), + DType::Bool(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + ) +} + +fn sum_state_scalar(partial: &SumPartial, nullability: Nullability) -> Scalar { + match &partial.sum { + SumState::Unsigned(v) => Scalar::primitive(*v, nullability), + SumState::Signed(v) => Scalar::primitive(*v, nullability), + SumState::Float(v) => Scalar::primitive(*v, nullability), + SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), + } +} + +fn sum_value_scalar(partial: &SumPartial) -> Scalar { + if partial.is_overflow || partial.is_empty { + return Scalar::null(partial.return_dtype.as_nullable()); + } + + sum_state_scalar(partial, Nullability::Nullable) +} + +/// Convert scalar legacy Sum partials read from storage (a single nullable primitive) to the struct partial shape. +/// +/// A legacy non-null scalar becomes a non-empty partial, while a legacy null becomes an overflowed +/// partial. Legacy Sum used zero for empty inputs, so a scalar partial cannot represent +/// `is_empty = true` and defaults to `false`. +pub fn normalize_legacy_partial_array(partials: ArrayRef) -> VortexResult { + if matches!(partials.dtype(), DType::Struct(..)) { + vortex_bail!( + "Expected scalar legacy Sum partials, found {}", + partials.dtype() + ); + } + + let len = partials.len(); + let sum_dtype = partials.dtype().as_nonnullable(); + let is_overflow = partials.is_null()?; + let sum = partials.fill_null(Scalar::zero_value(&sum_dtype))?; + let is_empty = ConstantArray::new(false, len).into_array(); + + Ok(StructArray::try_new( + FieldNames::from_iter([ + FieldName::from(SUM_FIELD), + FieldName::from(IS_OVERFLOW_FIELD), + FieldName::from(IS_EMPTY_FIELD), + ]), + vec![sum, is_overflow, is_empty], + len, + Validity::AllValid, + )? + .into_array()) +} + +fn sum_result_partial_scalar( + result: Scalar, + return_dtype: &DType, + is_empty: bool, +) -> VortexResult { + if !result.dtype().eq_ignore_nullability(return_dtype) { + vortex_bail!( + "Sum result has dtype {}, expected {}", + result.dtype(), + return_dtype + ); + } + + let is_overflow = result.is_null() && !is_empty; + let sum = if result.is_null() { + Scalar::zero_value(&return_dtype.as_nonnullable()) + } else { + result.cast(&return_dtype.as_nonnullable())? + }; + Ok(Scalar::struct_( + sum_partial_dtype(return_dtype.clone()), + vec![ + sum, + Scalar::bool(is_overflow, Nullability::NonNullable), + Scalar::bool(is_empty, Nullability::NonNullable), + ], + )) +} + +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); + }; + if sum.is_null() { + // A finalized null cannot distinguish an empty batch from overflow, so it is not a + // mergeable partial. Array stats cannot currently store nulls, but recompute if that + // representation changes in the future. + return Ok(false); + } + + let sum = sum.cast(&partial.return_dtype)?; + let partial_scalar = sum_result_partial_scalar(sum, &partial.return_dtype, false)?; + vtable.combine_partials(partial, partial_scalar)?; + Ok(true) +} + /// Checked add for u64, returning true if overflow occurred. #[inline(always)] fn checked_add_u64(acc: &mut u64, val: u64) -> bool { @@ -391,12 +639,16 @@ fn checked_add_i64(acc: &mut i64, val: i64) -> bool { } #[cfg(test)] -mod tests { +mod tests; + +#[cfg(test)] +mod arithmetic_tests { use num_traits::CheckedAdd; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use super::sum_result_partial_scalar; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; @@ -405,8 +657,8 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::BoolArray; @@ -431,6 +683,11 @@ mod tests { use crate::scalar::Scalar; use crate::validity::Validity; + fn partial_with_value(value: Scalar) -> VortexResult { + let return_dtype = value.dtype().as_nullable(); + sum_result_partial_scalar(value, &return_dtype, false) + } + /// Sum 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(); @@ -441,9 +698,11 @@ mod tests { return sum(array, &mut ctx); } - let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("Sum not supported for dtype: {}", array.dtype()) - })?; + let sum_dtype = Sum + .return_dtype(&SumAggregateOpts::default(), array.dtype()) + .ok_or_else(|| { + vortex_error::vortex_err!("Sum 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()) @@ -495,7 +754,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(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -512,7 +771,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(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -531,17 +790,24 @@ mod tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let options = SumAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); - Sum.combine_partials(&mut state, scalar1)?; + Sum.combine_partials(&mut state, partial_with_value(scalar1)?)?; let scalar2 = Scalar::primitive(50i64, Nullable); - Sum.combine_partials(&mut state, scalar2)?; + Sum.combine_partials(&mut state, partial_with_value(scalar2)?)?; let result = Sum.to_scalar(&state)?; Sum.reset(&mut state); - assert_eq!(result.as_primitive().typed_value::(), Some(150)); + assert_eq!( + result + .as_struct() + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(150) + ); Ok(()) } @@ -584,11 +850,8 @@ mod tests { // Grouped sum tests fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - elem_dtype.clone(), - )?; + let mut acc = + GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -652,7 +915,7 @@ 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(); + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -676,8 +939,7 @@ 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(Sum, NumericalAggregateOpts::default(), elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); @@ -745,7 +1007,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(); @@ -754,7 +1016,7 @@ mod tests { &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, 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 87d8da4b143..cb021023e15 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -7,11 +7,11 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; +use vortex_mask::Mask; 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; @@ -24,10 +24,9 @@ const SUM_CHUNK: usize = 1 << 16; pub(super) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, - ctx: &mut ExecutionCtx, + mask: &Mask, 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), @@ -183,12 +182,13 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; + use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::ConstantArray; @@ -260,7 +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())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -268,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::primitive(0f64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); Ok(()) } @@ -289,20 +289,20 @@ mod tests { } #[test] - fn sum_empty_produces_zero() -> VortexResult<()> { + fn sum_empty_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } #[test] - fn sum_empty_f64_produces_zero() -> VortexResult<()> { + fn sum_empty_f64_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + assert!(result.is_null()); Ok(()) } @@ -345,11 +345,8 @@ mod tests { Ok(()) } - /// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). - fn sum_with_options( - arr: &crate::ArrayRef, - options: NumericalAggregateOpts, - ) -> VortexResult { + /// Sum an array with explicit [`SumAggregateOpts`] (test-only helper). + fn sum_with_options(arr: &ArrayRef, options: SumAggregateOpts) -> VortexResult { let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; acc.accumulate(arr, &mut array_session().create_execution_ctx())?; acc.finish() @@ -359,7 +356,7 @@ mod tests { 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())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -368,7 +365,7 @@ mod tests { 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())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -381,7 +378,7 @@ mod tests { 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())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -395,7 +392,7 @@ mod tests { .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())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) } @@ -404,10 +401,10 @@ mod tests { 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())?; + let result = sum_with_options(&arr, SumAggregateOpts::default())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -425,7 +422,7 @@ mod tests { let mut acc = Accumulator::try_new( Sum, - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; @@ -444,7 +441,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(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; assert!(!acc.is_saturated()); let batch = diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs new file mode 100644 index 00000000000..c168b39e757 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -0,0 +1,992 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the default [`Sum`] behavior: the `{sum, is_overflow, is_empty}` state algebra, the +//! null-for-zero-valid-values rule across input kinds, NaN and overflow handling, +//! cached-statistic consumption, and grouped aggregation. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::Sum; +use super::SumAggregateOpts; +use super::sum; +use super::sum_result_partial_scalar; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::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::Scalar; +use crate::scalar::ScalarValue; +use crate::validity::Validity; + +/// Sum an array with explicit [`SumAggregateOpts`] (test-only helper). +fn sum_with_options(arr: &ArrayRef, options: SumAggregateOpts) -> VortexResult { + let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; + acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.finish() +} + +fn partial_with_value(value: Scalar) -> VortexResult { + let return_dtype = value.dtype().as_nullable(); + sum_result_partial_scalar(value, &return_dtype, false) +} + +#[test] +fn sum_uses_new_partial_shape_by_default() { + let options = SumAggregateOpts::default(); + let sum = Sum.bind(options); + assert_eq!(sum.id().as_ref(), "vortex.sum"); + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let partial_dtype = Sum.partial_dtype(&options, &input_dtype).unwrap(); + assert_eq!(partial_dtype.nullability(), Nullable); + let fields = partial_dtype.as_struct_fields(); + assert_eq!(fields.names().as_ref(), &["sum", "is_overflow", "is_empty"]); + assert_eq!( + fields.field("sum"), + Some(DType::Primitive(PType::I64, Nullability::NonNullable)) + ); + assert_eq!( + fields.field("is_overflow"), + Some(DType::Bool(Nullability::NonNullable)) + ); + assert_eq!( + fields.field("is_empty"), + Some(DType::Bool(Nullability::NonNullable)) + ); +} + +#[test] +fn legacy_options_only_describe_the_stored_partial() -> VortexResult<()> { + let options = SumAggregateOpts::deserialize(&NumericalAggregateOpts::skip_nans().serialize())?; + assert_eq!( + options, + SumAggregateOpts { + skip_nans: true, + struct_partial: false, + } + ); + + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!(matches!( + Sum.partial_dtype(&options, &input_dtype), + Some(DType::Struct(..)) + )); + + let mut acc = Accumulator::try_new(Sum, options, input_dtype)?; + assert_eq!( + acc.partial_scalar()? + .as_struct() + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(true) + ); + assert!(acc.finish()?.is_null()); + Ok(()) +} + +// State algebra: the `{sum, is_overflow, is_empty}` monoid. + +#[test] +fn sum_rejects_null_partial() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + let partial_dtype = Sum.to_scalar(&state)?.dtype().clone(); + + assert!( + Sum.combine_partials(&mut state, Scalar::null(partial_dtype)) + .is_err() + ); + Ok(()) +} + +#[test] +fn sum_rejects_partial_with_wrong_sum_dtype() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + let wrong_partial = partial_with_value(Scalar::primitive(1i32, Nullable))?; + + assert!(Sum.combine_partials(&mut state, wrong_partial).is_err()); + Ok(()) +} + +#[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(&SumAggregateOpts::default(), &dtype)?; + let empty = Sum.to_scalar(&state)?; + let fields = empty.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(0) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(false) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(true) + ); + Sum.combine_partials(&mut state, empty)?; + let partial = Sum.to_scalar(&state)?; + assert_eq!( + partial + .as_struct() + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(true) + ); + Ok(()) +} + +#[test] +fn sum_state_empty_is_identity() -> VortexResult<()> { + // Combining an empty state into a non-empty state changes nothing. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + Sum.combine_partials( + &mut state, + partial_with_value(Scalar::primitive(100i64, Nullable))?, + )?; + + let empty = Sum.to_scalar(&Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?)?; + Sum.combine_partials(&mut state, empty)?; + + let result = Sum.to_scalar(&state)?; + assert_eq!( + result + .as_struct() + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(100) + ); + Ok(()) +} + +#[test] +fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { + // Overflow sets the flag and poisons the merge even when combined with later values. + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut overflowed = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + Sum.combine_partials( + &mut overflowed, + partial_with_value(Scalar::primitive(i64::MAX, Nullable))?, + )?; + Sum.combine_partials( + &mut overflowed, + partial_with_value(Scalar::primitive(1i64, Nullable))?, + )?; + let overflowed = Sum.to_scalar(&overflowed)?; + let fields = overflowed.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(false) + ); + + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + Sum.combine_partials( + &mut state, + partial_with_value(Scalar::primitive(5i64, Nullable))?, + )?; + Sum.combine_partials(&mut state, overflowed)?; + Sum.combine_partials( + &mut state, + partial_with_value(Scalar::primitive(7i64, Nullable))?, + )?; + + let partial = Sum.to_scalar(&state)?; + assert_eq!( + partial + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + 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(Sum, SumAggregateOpts::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 = sum(&array, &mut array_session().create_execution_ctx())?; + assert!(result.is_null()); + Ok(()) +} + +#[test] +fn sum_all_nan_is_zero_not_null() -> VortexResult<()> { + // NaNs are valid values: with the default `skip_nans` they contribute nothing, but + // the sum is a genuine `0.0`, unlike an all-null array whose sum is null. + let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); + let result = sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) +} + +#[test] +fn stat_sum_remains_a_finalized_scalar() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + assert!(sum(&arr, &mut ctx)?.is_null()); + assert!(arr.statistics().get(Stat::Sum).is_absent()); + + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(0i64))); + assert_eq!( + sum(&arr, &mut ctx)?.as_primitive().typed_value::(), + Some(0) + ); + + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; + assert!( + Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable)) + .is_err(), + "live Sum only accepts canonical struct partials" + ); + 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 = Sum + .return_dtype(&SumAggregateOpts::default(), &input) + .unwrap(); + assert_eq!(dtype, expected); +} + +// One value smoke test per accumulate branch; summation arithmetic is pinned by the +// shared kernels' tests in the `sum` module. + +#[test] +fn sum_primitive_with_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); + let result = sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6)); + Ok(()) +} + +#[test] +fn sum_bool_with_nulls() -> VortexResult<()> { + let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); + let result = sum( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().typed_value::(), Some(2)); + Ok(()) +} + +#[test] +fn sum_decimal_with_nulls() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![100i32, 200i32, 300i32, 400i32], + DecimalDType::new(4, 2), + Validity::from_iter([true, false, true, true]), + ); + let result = 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 = sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result, 50u64.into()); + Ok(()) +} + +#[test] +fn sum_constant_false_is_zero_not_null() -> VortexResult<()> { + let array = ConstantArray::new(false, 10).into_array(); + let result = sum(&array, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) +} + +#[test] +fn sum_multi_batch_and_finish_resets() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + + let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + acc.accumulate(&batch1, &mut ctx)?; + let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(48)); + + // finish resets the state: an untouched accumulator is empty again. + assert!(acc.finish()?.is_null()); + let batch3 = PrimitiveArray::new(buffer![1i32], Validity::NonNullable).into_array(); + acc.accumulate(&batch3, &mut ctx)?; + assert_eq!(acc.finish()?.as_primitive().typed_value::(), Some(1)); + Ok(()) +} + +// Chunked accumulation: the nullable sum must merge across chunks. + +#[test] +fn sum_chunked_floats_with_nulls() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter(vec![Some(1.5f64), None, Some(3.2), Some(4.8)]); + let chunk2 = PrimitiveArray::from_option_iter(vec![Some(2.1f64), Some(5.7), None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(17.3)); + Ok(()) +} + +#[test] +fn sum_chunked_all_nulls_is_null() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter::(vec![None, None, None]); + let chunk2 = PrimitiveArray::from_option_iter::(vec![None, None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + let result = sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert!(result.is_null()); + Ok(()) +} + +#[test] +fn sum_chunked_empty_chunks() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter(vec![Some(10.5f64), Some(20.3)]); + let chunk2 = ConstantArray::new(Scalar::primitive(0f64, Nullable), 0); + let chunk3 = PrimitiveArray::from_option_iter(vec![Some(5.2f64)]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new( + vec![ + chunk1.into_array(), + chunk2.into_array(), + chunk3.into_array(), + ], + dtype, + )?; + + let result = sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(36.0)); + Ok(()) +} + +#[test] +fn sum_chunked_value_survives_empty_chunk() -> VortexResult<()> { + // One valid value in one chunk, followed by an all-null chunk: the value must survive merging + // with the second chunk's null identity. + let chunk1 = PrimitiveArray::from_option_iter::(vec![Some(1)]); + let chunk2 = PrimitiveArray::from_option_iter::(vec![None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = sum( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!(result.as_primitive().as_::(), Some(1)); + Ok(()) +} + +// NaN handling and its interplay with the nullable sum and cached statistics. + +#[test] +fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(f64::NAN), Some(3.0)]) + .into_array(); + let result = 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, SumAggregateOpts::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, SumAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_uses_cached_stat_sum() -> VortexResult<()> { + // The planted result differs from the actual data, proving the cache was consumed. + 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))); + let result = 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))); + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) +} + +#[test] +fn sum_constant_nan() -> VortexResult<()> { + let arr = ConstantArray::new(f64::NAN, 4).into_array(); + // NaN constants are skipped by default (a non-empty zero sum) and poison the sum otherwise. + let result = sum_with_options(&arr, SumAggregateOpts::default())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + + let result = sum_with_options(&arr, SumAggregateOpts::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 = 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( + Sum, + SumAggregateOpts::default(), + DType::Primitive(PType::F64, Nullability::NonNullable), + )?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + Ok(()) +} + +// Overflow: a null sum value plus an explicit saturation flag. + +#[test] +fn sum_checked_overflow_is_null_and_saturates() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::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 partial = acc.partial_scalar()?; + let fields = partial.as_struct(); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(false) + ); + + 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 = sum( + &decimal.into_array(), + &mut array_session().create_execution_ctx(), + )?; + assert_eq!( + result, + Scalar::null(DType::Decimal(decimal_dtype, Nullable)) + ); + Ok(()) +} + +#[test] +fn sum_decimal_near_precision_boundary() -> VortexResult<()> { + // Input precision 4 → return precision min(76, 4+10) = 14. + // Native type for precision 14 is I64 (max precision 18), so 14 < 18. + // Use combine_partials to push state near (but under) 10^14. + let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(99_999_999_999_990i64), + DecimalDType::new(14, 0), + Nullable, + ); + Sum.combine_partials(&mut state, partial_with_value(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); + Sum.combine_partials(&mut state, partial_with_value(small)?)?; + + let result = Sum.to_scalar(&state)?; + let fields = result.as_struct(); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(false) + ); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.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 = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(near_limit), + DecimalDType::new(14, 0), + Nullable, + ); + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; + + let one_more = Scalar::decimal( + DecimalValue::from(one_more), + DecimalDType::new(14, 0), + Nullable, + ); + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; + + let result = Sum.to_scalar(&state)?; + assert_eq!( + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + 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 = Sum.empty_partial(&SumAggregateOpts::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); + Sum.combine_partials(&mut state, partial_with_value(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(); + Sum.accumulate(&mut state, &columnar, &mut ctx)?; + + let result = Sum.to_scalar(&state)?; + assert_eq!( + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + 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(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish() +} + +#[test] +fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(5i64), None, Some(i64::MAX), Some(1)]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 0, 1, 2, 0].into_array(), + buffer![1i32, 0, 1, 2, 1].into_array(), + Validity::from_iter([true, true, true, true, false]), + )? + .into_array(); + let mut acc = GroupedAccumulator::try_new( + Sum, + SumAggregateOpts::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("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(false) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(false) + ); + + for index in [1, 2] { + let empty = partials.execute_scalar(index, &mut ctx)?; + let fields = empty.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(0) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(false) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(true) + ); + } + + let overflow = partials.execute_scalar(3, &mut ctx)?; + let fields = overflow.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(false) + ); + assert!(partials.execute_scalar(4, &mut ctx)?.is_null()); + Ok(()) +} + +#[test] +fn grouped_sum_fallback_empty_and_all_null_groups() -> VortexResult<()> { + // Bool elements are rejected by the primitive grouped kernel, forcing the generic + // per-group fallback: empty and all-null groups have null sums there too. + let mut ctx = array_session().create_execution_ctx(); + let elements = BoolArray::from_iter([Some(true), Some(true), None, None]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 2, 2].into_array(), + buffer![2i32, 0, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let result = run_grouped_sum(&groups, &DType::Bool(Nullable))?; + let expected = PrimitiveArray::from_option_iter([Some(2u64), None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_fixed_size_list() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); + let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_with_null_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) + .into_array(); + let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_with_null_group() -> VortexResult<()> { + // A null group must become a null row through the lazy finalize's validity handling. + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) + .into_array(); + let validity = Validity::from_iter([true, false, true]); + let groups = FixedSizeListArray::try_new(elements, 3, validity, 3)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::I32, Nullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + // The all-null group has a null sum + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_all_nan_is_zero_not_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![f64::NAN, f64::NAN, 3.0, 4.0], Validity::NonNullable) + .into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; + + let elem_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + + let expected = PrimitiveArray::from_option_iter([Some(0.0f64), Some(7.0)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_finish_resets() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype)?; + + let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; + acc.accumulate_list(&groups1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + + let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); + assert_arrays_eq!(&result1, &expected1, &mut ctx); + + let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; + acc.accumulate_list(&groups2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + + let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); + assert_arrays_eq!(&result2, &expected2, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); + let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new(buffer![1i32, 1, 1], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false, true]); + let groups = ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array(); + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups, &elem_dtype)?; + + // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. + let expected = + PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 92fac87892a..9580474d048 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -76,6 +76,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::session::AggregateFnSession; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::dtype::DType; @@ -173,17 +174,17 @@ mod tests { assert_eq!(deserialized, agg_fn); } - /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the - /// numeric aggregates, including the non-default NaN-including configuration. + /// Both Sum options survive a protobuf serialize/deserialize round-trip. #[rstest] - #[case(NumericalAggregateOpts::skip_nans())] - #[case(NumericalAggregateOpts::include_nans())] - fn numeric_aggregate_options_round_trip( - #[case] options: NumericalAggregateOpts, - ) -> VortexResult<()> { + #[case(SumAggregateOpts::skip_nans())] + #[case(SumAggregateOpts::include_nans())] + #[case(SumAggregateOpts { + skip_nans: true, + struct_partial: false, + })] + fn sum_aggregate_options_round_trip(#[case] options: SumAggregateOpts) -> VortexResult<()> { let session = crate::array_session(); let agg_fn = Sum.bind(options); - let proto = agg_fn.serialize_proto()?; let buf = proto.encode_to_vec(); let decoded = pb::AggregateFn::decode(buf.as_slice())?; @@ -193,6 +194,25 @@ mod tests { Ok(()) } + #[test] + fn legacy_sum_options_mark_a_stored_scalar_partial() -> VortexResult<()> { + let session = crate::array_session(); + let proto = pb::AggregateFn { + id: Sum.id().to_string(), + metadata: Some(NumericalAggregateOpts::skip_nans().serialize()), + }; + + let aggregate_fn = AggregateFnRef::from_proto(&proto, &session)?; + assert_eq!( + aggregate_fn.as_opt::(), + Some(&SumAggregateOpts { + skip_nans: true, + struct_partial: false, + }) + ); + Ok(()) + } + #[test] fn unknown_aggregate_fn_id_allow_unknown() { let session = VortexSession::empty().with::(); diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 149c72f499e..c82f42b98c1 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -46,8 +46,8 @@ mod tests { use crate::VortexSessionExecute; 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::sum::SumAggregateOpts; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -59,11 +59,8 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); - let mut acc = Accumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - batch.dtype().clone(), - )?; + let mut acc = + Accumulator::try_new(Sum, SumAggregateOpts::default(), batch.dtype().clone())?; acc.accumulate(batch, &mut ctx)?; acc.finish() } @@ -120,7 +117,7 @@ mod tests { DType::Primitive(PType::I32, Nullability::Nullable), )?; let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -158,7 +155,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/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 57e11135ef3..756d2eb4237 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -10,6 +10,7 @@ use enum_iterator::all; use num_enum::IntoPrimitive; use num_enum::TryFromPrimitive; +use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; @@ -189,7 +190,7 @@ impl Stat { Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. return aggregate_fn::fns::sum::Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); + .return_dtype(&SumAggregateOpts::skip_nans(), data_type); } }) } @@ -200,7 +201,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::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum.bind(SumAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 95b8ddbb55a..669f8344a96 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. /// @@ -83,7 +74,7 @@ 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) + Sum.return_dtype(&(*options).into(), elem_dtype) .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) } @@ -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 { @@ -134,60 +124,16 @@ 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 null-on-empty [`Sum`] behavior yields null for null, empty, and all-null lists. fn list_sum_impl( canonical: ArrayRef, elem_dtype: DType, options: &NumericalAggregateOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, (*options).into(), elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; - let sums = acc.finish()?; - - let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::() { - fsl.into_owned().into() - } else if let Some(lv) = canonical.as_opt::() { - lv.into_owned().into() - } else { - let dtype = canonical.dtype(); - vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}") - }; - - mask_empty_lists(grouped, sums, ctx) -} - -/// Applies a mask to `sums` that nullifies entries produced by lists without at least -/// one valid element. This is necessary because the grouped `Sum` aggregate only produces -/// nulls for null lists and sums that overflow. -fn mask_empty_lists( - grouped: GroupedArray, - sums: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let elements = grouped.elements(); - let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; - let ranges = grouped.group_ranges(ctx)?; - - let has_valid_element: BitBuffer = match elem_mask.bit_buffer() { - AllOr::All => match &ranges { - // fixed-size lists of non-zero width cannot have empty lists. - GroupRanges::FixedSizeList { size, .. } if *size > 0 => return Ok(sums), - GroupRanges::FixedSizeList { len, .. } => BitBuffer::full(false, *len), - GroupRanges::ListView { ranges } => ranges.iter().map(|&(_, size)| size > 0).collect(), - }, - AllOr::None => BitBuffer::full(false, ranges.len()), - AllOr::Some(bits) => ranges - .iter() - .map(|(offset, size)| size > 0 && bits.count_range(offset, offset + size) > 0) - .collect(), - }; - if has_valid_element.true_count() == has_valid_element.len() { - return Ok(sums); - } - - sums.mask(BoolArray::new(has_valid_element, Validity::NonNullable).into_array()) + acc.finish() } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index c961ac10d53..05a149d2a37 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -18,6 +18,7 @@ use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; +use crate::aggregate_fn::fns::sum::Sum; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::display::ExprDisplay; @@ -130,7 +131,13 @@ impl ScalarFnVTable for StatFn { } fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult { - let Some(dtype) = aggregate_fn.state_dtype(input_dtype) else { + let dtype = if aggregate_fn.is::() { + // Sum stats expose a scalar result even though canonical Sum state is a struct. + aggregate_fn.return_dtype(input_dtype) + } else { + aggregate_fn.state_dtype(input_dtype) + }; + let Some(dtype) = dtype else { vortex_bail!( "Aggregate function {} does not support input dtype {}", aggregate_fn, diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 1e0ceef02d3..92babebc772 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -17,6 +17,7 @@ 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::sum::SumAggregateOpts; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; @@ -50,11 +51,11 @@ fn bound_min_max(expr: BoundExpression) -> BoundExpression { /// 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, Sum.bind(SumAggregateOpts::skip_nans())) } fn bound_sum(expr: BoundExpression) -> BoundExpression { - bound_stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) + bound_stat(expr, Sum.bind(SumAggregateOpts::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 5ee407994de..d200756d475 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -18,6 +18,7 @@ use vortex::aggregate_fn::fns::max::Max; use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; use vortex::aggregate_fn::fns::sum::Sum; +use vortex::aggregate_fn::fns::sum::SumAggregateOpts; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -533,10 +534,14 @@ impl PushedAggregate { Ok(match self { Self::Min => Box::new(Accumulator::try_new(Min, opts, dtype)?), Self::Max => Box::new(Accumulator::try_new(Max, opts, dtype)?), - Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new( + Sum, + SumAggregateOpts::from(opts), + dtype, + )?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), - PairOptions(opts, opts), + PairOptions(opts.into(), opts), dtype, )?), Self::First => Box::new(Accumulator::try_new(First, AggregateEmptyOptions, dtype)?), diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 700aab86ce4..57a1a34c2c0 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -34,6 +34,7 @@ use crate::LazyReaderChildren; use crate::VTable; use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::zone_map::ZoneMap; +use crate::layouts::zoned::zone_map::normalize_sum_partial_fields; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = @@ -169,6 +170,7 @@ impl PruningState { async move { let mut ctx = session.create_execution_ctx(); let zones_array = zones_eval.await?.execute::(&mut ctx)?; + let zones_array = normalize_sum_partial_fields(zones_array, &aggregate_fns)?; // SAFETY: zoned layout validation checked that this zones child was // written from the same column dtype and aggregate stats-table schema. Ok(unsafe { diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index f35eb4a59a9..629387bea72 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -90,8 +91,21 @@ pub(crate) fn aggregate_stats_table_dtype( ) -> DType { DType::Struct( StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| { - aggregate_state_dtype(column_dtype, aggregate_fn) - .map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable())) + let dtype = if aggregate_fn + .as_opt::() + .is_some_and(|options| !options.struct_partial) + { + aggregate_fn.return_dtype(column_dtype).or_else(|| { + if let DType::Extension(ext) = column_dtype { + aggregate_fn.return_dtype(ext.storage_dtype()) + } else { + None + } + }) + } else { + aggregate_state_dtype(column_dtype, aggregate_fn) + }; + dtype.map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable())) })), Nullability::NonNullable, ) @@ -194,6 +208,7 @@ mod tests { 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::sum::SumAggregateOpts; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -260,7 +275,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - Sum.bind(NumericalAggregateOpts::skip_nans()), + Sum.bind(SumAggregateOpts::skip_nans()), ], ); @@ -269,8 +284,29 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - Sum.bind(NumericalAggregateOpts::skip_nans()).to_string(), + Sum.bind(SumAggregateOpts::skip_nans()).to_string(), ] ); } + + #[test] + fn stored_sum_dtype_uses_the_persisted_partial_shape() -> VortexResult<()> { + let column_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let legacy_options = + SumAggregateOpts::deserialize(&NumericalAggregateOpts::default().serialize())?; + let legacy = aggregate_stats_table_dtype(&column_dtype, &[Sum.bind(legacy_options)]); + let current = + aggregate_stats_table_dtype(&column_dtype, &[Sum.bind(SumAggregateOpts::default())]); + + assert!(matches!( + legacy.as_struct_fields().field("vortex.sum()"), + Some(DType::Primitive(PType::I64, Nullability::Nullable)) + )); + assert!(matches!( + current.as_struct_fields().field("vortex.sum()"), + Some(DType::Struct(..)) + )); + + Ok(()) + } } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index cca134c0759..92c4b4c9870 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,6 +25,7 @@ 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::sum::SumAggregateOpts; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; @@ -209,10 +210,10 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A let mut aggregate_fns = vec![max, min]; if Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) + .return_dtype(&SumAggregateOpts::skip_nans(), dtype) .is_some() { - aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); + aggregate_fns.push(Sum.bind(SumAggregateOpts::skip_nans())); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c84c0b443dd..2248aaed97c 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,6 +16,8 @@ use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum::normalize_legacy_partial_array; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -24,8 +26,12 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::eq; +use vortex_array::expr::fill_null; use vortex_array::expr::get_item; use vortex_array::expr::lit; +use vortex_array::expr::mask; +use vortex_array::expr::not; +use vortex_array::expr::or; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar_fn::EmptyOptions; @@ -79,8 +85,10 @@ impl ZoneMap { if &expected_dtype != array.dtype() { vortex_bail!("Array dtype does not match expected zone map dtype: {expected_dtype}"); } + let array = normalize_sum_partial_fields(array, &aggregate_fns)?; - // SAFETY: We checked that the array matches the expected stats-table schema. + // SAFETY: We checked the stored stats-table schema, then normalized legacy Sum fields to + // the canonical runtime representation. Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) }) } @@ -165,6 +173,28 @@ impl ZoneMap { } } +pub(super) fn normalize_sum_partial_fields( + array: StructArray, + aggregate_fns: &[AggregateFnRef], +) -> VortexResult { + let names = array.names().clone(); + let mut fields = array.iter_unmasked_fields().cloned().collect::>(); + + for aggregate_fn in aggregate_fns { + if !aggregate_fn.is::() { + continue; + } + let Some(index) = names.find(aggregate_fn.to_string()) else { + continue; + }; + if !matches!(fields[index].dtype(), DType::Struct(..)) { + fields[index] = normalize_legacy_partial_array(fields[index].clone())?; + } + } + + StructArray::try_new(names, fields, array.len(), array.struct_validity()) +} + struct ZoneMapStatsBinder<'a> { zone_map: &'a ZoneMap, } @@ -293,6 +323,17 @@ impl ZoneMap { fn aggregate_result_expr(stored: &AggregateFnRef, state_expr: Expression) -> Expression { if stored.is::() { get_item(BOUNDED_MAX_BOUND, state_expr) + } else if stored.is::() { + // The sum is null if either there was an overflow or the underlying array was empty + // (no valid elements). + let is_invalid = fill_null( + or( + get_item("is_overflow", state_expr.clone()), + get_item("is_empty", state_expr.clone()), + ), + lit(true), + ); + mask(get_item("sum", state_expr), not(is_invalid)) } else { state_expr } @@ -356,6 +397,8 @@ mod tests { 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::sum::SumAggregateOpts; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -406,6 +449,31 @@ mod tests { unsafe { NonZeroUsize::new_unchecked(64) } } + #[test] + fn legacy_scalar_sum_field_is_normalized_on_read() -> VortexResult<()> { + let options = + SumAggregateOpts::deserialize(&NumericalAggregateOpts::default().serialize())?; + let sum = Sum.bind(options); + let legacy_partials = + PrimitiveArray::from_option_iter([Some(0i64), Some(5i64), None]).into_array(); + let legacy_table = StructArray::from_fields(&[(sum.to_string(), legacy_partials)])?; + let zone_map = ZoneMap::try_new( + PType::I32.into(), + legacy_table, + Arc::new([sum.clone()]), + 1, + 3, + )?; + let result_expr = zone_map + .aggregate_field_expr(&sum) + .expect("normalized Sum field"); + let result = zone_map.array.into_array().apply(&result_expr)?; + let expected = + PrimitiveArray::from_option_iter([Some(0i64), Some(5i64), None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut SESSION.create_execution_ctx()); + Ok(()) + } + #[test] fn test_zone_map_prunes() { // Construct a zone map with 3 zones: diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index 00ffaac433c..745017e8e64 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -32,6 +32,14 @@ message NumericalAggregateOpts { bool skip_nans = 1; } +// Options for `vortex.sum`. The optional partial-shape marker distinguishes newly serialized +// options from the historical `NumericalAggregateOpts` encoding, which used the same tag for +// `skip_nans` but had no second field. +message SumAggregateOpts { + bool skip_nans = 1; + optional bool struct_partial = 2; +} + // Options for `vortex.literal` message LiteralOpts { vortex.scalar.Scalar value = 1; diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index a44328623e3..afbb82cb1d4 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -26,6 +26,16 @@ pub struct NumericalAggregateOpts { #[prost(bool, tag = "1")] pub skip_nans: bool, } +/// Options for `vortex.sum`. The optional partial-shape marker distinguishes newly serialized +/// options from the historical `NumericalAggregateOpts` encoding, which used the same tag for +/// `skip_nans` but had no second field. +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct SumAggregateOpts { + #[prost(bool, tag = "1")] + pub skip_nans: bool, + #[prost(bool, optional, tag = "2")] + pub struct_partial: ::core::option::Option, +} /// Options for `vortex.literal` #[derive(Clone, PartialEq, ::prost::Message)] pub struct LiteralOpts { diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index aece56aa9d6..096db5adedf 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -314,7 +314,7 @@ impl PyVortexWriteOptions { /// >>> vx.io.VortexWriteOptions.default().write(sprl, "chonky.vortex") /// >>> import os /// >>> os.path.getsize('chonky.vortex') - /// 215932 + /// 216108 /// /// Wow, Vortex manages to use about two bytes per integer! So advanced. So tiny. /// @@ -324,7 +324,7 @@ impl PyVortexWriteOptions { /// /// >>> vx.io.VortexWriteOptions.compact().write(sprl, "tiny.vortex") /// >>> os.path.getsize('tiny.vortex') - /// 55060 + /// 55232 /// /// Random numbers are not (usually) composed of random bytes! #[staticmethod] diff --git a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt index fbf9d2b74fe..c0c7f9128e9 100644 --- a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt +++ b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt @@ -13,8 +13,7 @@ statement ok COPY (SELECT CAST(x AS INTEGER) AS x FROM (VALUES (CAST(NULL AS INTEGER)),(CAST(NULL AS INTEGER))) t(x)) TO '${WORK_DIR}/i-null.vortex'; -# TODO(https://github.com/vortex-data/vortex/issues/9084): this should be fixed (expected NULL) query I SELECT sum(x) FROM '${WORK_DIR}/i-null.vortex'; ---- -0 +NULL diff --git a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt index 59931886824..b6640cf0f39 100644 --- a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt +++ b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt @@ -123,7 +123,7 @@ COPY ( query IRR SELECT count(x), sum(x), max(x) FROM '${WORK_DIR}/all-null.vortex'; ---- -0 0 NULL +0 NULL NULL query R SELECT min(x) FROM '${WORK_DIR}/all-null.vortex';