diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 312c5c9c184..11477e57503 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,12 +154,14 @@ fn list_element_dtype(list_view: &ArrayRef) -> DType { fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef where - V: AggregateFnVTable + Clone, - V::Options: Default, + V: AggregateFnVTable + Clone, { - let mut acc = - GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) - .unwrap(); + let mut acc = GroupedAccumulator::try_new( + vtable, + NumericalAggregateOpts::default(), + list_element_dtype(list_view), + ) + .unwrap(); acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); divan::black_box(acc.finish().unwrap()) @@ -197,59 +199,6 @@ 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 52255481a43..69bae4e1053 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -119,18 +119,9 @@ impl DynAccumulator for Accumulator { batch.dtype() ); - // 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. + // 0. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot, + // consume that exact stat before kernel dispatch or decode. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) - && stat != Stat::Sum && let Precision::Exact(partial) = batch.statistics().get(stat) { let partial = if partial.dtype() == &self.partial_dtype { @@ -176,7 +167,7 @@ impl DynAccumulator for Accumulator { } // 2. Allow the vtable to short-circuit on the raw array before decompression. - if !checked_cached_sum && self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { + if self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { return Ok(()); } @@ -289,7 +280,6 @@ 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; @@ -298,10 +288,7 @@ 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` @@ -333,7 +320,7 @@ mod tests { } } - /// Sum partial sentinel `{sum: 42.0, is_overflow: false, is_empty: false}` — distinguishable from the natural Sum of + /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -344,7 +331,7 @@ mod tests { _batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - Ok(Some(sum_partial(42.0))) + Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) } } @@ -363,7 +350,7 @@ mod tests { Accumulator::try_new( Mean::combined(), PairOptions( - SumAggregateOpts::default(), + NumericalAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -372,24 +359,11 @@ mod tests { fn sentinel_partial() -> Scalar { let acc = mean_f64_accumulator().expect("build accumulator"); - let sum = sum_partial(42.0); + let sum = Scalar::primitive(42.0f64, Nullability::Nullable); 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] @@ -407,13 +381,7 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum") - .unwrap() - .as_struct() - .field("sum") - .unwrap() - .as_primitive() - .as_::(), + s.field("sum").unwrap().as_primitive().as_::(), Some(42.0) ); assert_eq!( @@ -440,13 +408,7 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum") - .unwrap() - .as_struct() - .field("sum") - .unwrap() - .as_primitive() - .as_::(), + s.field("sum").unwrap().as_primitive().as_::(), Some(7.0) ); assert_eq!( @@ -477,13 +439,7 @@ 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_struct() - .field("sum") - .unwrap() - .as_primitive() - .as_::(), + s.field("sum").unwrap().as_primitive().as_::(), Some(42.0) ); assert_eq!( @@ -492,26 +448,4 @@ 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 8bd26dad961..286f4194cd5 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -273,10 +273,7 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn flush(&mut self) -> VortexResult { - let mut states = std::mem::take(&mut self.partials); - if states.len() == 1 { - return Ok(states.pop().vortex_expect("checked one partial")); - } + let states = std::mem::take(&mut self.partials); 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 16a1b5431fe..da93d2352a3 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -19,7 +19,6 @@ 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; @@ -41,7 +40,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - SumAggregateOpts::default(), + NumericalAggregateOpts::default(), NumericalAggregateOpts::default(), ), array.dtype().clone(), @@ -320,7 +319,7 @@ mod tests { let keep_nans = NumericalAggregateOpts::include_nans(); let mut acc = Accumulator::try_new( Mean::combined(), - PairOptions(keep_nans.into(), keep_nans), + PairOptions(keep_nans, keep_nans), array.dtype().clone(), )?; acc.accumulate(&array, &mut ctx)?; @@ -412,7 +411,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - SumAggregateOpts::default(), + NumericalAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -447,7 +446,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - SumAggregateOpts::default(), + NumericalAggregateOpts::default(), NumericalAggregateOpts::default(), ), DType::Primitive(PType::F64, Nullability::Nullable), @@ -475,7 +474,7 @@ mod tests { let mut acc = GroupedAccumulator::try_new( Mean::combined(), PairOptions( - SumAggregateOpts::default(), + NumericalAggregateOpts::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 4e33a567dd2..b3993840586 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -6,22 +6,23 @@ 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, - mask: &Mask, + ctx: &mut ExecutionCtx, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; + let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, @@ -39,8 +40,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; @@ -100,16 +101,16 @@ mod tests { &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } #[test] - fn sum_bool_empty_produces_null() -> VortexResult<()> { + fn sum_bool_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -117,7 +118,7 @@ mod tests { fn sum_bool_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -135,7 +136,7 @@ mod tests { fn sum_bool_return_dtype() -> VortexResult<()> { let dtype = Sum .return_dtype( - &SumAggregateOpts::default(), + &NumericalAggregateOpts::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 14da9d60799..0f366620e5c 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::null(DType::Primitive(PType::U64, Nullable))); + assert_eq!(result, Scalar::primitive(0u64, 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::null(DType::Primitive(PType::U64, Nullable))); + assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -187,7 +187,11 @@ mod tests { let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, - Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) + Scalar::decimal( + DecimalValue::I256(i256::ZERO), + DecimalDType::new(20, 2), + Nullable + ) ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 03c3f116fc5..872e5769a01 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -13,6 +13,7 @@ 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; @@ -25,9 +26,10 @@ use crate::scalar::DecimalValue; pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, - mask: &Mask, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = match mask { + let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + let validity = match &mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), Mask::AllFalse(_) => { @@ -109,12 +111,11 @@ 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; @@ -128,11 +129,6 @@ 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( @@ -361,31 +357,23 @@ 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(&SumAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::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)?)?; + Sum.combine_partials(&mut state, near_limit)?; // Add a small value that keeps us just under 10^14. let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, partial_with_value(small)?)?; + Sum.combine_partials(&mut state, 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!(!result.is_null()); assert_eq!( - fields - .field("sum") - .and_then(|sum| sum.as_decimal().decimal_value()), + result.as_decimal().decimal_value(), Some(DecimalValue::I256(i256::from_i128(99_999_999_999_999))) ); Ok(()) @@ -399,27 +387,25 @@ 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(&SumAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::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, partial_with_value(near_limit)?)?; + Sum.combine_partials(&mut state, near_limit)?; // Push the sum to exactly 10^14, exceeding precision 14. let one_more = Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; + Sum.combine_partials(&mut state, 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) + result.dtype(), + &DType::Decimal(DecimalDType::new(14, 0), Nullable) ); Ok(()) } @@ -428,30 +414,24 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::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, partial_with_value(near_limit)?)?; + Sum.combine_partials(&mut state, near_limit)?; let one_more = Scalar::decimal( DecimalValue::from(-1i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; + Sum.combine_partials(&mut state, 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) - ); + assert!(result.is_null()); Ok(()) } @@ -466,13 +446,13 @@ mod tests { // a real array that pushes it over. let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); let return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; let near_limit = Scalar::decimal(DecimalValue::from(near_limit_val), return_dtype, Nullable); - Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; + Sum.combine_partials(&mut state, near_limit)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = @@ -484,13 +464,7 @@ mod tests { 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) - ); + assert!(result.is_null()); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 64b0ce81995..efe0825d4d6 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,16 +1,14 @@ // 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; @@ -18,14 +16,10 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; use crate::dtype::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)] @@ -84,47 +78,31 @@ fn grouped_sum( .as_ref() .validity()? .execute_mask(elements.as_ref().len(), ctx)?; - let all_valid = elem_mask.all_true(); + let all_valid = matches!(elem_mask.slices(), AllOr::All); - let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), + let result = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); - collect_sums::( - values, group_ranges, group_validity, &elem_mask, all_valid, 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 }) } ); - 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()) + Ok(result.into_array()) } -/// Reduce each group's element slice into a non-null sum, overflow bitmap, and empty bitmap. +/// Reduce each group's element slice into a nullable sum. A group is null when the group +/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -132,37 +110,24 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> (PrimitiveArray, BitBuffer, BitBuffer) { - let group_count = group_ranges.len(); - let mut is_overflow = BitBufferMut::new_unset(group_count); - let mut is_empty = BitBufferMut::new_unset(group_count); +) -> PrimitiveArray { let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - return A::default(); + return None; } let mut acc = A::default(); - let (overflow, any_valid) = if all_valid { - (sum_run(&mut acc, &values[offset..offset + size]), size > 0) + let overflow = if all_valid { + sum_run(&mut acc, &values[offset..offset + size]) } else { sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) }; - 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 + (!overflow).then_some(acc) }); - let sums = PrimitiveArray::from_iter(sums); - (sums, is_overflow.freeze(), is_empty.freeze()) + PrimitiveArray::from_option_iter(sums) } -/// Sum valid runs in one group, returning `(overflow, any_valid)`. +/// 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. fn sum_masked_group( acc: &mut A, values: &[T], @@ -170,23 +135,17 @@ fn sum_masked_group( size: usize, elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> (bool, bool) { - match elem_mask { - Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), - Mask::AllFalse(_) => (false, false), - Mask::Values(mask_values) => { - let validity = mask_values - .bit_buffer() - .as_view() - .slice(offset..offset + size); - let mut any_valid = false; - for (start, end) in validity.set_slices() { - any_valid = true; +) -> bool { + match elem_mask.slice(offset..offset + size).slices() { + AllOr::All => sum_run(acc, &values[offset..offset + size]), + AllOr::None => false, + AllOr::Some(runs) => { + for &(start, end) in runs { if sum_run(acc, &values[offset + start..offset + end]) { - return (true, true); + return true; } } - (false, any_valid) + false } } } @@ -203,8 +162,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; @@ -220,8 +179,11 @@ 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, SumAggregateOpts::default(), elem_dtype.clone())?; + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -238,8 +200,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .return_dtype(&SumAggregateOpts::default(), elem_dtype) - .expect("sum return dtype"); + .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .expect("sum partial dtype"); let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); for (i, &(offset, size)) in ranges.iter().enumerate() { if group_valid[i] { @@ -329,7 +291,8 @@ mod tests { let actual = grouped_sum_actual(&groups, &elem_dtype)?; let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - let direct = PrimitiveArray::from_option_iter([Some(4i64), None, None, Some(9i64)]); + let direct = + PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) @@ -406,7 +369,7 @@ mod tests { let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; let mut acc = - GroupedAccumulator::try_new(Sum, SumAggregateOpts::include_nans(), elem_dtype)?; + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::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 89b19b5575f..7ad62eacb87 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -6,18 +6,12 @@ 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 prost::Message; +use vortex_error::VortexExpect; 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; @@ -29,48 +23,43 @@ 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 { - if let Precision::Exact(sum) = array.statistics().get(Stat::Sum) { - return Ok(sum); + // Short-circuit using cached array statistics. + if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { + return Ok(sum_scalar); } - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), array.dtype().clone())?; + // 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(), + )?; 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)); } @@ -78,95 +67,15 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// 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`]. +/// Sum an array, starting from zero. /// -/// 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(()) - } -} +/// 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)] +pub struct Sum; // Both Spark and DataFusion use this heuristic. // - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 @@ -179,7 +88,7 @@ pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { } impl AggregateFnVTable for Sum { - type Options = SumAggregateOpts; + type Options = NumericalAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { @@ -196,12 +105,12 @@ impl AggregateFnVTable for Sum { metadata: &[u8], _session: &VortexSession, ) -> VortexResult { - SumAggregateOpts::deserialize(metadata) + NumericalAggregateOpts::deserialize(metadata) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - // When a sum overflows, we return a null sum value. Therefore, all return dtypes are - // nullable. + // When a sum overflows, we return a sum _value_ of null. Therefore, we all return dtypes + // are nullable. use Nullability::Nullable; Some(match input_dtype { @@ -227,8 +136,7 @@ impl AggregateFnVTable for Sum { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - let return_dtype = self.return_dtype(options, input_dtype)?; - Some(sum_partial_dtype(return_dtype)) + self.return_dtype(options, input_dtype) } fn empty_partial( @@ -239,57 +147,94 @@ impl AggregateFnVTable for Sum { let return_dtype = self .return_dtype(options, input_dtype) .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; - let sum = make_zero_state(&return_dtype); + let initial = make_zero_state(&return_dtype); + Ok(SumPartial { return_dtype, - sum, - is_overflow: false, - is_empty: true, + current: Some(initial), skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let (other_sum, other_is_overflow, other_is_empty) = decode_sum_partial_scalar(other)?; - validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; - - if partial.is_overflow { + if other.is_null() { + // A null partial means the sub-accumulator saturated (overflow). + partial.current = None; return Ok(()); } - if other_is_overflow { - partial.is_overflow = true; - partial.is_empty = false; - return Ok(()); - } - if other_is_empty { + let Some(ref mut inner) = partial.current else { return Ok(()); + }; + let saturated = match inner { + SumState::Unsigned(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + checked_add_u64(acc, val) + } + SumState::Signed(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + checked_add_i64(acc, val) + } + SumState::Float(acc) => { + let val = other + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + *acc += val; + false + } + SumState::Decimal { value, dtype } => { + let val = other + .as_decimal() + .decimal_value() + .vortex_expect("checked non-null"); + match value.checked_add(&val) { + Some(r) => { + *value = r; + !value.fits_in_precision(*dtype) + } + None => true, + } + } + }; + if saturated { + partial.current = None; } - - partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; - partial.is_empty = false; Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::struct_( - sum_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), - ], - )) + Ok(match &partial.current { + None => Scalar::null(partial.return_dtype.as_nullable()), + Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable), + Some(SumState::Decimal { value, .. }) => { + let decimal_dtype = *partial + .return_dtype + .as_decimal_opt() + .vortex_expect("return dtype must be decimal"); + Scalar::decimal(*value, decimal_dtype, Nullability::Nullable) + } + }) } fn reset(&self, partial: &mut Self::Partial) { - partial.sum = make_zero_state(&partial.return_dtype); - partial.is_overflow = false; - partial.is_empty = true; + partial.current = Some(make_zero_state(&partial.return_dtype)); } #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { - partial.is_overflow || matches!(&partial.sum, SumState::Float(v) if v.is_nan()) + match partial.current.as_ref() { + None => true, + Some(SumState::Float(v)) => v.is_nan(), + Some(_) => false, + } } fn try_accumulate( @@ -298,27 +243,31 @@ impl AggregateFnVTable for Sum { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - 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(_)) { + // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes + // the default dispatch path. + if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { Precision::Exact(0) => { // NaN-free batch: the cached NaN-skipping sum (if any) equals the // NaN-including sum. - try_accumulate_cached_sum(self, partial, batch) + if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + self.combine_partials(partial, sum)?; + return Ok(true); + } + Ok(false) } Precision::Exact(_) => { // At least one NaN value: the sum is NaN without scanning the batch. - let SumState::Float(acc) = &mut partial.sum else { - unreachable!("checked float sum state") - }; - *acc = f64::NAN; - partial.is_empty = false; + if let Some(SumState::Float(acc)) = partial.current.as_mut() { + *acc = f64::NAN; + } Ok(true) } _ => Ok(false), @@ -331,85 +280,60 @@ 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 (result, any_valid) = match batch { + let mut inner = match partial.current.take() { + Some(inner) => inner, + None => return Ok(()), + }; + + let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => { - 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) - } + Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans), + Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), + Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), }; - if any_valid { - partial.is_empty = false; - } - if result? { - partial.is_overflow = true; - partial.is_empty = false; + match result { + Ok(false) => partial.current = Some(inner), + Ok(true) => {} // saturated: current stays None + Err(e) => { + partial.current = Some(inner); + return Err(e); + } } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - let sum = partials.get_item(SUM_FIELD)?; - let is_invalid = partials - .get_item(IS_OVERFLOW_FIELD)? - .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? - .fill_null(true)?; - sum.mask(is_invalid.not()?) + Ok(partials) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(sum_value_scalar(partial)) + self.to_scalar(partial) } } -/// In-memory state for sum accumulation. +/// The group state for a sum aggregate, containing the accumulated value and configuration +/// needed for reset/result without external context. pub struct SumPartial { return_dtype: DType, - /// 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, + /// The current accumulated state, or `None` if saturated (checked overflow). + current: Option, /// Whether NaN values in float inputs are skipped. skip_nans: bool, } @@ -427,60 +351,6 @@ 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 { @@ -496,124 +366,6 @@ 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 { @@ -639,16 +391,12 @@ fn checked_add_i64(acc: &mut i64, val: i64) -> bool { } #[cfg(test)] -mod tests; - -#[cfg(test)] -mod arithmetic_tests { +mod 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; @@ -657,8 +405,8 @@ mod arithmetic_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; @@ -683,11 +431,6 @@ mod arithmetic_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(); @@ -698,11 +441,9 @@ mod arithmetic_tests { return sum(array, &mut ctx); } - let sum_dtype = Sum - .return_dtype(&SumAggregateOpts::default(), array.dtype()) - .ok_or_else(|| { - vortex_error::vortex_err!("Sum not supported for dtype: {}", array.dtype()) - })?; + let sum_dtype = Stat::Sum.dtype(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()) @@ -754,7 +495,7 @@ mod arithmetic_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, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -771,7 +512,7 @@ mod arithmetic_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, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -790,24 +531,17 @@ mod arithmetic_tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let options = SumAggregateOpts::default(); - let mut state = Sum.empty_partial(&options, &dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); - Sum.combine_partials(&mut state, partial_with_value(scalar1)?)?; + Sum.combine_partials(&mut state, scalar1)?; let scalar2 = Scalar::primitive(50i64, Nullable); - Sum.combine_partials(&mut state, partial_with_value(scalar2)?)?; + Sum.combine_partials(&mut state, scalar2)?; let result = Sum.to_scalar(&state)?; Sum.reset(&mut state); - assert_eq!( - result - .as_struct() - .field("sum") - .and_then(|sum| sum.as_primitive().typed_value::()), - Some(150) - ); + assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) } @@ -850,8 +584,11 @@ mod arithmetic_tests { // Grouped sum tests fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = - GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -915,7 +652,7 @@ mod arithmetic_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([None, Some(7i64)]).into_array(); + let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -939,7 +676,8 @@ mod arithmetic_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, SumAggregateOpts::default(), elem_dtype)?; + let mut acc = + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); @@ -1007,7 +745,7 @@ mod arithmetic_tests { } #[test] - fn sum_chunked_floats_all_nulls_is_null() -> VortexResult<()> { + fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> { let chunk1 = PrimitiveArray::from_option_iter::(vec![None, None, None]); let chunk2 = PrimitiveArray::from_option_iter::(vec![None, None]); let dtype = chunk1.dtype().clone(); @@ -1016,7 +754,7 @@ mod arithmetic_tests { &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); + assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index cb021023e15..87d8da4b143 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,9 +24,10 @@ const SUM_CHUNK: usize = 1 << 16; pub(super) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, - mask: &Mask, + ctx: &mut ExecutionCtx, skip_nans: bool, ) -> VortexResult { + let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), AllOr::All => accumulate_primitive_all(inner, p, skip_nans), @@ -182,13 +183,12 @@ 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!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); 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::null(DType::Primitive(PType::F64, Nullable))); + assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } @@ -289,20 +289,20 @@ mod tests { } #[test] - fn sum_empty_produces_null() -> VortexResult<()> { + fn sum_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } #[test] - fn sum_empty_f64_produces_null() -> VortexResult<()> { + fn sum_empty_f64_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) } @@ -345,8 +345,11 @@ mod tests { Ok(()) } - /// Sum an array with explicit [`SumAggregateOpts`] (test-only helper). - fn sum_with_options(arr: &ArrayRef, options: SumAggregateOpts) -> VortexResult { + /// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + fn sum_with_options( + arr: &crate::ArrayRef, + options: NumericalAggregateOpts, + ) -> VortexResult { let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; acc.accumulate(arr, &mut array_session().create_execution_ctx())?; acc.finish() @@ -356,7 +359,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, SumAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -365,7 +368,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, SumAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -378,7 +381,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, SumAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -392,7 +395,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, SumAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) } @@ -401,10 +404,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, SumAggregateOpts::default())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); - let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -422,7 +425,7 @@ mod tests { let mut acc = Accumulator::try_new( Sum, - SumAggregateOpts::default(), + NumericalAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; @@ -441,7 +444,7 @@ mod tests { #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::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 deleted file mode 100644 index c168b39e757..00000000000 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ /dev/null @@ -1,992 +0,0 @@ -// 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 9580474d048..92fac87892a 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -76,7 +76,6 @@ 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; @@ -174,17 +173,17 @@ mod tests { assert_eq!(deserialized, agg_fn); } - /// Both Sum options survive a protobuf serialize/deserialize round-trip. + /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the + /// numeric aggregates, including the non-default NaN-including configuration. #[rstest] - #[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<()> { + #[case(NumericalAggregateOpts::skip_nans())] + #[case(NumericalAggregateOpts::include_nans())] + fn numeric_aggregate_options_round_trip( + #[case] options: NumericalAggregateOpts, + ) -> 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())?; @@ -194,25 +193,6 @@ 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 c82f42b98c1..149c72f499e 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,8 +59,11 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); - let mut acc = - Accumulator::try_new(Sum, SumAggregateOpts::default(), batch.dtype().clone())?; + let mut acc = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + batch.dtype().clone(), + )?; acc.accumulate(batch, &mut ctx)?; acc.finish() } @@ -117,7 +120,7 @@ mod tests { DType::Primitive(PType::I32, Nullability::Nullable), )?; let result = run_sum(&chunked.into_array())?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -155,7 +158,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let chunked = ChunkedArray::try_new(vec![], dtype)?; let result = run_sum(&chunked.into_array())?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 756d2eb4237..57e11135ef3 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -10,7 +10,6 @@ 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; @@ -190,7 +189,7 @@ impl Stat { Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. return aggregate_fn::fns::sum::Sum - .return_dtype(&SumAggregateOpts::skip_nans(), data_type); + .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); } }) } @@ -201,7 +200,7 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::sum::Sum.bind(SumAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 669f8344a96..95b8ddbb55a 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -1,9 +1,11 @@ // 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; @@ -14,16 +16,23 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; +use crate::arrays::BoolArray; use crate::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. /// @@ -74,7 +83,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).into(), elem_dtype) + Sum.return_dtype(options, elem_dtype) .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) } @@ -91,6 +100,7 @@ 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 { @@ -124,16 +134,60 @@ impl ScalarFnVTable for ListSum { /// Sum each list of a canonical `array` into one value per list. /// -/// The null-on-empty [`Sum`] behavior yields null for null, empty, and all-null lists. +/// Note that we need to nullify sums produced by empty or all-null lists, +/// since grouped sum kernels default to 0 for these. fn list_sum_impl( canonical: ArrayRef, elem_dtype: DType, options: &NumericalAggregateOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, (*options).into(), elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; - acc.finish() + 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()) } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 05a149d2a37..c961ac10d53 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -18,7 +18,6 @@ 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; @@ -131,13 +130,7 @@ impl ScalarFnVTable for StatFn { } fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult { - 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 { + let Some(dtype) = aggregate_fn.state_dtype(input_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 92babebc772..1e0ceef02d3 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -17,7 +17,6 @@ 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; @@ -51,11 +50,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(SumAggregateOpts::skip_nans())) + stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } fn bound_sum(expr: BoundExpression) -> BoundExpression { - bound_stat(expr, Sum.bind(SumAggregateOpts::skip_nans())) + bound_stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index d200756d475..5ee407994de 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -18,7 +18,6 @@ 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; @@ -534,14 +533,10 @@ 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, - SumAggregateOpts::from(opts), - dtype, - )?), + Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), - PairOptions(opts.into(), opts), + PairOptions(opts, 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 57a1a34c2c0..700aab86ce4 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -34,7 +34,6 @@ 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 = @@ -170,7 +169,6 @@ 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 629387bea72..f35eb4a59a9 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -7,7 +7,6 @@ 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; @@ -91,21 +90,8 @@ pub(crate) fn aggregate_stats_table_dtype( ) -> DType { DType::Struct( StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| { - 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())) + aggregate_state_dtype(column_dtype, aggregate_fn) + .map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable())) })), Nullability::NonNullable, ) @@ -208,7 +194,6 @@ 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; @@ -275,7 +260,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - Sum.bind(SumAggregateOpts::skip_nans()), + Sum.bind(NumericalAggregateOpts::skip_nans()), ], ); @@ -284,29 +269,8 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - Sum.bind(SumAggregateOpts::skip_nans()).to_string(), + Sum.bind(NumericalAggregateOpts::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 92c4b4c9870..cca134c0759 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,7 +25,6 @@ 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; @@ -210,10 +209,10 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A let mut aggregate_fns = vec![max, min]; if Sum - .return_dtype(&SumAggregateOpts::skip_nans(), dtype) + .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) .is_some() { - aggregate_fns.push(Sum.bind(SumAggregateOpts::skip_nans())); + aggregate_fns.push(Sum.bind(NumericalAggregateOpts::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 2248aaed97c..c84c0b443dd 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,8 +16,6 @@ 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; @@ -26,12 +24,8 @@ 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; @@ -85,10 +79,8 @@ 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 the stored stats-table schema, then normalized legacy Sum fields to - // the canonical runtime representation. + // SAFETY: We checked that the array matches the expected stats-table schema. Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) }) } @@ -173,28 +165,6 @@ 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, } @@ -323,17 +293,6 @@ 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 } @@ -397,8 +356,6 @@ 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; @@ -449,31 +406,6 @@ 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 745017e8e64..00ffaac433c 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -32,14 +32,6 @@ 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 afbb82cb1d4..a44328623e3 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -26,16 +26,6 @@ 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 096db5adedf..aece56aa9d6 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') - /// 216108 + /// 215932 /// /// 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') - /// 55232 + /// 55060 /// /// 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 c0c7f9128e9..fbf9d2b74fe 100644 --- a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt +++ b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt @@ -13,7 +13,8 @@ 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'; ---- -NULL +0 diff --git a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt index b6640cf0f39..59931886824 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 NULL NULL +0 0 NULL query R SELECT min(x) FROM '${WORK_DIR}/all-null.vortex';