From 61bf257e9aec6b816ce91b2b5ab999d704e317d7 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 31 Jul 2026 13:30:52 +0100 Subject: [PATCH 01/18] Use canonical partial state for Sum Signed-off-by: Matt Katz --- encodings/sparse/src/compute/sum.rs | 1 + vortex-array/benches/aggregate_grouped.rs | 125 ++- vortex-array/src/aggregate_fn/accumulator.rs | 54 +- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 8 +- .../src/aggregate_fn/fns/sum/constant.rs | 12 +- .../src/aggregate_fn/fns/sum/decimal.rs | 22 +- .../src/aggregate_fn/fns/sum/grouped.rs | 82 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 367 +++++-- .../src/aggregate_fn/fns/sum/primitive.rs | 20 +- .../src/aggregate_fn/fns/sum/tests.rs | 899 ++++++++++++++++++ vortex-array/src/aggregate_fn/proto.rs | 1 - .../src/arrays/chunked/compute/aggregate.rs | 4 +- vortex-array/src/expr/stats/mod.rs | 4 +- vortex-array/src/scalar_fn/fns/list_sum.rs | 58 +- vortex-array/src/scalar_fn/fns/stat.rs | 4 +- vortex-layout/src/layouts/zoned/mod.rs | 50 +- vortex-layout/src/layouts/zoned/pruning.rs | 2 + vortex-layout/src/layouts/zoned/schema.rs | 47 + vortex-layout/src/layouts/zoned/writer.rs | 8 +- vortex-layout/src/layouts/zoned/zone_map.rs | 57 ++ 21 files changed, 1574 insertions(+), 253 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/sum/tests.rs diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index 8f85124b3b5..223589e8a7f 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -100,6 +100,7 @@ mod tests { kernel_result, canonical_result, "kernel and canonical sum paths disagree" ); + Ok(kernel_result) } diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..739b51dcee1 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -11,9 +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::Accumulator; use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; @@ -23,6 +26,9 @@ use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -152,27 +158,63 @@ fn list_element_dtype(list_view: &ArrayRef) -> DType { } } -fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef +fn grouped_accumulator(list_view: &ArrayRef, vtable: V, options: V::Options) -> ArrayRef where - V: AggregateFnVTable + Clone, + V: AggregateFnVTable + Clone, { - let mut acc = GroupedAccumulator::try_new( - vtable, - NumericalAggregateOpts::default(), - list_element_dtype(list_view), - ) - .unwrap(); + let mut acc = + GroupedAccumulator::try_new(vtable, options, list_element_dtype(list_view)).unwrap(); acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); divan::black_box(acc.finish().unwrap()) } +#[divan::bench] +fn sum_legacy_scalar_partial_merge(bencher: Bencher) { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let partial = Scalar::primitive(1i64, Nullability::Nullable); + bencher + .with_inputs(|| partial.clone()) + .bench_refs(|partial| { + let mut acc = + Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()) + .unwrap(); + for _ in 0..GROUP_COUNT { + acc.combine_partials(partial.clone()).unwrap(); + } + divan::black_box(acc.finish().unwrap()) + }); +} + +#[divan::bench] +fn sum_canonical_partial_merge(bencher: Bencher) { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut source = + Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()).unwrap(); + source + .combine_partials(Scalar::primitive(1i64, Nullability::Nullable)) + .unwrap(); + let partial = source.flush().unwrap(); + + bencher + .with_inputs(|| partial.clone()) + .bench_refs(|partial| { + let mut acc = + Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()) + .unwrap(); + for _ in 0..GROUP_COUNT { + acc.combine_partials(partial.clone()).unwrap(); + } + divan::black_box(acc.finish().unwrap()) + }); +} + #[divan::bench] fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); } #[divan::bench] @@ -180,7 +222,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); } #[divan::bench] @@ -188,7 +230,7 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); } #[divan::bench] @@ -196,7 +238,62 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); +} + +/// 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, + options: V::Options, +) -> ArrayRef +where + V: AggregateFnVTable + Clone, +{ + let mut acc = + GroupedAccumulator::try_new(vtable, options, 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, NumericalAggregateOpts::default()) + }); +} + +#[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, NumericalAggregateOpts::default()) + }); +} + +#[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, NumericalAggregateOpts::default()) + }); +} + +#[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, NumericalAggregateOpts::default()) + }); } #[divan::bench] @@ -204,7 +301,7 @@ fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Count)); + .bench_refs(|input| grouped_accumulator(input, Count, NumericalAggregateOpts::default())); } #[divan::bench] @@ -212,5 +309,5 @@ fn count_varbinview(bencher: Bencher) { let input = varbinview_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Count)); + .bench_refs(|input| grouped_accumulator(input, Count, NumericalAggregateOpts::default())); } diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 69bae4e1053..0499b5d1ecb 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -124,19 +124,8 @@ impl DynAccumulator for Accumulator { if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - let partial = if partial.dtype() == &self.partial_dtype { - partial - } else { - vortex_ensure!( - partial.dtype().eq_ignore_nullability(&self.partial_dtype), - "Aggregate {} read legacy stat {} with dtype {}, expected {}", - self.aggregate_fn, - stat, - partial.dtype(), - self.partial_dtype, - ); - partial.cast(&self.partial_dtype)? - }; + // Legacy stat slots can use an older partial shape. The aggregate vtable owns that + // compatibility logic (for example, Sum accepts both scalar and struct partials). self.vtable.combine_partials(&mut self.partial, partial)?; return Ok(()); } @@ -320,7 +309,7 @@ mod tests { } } - /// Sum partial sentinel `42.0` — distinguishable from the natural Sum of + /// Sum partial sentinel `{sum: 42.0, is_overflow: false, is_empty: false}` — distinguishable from the natural Sum of /// `dict_of_seven()` which is `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -331,7 +320,7 @@ mod tests { _batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) + Ok(Some(sum_partial(42.0))) } } @@ -359,11 +348,20 @@ mod tests { fn sentinel_partial() -> Scalar { let acc = mean_f64_accumulator().expect("build accumulator"); - let sum = Scalar::primitive(42.0f64, Nullability::Nullable); + let sum = sum_partial(42.0); let count = Scalar::primitive(1u64, Nullability::NonNullable); Scalar::struct_(acc.partial_dtype, vec![sum, count]) } + fn sum_partial(value: f64) -> Scalar { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = + Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype).expect("sum"); + acc.combine_partials(Scalar::primitive(value, Nullability::Nullable)) + .expect("legacy scalar partial"); + acc.flush().expect("sum partial") + } + /// Kernel registered for `(Dict, Combined)` fires in preference to /// `Combined::try_accumulate`'s fan-out path — proves the dispatch reorder. #[test] @@ -381,7 +379,13 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(42.0) ); assert_eq!( @@ -408,7 +412,13 @@ mod tests { let s = partial.as_struct(); assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(7.0) ); assert_eq!( @@ -439,7 +449,13 @@ mod tests { // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the // batch's valid_count, so count is the real 1. assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), + s.field("sum") + .unwrap() + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), Some(42.0) ); assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index da93d2352a3..d2c1ac78604 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -333,7 +333,7 @@ mod tests { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; - assert_eq!(result.as_primitive().as_::(), None); + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3993840586..80aa62831bd 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -13,7 +13,7 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; -pub(super) fn accumulate_bool( +pub(crate) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, @@ -101,16 +101,16 @@ mod tests { &arr.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } #[test] - fn sum_bool_empty_produces_zero() -> VortexResult<()> { + fn sum_bool_empty_produces_null() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..1e64e5348d8 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -15,7 +15,7 @@ use crate::scalar::Scalar; /// /// Returns `Ok(None)` if the scalar is null (no contribution to the sum). /// Returns a null scalar on overflow (saturation). -pub(super) fn multiply_constant( +pub(crate) fn multiply_constant( scalar: &Scalar, len: usize, return_dtype: &DType, @@ -127,7 +127,7 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -151,7 +151,7 @@ mod tests { fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); let result = sum(&array, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0u64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::U64, Nullable))); Ok(()) } @@ -187,11 +187,7 @@ mod tests { let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, - Scalar::decimal( - DecimalValue::I256(i256::ZERO), - DecimalDType::new(20, 2), - Nullable - ) + Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullable)) ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 872e5769a01..60c4e20de33 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -23,7 +23,7 @@ use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. /// Returns Ok(true) if saturated (overflow), Ok(false) if not. -pub(super) fn accumulate_decimal( +pub(crate) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, @@ -357,7 +357,8 @@ mod tests { // Native type for precision 14 is I64 (max precision 18), so 14 < 18. // Use combine_partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let options = NumericalAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), @@ -370,7 +371,7 @@ mod tests { let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, small)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(!result.is_null()); assert_eq!( result.as_decimal().decimal_value(), @@ -387,7 +388,8 @@ mod tests { // i256 arithmetic does not overflow. This tests the precision-based // saturation path in combine_partials. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let options = NumericalAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), @@ -401,7 +403,7 @@ mod tests { Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); assert_eq!( result.dtype(), @@ -414,7 +416,8 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let options = NumericalAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), @@ -430,7 +433,7 @@ mod tests { ); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } @@ -446,7 +449,8 @@ mod tests { // a real array that pushes it over. let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); let return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let options = NumericalAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; @@ -463,7 +467,7 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index efe0825d4d6..a8630dee99b 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,10 +1,15 @@ // 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::IS_EMPTY_FIELD; +use super::IS_OVERFLOW_FIELD; +use super::SUM_FIELD; use super::Sum; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; @@ -16,10 +21,16 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; use crate::dtype::NativePType; +use crate::dtype::Nullability; use crate::match_each_native_ptype; +use crate::validity::Validity; /// Encoding-specific grouped [`Sum`] kernel for primitive element arrays. #[derive(Debug)] @@ -80,29 +91,43 @@ fn grouped_sum( .execute_mask(elements.as_ref().len(), ctx)?; let all_valid = matches!(elem_mask.slices(), AllOr::All); - let result = match_each_native_ptype!(elements.ptype(), + let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), unsigned: |T| { let values = elements.as_slice::(); - collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, - sum_unsigned_all) + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, sum_unsigned_all) }, signed: |T| { let values = elements.as_slice::(); - collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, - sum_signed_all) + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, sum_signed_all) }, floating: |T| { let values = elements.as_slice::(); - collect_sums::(values, group_ranges, group_validity, &elem_mask, all_valid, + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, |acc, slice| { sum_float_all(acc, slice, skip_nans); false }) } ); - Ok(result.into_array()) + Ok(StructArray::try_new( + FieldNames::from_iter([ + FieldName::from(SUM_FIELD), + FieldName::from(IS_OVERFLOW_FIELD), + FieldName::from(IS_EMPTY_FIELD), + ]), + vec![ + sums.into_array(), + BoolArray::new(is_overflow, Validity::NonNullable).into_array(), + BoolArray::new(is_empty, Validity::NonNullable).into_array(), + ], + group_validity.len(), + Validity::from_mask(group_validity.clone(), Nullability::Nullable), + )? + .into_array()) } -/// Reduce each group's element slice into a nullable sum. A group is null when the group -/// itself is invalid, or when summing it overflows (`sum_run` returns `true`). +/// Reduce each group's element slice into a non-null sum, overflow bitmap, and empty bitmap. fn collect_sums( values: &[T], group_ranges: &GroupRanges, @@ -110,24 +135,30 @@ fn collect_sums( elem_mask: &Mask, all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> PrimitiveArray { +) -> (PrimitiveArray, BitBuffer, BitBuffer) { + let mut is_overflow = BitBufferMut::with_capacity(group_ranges.len()); + let mut is_empty = BitBufferMut::with_capacity(group_ranges.len()); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - return None; + is_overflow.append(false); + is_empty.append(true); + return A::default(); } let mut acc = A::default(); - let overflow = if all_valid { - sum_run(&mut acc, &values[offset..offset + size]) + let (overflow, any_valid) = if all_valid { + (sum_run(&mut acc, &values[offset..offset + size]), size > 0) } else { sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) }; - (!overflow).then_some(acc) + is_overflow.append(overflow); + is_empty.append(!any_valid); + acc }); - PrimitiveArray::from_option_iter(sums) + let sums = PrimitiveArray::from_iter(sums); + (sums, is_overflow.freeze(), is_empty.freeze()) } -/// Sum the valid elements of a single group, using the contiguous valid runs of the element mask -/// intersected with the group's `[offset, offset + size)` range. +/// Sum valid runs in one group, returning `(overflow, any_valid)`. fn sum_masked_group( acc: &mut A, values: &[T], @@ -135,17 +166,17 @@ fn sum_masked_group( size: usize, elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> bool { +) -> (bool, bool) { match elem_mask.slice(offset..offset + size).slices() { - AllOr::All => sum_run(acc, &values[offset..offset + size]), - AllOr::None => false, + AllOr::All => (sum_run(acc, &values[offset..offset + size]), size > 0), + AllOr::None => (false, false), AllOr::Some(runs) => { for &(start, end) in runs { if sum_run(acc, &values[offset + start..offset + end]) { - return true; + return (true, true); } } - false + (false, !runs.is_empty()) } } } @@ -200,8 +231,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) - .expect("sum partial dtype"); + .return_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .expect("sum return dtype"); let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); for (i, &(offset, size)) in ranges.iter().enumerate() { if group_valid[i] { @@ -291,8 +322,7 @@ mod tests { let actual = grouped_sum_actual(&groups, &elem_dtype)?; let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - let direct = - PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); + let direct = PrimitiveArray::from_option_iter([Some(4i64), None, None, Some(9i64)]); assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 7ad62eacb87..a93b3fd910c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -23,34 +23,40 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::arrays::ConstantArray; +use crate::arrays::StructArray; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; use crate::dtype::MAX_PRECISION; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::StructFields; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +const SUM_FIELD: &str = "sum"; +const IS_OVERFLOW_FIELD: &str = "is_overflow"; +const IS_EMPTY_FIELD: &str = "is_empty"; /// Return the sum of an array. /// /// See [`Sum`] for details. pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // Short-circuit using cached array statistics. - if let Precision::Exact(sum_scalar) = array.statistics().get(Stat::Sum) { - return Ok(sum_scalar); - } - - // 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(), @@ -59,7 +65,6 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { acc.accumulate(array, ctx)?; let result = acc.finish()?; - // Cache the computed sum as a statistic (only if non-null, i.e. no overflow). if let Some(val) = result.value().cloned() { array.statistics().set(Stat::Sum, Precision::Exact(val)); } @@ -67,14 +72,11 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// Sum an array, starting from zero. -/// -/// If the sum overflows, a null scalar will be returned. -/// If the array is all-invalid, the sum will be zero. +/// Sum an array, returning null when it has no valid values. /// -/// 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)] +/// If the sum overflows, a null scalar is returned. Legacy scalar partials remain supported; their +/// zero identity preserves the historical zero-on-empty behavior when encountered. +#[derive(Clone, Copy, Debug)] pub struct Sum; // Both Spark and DataFusion use this heuristic. @@ -109,8 +111,8 @@ impl AggregateFnVTable for Sum { } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - // When a sum overflows, we return a sum _value_ of null. Therefore, we all return dtypes - // are nullable. + // When a sum overflows, we return a null sum value. Therefore, all return dtypes are + // nullable. use Nullability::Nullable; Some(match input_dtype { @@ -136,7 +138,8 @@ impl AggregateFnVTable for Sum { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + let return_dtype = self.return_dtype(options, input_dtype)?; + Some(sum_partial_dtype(return_dtype)) } fn empty_partial( @@ -147,25 +150,41 @@ impl AggregateFnVTable for Sum { let return_dtype = self .return_dtype(options, input_dtype) .ok_or_else(|| vortex_err!("Unsupported sum dtype: {}", input_dtype))?; - let initial = make_zero_state(&return_dtype); - + let sum = make_zero_state(&return_dtype); Ok(SumPartial { return_dtype, - current: Some(initial), + sum, + is_overflow: false, + is_empty: true, skip_nans: options.skip_nans, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - // A null partial means the sub-accumulator saturated (overflow). - partial.current = None; + let other = normalize_partial_scalar(other, &partial.return_dtype)?; + let fields = other.as_struct(); + let other = fields + .field(SUM_FIELD) + .ok_or_else(|| vortex_err!("Sum partial is missing the `{SUM_FIELD}` field"))?; + let other_is_overflow = fields + .field(IS_OVERFLOW_FIELD) + .and_then(|is_overflow| is_overflow.as_bool().value()) + .ok_or_else(|| vortex_err!("Sum partial has an invalid `{IS_OVERFLOW_FIELD}` field"))?; + let other_is_empty = fields + .field(IS_EMPTY_FIELD) + .and_then(|is_empty| is_empty.as_bool().value()) + .ok_or_else(|| vortex_err!("Sum partial has an invalid `{IS_EMPTY_FIELD}` field"))?; + + partial.is_empty &= other_is_empty; + if partial.is_overflow || other_is_overflow { + partial.is_overflow = true; return Ok(()); } - let Some(ref mut inner) = partial.current else { + if other_is_empty { return Ok(()); - }; - let saturated = match inner { + } + + let saturated = match &mut partial.sum { SumState::Unsigned(acc) => { let val = other .as_primitive() @@ -203,38 +222,33 @@ impl AggregateFnVTable for Sum { } }; if saturated { - partial.current = None; + partial.is_overflow = true; + } else { + partial.is_empty = false; } Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(match &partial.current { - None => Scalar::null(partial.return_dtype.as_nullable()), - Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable), - Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable), - Some(SumState::Float(v)) => Scalar::primitive(*v, Nullability::Nullable), - Some(SumState::Decimal { value, .. }) => { - let decimal_dtype = *partial - .return_dtype - .as_decimal_opt() - .vortex_expect("return dtype must be decimal"); - Scalar::decimal(*value, decimal_dtype, Nullability::Nullable) - } - }) + Ok(Scalar::struct_( + sum_partial_dtype(partial.return_dtype.clone()), + vec![ + sum_state_scalar(partial), + Scalar::bool(partial.is_overflow, Nullability::NonNullable), + Scalar::bool(partial.is_empty, Nullability::NonNullable), + ], + )) } fn reset(&self, partial: &mut Self::Partial) { - partial.current = Some(make_zero_state(&partial.return_dtype)); + partial.sum = make_zero_state(&partial.return_dtype); + partial.is_overflow = false; + partial.is_empty = true; } #[inline] fn is_saturated(&self, partial: &Self::Partial) -> bool { - match partial.current.as_ref() { - None => true, - Some(SumState::Float(v)) => v.is_nan(), - Some(_) => false, - } + partial.is_overflow || matches!(&partial.sum, SumState::Float(v) if v.is_nan()) } fn try_accumulate( @@ -243,31 +257,27 @@ impl AggregateFnVTable for Sum { batch: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult { - // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes - // the default dispatch path. - if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + if partial.skip_nans { + return try_accumulate_cached_sum(self, partial, batch); + } + + // NaN-aware short-circuits only apply to NaN-including float sums. + if !matches!(&partial.sum, SumState::Float(_)) { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { Precision::Exact(0) => { // NaN-free batch: the cached NaN-skipping sum (if any) equals the // NaN-including sum. - if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { - let sum = if sum.dtype() == &partial.return_dtype { - sum - } else { - sum.cast(&partial.return_dtype)? - }; - self.combine_partials(partial, sum)?; - return Ok(true); - } - Ok(false) + try_accumulate_cached_sum(self, partial, batch) } Precision::Exact(_) => { // At least one NaN value: the sum is NaN without scanning the batch. - if let Some(SumState::Float(acc)) = partial.current.as_mut() { - *acc = f64::NAN; - } + let SumState::Float(acc) = &mut partial.sum else { + unreachable!("checked float sum state") + }; + *acc = f64::NAN; + partial.is_empty = false; Ok(true) } _ => Ok(false), @@ -280,8 +290,15 @@ 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(()); @@ -293,47 +310,72 @@ impl AggregateFnVTable for Sum { } let skip_nans = partial.skip_nans; - let mut inner = match partial.current.take() { - Some(inner) => inner, - None => return Ok(()), + let any_valid = if partial.is_empty { + match batch { + Columnar::Canonical(c) => match c { + Canonical::Primitive(p) => { + any_valid(p.as_ref().validity()?, p.as_ref().len(), ctx)? + } + Canonical::Bool(b) => any_valid(b.as_ref().validity()?, b.as_ref().len(), ctx)?, + Canonical::Decimal(d) => { + any_valid(d.as_ref().validity()?, d.as_ref().len(), ctx)? + } + _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), + }, + Columnar::Constant(_) => unreachable!(), + } + } else { + false }; let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans), - Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), - Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx), + Canonical::Primitive(p) => { + accumulate_primitive(&mut partial.sum, p, ctx, skip_nans) + } + Canonical::Bool(b) => accumulate_bool(&mut partial.sum, b, ctx), + Canonical::Decimal(d) => accumulate_decimal(&mut partial.sum, d, ctx), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), }; match result { - Ok(false) => partial.current = Some(inner), - Ok(true) => {} // saturated: current stays None - Err(e) => { - partial.current = Some(inner); - return Err(e); + Ok(false) => { + if any_valid { + partial.is_empty = false; + } } + Ok(true) => partial.is_overflow = true, + Err(e) => return Err(e), } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) + let partials = normalize_partial_array(partials)?; + let sum = partials.get_item(SUM_FIELD)?; + let is_invalid = partials + .get_item(IS_OVERFLOW_FIELD)? + .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? + .fill_null(true)?; + sum.mask(is_invalid.not()?) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + Ok(sum_value_scalar(partial)) } } -/// The group state for a sum aggregate, containing the accumulated value and configuration -/// needed for reset/result without external context. +/// In-memory state for sum accumulation. pub struct SumPartial { return_dtype: DType, - /// The current accumulated state, or `None` if saturated (checked overflow). - current: Option, + /// The non-null running sum, initialized to zero. + sum: SumState, + /// Whether checked arithmetic overflowed. + is_overflow: bool, + /// Whether no valid value has been accumulated. + is_empty: bool, /// Whether NaN values in float inputs are skipped. skip_nans: bool, } @@ -351,7 +393,7 @@ pub enum SumState { }, } -fn make_zero_state(return_dtype: &DType) -> SumState { +pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0), @@ -366,9 +408,152 @@ fn make_zero_state(return_dtype: &DType) -> SumState { } } +fn sum_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct( + 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), + ], + ), + Nullability::Nullable, + ) +} + +fn sum_state_scalar(partial: &SumPartial) -> Scalar { + match &partial.sum { + SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::NonNullable), + SumState::Signed(v) => Scalar::primitive(*v, Nullability::NonNullable), + SumState::Float(v) => Scalar::primitive(*v, Nullability::NonNullable), + 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::NonNullable) + } + } +} + +fn sum_value_scalar(partial: &SumPartial) -> Scalar { + if partial.is_overflow || partial.is_empty { + return Scalar::null(partial.return_dtype.as_nullable()); + } + + match &partial.sum { + SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::Nullable), + SumState::Signed(v) => Scalar::primitive(*v, Nullability::Nullable), + SumState::Float(v) => Scalar::primitive(*v, Nullability::Nullable), + 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) + } + } +} + +/// Normalize an array of scalar legacy Sum partials into the canonical struct partial shape. +/// +/// Canonical partial arrays are returned unchanged. 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`. +pub fn normalize_partial_array(partials: ArrayRef) -> VortexResult { + if matches!(partials.dtype(), DType::Struct(..)) { + return Ok(partials); + } + + 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 normalize_partial_scalar(partial: Scalar, return_dtype: &DType) -> VortexResult { + let partial_dtype = sum_partial_dtype(return_dtype.clone()); + if matches!(partial.dtype(), DType::Struct(..)) { + if partial.is_null() { + return Ok(Scalar::struct_( + partial_dtype, + vec![ + Scalar::zero_value(&return_dtype.as_nonnullable()), + Scalar::bool(true, Nullability::NonNullable), + Scalar::bool(false, Nullability::NonNullable), + ], + )); + } + return partial.cast(&partial_dtype); + } + + if !partial.dtype().eq_ignore_nullability(return_dtype) { + vortex_bail!( + "Legacy Sum partial has dtype {}, expected {}", + partial.dtype(), + return_dtype + ); + } + + let is_overflow = partial.is_null(); + let sum = if is_overflow { + Scalar::zero_value(&return_dtype.as_nonnullable()) + } else { + partial.cast(&return_dtype.as_nonnullable())? + }; + Ok(Scalar::struct_( + partial_dtype, + vec![ + sum, + Scalar::bool(is_overflow, Nullability::NonNullable), + Scalar::bool(false, 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); + }; + + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + vtable.combine_partials(partial, sum)?; + Ok(true) +} + +fn any_valid(validity: Validity, len: usize, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(validity.execute_mask(len, ctx)?.true_count() > 0) +} + /// Checked add for u64, returning true if overflow occurred. #[inline(always)] -fn checked_add_u64(acc: &mut u64, val: u64) -> bool { +pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; @@ -380,7 +565,7 @@ fn checked_add_u64(acc: &mut u64, val: u64) -> bool { /// Checked add for i64, returning true if overflow occurred. #[inline(always)] -fn checked_add_i64(acc: &mut i64, val: i64) -> bool { +pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; @@ -391,7 +576,10 @@ fn checked_add_i64(acc: &mut i64, val: i64) -> bool { } #[cfg(test)] -mod tests { +mod tests; + +#[cfg(test)] +mod arithmetic_tests { use num_traits::CheckedAdd; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -531,7 +719,8 @@ mod tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let options = NumericalAggregateOpts::default(); + let mut state = Sum.empty_partial(&options, &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); Sum.combine_partials(&mut state, scalar1)?; @@ -539,7 +728,7 @@ mod tests { let scalar2 = Scalar::primitive(50i64, Nullable); Sum.combine_partials(&mut state, scalar2)?; - let result = Sum.to_scalar(&state)?; + let result = Sum.finalize_scalar(&state)?; Sum.reset(&mut state); assert_eq!(result.as_primitive().typed_value::(), Some(150)); Ok(()) @@ -652,7 +841,7 @@ mod tests { let elem_dtype = DType::Primitive(PType::I32, Nullable); let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; - let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); + let expected = PrimitiveArray::from_option_iter([None, Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -745,7 +934,7 @@ mod tests { } #[test] - fn sum_chunked_floats_all_nulls_is_zero() -> VortexResult<()> { + fn sum_chunked_floats_all_nulls_is_null() -> VortexResult<()> { let chunk1 = PrimitiveArray::from_option_iter::(vec![None, None, None]); let chunk2 = PrimitiveArray::from_option_iter::(vec![None, None]); let dtype = chunk1.dtype().clone(); @@ -754,7 +943,7 @@ mod tests { &chunked.into_array(), &mut array_session().create_execution_ctx(), )?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..8728ebe88eb 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -21,7 +21,7 @@ use crate::match_each_native_ptype; /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. const SUM_CHUNK: usize = 1 << 16; -pub(super) fn accumulate_primitive( +pub(crate) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -66,7 +66,7 @@ fn accumulate_primitive_all( /// Sum the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values /// are skipped to match the scalar `sum` semantics; otherwise any NaN poisons the accumulator to /// NaN. Floats cannot overflow the accumulator, so this never reports saturation. -pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { +pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { if skip_nans { for &v in slice { if !v.is_nan() { @@ -84,7 +84,7 @@ pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan /// chunks of [`SUM_CHUNK`] with a single checked add per chunk, which lets the inner loop vectorize /// to packed widening adds. `u64` input keeps a per-element checked add since a chunk of `u64`s /// could itself overflow. Returns `true` on overflow. -pub(super) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool +pub(crate) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { @@ -106,7 +106,7 @@ where } /// Signed counterpart of [`sum_unsigned_all`]. -pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool +pub(crate) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { @@ -260,7 +260,7 @@ mod tests { fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -268,7 +268,7 @@ mod tests { fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let result = sum(&arr, &mut array_session().create_execution_ctx())?; - assert_eq!(result, Scalar::primitive(0f64, Nullable)); + assert_eq!(result, Scalar::null(DType::Primitive(PType::F64, Nullable))); Ok(()) } @@ -289,20 +289,20 @@ mod tests { } #[test] - fn sum_empty_produces_zero() -> VortexResult<()> { + fn sum_empty_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } #[test] - fn sum_empty_f64_produces_zero() -> VortexResult<()> { + fn sum_empty_f64_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; - assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs new file mode 100644 index 00000000000..68683071ee0 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -0,0 +1,899 @@ +// 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::sum; +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::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; +use crate::validity::Validity; + +/// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). +fn sum_with_options(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { + let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; + acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.finish() +} + +#[test] +fn sum_uses_new_partial_shape_by_default() { + let options = NumericalAggregateOpts::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)) + ); +} + +// State algebra: the `{sum, is_overflow, is_empty}` monoid. + +#[test] +fn sum_state_empty_is_null() -> VortexResult<()> { + // A state that never saw a valid value finalizes to null, and combining empty states + // stays empty. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&NumericalAggregateOpts::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)?; + assert!(Sum.finalize_scalar(&state)?.is_null()); + Ok(()) +} + +#[test] +fn sum_state_empty_is_identity() -> VortexResult<()> { + // Combining an empty state into a non-empty state changes nothing. + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut state, Scalar::primitive(100i64, Nullable))?; + + let empty = Sum.to_scalar(&Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?)?; + Sum.combine_partials(&mut state, empty)?; + + let result = Sum.finalize_scalar(&state)?; + assert_eq!(result.as_primitive().typed_value::(), Some(100)); + Ok(()) +} + +#[test] +fn sum_state_overflow_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(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut overflowed, Scalar::primitive(i64::MAX, Nullable))?; + Sum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; + let overflowed = Sum.to_scalar(&overflowed)?; + let 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(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; + Sum.combine_partials(&mut state, overflowed)?; + Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; + + assert!(Sum.finalize_scalar(&state)?.is_null()); + Ok(()) +} + +// 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, NumericalAggregateOpts::default(), dtype)?; + assert!(acc.finish()?.is_null()); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::from_option_iter([None::, None, None]).into_array())] +#[case::float(PrimitiveArray::from_option_iter::([None, None, None]).into_array())] +#[case::bool(BoolArray::from_iter([None::, None, None]).into_array())] +#[case::constant_primitive( + ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10).into_array() +)] +#[case::constant_bool(ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array())] +#[case::constant_decimal( + ConstantArray::new(Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullable)), 10) + .into_array() +)] +fn sum_all_null_is_null(#[case] array: ArrayRef) -> VortexResult<()> { + let result = 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 legacy_scalar_partial_preserves_zero_on_empty() -> 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()); + + // A scalar `Stat::Sum` is an old partial. Its zero identity cannot encode emptiness, so its + // historical zero-on-empty result is preserved when it is encountered. + 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(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable))?; + assert_eq!( + Sum.finalize_scalar(&state)? + .as_primitive() + .typed_value::(), + Some(0) + ); + + let mut overflowed = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + Sum.combine_partials( + &mut overflowed, + Scalar::null(DType::Primitive(PType::I64, 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(0) + ); + 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) + ); + 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(&NumericalAggregateOpts::default(), &input) + .unwrap(); + assert_eq!(dtype, expected); +} + +// One value smoke test per accumulate branch; summation arithmetic is pinned by the +// shared kernels' tests in the `sum` module. + +#[test] +fn sum_primitive_with_nulls() -> VortexResult<()> { + let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); + let result = 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, NumericalAggregateOpts::default(), dtype)?; + + let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + acc.accumulate(&batch1, &mut ctx)?; + let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().typed_value::(), Some(48)); + + // finish resets the state: an untouched accumulator is empty again. + assert!(acc.finish()?.is_null()); + let batch3 = PrimitiveArray::new(buffer![1i32], Validity::NonNullable).into_array(); + acc.accumulate(&batch3, &mut ctx)?; + assert_eq!(acc.finish()?.as_primitive().typed_value::(), Some(1)); + Ok(()) +} + +// Chunked accumulation: the nullable sum must merge across chunks. + +#[test] +fn sum_chunked_floats_with_nulls() -> VortexResult<()> { + let chunk1 = PrimitiveArray::from_option_iter(vec![Some(1.5f64), None, Some(3.2), Some(4.8)]); + let chunk2 = PrimitiveArray::from_option_iter(vec![Some(2.1f64), Some(5.7), None]); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; + + let result = 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, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + // The array has no NaNs; a planted exact NaNCount stat proves the NaN poisoning came + // from the stat rather than a scan. + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_uses_cached_stat_sum() -> VortexResult<()> { + // A planted exact `Stat::Sum` with a known null count is consumed instead of a scan + // (the planted value differs from the actual data to prove it). + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); + arr.statistics() + .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); + let result = sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) +} + +#[test] +fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact Sum stat is usable as-is. + let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); + arr.statistics() + .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) +} + +#[test] +fn sum_constant_nan() -> VortexResult<()> { + let arr = ConstantArray::new(f64::NAN, 4).into_array(); + // NaN constants are skipped by default (a non-empty zero sum) and poison the sum otherwise. + let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) +} + +#[test] +fn sum_f64_with_infinity() -> VortexResult<()> { + let batch = PrimitiveArray::new( + buffer![1.0f64, f64::INFINITY, f64::NEG_INFINITY, 2.0], + Validity::NonNullable, + ) + .into_array(); + let acc = 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, + NumericalAggregateOpts::default(), + DType::Primitive(PType::F64, Nullability::NonNullable), + )?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + Ok(()) +} + +// Overflow: a null sum value plus an explicit saturation flag. + +#[test] +fn sum_checked_overflow_is_null_and_saturates() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + assert!(!acc.is_saturated()); + + let batch = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(acc.is_saturated()); + let result = acc.finish()?; + assert!(result.is_null()); + + // finish resets state, clearing saturation + assert!(!acc.is_saturated()); + Ok(()) +} + +#[test] +fn sum_decimal_i256_overflow() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(76, 0); + let decimal = DecimalArray::new( + buffer![i256::MAX, i256::MAX, i256::MAX], + decimal_dtype, + Validity::AllValid, + ); + + let result = 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(&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, near_limit)?; + + // Add a small value that keeps us just under 10^14. + let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); + Sum.combine_partials(&mut state, small)?; + + let result = Sum.finalize_scalar(&state)?; + assert!(!result.is_null()); + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(99_999_999_999_999))) + ); + Ok(()) +} + +#[rstest] +#[case::positive(99_999_999_999_999i64, 1i64)] +#[case::negative(-99_999_999_999_999i64, -1i64)] +fn sum_decimal_precision_overflow_within_i256( + #[case] near_limit: i64, + #[case] one_more: i64, +) -> VortexResult<()> { + // Input precision 4 → return precision 14. Native I64 (max 18). + // The max representable magnitude for precision 14 is 10^14 - 1: pushing the sum to + // exactly ±10^14 fails fits_in_precision even though i256 arithmetic does not + // overflow. This tests the precision-based saturation path in combine_partials. + let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + + let near_limit = Scalar::decimal( + DecimalValue::from(near_limit), + DecimalDType::new(14, 0), + Nullable, + ); + Sum.combine_partials(&mut state, near_limit)?; + + let one_more = Scalar::decimal( + DecimalValue::from(one_more), + DecimalDType::new(14, 0), + Nullable, + ); + Sum.combine_partials(&mut state, one_more)?; + + let result = Sum.finalize_scalar(&state)?; + assert!(result.is_null()); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(14, 0), Nullable) + ); + Ok(()) +} + +#[test] +fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { + // Test precision overflow via the accumulate_decimal path (not combine_partials). + // Input precision 27 → return precision 37. Native for 37 is I128 (max 38), so 37 < 38. + // Use combine_partials to get the state close to 10^37, then accumulate a real array + // that pushes it over. + let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); + let return_dtype = DecimalDType::new(37, 0); + let mut state = 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, 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.finalize_scalar(&state)?; + assert!(result.is_null()); + Ok(()) +} + +// Grouped aggregation: empty, all-null, and null groups through the lazy finalize. + +fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { + let mut acc = + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype.clone())?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(groups, &mut ctx)?; + acc.finish() +} + +#[test] +fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(5i64), None, Some(i64::MAX), Some(1)]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 0, 1, 2, 0].into_array(), + buffer![1i32, 0, 1, 2, 1].into_array(), + Validity::from_iter([true, true, true, true, false]), + )? + .into_array(); + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + DType::Primitive(PType::I64, Nullable), + )?; + acc.accumulate_list(&groups, &mut ctx)?; + let partials = acc.flush()?; + + let value = partials.execute_scalar(0, &mut ctx)?; + let fields = value.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(5) + ); + assert_eq!( + fields + .field("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, NumericalAggregateOpts::default(), elem_dtype)?; + + let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; + acc.accumulate_list(&groups1.into_array(), &mut ctx)?; + let result1 = acc.finish()?; + + let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); + assert_arrays_eq!(&result1, &expected1, &mut ctx); + + let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; + acc.accumulate_list(&groups2.into_array(), &mut ctx)?; + let result2 = acc.finish()?; + + let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); + assert_arrays_eq!(&result2, &expected2, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); + let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new(buffer![1i32, 1, 1], Validity::NonNullable).into_array(); + let validity = Validity::from_iter([true, false, true]); + let groups = ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array(); + + let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = run_grouped_sum(&groups, &elem_dtype)?; + + // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. + let expected = + PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 92fac87892a..00c6600d12b 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -183,7 +183,6 @@ mod tests { ) -> 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())?; diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 149c72f499e..066654a44f3 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -120,7 +120,7 @@ mod tests { DType::Primitive(PType::I32, Nullability::Nullable), )?; let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } @@ -158,7 +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_eq!(result.as_primitive().typed_value::(), Some(0)); + assert!(result.is_null()); Ok(()) } diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 57e11135ef3..8b4f5e4c60a 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -188,8 +188,8 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - return aggregate_fn::fns::sum::Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); + let options = NumericalAggregateOpts::skip_nans(); + return aggregate_fn::fns::sum::Sum.return_dtype(&options, data_type); } }) } diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 95b8ddbb55a..2ad3369d37c 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -1,11 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -16,23 +14,16 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; -use crate::aggregate_fn::GroupRanges; use crate::aggregate_fn::GroupedAccumulator; -use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; -use crate::arrays::BoolArray; use crate::arrays::ConstantArray; -use crate::arrays::FixedSizeList; -use crate::arrays::ListView; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::validity::Validity; /// Sum of the elements in each list of a `List` or `FixedSizeList` typed array. /// @@ -100,7 +91,6 @@ impl ScalarFnVTable for ListSum { other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), }; - // `mask_empty_lists` needs access to list elements validity and sizes let columnar = input.execute::(ctx)?; match columnar { @@ -134,8 +124,7 @@ impl ScalarFnVTable for ListSum { /// Sum each list of a canonical `array` into one value per list. /// -/// Note that we need to nullify sums produced by empty or all-null lists, -/// since grouped sum kernels default to 0 for these. +/// The null-on-empty [`Sum`] behavior yields null for null, empty, and all-null lists. fn list_sum_impl( canonical: ArrayRef, elem_dtype: DType, @@ -144,50 +133,7 @@ fn list_sum_impl( ) -> VortexResult { let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; - let sums = acc.finish()?; - - let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::() { - fsl.into_owned().into() - } else if let Some(lv) = canonical.as_opt::() { - lv.into_owned().into() - } else { - let dtype = canonical.dtype(); - vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}") - }; - - mask_empty_lists(grouped, sums, ctx) -} - -/// Applies a mask to `sums` that nullifies entries produced by lists without at least -/// one valid element. This is necessary because the grouped `Sum` aggregate only produces -/// nulls for null lists and sums that overflow. -fn mask_empty_lists( - grouped: GroupedArray, - sums: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let elements = grouped.elements(); - let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; - let ranges = grouped.group_ranges(ctx)?; - - let has_valid_element: BitBuffer = match elem_mask.bit_buffer() { - AllOr::All => match &ranges { - // fixed-size lists of non-zero width cannot have empty lists. - GroupRanges::FixedSizeList { size, .. } if *size > 0 => return Ok(sums), - GroupRanges::FixedSizeList { len, .. } => BitBuffer::full(false, *len), - GroupRanges::ListView { ranges } => ranges.iter().map(|&(_, size)| size > 0).collect(), - }, - AllOr::None => BitBuffer::full(false, ranges.len()), - AllOr::Some(bits) => ranges - .iter() - .map(|(offset, size)| size > 0 && bits.count_range(offset, offset + size) > 0) - .collect(), - }; - if has_valid_element.true_count() == has_valid_element.len() { - return Ok(sums); - } - - sums.mask(BoolArray::new(has_valid_element, Validity::NonNullable).into_array()) + acc.finish() } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index c961ac10d53..843657d8a41 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -57,7 +57,7 @@ impl Display for StatOptions { } } -/// Scalar function that broadcasts a stored aggregate partial over the input rows. +/// Scalar function that broadcasts a stored aggregate result over the input rows. /// /// The only current consumer is **row-wise pruning**: substituting `stat(col, agg)` into a /// predicate produces a cheap, row-aligned approximation whose constant runs let downstream @@ -130,7 +130,7 @@ impl ScalarFnVTable for StatFn { } fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult { - let Some(dtype) = aggregate_fn.state_dtype(input_dtype) else { + let Some(dtype) = aggregate_fn.return_dtype(input_dtype) else { vortex_bail!( "Aggregate function {} does not support input dtype {}", aggregate_fn, diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..3efc718de9f 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -58,6 +58,7 @@ use crate::layouts::zoned::schema::AggregateSpecProto; use crate::layouts::zoned::schema::aggregate_specs_from_fns; use crate::layouts::zoned::schema::aggregate_stats_table_dtype; use crate::layouts::zoned::schema::legacy_stats_table_dtype; +use crate::layouts::zoned::schema::legacy_sum_aggregate_stats_table_dtype; use crate::layouts::zoned::schema::try_aggregate_fns_from_specs; use crate::segments::SegmentSource; @@ -79,6 +80,7 @@ pub struct ZonedData { zone_len: usize, zone_map_schema: ZoneMapSchema, stats_table_dtype: DType, + legacy_sum_partials: bool, } /// A layout annotating a data child with per-zone statistics. @@ -108,6 +110,7 @@ impl VTable for Zoned { vortex_panic!("Cannot serialize legacy stats schema as vortex.zoned") } }, + legacy_sum_partials: layout.legacy_sum_partials, } } @@ -130,16 +133,22 @@ impl VTable for Zoned { zone_len: 0, zone_map_schema: ZoneMapSchema::AggregateFns(Arc::new([])), stats_table_dtype: aggregate_stats_table_dtype(args.dtype, &[]), + legacy_sum_partials: metadata.legacy_sum_partials, }); }; aggregate_specs_from_fns(&aggregate_fns)?; - let stats_table_dtype = aggregate_stats_table_dtype(args.dtype, &aggregate_fns); + let stats_table_dtype = if metadata.legacy_sum_partials { + legacy_sum_aggregate_stats_table_dtype(args.dtype, &aggregate_fns) + } else { + aggregate_stats_table_dtype(args.dtype, &aggregate_fns) + }; args.children.child(0, args.dtype)?; args.children.child(1, &stats_table_dtype)?; Ok(ZonedData { zone_len: metadata.zone_len as usize, zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), stats_table_dtype, + legacy_sum_partials: metadata.legacy_sum_partials, }) } @@ -208,6 +217,7 @@ impl VTable for LegacyStats { zone_len: metadata.zone_len as usize, zone_map_schema: metadata.zone_map_schema.clone(), stats_table_dtype, + legacy_sum_partials: false, }) } @@ -267,7 +277,6 @@ impl LegacyStatsLayout { usize::try_from(self.children().child_row_count(1)) .vortex_expect("Invalid number of zones, cannot handle more than usize zones") } - /// Returns display names for the zone-map aggregates stored by this layout. pub fn present_aggregates(&self) -> Arc<[String]> { present_aggregates(&self.zone_map_schema) @@ -309,6 +318,7 @@ impl ZonedLayout { zone_len: zone_len.get(), zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), stats_table_dtype: expected_dtype, + legacy_sum_partials: false, }, ) .into_typed()) @@ -394,6 +404,7 @@ fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { pub struct ZonedMetadata { pub(super) zone_len: u32, pub(super) aggregate_specs: Arc<[AggregateSpecProto]>, + legacy_sum_partials: bool, } /// Serialized metadata for legacy `vortex.stats` layouts. @@ -403,7 +414,8 @@ pub struct LegacyStatsMetadata { pub(crate) zone_map_schema: ZoneMapSchema, } -const ZONED_METADATA_PROTO_VERSION: u8 = 1; +const LEGACY_SUM_PARTIAL_METADATA_VERSION: u8 = 1; +const ZONED_METADATA_PROTO_VERSION: u8 = 2; #[derive(Clone, PartialEq, Message)] struct ZonedMetadataProto { @@ -422,7 +434,10 @@ impl DeserializeMetadata for ZonedMetadata { }; vortex_ensure!( - version == ZONED_METADATA_PROTO_VERSION, + matches!( + version, + LEGACY_SUM_PARTIAL_METADATA_VERSION | ZONED_METADATA_PROTO_VERSION + ), "Unsupported zoned metadata version: {}", version ); @@ -432,6 +447,7 @@ impl DeserializeMetadata for ZonedMetadata { Ok(Self { zone_len: proto.zone_len, aggregate_specs: proto.aggregate_specs.into(), + legacy_sum_partials: version == LEGACY_SUM_PARTIAL_METADATA_VERSION, }) } } @@ -442,7 +458,12 @@ impl SerializeMetadata for ZonedMetadata { zone_len: self.zone_len, aggregate_specs: self.aggregate_specs.to_vec(), }; - let mut metadata = vec![ZONED_METADATA_PROTO_VERSION]; + let version = if self.legacy_sum_partials { + LEGACY_SUM_PARTIAL_METADATA_VERSION + } else { + ZONED_METADATA_PROTO_VERSION + }; + let mut metadata = vec![version]; metadata.extend(proto.encode_to_vec()); metadata } @@ -520,6 +541,7 @@ mod tests { #[case(ZonedMetadata { zone_len: u32::MAX, aggregate_specs: Arc::new([]), + legacy_sum_partials: false, })] #[case::min_max(ZonedMetadata { zone_len: 314, @@ -527,6 +549,7 @@ mod tests { aggregate_spec(Max.bind(NumericalAggregateOpts::skip_nans())), aggregate_spec(Min.bind(NumericalAggregateOpts::skip_nans())), ]), + legacy_sum_partials: false, })] fn test_metadata_serialization(#[case] metadata: ZonedMetadata) { let serialized = metadata.clone().serialize(); @@ -544,6 +567,7 @@ mod tests { let metadata = ZonedMetadata { zone_len: 314, aggregate_specs: Arc::new([AggregateSpecProto::try_from_aggregate_fn(&aggregate_fn)?]), + legacy_sum_partials: false, }; let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?; @@ -555,6 +579,19 @@ mod tests { Ok(()) } + #[test] + fn test_legacy_sum_partial_metadata_version_round_trip() { + let metadata = ZonedMetadata { + zone_len: 314, + aggregate_specs: Arc::new([]), + legacy_sum_partials: true, + }; + + let serialized = metadata.clone().serialize(); + assert_eq!(serialized[0], LEGACY_SUM_PARTIAL_METADATA_VERSION); + assert_eq!(ZonedMetadata::deserialize(&serialized).unwrap(), metadata); + } + #[test] fn test_deserialize_legacy_stat_bitset_as_legacy_stats() { let mut serialized = u32::MAX.to_le_bytes().to_vec(); @@ -648,6 +685,7 @@ mod tests { let metadata = ZonedMetadata { zone_len: 3, aggregate_specs: Arc::new([]), + legacy_sum_partials: false, }; let children = OwnedLayoutChildren::layout_children(vec![]); let session = vortex_array::array_session(); @@ -689,6 +727,7 @@ mod tests { let metadata = ZonedMetadata { zone_len: 8, aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), + legacy_sum_partials: false, }; let layout = ::build( @@ -725,6 +764,7 @@ mod tests { let metadata = ZonedMetadata { zone_len: 8, aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), + legacy_sum_partials: false, }; let result = ::build( diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 700aab86ce4..57a1a34c2c0 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -34,6 +34,7 @@ use crate::LazyReaderChildren; use crate::VTable; use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::zone_map::ZoneMap; +use crate::layouts::zoned::zone_map::normalize_sum_partial_fields; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = @@ -169,6 +170,7 @@ impl PruningState { async move { let mut ctx = session.create_execution_ctx(); let zones_array = zones_eval.await?.execute::(&mut ctx)?; + let zones_array = normalize_sum_partial_fields(zones_array, &aggregate_fns)?; // SAFETY: zoned layout validation checked that this zones child was // written from the same column dtype and aggregate stats-table schema. Ok(unsafe { diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index f35eb4a59a9..3b17ac49831 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -97,6 +98,24 @@ pub(crate) fn aggregate_stats_table_dtype( ) } +/// Return the auxiliary stats-table schema written before Sum adopted a struct partial. +pub(crate) fn legacy_sum_aggregate_stats_table_dtype( + column_dtype: &DType, + aggregate_fns: &[AggregateFnRef], +) -> DType { + DType::Struct( + StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| { + let state_dtype = if aggregate_fn.is::() { + aggregate_return_dtype(column_dtype, aggregate_fn) + } else { + aggregate_state_dtype(column_dtype, aggregate_fn) + }?; + Some((aggregate_fn.to_string(), state_dtype.as_nullable())) + })), + Nullability::NonNullable, + ) +} + pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[Stat]) -> DType { assert!(present_stats.is_sorted(), "Stats must be sorted"); DType::Struct( @@ -182,6 +201,16 @@ pub(crate) fn aggregate_state_dtype( }) } +fn aggregate_return_dtype(column_dtype: &DType, aggregate_fn: &AggregateFnRef) -> Option { + aggregate_fn.return_dtype(column_dtype).or_else(|| { + if let DType::Extension(ext) = column_dtype { + aggregate_fn.return_dtype(ext.storage_dtype()) + } else { + None + } + }) +} + pub(crate) fn default_bounded_stat_max_bytes() -> std::num::NonZeroUsize { // SAFETY: 64 is non-zero. unsafe { std::num::NonZeroUsize::new_unchecked(64) } @@ -273,4 +302,22 @@ mod tests { ] ); } + + #[test] + fn legacy_sum_stats_table_dtype_uses_scalar_partial() { + let column_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let sum = Sum.bind(NumericalAggregateOpts::default()); + let legacy = + legacy_sum_aggregate_stats_table_dtype(&column_dtype, std::slice::from_ref(&sum)); + let current = aggregate_stats_table_dtype(&column_dtype, &[sum]); + + 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(..)) + )); + } } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index cca134c0759..f2924c36e6c 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -208,11 +208,9 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A }; let mut aggregate_fns = vec![max, min]; - if Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) - .is_some() - { - aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); + let sum_options = NumericalAggregateOpts::skip_nans(); + if Sum.return_dtype(&sum_options, dtype).is_some() { + aggregate_fns.push(Sum.bind(sum_options)); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c84c0b443dd..6bda75acc96 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,6 +16,8 @@ use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum::normalize_partial_array; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -24,8 +26,12 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::eq; +use vortex_array::expr::fill_null; use vortex_array::expr::get_item; use vortex_array::expr::lit; +use vortex_array::expr::mask; +use vortex_array::expr::not; +use vortex_array::expr::or; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar_fn::EmptyOptions; @@ -165,6 +171,26 @@ 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; + }; + fields[index] = normalize_partial_array(fields[index].clone())?; + } + + StructArray::try_new(names, fields, array.len(), array.struct_validity()) +} + struct ZoneMapStatsBinder<'a> { zone_map: &'a ZoneMap, } @@ -293,6 +319,15 @@ 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::() { + let is_invalid = fill_null( + or( + get_item("is_overflow", state_expr.clone()), + get_item("is_empty", state_expr.clone()), + ), + lit(true), + ); + mask(get_item("sum", state_expr), not(is_invalid)) } else { state_expr } @@ -356,6 +391,7 @@ 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::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -387,6 +423,7 @@ mod tests { use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ZoneMap; + use crate::layouts::zoned::zone_map::normalize_sum_partial_fields; use crate::test::SESSION; fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { @@ -406,6 +443,26 @@ mod tests { unsafe { NonZeroUsize::new_unchecked(64) } } + #[test] + fn legacy_scalar_sum_field_is_normalized_once() -> VortexResult<()> { + let sum = Sum.bind(NumericalAggregateOpts::default()); + 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 normalized = normalize_sum_partial_fields(legacy_table, std::slice::from_ref(&sum))?; + let zone_map = + ZoneMap::try_new(PType::I32.into(), normalized, 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: From a3e5f3370f4e00839048dda11a2e8c73c063509f Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 31 Jul 2026 16:41:37 +0100 Subject: [PATCH 02/18] Encode Sum partial shape in aggregate options Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 42 +++-- vortex-array/src/aggregate_fn/accumulator.rs | 6 +- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 11 +- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 8 +- .../src/aggregate_fn/fns/sum/decimal.rs | 10 +- .../src/aggregate_fn/fns/sum/grouped.rs | 16 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 145 +++++++++++++++--- .../src/aggregate_fn/fns/sum/primitive.rs | 29 ++-- .../src/aggregate_fn/fns/sum/tests.rs | 79 ++++++---- vortex-array/src/aggregate_fn/proto.rs | 35 ++++- .../src/arrays/chunked/compute/aggregate.rs | 9 +- vortex-array/src/expr/stats/mod.rs | 5 +- vortex-array/src/scalar_fn/fns/list_sum.rs | 4 +- vortex-array/src/stats/expr.rs | 5 +- vortex-duckdb/src/convert/expr.rs | 9 +- vortex-layout/src/layouts/zoned/mod.rs | 53 +------ vortex-layout/src/layouts/zoned/schema.rs | 46 ++---- vortex-layout/src/layouts/zoned/writer.rs | 3 +- vortex-layout/src/layouts/zoned/zone_map.rs | 20 ++- vortex-proto/proto/expr.proto | 8 + vortex-proto/src/generated/vortex.expr.rs | 10 ++ 21 files changed, 331 insertions(+), 222 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 739b51dcee1..8a041b2275e 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -22,6 +22,7 @@ use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum::SumAggregateOpts; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -177,8 +178,7 @@ fn sum_legacy_scalar_partial_merge(bencher: Bencher) { .with_inputs(|| partial.clone()) .bench_refs(|partial| { let mut acc = - Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()) - .unwrap(); + Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); for _ in 0..GROUP_COUNT { acc.combine_partials(partial.clone()).unwrap(); } @@ -189,8 +189,7 @@ fn sum_legacy_scalar_partial_merge(bencher: Bencher) { #[divan::bench] fn sum_canonical_partial_merge(bencher: Bencher) { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut source = - Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()).unwrap(); + let mut source = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); source .combine_partials(Scalar::primitive(1i64, Nullability::Nullable)) .unwrap(); @@ -200,8 +199,7 @@ fn sum_canonical_partial_merge(bencher: Bencher) { .with_inputs(|| partial.clone()) .bench_refs(|partial| { let mut acc = - Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone()) - .unwrap(); + Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); for _ in 0..GROUP_COUNT { acc.combine_partials(partial.clone()).unwrap(); } @@ -214,7 +212,7 @@ fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); } #[divan::bench] @@ -222,7 +220,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); } #[divan::bench] @@ -230,7 +228,7 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); } #[divan::bench] @@ -238,7 +236,7 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); } /// Like [`grouped_accumulator`], but executes the lazy finalize result to canonical so the @@ -267,33 +265,33 @@ where #[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, NumericalAggregateOpts::default()) - }); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum, SumAggregateOpts::default())); } #[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, NumericalAggregateOpts::default()) - }); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum, SumAggregateOpts::default())); } #[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, NumericalAggregateOpts::default()) - }); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum, SumAggregateOpts::default())); } #[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, NumericalAggregateOpts::default()) - }); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum, SumAggregateOpts::default())); } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 0499b5d1ecb..884fac3c4b0 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -269,6 +269,7 @@ mod tests { use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::session::AggregateFnSession; use crate::array::VTable; @@ -339,7 +340,7 @@ mod tests { Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -355,8 +356,7 @@ mod tests { fn sum_partial(value: f64) -> Scalar { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = - Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype).expect("sum"); + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype).expect("sum"); acc.combine_partials(Scalar::primitive(value, Nullability::Nullable)) .expect("legacy scalar partial"); acc.flush().expect("sum partial") diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index d2c1ac78604..cf75d52ddd2 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -19,6 +19,7 @@ use crate::aggregate_fn::combined::CombinedOptions; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum_decimal_dtype; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; @@ -40,7 +41,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), array.dtype().clone(), @@ -319,7 +320,7 @@ mod tests { let keep_nans = NumericalAggregateOpts::include_nans(); let mut acc = Accumulator::try_new( Mean::combined(), - PairOptions(keep_nans, keep_nans), + PairOptions(keep_nans.into(), keep_nans), array.dtype().clone(), )?; acc.accumulate(&array, &mut ctx)?; @@ -411,7 +412,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), dtype, @@ -446,7 +447,7 @@ mod tests { let mut acc = Accumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), DType::Primitive(PType::F64, Nullability::Nullable), @@ -474,7 +475,7 @@ mod tests { let mut acc = GroupedAccumulator::try_new( Mean::combined(), PairOptions( - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), NumericalAggregateOpts::default(), ), DType::Primitive(PType::F64, Nullability::Nullable), diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index 80aa62831bd..d74568df356 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -40,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; @@ -108,7 +108,7 @@ mod tests { #[test] fn sum_bool_empty_produces_null() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert!(result.is_null()); Ok(()) @@ -118,7 +118,7 @@ mod tests { fn sum_bool_finish_resets_state() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -136,7 +136,7 @@ mod tests { fn sum_bool_return_dtype() -> VortexResult<()> { let dtype = Sum .return_dtype( - &NumericalAggregateOpts::default(), + &SumAggregateOpts::default(), &DType::Bool(Nullability::NonNullable), ) .unwrap(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 60c4e20de33..ad517dbfdfc 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -114,8 +114,8 @@ mod tests { 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; @@ -357,7 +357,7 @@ 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 options = NumericalAggregateOpts::default(); + let options = SumAggregateOpts::default(); let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( @@ -388,7 +388,7 @@ 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 options = NumericalAggregateOpts::default(); + let options = SumAggregateOpts::default(); let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( @@ -416,7 +416,7 @@ 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 options = NumericalAggregateOpts::default(); + let options = SumAggregateOpts::default(); let mut state = Sum.empty_partial(&options, &input_dtype)?; let near_limit = Scalar::decimal( @@ -449,7 +449,7 @@ 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 options = NumericalAggregateOpts::default(); + let options = SumAggregateOpts::default(); let mut state = Sum.empty_partial(&options, &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index a8630dee99b..38c3c103567 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -46,6 +46,9 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); }; + if !options.struct_partial { + return Ok(None); + } try_grouped_sum(groups, ctx, options.skip_nans) } } @@ -193,8 +196,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; @@ -210,11 +213,8 @@ mod tests { /// Run a grouped sum through the accumulator. fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - elem_dtype.clone(), - )?; + let mut acc = + GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -231,7 +231,7 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .return_dtype(&NumericalAggregateOpts::default(), elem_dtype) + .return_dtype(&SumAggregateOpts::default(), elem_dtype) .expect("sum return dtype"); let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); for (i, &(offset, size)) in ranges.iter().enumerate() { @@ -399,7 +399,7 @@ mod tests { let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; let mut acc = - GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; + GroupedAccumulator::try_new(Sum, SumAggregateOpts::include_nans(), elem_dtype)?; acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; let actual = acc.finish()?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index a93b3fd910c..fdb15a5c24a 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -6,12 +6,18 @@ mod constant; mod decimal; mod grouped; mod primitive; +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; +use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_proto::expr as pb; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -57,11 +63,7 @@ const IS_EMPTY_FIELD: &str = "is_empty"; /// /// See [`Sum`] for details. pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let mut acc = Accumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - array.dtype().clone(), - )?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), array.dtype().clone())?; acc.accumulate(array, ctx)?; let result = acc.finish()?; @@ -79,6 +81,86 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { #[derive(Clone, Copy, Debug)] pub struct Sum; +/// Options for [`Sum`]. +/// +/// New sums use a struct partial that can distinguish an empty input from a zero sum. The +/// `struct_partial` field exists for deserializing aggregates written before that partial was +/// introduced; 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 partials use the `{ sum, is_overflow, is_empty }` struct representation. + 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. + pub fn serialize(&self) -> Vec { + pb::SumAggregateOpts { + skip_nans: self.skip_nans, + struct_partial: Some(self.struct_partial), + } + .encode_to_vec() + } + + /// Deserialize these options from protobuf-encoded metadata bytes. + /// + /// Historical Sum options were serialized as [`NumericalAggregateOpts`]. They use the same + /// wire representation for `skip_nans` and omit `struct_partial`, which selects the legacy + /// scalar partial representation. + pub fn deserialize(metadata: &[u8]) -> VortexResult { + let options = pb::SumAggregateOpts::decode(metadata)?; + Ok(Self { + skip_nans: options.skip_nans, + struct_partial: options.struct_partial.unwrap_or(false), + }) + } +} + +impl Default for SumAggregateOpts { + fn default() -> Self { + Self::skip_nans() + } +} + +impl From for SumAggregateOpts { + fn from(options: NumericalAggregateOpts) -> Self { + Self { + skip_nans: options.skip_nans, + struct_partial: true, + } + } +} + +impl Display for SumAggregateOpts { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + // The partial representation is a storage-compatibility detail. Keeping it out of the + // display preserves the stats-table field name across old and new Sum partials. + if !self.skip_nans { + write!(f, "skip_nans=false")?; + } + Ok(()) + } +} + // Both Spark and DataFusion use this heuristic. // - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 // - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 @@ -90,7 +172,7 @@ pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { } impl AggregateFnVTable for Sum { - type Options = NumericalAggregateOpts; + type Options = SumAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { @@ -107,7 +189,7 @@ impl AggregateFnVTable for Sum { metadata: &[u8], _session: &VortexSession, ) -> VortexResult { - NumericalAggregateOpts::deserialize(metadata) + SumAggregateOpts::deserialize(metadata) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -139,7 +221,11 @@ 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)) + if options.struct_partial { + Some(sum_partial_dtype(return_dtype)) + } else { + Some(return_dtype) + } } fn empty_partial( @@ -157,6 +243,7 @@ impl AggregateFnVTable for Sum { is_overflow: false, is_empty: true, skip_nans: options.skip_nans, + struct_partial: options.struct_partial, }) } @@ -230,6 +317,10 @@ impl AggregateFnVTable for Sum { } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + if !partial.struct_partial { + return Ok(legacy_sum_value_scalar(partial)); + } + Ok(Scalar::struct_( sum_partial_dtype(partial.return_dtype.clone()), vec![ @@ -363,7 +454,11 @@ impl AggregateFnVTable for Sum { } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(sum_value_scalar(partial)) + if partial.struct_partial { + Ok(sum_value_scalar(partial)) + } else { + Ok(legacy_sum_value_scalar(partial)) + } } } @@ -378,6 +473,8 @@ pub struct SumPartial { is_empty: bool, /// Whether NaN values in float inputs are skipped. skip_nans: bool, + /// Whether this accumulator emits the canonical struct partial. + struct_partial: bool, } /// The accumulated sum value. @@ -446,6 +543,18 @@ fn sum_value_scalar(partial: &SumPartial) -> Scalar { return Scalar::null(partial.return_dtype.as_nullable()); } + nullable_sum_state_scalar(partial) +} + +fn legacy_sum_value_scalar(partial: &SumPartial) -> Scalar { + if partial.is_overflow { + return Scalar::null(partial.return_dtype.as_nullable()); + } + + nullable_sum_state_scalar(partial) +} + +fn nullable_sum_state_scalar(partial: &SumPartial) -> Scalar { match &partial.sum { SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::Nullable), SumState::Signed(v) => Scalar::primitive(*v, Nullability::Nullable), @@ -593,8 +702,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,7 +792,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, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -700,7 +809,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, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -719,7 +828,7 @@ mod arithmetic_tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let options = NumericalAggregateOpts::default(); + let options = SumAggregateOpts::default(); let mut state = Sum.empty_partial(&options, &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); @@ -773,11 +882,8 @@ mod arithmetic_tests { // Grouped sum tests fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - elem_dtype.clone(), - )?; + let mut acc = + GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -865,8 +971,7 @@ 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, NumericalAggregateOpts::default(), elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 8728ebe88eb..a58707c09d1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -187,8 +187,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::aggregate_fn::fns::sum::sum; use crate::array_session; use crate::arrays::ConstantArray; @@ -291,7 +291,7 @@ mod tests { #[test] fn sum_empty_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert!(result.is_null()); Ok(()) @@ -300,7 +300,7 @@ mod tests { #[test] fn sum_empty_f64_produces_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert!(result.is_null()); Ok(()) @@ -345,11 +345,8 @@ mod tests { Ok(()) } - /// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). - fn sum_with_options( - arr: &crate::ArrayRef, - options: NumericalAggregateOpts, - ) -> VortexResult { + /// Sum an array with explicit [`SumAggregateOpts`] (test-only helper). + fn sum_with_options(arr: &crate::ArrayRef, options: SumAggregateOpts) -> VortexResult { let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; acc.accumulate(arr, &mut array_session().create_execution_ctx())?; acc.finish() @@ -359,7 +356,7 @@ mod tests { fn sum_f64_with_nan_not_skipping() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f64, f64::NAN, 2.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -368,7 +365,7 @@ mod tests { fn sum_f64_without_nan_not_skipping() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -381,7 +378,7 @@ mod tests { PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); arr.statistics() .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -395,7 +392,7 @@ mod tests { .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); arr.statistics() .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) } @@ -404,10 +401,10 @@ mod tests { fn sum_constant_nan() -> VortexResult<()> { let arr = ConstantArray::new(f64::NAN, 4).into_array(); // NaN constants are skipped by default and poison the sum otherwise. - let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; + let result = sum_with_options(&arr, SumAggregateOpts::default())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -425,7 +422,7 @@ mod tests { let mut acc = Accumulator::try_new( Sum, - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; @@ -444,7 +441,7 @@ mod tests { #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; assert!(!acc.is_saturated()); let batch = diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs index 68683071ee0..f1c16e21295 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -10,6 +10,7 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use super::Sum; +use super::SumAggregateOpts; use super::sum; use crate::ArrayRef; use crate::IntoArray; @@ -43,8 +44,8 @@ use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::validity::Validity; -/// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). -fn sum_with_options(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { +/// 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() @@ -52,7 +53,7 @@ fn sum_with_options(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexRe #[test] fn sum_uses_new_partial_shape_by_default() { - let options = NumericalAggregateOpts::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); @@ -74,6 +75,32 @@ fn sum_uses_new_partial_shape_by_default() { ); } +#[test] +fn legacy_options_use_scalar_partial_and_zero_on_empty() -> 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_eq!( + Sum.partial_dtype(&options, &input_dtype), + Some(DType::Primitive(PType::I64, Nullable)) + ); + + let mut acc = Accumulator::try_new(Sum, options, input_dtype)?; + assert_eq!( + acc.partial_scalar()?.as_primitive().typed_value::(), + Some(0) + ); + assert_eq!(acc.finish()?.as_primitive().typed_value::(), Some(0)); + Ok(()) +} + // State algebra: the `{sum, is_overflow, is_empty}` monoid. #[test] @@ -81,7 +108,7 @@ fn sum_state_empty_is_null() -> VortexResult<()> { // A state that never saw a valid value finalizes to null, and combining empty states // stays empty. let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; let empty = Sum.to_scalar(&state)?; let fields = empty.as_struct(); assert_eq!( @@ -111,10 +138,10 @@ fn sum_state_empty_is_null() -> VortexResult<()> { 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(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials(&mut state, Scalar::primitive(100i64, Nullable))?; - let empty = Sum.to_scalar(&Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?)?; + let empty = Sum.to_scalar(&Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?)?; Sum.combine_partials(&mut state, empty)?; let result = Sum.finalize_scalar(&state)?; @@ -126,7 +153,7 @@ fn sum_state_empty_is_identity() -> VortexResult<()> { 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(&NumericalAggregateOpts::default(), &dtype)?; + let mut overflowed = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials(&mut overflowed, Scalar::primitive(i64::MAX, Nullable))?; Sum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; let overflowed = Sum.to_scalar(&overflowed)?; @@ -150,7 +177,7 @@ fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { Some(false) ); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; Sum.combine_partials(&mut state, overflowed)?; Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; @@ -166,7 +193,7 @@ fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { #[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, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; assert!(acc.finish()?.is_null()); Ok(()) } @@ -215,7 +242,7 @@ fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { ); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable))?; assert_eq!( Sum.finalize_scalar(&state)? @@ -224,7 +251,7 @@ fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { Some(0) ); - let mut overflowed = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; + let mut overflowed = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials( &mut overflowed, Scalar::null(DType::Primitive(PType::I64, Nullable)), @@ -277,7 +304,7 @@ fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { )] fn sum_return_dtype_widens(#[case] input: DType, #[case] expected: DType) { let dtype = Sum - .return_dtype(&NumericalAggregateOpts::default(), &input) + .return_dtype(&SumAggregateOpts::default(), &input) .unwrap(); assert_eq!(dtype, expected); } @@ -343,7 +370,7 @@ fn sum_constant_false_is_zero_not_null() -> VortexResult<()> { 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, NumericalAggregateOpts::default(), dtype)?; + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -446,7 +473,7 @@ fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { fn sum_f64_with_nan_not_skipping() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f64, f64::NAN, 2.0], Validity::NonNullable).into_array(); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -458,7 +485,7 @@ fn sum_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> let arr = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); arr.statistics() .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -487,7 +514,7 @@ fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); arr.statistics() .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) } @@ -496,10 +523,10 @@ fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { fn sum_constant_nan() -> VortexResult<()> { let arr = ConstantArray::new(f64::NAN, 4).into_array(); // NaN constants are skipped by default (a non-empty zero sum) and poison the sum otherwise. - let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; + let result = sum_with_options(&arr, SumAggregateOpts::default())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); - let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -517,7 +544,7 @@ fn sum_f64_with_infinity() -> VortexResult<()> { let mut acc = Accumulator::try_new( Sum, - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; @@ -530,7 +557,7 @@ fn sum_f64_with_infinity() -> VortexResult<()> { #[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, NumericalAggregateOpts::default(), dtype)?; + 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(); @@ -570,7 +597,7 @@ fn sum_decimal_near_precision_boundary() -> VortexResult<()> { // Native type for precision 14 is I64 (max precision 18), so 14 < 18. // Use combine_partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), @@ -604,7 +631,7 @@ fn sum_decimal_precision_overflow_within_i256( // 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(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(near_limit), @@ -637,7 +664,7 @@ fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { // that pushes it over. let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); let return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; @@ -659,7 +686,7 @@ fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = - GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype.clone())?; + 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() @@ -679,7 +706,7 @@ fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexRe .into_array(); let mut acc = GroupedAccumulator::try_new( Sum, - NumericalAggregateOpts::default(), + SumAggregateOpts::default(), DType::Primitive(PType::I64, Nullable), )?; acc.accumulate_list(&groups, &mut ctx)?; @@ -858,7 +885,7 @@ fn grouped_sum_all_nan_is_zero_not_null() -> VortexResult<()> { fn grouped_sum_finish_resets() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 00c6600d12b..a3e72ba8e0c 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -76,6 +76,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::session::AggregateFnSession; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::dtype::DType; @@ -173,14 +174,15 @@ mod tests { assert_eq!(deserialized, agg_fn); } - /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the - /// numeric aggregates, including the non-default NaN-including configuration. + /// Both Sum options survive a protobuf serialize/deserialize round-trip. #[rstest] - #[case(NumericalAggregateOpts::skip_nans())] - #[case(NumericalAggregateOpts::include_nans())] - fn numeric_aggregate_options_round_trip( - #[case] options: NumericalAggregateOpts, - ) -> VortexResult<()> { + #[case(SumAggregateOpts::skip_nans())] + #[case(SumAggregateOpts::include_nans())] + #[case(SumAggregateOpts { + skip_nans: true, + struct_partial: false, + })] + fn sum_aggregate_options_round_trip(#[case] options: SumAggregateOpts) -> VortexResult<()> { let session = crate::array_session(); let agg_fn = Sum.bind(options); let proto = agg_fn.serialize_proto()?; @@ -192,6 +194,25 @@ mod tests { Ok(()) } + #[test] + fn legacy_sum_options_select_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 066654a44f3..c82f42b98c1 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -46,8 +46,8 @@ mod tests { use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -59,11 +59,8 @@ mod tests { fn run_sum(batch: &crate::ArrayRef) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); - let mut acc = Accumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - batch.dtype().clone(), - )?; + let mut acc = + Accumulator::try_new(Sum, SumAggregateOpts::default(), batch.dtype().clone())?; acc.accumulate(batch, &mut ctx)?; acc.finish() } diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 8b4f5e4c60a..b1100fe2eb3 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -188,7 +188,7 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - let options = NumericalAggregateOpts::skip_nans(); + let options = aggregate_fn::fns::sum::SumAggregateOpts::skip_nans(); return aggregate_fn::fns::sum::Sum.return_dtype(&options, data_type); } }) @@ -200,7 +200,8 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum + .bind(aggregate_fn::fns::sum::SumAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 2ad3369d37c..669f8344a96 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -74,7 +74,7 @@ impl ScalarFnVTable for ListSum { DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref(), other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), }; - Sum.return_dtype(options, elem_dtype) + Sum.return_dtype(&(*options).into(), elem_dtype) .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) } @@ -131,7 +131,7 @@ fn list_sum_impl( options: &NumericalAggregateOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; + let mut acc = GroupedAccumulator::try_new(Sum, (*options).into(), elem_dtype)?; acc.accumulate_list(&canonical, ctx)?; acc.finish() } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 1e0ceef02d3..92babebc772 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -17,6 +17,7 @@ use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; @@ -50,11 +51,11 @@ fn bound_min_max(expr: BoundExpression) -> BoundExpression { /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. - stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) + stat(expr, Sum.bind(SumAggregateOpts::skip_nans())) } fn bound_sum(expr: BoundExpression) -> BoundExpression { - bound_stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) + bound_stat(expr, Sum.bind(SumAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 5ee407994de..d200756d475 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -18,6 +18,7 @@ use vortex::aggregate_fn::fns::max::Max; use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; use vortex::aggregate_fn::fns::sum::Sum; +use vortex::aggregate_fn::fns::sum::SumAggregateOpts; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -533,10 +534,14 @@ impl PushedAggregate { Ok(match self { Self::Min => Box::new(Accumulator::try_new(Min, opts, dtype)?), Self::Max => Box::new(Accumulator::try_new(Max, opts, dtype)?), - Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new( + Sum, + SumAggregateOpts::from(opts), + dtype, + )?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), - PairOptions(opts, opts), + PairOptions(opts.into(), opts), dtype, )?), Self::First => Box::new(Accumulator::try_new(First, AggregateEmptyOptions, dtype)?), diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 3efc718de9f..01de0b2b452 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -58,7 +58,6 @@ use crate::layouts::zoned::schema::AggregateSpecProto; use crate::layouts::zoned::schema::aggregate_specs_from_fns; use crate::layouts::zoned::schema::aggregate_stats_table_dtype; use crate::layouts::zoned::schema::legacy_stats_table_dtype; -use crate::layouts::zoned::schema::legacy_sum_aggregate_stats_table_dtype; use crate::layouts::zoned::schema::try_aggregate_fns_from_specs; use crate::segments::SegmentSource; @@ -80,7 +79,6 @@ pub struct ZonedData { zone_len: usize, zone_map_schema: ZoneMapSchema, stats_table_dtype: DType, - legacy_sum_partials: bool, } /// A layout annotating a data child with per-zone statistics. @@ -110,7 +108,6 @@ impl VTable for Zoned { vortex_panic!("Cannot serialize legacy stats schema as vortex.zoned") } }, - legacy_sum_partials: layout.legacy_sum_partials, } } @@ -133,22 +130,16 @@ impl VTable for Zoned { zone_len: 0, zone_map_schema: ZoneMapSchema::AggregateFns(Arc::new([])), stats_table_dtype: aggregate_stats_table_dtype(args.dtype, &[]), - legacy_sum_partials: metadata.legacy_sum_partials, }); }; aggregate_specs_from_fns(&aggregate_fns)?; - let stats_table_dtype = if metadata.legacy_sum_partials { - legacy_sum_aggregate_stats_table_dtype(args.dtype, &aggregate_fns) - } else { - aggregate_stats_table_dtype(args.dtype, &aggregate_fns) - }; + let stats_table_dtype = aggregate_stats_table_dtype(args.dtype, &aggregate_fns); args.children.child(0, args.dtype)?; args.children.child(1, &stats_table_dtype)?; Ok(ZonedData { zone_len: metadata.zone_len as usize, zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), stats_table_dtype, - legacy_sum_partials: metadata.legacy_sum_partials, }) } @@ -217,7 +208,6 @@ impl VTable for LegacyStats { zone_len: metadata.zone_len as usize, zone_map_schema: metadata.zone_map_schema.clone(), stats_table_dtype, - legacy_sum_partials: false, }) } @@ -318,7 +308,6 @@ impl ZonedLayout { zone_len: zone_len.get(), zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), stats_table_dtype: expected_dtype, - legacy_sum_partials: false, }, ) .into_typed()) @@ -404,7 +393,6 @@ fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { pub struct ZonedMetadata { pub(super) zone_len: u32, pub(super) aggregate_specs: Arc<[AggregateSpecProto]>, - legacy_sum_partials: bool, } /// Serialized metadata for legacy `vortex.stats` layouts. @@ -414,8 +402,7 @@ pub struct LegacyStatsMetadata { pub(crate) zone_map_schema: ZoneMapSchema, } -const LEGACY_SUM_PARTIAL_METADATA_VERSION: u8 = 1; -const ZONED_METADATA_PROTO_VERSION: u8 = 2; +const ZONED_METADATA_PROTO_VERSION: u8 = 1; #[derive(Clone, PartialEq, Message)] struct ZonedMetadataProto { @@ -433,21 +420,13 @@ impl DeserializeMetadata for ZonedMetadata { vortex_bail!("Zoned metadata missing protobuf version"); }; - vortex_ensure!( - matches!( - version, - LEGACY_SUM_PARTIAL_METADATA_VERSION | ZONED_METADATA_PROTO_VERSION - ), - "Unsupported zoned metadata version: {}", - version - ); + vortex_ensure_eq!(version, ZONED_METADATA_PROTO_VERSION); vortex_ensure!(!proto_bytes.is_empty(), "Zoned metadata missing protobuf"); let proto = ZonedMetadataProto::decode(proto_bytes)?; Ok(Self { zone_len: proto.zone_len, aggregate_specs: proto.aggregate_specs.into(), - legacy_sum_partials: version == LEGACY_SUM_PARTIAL_METADATA_VERSION, }) } } @@ -458,12 +437,7 @@ impl SerializeMetadata for ZonedMetadata { zone_len: self.zone_len, aggregate_specs: self.aggregate_specs.to_vec(), }; - let version = if self.legacy_sum_partials { - LEGACY_SUM_PARTIAL_METADATA_VERSION - } else { - ZONED_METADATA_PROTO_VERSION - }; - let mut metadata = vec![version]; + let mut metadata = vec![ZONED_METADATA_PROTO_VERSION]; metadata.extend(proto.encode_to_vec()); metadata } @@ -541,7 +515,6 @@ mod tests { #[case(ZonedMetadata { zone_len: u32::MAX, aggregate_specs: Arc::new([]), - legacy_sum_partials: false, })] #[case::min_max(ZonedMetadata { zone_len: 314, @@ -549,7 +522,6 @@ mod tests { aggregate_spec(Max.bind(NumericalAggregateOpts::skip_nans())), aggregate_spec(Min.bind(NumericalAggregateOpts::skip_nans())), ]), - legacy_sum_partials: false, })] fn test_metadata_serialization(#[case] metadata: ZonedMetadata) { let serialized = metadata.clone().serialize(); @@ -567,7 +539,6 @@ mod tests { let metadata = ZonedMetadata { zone_len: 314, aggregate_specs: Arc::new([AggregateSpecProto::try_from_aggregate_fn(&aggregate_fn)?]), - legacy_sum_partials: false, }; let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?; @@ -579,19 +550,6 @@ mod tests { Ok(()) } - #[test] - fn test_legacy_sum_partial_metadata_version_round_trip() { - let metadata = ZonedMetadata { - zone_len: 314, - aggregate_specs: Arc::new([]), - legacy_sum_partials: true, - }; - - let serialized = metadata.clone().serialize(); - assert_eq!(serialized[0], LEGACY_SUM_PARTIAL_METADATA_VERSION); - assert_eq!(ZonedMetadata::deserialize(&serialized).unwrap(), metadata); - } - #[test] fn test_deserialize_legacy_stat_bitset_as_legacy_stats() { let mut serialized = u32::MAX.to_le_bytes().to_vec(); @@ -685,7 +643,6 @@ mod tests { let metadata = ZonedMetadata { zone_len: 3, aggregate_specs: Arc::new([]), - legacy_sum_partials: false, }; let children = OwnedLayoutChildren::layout_children(vec![]); let session = vortex_array::array_session(); @@ -727,7 +684,6 @@ mod tests { let metadata = ZonedMetadata { zone_len: 8, aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), - legacy_sum_partials: false, }; let layout = ::build( @@ -764,7 +720,6 @@ mod tests { let metadata = ZonedMetadata { zone_len: 8, aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), - legacy_sum_partials: false, }; let result = ::build( diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 3b17ac49831..5ffa94e2eb6 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; @@ -98,24 +97,6 @@ pub(crate) fn aggregate_stats_table_dtype( ) } -/// Return the auxiliary stats-table schema written before Sum adopted a struct partial. -pub(crate) fn legacy_sum_aggregate_stats_table_dtype( - column_dtype: &DType, - aggregate_fns: &[AggregateFnRef], -) -> DType { - DType::Struct( - StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| { - let state_dtype = if aggregate_fn.is::() { - aggregate_return_dtype(column_dtype, aggregate_fn) - } else { - aggregate_state_dtype(column_dtype, aggregate_fn) - }?; - Some((aggregate_fn.to_string(), state_dtype.as_nullable())) - })), - Nullability::NonNullable, - ) -} - pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[Stat]) -> DType { assert!(present_stats.is_sorted(), "Stats must be sorted"); DType::Struct( @@ -201,16 +182,6 @@ pub(crate) fn aggregate_state_dtype( }) } -fn aggregate_return_dtype(column_dtype: &DType, aggregate_fn: &AggregateFnRef) -> Option { - aggregate_fn.return_dtype(column_dtype).or_else(|| { - if let DType::Extension(ext) = column_dtype { - aggregate_fn.return_dtype(ext.storage_dtype()) - } else { - None - } - }) -} - pub(crate) fn default_bounded_stat_max_bytes() -> std::num::NonZeroUsize { // SAFETY: 64 is non-zero. unsafe { std::num::NonZeroUsize::new_unchecked(64) } @@ -223,6 +194,7 @@ mod tests { use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::sum::Sum; + use vortex_array::aggregate_fn::fns::sum::SumAggregateOpts; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -289,7 +261,7 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()), Min.bind(NumericalAggregateOpts::skip_nans()), - Sum.bind(NumericalAggregateOpts::skip_nans()), + Sum.bind(SumAggregateOpts::skip_nans()), ], ); @@ -298,18 +270,19 @@ mod tests { &[ Max.bind(NumericalAggregateOpts::skip_nans()).to_string(), Min.bind(NumericalAggregateOpts::skip_nans()).to_string(), - Sum.bind(NumericalAggregateOpts::skip_nans()).to_string(), + Sum.bind(SumAggregateOpts::skip_nans()).to_string(), ] ); } #[test] - fn legacy_sum_stats_table_dtype_uses_scalar_partial() { + fn sum_stats_table_dtype_uses_option_partial_shape() -> VortexResult<()> { let column_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let sum = Sum.bind(NumericalAggregateOpts::default()); - let legacy = - legacy_sum_aggregate_stats_table_dtype(&column_dtype, std::slice::from_ref(&sum)); - let current = aggregate_stats_table_dtype(&column_dtype, &[sum]); + 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()"), @@ -319,5 +292,6 @@ mod tests { 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 f2924c36e6c..c9b563900a1 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -25,6 +25,7 @@ use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum::SumAggregateOpts; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; @@ -208,7 +209,7 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A }; let mut aggregate_fns = vec![max, min]; - let sum_options = NumericalAggregateOpts::skip_nans(); + let sum_options = SumAggregateOpts::skip_nans(); if Sum.return_dtype(&sum_options, dtype).is_some() { aggregate_fns.push(Sum.bind(sum_options)); } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 6bda75acc96..f15e54237b1 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -85,8 +85,10 @@ impl ZoneMap { if &expected_dtype != array.dtype() { vortex_bail!("Array dtype does not match expected zone map dtype: {expected_dtype}"); } + let array = normalize_sum_partial_fields(array, &aggregate_fns)?; - // SAFETY: We checked that the array matches the expected stats-table schema. + // SAFETY: We checked the stored stats-table schema, then normalized legacy Sum fields to + // the canonical runtime representation. Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) }) } @@ -392,6 +394,7 @@ mod tests { 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; @@ -423,7 +426,6 @@ mod tests { use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ZoneMap; - use crate::layouts::zoned::zone_map::normalize_sum_partial_fields; use crate::test::SESSION; fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { @@ -445,13 +447,19 @@ mod tests { #[test] fn legacy_scalar_sum_field_is_normalized_once() -> VortexResult<()> { - let sum = Sum.bind(NumericalAggregateOpts::default()); + 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 normalized = normalize_sum_partial_fields(legacy_table, std::slice::from_ref(&sum))?; - let zone_map = - ZoneMap::try_new(PType::I32.into(), normalized, Arc::new([sum.clone()]), 1, 3)?; + 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) diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index 00ffaac433c..745017e8e64 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -32,6 +32,14 @@ message NumericalAggregateOpts { bool skip_nans = 1; } +// Options for `vortex.sum`. The optional partial-shape marker distinguishes newly serialized +// options from the historical `NumericalAggregateOpts` encoding, which used the same tag for +// `skip_nans` but had no second field. +message SumAggregateOpts { + bool skip_nans = 1; + optional bool struct_partial = 2; +} + // Options for `vortex.literal` message LiteralOpts { vortex.scalar.Scalar value = 1; diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index a44328623e3..afbb82cb1d4 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -26,6 +26,16 @@ pub struct NumericalAggregateOpts { #[prost(bool, tag = "1")] pub skip_nans: bool, } +/// Options for `vortex.sum`. The optional partial-shape marker distinguishes newly serialized +/// options from the historical `NumericalAggregateOpts` encoding, which used the same tag for +/// `skip_nans` but had no second field. +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct SumAggregateOpts { + #[prost(bool, tag = "1")] + pub skip_nans: bool, + #[prost(bool, optional, tag = "2")] + pub struct_partial: ::core::option::Option, +} /// Options for `vortex.literal` #[derive(Clone, PartialEq, ::prost::Message)] pub struct LiteralOpts { From 09a1b4c7b8ccbba1c0602c5d87dbfd7c240aef5c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 31 Jul 2026 16:50:55 +0100 Subject: [PATCH 03/18] Simplify grouped aggregate benchmark helpers Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 37 +++++++++++------------ 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 8a041b2275e..83f8db70796 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -19,7 +19,6 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; 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::aggregate_fn::fns::sum::SumAggregateOpts; @@ -159,12 +158,14 @@ fn list_element_dtype(list_view: &ArrayRef) -> DType { } } -fn grouped_accumulator(list_view: &ArrayRef, vtable: V, options: V::Options) -> ArrayRef +fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef where V: AggregateFnVTable + Clone, + V::Options: Default, { let mut acc = - GroupedAccumulator::try_new(vtable, options, list_element_dtype(list_view)).unwrap(); + GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) + .unwrap(); acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); divan::black_box(acc.finish().unwrap()) @@ -212,7 +213,7 @@ fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -220,7 +221,7 @@ fn sum_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -228,7 +229,7 @@ fn sum_f64_all_valid(bencher: Bencher) { let input = f64_all_valid_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Sum)); } #[divan::bench] @@ -236,21 +237,19 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { let input = f64_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Sum, SumAggregateOpts::default())); + .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, - options: V::Options, -) -> ArrayRef +fn grouped_accumulator_canonical(list_view: &ArrayRef, vtable: V) -> ArrayRef where V: AggregateFnVTable + Clone, + V::Options: Default, { let mut acc = - GroupedAccumulator::try_new(vtable, options, list_element_dtype(list_view)).unwrap(); + 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 @@ -267,7 +266,7 @@ 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, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); } #[divan::bench] @@ -275,7 +274,7 @@ 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, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); } #[divan::bench] @@ -283,7 +282,7 @@ 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, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); } #[divan::bench] @@ -291,7 +290,7 @@ 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, SumAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); } #[divan::bench] @@ -299,7 +298,7 @@ fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Count, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Count)); } #[divan::bench] @@ -307,5 +306,5 @@ fn count_varbinview(bencher: Bencher) { let input = varbinview_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Count, NumericalAggregateOpts::default())); + .bench_refs(|input| grouped_accumulator(input, Count)); } From 1a78215c700286c07a5edc1f5ce8176ddc71b4a0 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 31 Jul 2026 17:12:22 +0100 Subject: [PATCH 04/18] Preserve cached stat validation outside Sum Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/accumulator.rs | 24 ++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 884fac3c4b0..4fb2b330017 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -119,13 +119,33 @@ impl DynAccumulator for Accumulator { batch.dtype() ); + // Sum's legacy Stat slot stores a scalar, while its canonical partial is a struct. Let + // Sum normalize either shape before the generic legacy-stat cast path. + if Stat::from_aggregate_fn(&self.aggregate_fn) == Some(Stat::Sum) + && let Precision::Exact(partial) = batch.statistics().get(Stat::Sum) + { + self.vtable.combine_partials(&mut self.partial, partial)?; + return Ok(()); + } + // 0. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot, // consume that exact stat before kernel dispatch or decode. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) && let Precision::Exact(partial) = batch.statistics().get(stat) { - // Legacy stat slots can use an older partial shape. The aggregate vtable owns that - // compatibility logic (for example, Sum accepts both scalar and struct partials). + let partial = if partial.dtype() == &self.partial_dtype { + partial + } else { + vortex_ensure!( + partial.dtype().eq_ignore_nullability(&self.partial_dtype), + "Aggregate {} read legacy stat {} with dtype {}, expected {}", + self.aggregate_fn, + stat, + partial.dtype(), + self.partial_dtype, + ); + partial.cast(&self.partial_dtype)? + }; self.vtable.combine_partials(&mut self.partial, partial)?; return Ok(()); } From 784c690c7c15ba7120b7bd8ee1d968ada699ed27 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 08:58:57 -0400 Subject: [PATCH 05/18] Special-case Sum stat result dtype Signed-off-by: Matt Katz --- vortex-array/src/scalar_fn/fns/stat.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 843657d8a41..05a149d2a37 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -18,6 +18,7 @@ use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; +use crate::aggregate_fn::fns::sum::Sum; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::display::ExprDisplay; @@ -57,7 +58,7 @@ impl Display for StatOptions { } } -/// Scalar function that broadcasts a stored aggregate result over the input rows. +/// Scalar function that broadcasts a stored aggregate partial over the input rows. /// /// The only current consumer is **row-wise pruning**: substituting `stat(col, agg)` into a /// predicate produces a cheap, row-aligned approximation whose constant runs let downstream @@ -130,7 +131,13 @@ impl ScalarFnVTable for StatFn { } fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult { - let Some(dtype) = aggregate_fn.return_dtype(input_dtype) else { + let dtype = if aggregate_fn.is::() { + // Sum stats expose a scalar result even though canonical Sum state is a struct. + aggregate_fn.return_dtype(input_dtype) + } else { + aggregate_fn.state_dtype(input_dtype) + }; + let Some(dtype) = dtype else { vortex_bail!( "Aggregate function {} does not support input dtype {}", aggregate_fn, From 31ca2c29a0a79e4a3c0a35aa42f8fe5e84c087b4 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 10:20:41 -0400 Subject: [PATCH 06/18] Restore cached Sum fast path Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index fdb15a5c24a..ed75774b9be 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -63,6 +63,10 @@ const IS_EMPTY_FIELD: &str = "is_empty"; /// /// 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); + } + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), array.dtype().clone())?; acc.accumulate(array, ctx)?; let result = acc.finish()?; From 8dadbdd28c69f82df44ab13ac000a9ff8e160a3a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 10:33:57 -0400 Subject: [PATCH 07/18] clean Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum/primitive.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index a58707c09d1..fd621581611 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -183,6 +183,7 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; + use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; @@ -346,7 +347,7 @@ mod tests { } /// Sum an array with explicit [`SumAggregateOpts`] (test-only helper). - fn sum_with_options(arr: &crate::ArrayRef, options: SumAggregateOpts) -> VortexResult { + 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() From 7dc9887f983a6811d81fba59ad1e6e419dc5f70e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 11:22:01 -0400 Subject: [PATCH 08/18] Validate canonical Sum partial state Signed-off-by: Matt Katz --- encodings/sparse/src/compute/sum.rs | 1 - .../src/aggregate_fn/fns/sum/decimal.rs | 56 ++++++++++----- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 16 +++-- .../src/aggregate_fn/fns/sum/tests.rs | 70 ++++++++++++++----- vortex-layout/src/layouts/zoned/zone_map.rs | 2 + 5 files changed, 104 insertions(+), 41 deletions(-) diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index 223589e8a7f..8f85124b3b5 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -100,7 +100,6 @@ mod tests { kernel_result, canonical_result, "kernel and canonical sum paths disagree" ); - Ok(kernel_result) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index ad517dbfdfc..c6aa03f7d0a 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -357,8 +357,7 @@ 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 options = SumAggregateOpts::default(); - let mut state = Sum.empty_partial(&options, &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), @@ -371,10 +370,18 @@ mod tests { let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, small)?; - let result = Sum.finalize_scalar(&state)?; - assert!(!result.is_null()); + let result = Sum.to_scalar(&state)?; + let fields = result.as_struct(); assert_eq!( - result.as_decimal().decimal_value(), + 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(()) @@ -388,8 +395,7 @@ 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 options = SumAggregateOpts::default(); - let mut state = Sum.empty_partial(&options, &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), @@ -403,11 +409,13 @@ mod tests { Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.finalize_scalar(&state)?; - assert!(result.is_null()); + let result = Sum.to_scalar(&state)?; assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(14, 0), Nullable) + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) ); Ok(()) } @@ -416,8 +424,7 @@ 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 options = SumAggregateOpts::default(); - let mut state = Sum.empty_partial(&options, &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), @@ -433,8 +440,14 @@ mod tests { ); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.finalize_scalar(&state)?; - assert!(result.is_null()); + 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(()) } @@ -449,8 +462,7 @@ 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 options = SumAggregateOpts::default(); - let mut state = Sum.empty_partial(&options, &input_dtype)?; + let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; @@ -467,8 +479,14 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Sum.finalize_scalar(&state)?; - assert!(result.is_null()); + 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(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index ed75774b9be..a34d3f017ef 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -304,11 +304,11 @@ impl AggregateFnVTable for Sum { .decimal_value() .vortex_expect("checked non-null"); match value.checked_add(&val) { - Some(r) => { + Some(r) if r.fits_in_precision(*dtype) => { *value = r; - !value.fits_in_precision(*dtype) + false } - None => true, + Some(_) | None => true, } } }; @@ -841,9 +841,15 @@ mod arithmetic_tests { let scalar2 = Scalar::primitive(50i64, Nullable); Sum.combine_partials(&mut state, scalar2)?; - let result = Sum.finalize_scalar(&state)?; + let result = Sum.to_scalar(&state)?; Sum.reset(&mut state); - assert_eq!(result.as_primitive().typed_value::(), Some(150)); + assert_eq!( + result + .as_struct() + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(150) + ); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs index f1c16e21295..5c2fd65dd71 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -130,7 +130,14 @@ fn sum_state_empty_is_null() -> VortexResult<()> { Some(true) ); Sum.combine_partials(&mut state, empty)?; - assert!(Sum.finalize_scalar(&state)?.is_null()); + 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(()) } @@ -144,8 +151,14 @@ fn sum_state_empty_is_identity() -> VortexResult<()> { let empty = Sum.to_scalar(&Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?)?; Sum.combine_partials(&mut state, empty)?; - let result = Sum.finalize_scalar(&state)?; - assert_eq!(result.as_primitive().typed_value::(), Some(100)); + let result = Sum.to_scalar(&state)?; + assert_eq!( + result + .as_struct() + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(100) + ); Ok(()) } @@ -182,7 +195,14 @@ fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { Sum.combine_partials(&mut state, overflowed)?; Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; - assert!(Sum.finalize_scalar(&state)?.is_null()); + 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(()) } @@ -244,10 +264,12 @@ fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable))?; + let partial = Sum.to_scalar(&state)?; assert_eq!( - Sum.finalize_scalar(&state)? - .as_primitive() - .typed_value::(), + partial + .as_struct() + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), Some(0) ); @@ -610,10 +632,18 @@ fn sum_decimal_near_precision_boundary() -> VortexResult<()> { let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); Sum.combine_partials(&mut state, small)?; - let result = Sum.finalize_scalar(&state)?; - assert!(!result.is_null()); + 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!( - result.as_decimal().decimal_value(), + fields + .field("sum") + .and_then(|sum| sum.as_decimal().decimal_value()), Some(DecimalValue::I256(i256::from_i128(99_999_999_999_999))) ); Ok(()) @@ -647,11 +677,13 @@ fn sum_decimal_precision_overflow_within_i256( ); Sum.combine_partials(&mut state, one_more)?; - let result = Sum.finalize_scalar(&state)?; - assert!(result.is_null()); + let result = Sum.to_scalar(&state)?; assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(14, 0), Nullable) + result + .as_struct() + .field("is_overflow") + .and_then(|is_overflow| is_overflow.as_bool().value()), + Some(true) ); Ok(()) } @@ -677,8 +709,14 @@ fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; - let result = Sum.finalize_scalar(&state)?; - assert!(result.is_null()); + 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(()) } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index f15e54237b1..629148bde68 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -322,6 +322,8 @@ fn aggregate_result_expr(stored: &AggregateFnRef, state_expr: Expression) -> Exp 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()), From 4f60d674fd761f6cd60fcc7d09d4f34ecdcbc8ab Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 11:28:07 -0400 Subject: [PATCH 09/18] Narrow Sum helper visibility Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/constant.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/decimal.rs | 2 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 6 +++--- vortex-array/src/aggregate_fn/fns/sum/primitive.rs | 8 ++++---- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index d74568df356..b3d7d742292 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -13,7 +13,7 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; -pub(crate) fn accumulate_bool( +pub(super) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 1e64e5348d8..14da9d60799 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -15,7 +15,7 @@ use crate::scalar::Scalar; /// /// Returns `Ok(None)` if the scalar is null (no contribution to the sum). /// Returns a null scalar on overflow (saturation). -pub(crate) fn multiply_constant( +pub(super) fn multiply_constant( scalar: &Scalar, len: usize, return_dtype: &DType, diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index c6aa03f7d0a..63bbf1b8693 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -23,7 +23,7 @@ use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. /// Returns Ok(true) if saturated (overflow), Ok(false) if not. -pub(crate) fn accumulate_decimal( +pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index a34d3f017ef..dcb560dce26 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -494,7 +494,7 @@ pub enum SumState { }, } -pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { +fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0), @@ -666,7 +666,7 @@ fn any_valid(validity: Validity, len: usize, ctx: &mut ExecutionCtx) -> VortexRe /// Checked add for u64, returning true if overflow occurred. #[inline(always)] -pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { +fn checked_add_u64(acc: &mut u64, val: u64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; @@ -678,7 +678,7 @@ pub(crate) fn checked_add_u64(acc: &mut u64, val: u64) -> bool { /// Checked add for i64, returning true if overflow occurred. #[inline(always)] -pub(crate) fn checked_add_i64(acc: &mut i64, val: i64) -> bool { +fn checked_add_i64(acc: &mut i64, val: i64) -> bool { match acc.checked_add(val) { Some(r) => { *acc = r; diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index fd621581611..bc6d9b5cba7 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -21,7 +21,7 @@ use crate::match_each_native_ptype; /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. const SUM_CHUNK: usize = 1 << 16; -pub(crate) fn accumulate_primitive( +pub(super) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -66,7 +66,7 @@ fn accumulate_primitive_all( /// Sum the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values /// are skipped to match the scalar `sum` semantics; otherwise any NaN poisons the accumulator to /// NaN. Floats cannot overflow the accumulator, so this never reports saturation. -pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { +pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { if skip_nans { for &v in slice { if !v.is_nan() { @@ -84,7 +84,7 @@ pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan /// chunks of [`SUM_CHUNK`] with a single checked add per chunk, which lets the inner loop vectorize /// to packed widening adds. `u64` input keeps a per-element checked add since a chunk of `u64`s /// could itself overflow. Returns `true` on overflow. -pub(crate) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool +pub(super) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { @@ -106,7 +106,7 @@ where } /// Signed counterpart of [`sum_unsigned_all`]. -pub(crate) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool +pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { From b36dab04f35b5ff40a69521f77f8535c38eaf027 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 11:39:03 -0400 Subject: [PATCH 10/18] fix Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 2 +- vortex-array/src/expr/stats/mod.rs | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index cf75d52ddd2..16a1b5431fe 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -334,7 +334,7 @@ mod tests { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; - assert!(result.is_null()); + assert_eq!(result.as_primitive().as_::(), None); Ok(()) } diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index b1100fe2eb3..756d2eb4237 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -10,6 +10,7 @@ use enum_iterator::all; use num_enum::IntoPrimitive; use num_enum::TryFromPrimitive; +use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; @@ -188,8 +189,8 @@ impl Stat { } Self::Sum => { // Statistics follow NaN-skipping semantics; request it explicitly. - let options = aggregate_fn::fns::sum::SumAggregateOpts::skip_nans(); - return aggregate_fn::fns::sum::Sum.return_dtype(&options, data_type); + return aggregate_fn::fns::sum::Sum + .return_dtype(&SumAggregateOpts::skip_nans(), data_type); } }) } @@ -200,8 +201,7 @@ impl Stat { Some(match self { Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::sum::Sum - .bind(aggregate_fn::fns::sum::SumAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum.bind(SumAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { From daaad48cc786811a7a5721de1c5a10ce1dd05e43 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 11:45:02 -0400 Subject: [PATCH 11/18] fix Signed-off-by: Matt Katz --- vortex-array/benches/aggregate_grouped.rs | 43 ----------------------- 1 file changed, 43 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 83f8db70796..312c5c9c184 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -14,21 +14,15 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum::SumAggregateOpts; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -171,43 +165,6 @@ where divan::black_box(acc.finish().unwrap()) } -#[divan::bench] -fn sum_legacy_scalar_partial_merge(bencher: Bencher) { - let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let partial = Scalar::primitive(1i64, Nullability::Nullable); - bencher - .with_inputs(|| partial.clone()) - .bench_refs(|partial| { - let mut acc = - Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); - for _ in 0..GROUP_COUNT { - acc.combine_partials(partial.clone()).unwrap(); - } - divan::black_box(acc.finish().unwrap()) - }); -} - -#[divan::bench] -fn sum_canonical_partial_merge(bencher: Bencher) { - let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut source = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); - source - .combine_partials(Scalar::primitive(1i64, Nullability::Nullable)) - .unwrap(); - let partial = source.flush().unwrap(); - - bencher - .with_inputs(|| partial.clone()) - .bench_refs(|partial| { - let mut acc = - Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype.clone()).unwrap(); - for _ in 0..GROUP_COUNT { - acc.combine_partials(partial.clone()).unwrap(); - } - divan::black_box(acc.finish().unwrap()) - }); -} - #[divan::bench] fn sum_i32_nullable_all_valid(bencher: Bencher) { let input = i32_nullable_all_valid_input(); From 37dbbb6ada10cdae5e180ff4fb8bee7811a1a81c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 16:26:13 -0400 Subject: [PATCH 12/18] Use canonical Sum partials at runtime Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/accumulator.rs | 56 +++-- vortex-array/src/aggregate_fn/fns/sum/bool.rs | 5 +- .../src/aggregate_fn/fns/sum/decimal.rs | 26 ++- .../src/aggregate_fn/fns/sum/grouped.rs | 3 - vortex-array/src/aggregate_fn/fns/sum/mod.rs | 192 ++++++++---------- .../src/aggregate_fn/fns/sum/primitive.rs | 5 +- .../src/aggregate_fn/fns/sum/tests.rs | 125 ++++++------ vortex-array/src/aggregate_fn/proto.rs | 2 +- vortex-layout/src/layouts/zoned/schema.rs | 21 +- vortex-layout/src/layouts/zoned/zone_map.rs | 9 +- 10 files changed, 233 insertions(+), 211 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 4fb2b330017..52255481a43 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -119,18 +119,18 @@ impl DynAccumulator for Accumulator { batch.dtype() ); - // Sum's legacy Stat slot stores a scalar, while its canonical partial is a struct. Let - // Sum normalize either shape before the generic legacy-stat cast path. - if Stat::from_aggregate_fn(&self.aggregate_fn) == Some(Stat::Sum) - && let Precision::Exact(partial) = batch.statistics().get(Stat::Sum) - { - self.vtable.combine_partials(&mut self.partial, partial)?; + // 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. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot, - // consume that exact stat before kernel dispatch or decode. + // 0. Cached stats bridge: consume an exact partial from the aggregate's Stat slot before + // kernel dispatch or decode. Sum is handled above because its cache stores a result. if let Some(stat) = Stat::from_aggregate_fn(&self.aggregate_fn) + && stat != Stat::Sum && let Precision::Exact(partial) = batch.statistics().get(stat) { let partial = if partial.dtype() == &self.partial_dtype { @@ -176,7 +176,7 @@ impl DynAccumulator for Accumulator { } // 2. Allow the vtable to short-circuit on the raw array before decompression. - if self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { + if !checked_cached_sum && self.vtable.try_accumulate(&mut self.partial, batch, ctx)? { return Ok(()); } @@ -298,7 +298,10 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; use crate::scalar::Scalar; + use crate::scalar::ScalarValue; /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` @@ -376,10 +379,15 @@ mod tests { fn sum_partial(value: f64) -> Scalar { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype).expect("sum"); - acc.combine_partials(Scalar::primitive(value, Nullability::Nullable)) - .expect("legacy scalar partial"); - acc.flush().expect("sum partial") + 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 @@ -484,4 +492,26 @@ mod tests { ); Ok(()) } + + #[test] + fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { + static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let batch = dict_of_seven(); + batch + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(11.0f64))); + + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, SumAggregateOpts::default(), dtype)?; + acc.accumulate(&batch, &mut ctx)?; + + assert_eq!(acc.finish()?.as_primitive().as_::(), Some(11.0)); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3d7d742292..4e33a567dd2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -6,23 +6,22 @@ use std::ops::BitAnd; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; +use vortex_mask::Mask; use super::SumState; use super::checked_add_u64; -use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; pub(super) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, - ctx: &mut ExecutionCtx, + mask: &Mask, ) -> VortexResult { let SumState::Unsigned(acc) = inner else { vortex_panic!("expected unsigned sum state for bool input"); }; - let mask = b.as_ref().validity()?.execute_mask(b.as_ref().len(), ctx)?; let true_count = match mask.bit_buffer() { AllOr::None => return Ok(false), AllOr::All => b.bit_buffer_view().true_count() as u64, diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 63bbf1b8693..03c3f116fc5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -13,7 +13,6 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use super::SumState; -use crate::ExecutionCtx; use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; @@ -26,10 +25,9 @@ use crate::scalar::DecimalValue; pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, - ctx: &mut ExecutionCtx, + mask: &Mask, ) -> VortexResult { - let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - let validity = match &mask { + let validity = match mask { Mask::AllTrue(_) => None, Mask::Values(mask_values) => Some(mask_values.bit_buffer()), Mask::AllFalse(_) => { @@ -111,6 +109,7 @@ 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; @@ -129,6 +128,11 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; + fn partial_with_value(value: Scalar) -> VortexResult { + let return_dtype = value.dtype().as_nullable(); + sum_result_partial_scalar(value, &return_dtype, false) + } + #[test] fn sum_decimal_basic() -> VortexResult<()> { let decimal = DecimalArray::new( @@ -364,11 +368,11 @@ mod tests { DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Add a small value that keeps us just under 10^14. let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, small)?; + Sum.combine_partials(&mut state, partial_with_value(small)?)?; let result = Sum.to_scalar(&state)?; let fields = result.as_struct(); @@ -402,12 +406,12 @@ mod tests { DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Push the sum to exactly 10^14, exceeding precision 14. let one_more = Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, one_more)?; + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; let result = Sum.to_scalar(&state)?; assert_eq!( @@ -431,14 +435,14 @@ mod tests { DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; let one_more = Scalar::decimal( DecimalValue::from(-1i64), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, one_more)?; + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; let result = Sum.to_scalar(&state)?; assert_eq!( @@ -468,7 +472,7 @@ mod tests { let near_limit_val: i128 = 10i128.pow(37) - 1; let near_limit = Scalar::decimal(DecimalValue::from(near_limit_val), return_dtype, Nullable); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 38c3c103567..af655df0db9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -46,9 +46,6 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); }; - if !options.struct_partial { - return Ok(None); - } try_grouped_sum(groups, ctx, options.skip_nans) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index dcb560dce26..9416f5dfc4b 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -78,24 +78,24 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(result) } -/// Sum an array, returning null when it has no valid values. -/// -/// If the sum overflows, a null scalar is returned. Legacy scalar partials remain supported; their -/// zero identity preserves the historical zero-on-empty behavior when encountered. +/// 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`]. /// -/// New sums use a struct partial that can distinguish an empty input from a zero sum. The -/// `struct_partial` field exists for deserializing aggregates written before that partial was -/// introduced; callers should normally construct these options with [`Default`], +/// 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 partials use the `{ sum, is_overflow, is_empty }` struct representation. + /// 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, } @@ -117,6 +117,9 @@ impl SumAggregateOpts { } /// 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, @@ -225,11 +228,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)?; - if options.struct_partial { - Some(sum_partial_dtype(return_dtype)) - } else { - Some(return_dtype) - } + Some(sum_partial_dtype(return_dtype)) } fn empty_partial( @@ -247,12 +246,17 @@ impl AggregateFnVTable for Sum { is_overflow: false, is_empty: true, skip_nans: options.skip_nans, - struct_partial: options.struct_partial, }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let other = normalize_partial_scalar(other, &partial.return_dtype)?; + let other = other.cast(&sum_partial_dtype(partial.return_dtype.clone()))?; + if other.is_null() { + partial.is_empty = false; + partial.is_overflow = true; + return Ok(()); + } + let fields = other.as_struct(); let other = fields .field(SUM_FIELD) @@ -321,10 +325,6 @@ impl AggregateFnVTable for Sum { } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - if !partial.struct_partial { - return Ok(legacy_sum_value_scalar(partial)); - } - Ok(Scalar::struct_( sum_partial_dtype(partial.return_dtype.clone()), vec![ @@ -399,56 +399,49 @@ impl AggregateFnVTable for Sum { 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 any_valid = if partial.is_empty { - match batch { - Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => { - any_valid(p.as_ref().validity()?, p.as_ref().len(), ctx)? - } - Canonical::Bool(b) => any_valid(b.as_ref().validity()?, b.as_ref().len(), ctx)?, - Canonical::Decimal(d) => { - any_valid(d.as_ref().validity()?, d.as_ref().len(), ctx)? - } - _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), - }, - Columnar::Constant(_) => unreachable!(), - } - } else { - false - }; - - let result = match batch { + let (result, any_valid) = match batch { Columnar::Canonical(c) => match c { Canonical::Primitive(p) => { - accumulate_primitive(&mut partial.sum, p, ctx, skip_nans) + 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::Bool(b) => accumulate_bool(&mut partial.sum, b, ctx), - Canonical::Decimal(d) => accumulate_decimal(&mut partial.sum, d, ctx), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), }, Columnar::Constant(_) => unreachable!(), }; - match result { - Ok(false) => { - if any_valid { - partial.is_empty = false; - } - } - Ok(true) => partial.is_overflow = true, - Err(e) => return Err(e), + if any_valid { + partial.is_empty = false; + } + if result? { + partial.is_overflow = true; + partial.is_empty = false; } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - let partials = normalize_partial_array(partials)?; let sum = partials.get_item(SUM_FIELD)?; let is_invalid = partials .get_item(IS_OVERFLOW_FIELD)? @@ -458,11 +451,7 @@ impl AggregateFnVTable for Sum { } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - if partial.struct_partial { - Ok(sum_value_scalar(partial)) - } else { - Ok(legacy_sum_value_scalar(partial)) - } + Ok(sum_value_scalar(partial)) } } @@ -477,8 +466,6 @@ pub struct SumPartial { is_empty: bool, /// Whether NaN values in float inputs are skipped. skip_nans: bool, - /// Whether this accumulator emits the canonical struct partial. - struct_partial: bool, } /// The accumulated sum value. @@ -550,14 +537,6 @@ fn sum_value_scalar(partial: &SumPartial) -> Scalar { nullable_sum_state_scalar(partial) } -fn legacy_sum_value_scalar(partial: &SumPartial) -> Scalar { - if partial.is_overflow { - return Scalar::null(partial.return_dtype.as_nullable()); - } - - nullable_sum_state_scalar(partial) -} - fn nullable_sum_state_scalar(partial: &SumPartial) -> Scalar { match &partial.sum { SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::Nullable), @@ -573,14 +552,17 @@ fn nullable_sum_state_scalar(partial: &SumPartial) -> Scalar { } } -/// Normalize an array of scalar legacy Sum partials into the canonical struct partial shape. +/// Convert scalar legacy Sum partials read from storage (a single nullable primitive) to the struct partial shape. /// -/// Canonical partial arrays are returned unchanged. 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`. -pub fn normalize_partial_array(partials: ArrayRef) -> VortexResult { +/// 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(..)) { - return Ok(partials); + vortex_bail!( + "Expected scalar legacy Sum partials, found {}", + partials.dtype() + ); } let len = partials.len(); @@ -602,42 +584,31 @@ pub fn normalize_partial_array(partials: ArrayRef) -> VortexResult { .into_array()) } -fn normalize_partial_scalar(partial: Scalar, return_dtype: &DType) -> VortexResult { - let partial_dtype = sum_partial_dtype(return_dtype.clone()); - if matches!(partial.dtype(), DType::Struct(..)) { - if partial.is_null() { - return Ok(Scalar::struct_( - partial_dtype, - vec![ - Scalar::zero_value(&return_dtype.as_nonnullable()), - Scalar::bool(true, Nullability::NonNullable), - Scalar::bool(false, Nullability::NonNullable), - ], - )); - } - return partial.cast(&partial_dtype); - } - - if !partial.dtype().eq_ignore_nullability(return_dtype) { +fn sum_result_partial_scalar( + result: Scalar, + return_dtype: &DType, + is_empty: bool, +) -> VortexResult { + if !result.dtype().eq_ignore_nullability(return_dtype) { vortex_bail!( - "Legacy Sum partial has dtype {}, expected {}", - partial.dtype(), + "Sum result has dtype {}, expected {}", + result.dtype(), return_dtype ); } - let is_overflow = partial.is_null(); - let sum = if is_overflow { + let is_overflow = result.is_null() && !is_empty; + let sum = if result.is_null() { Scalar::zero_value(&return_dtype.as_nonnullable()) } else { - partial.cast(&return_dtype.as_nonnullable())? + result.cast(&return_dtype.as_nonnullable())? }; Ok(Scalar::struct_( - partial_dtype, + sum_partial_dtype(return_dtype.clone()), vec![ sum, Scalar::bool(is_overflow, Nullability::NonNullable), - Scalar::bool(false, Nullability::NonNullable), + Scalar::bool(is_empty, Nullability::NonNullable), ], )) } @@ -651,19 +622,12 @@ fn try_accumulate_cached_sum( return Ok(false); }; - let sum = if sum.dtype() == &partial.return_dtype { - sum - } else { - sum.cast(&partial.return_dtype)? - }; - vtable.combine_partials(partial, sum)?; + 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) } -fn any_valid(validity: Validity, len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(validity.execute_mask(len, ctx)?.true_count() > 0) -} - /// Checked add for u64, returning true if overflow occurred. #[inline(always)] fn checked_add_u64(acc: &mut u64, val: u64) -> bool { @@ -698,6 +662,7 @@ mod arithmetic_tests { use vortex_error::VortexExpect; use vortex_error::VortexResult; + use super::sum_result_partial_scalar; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; @@ -732,6 +697,11 @@ 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(); @@ -742,9 +712,11 @@ mod arithmetic_tests { return sum(array, &mut ctx); } - let sum_dtype = Stat::Sum.dtype(array.dtype()).ok_or_else(|| { - vortex_error::vortex_err!("Sum not supported for dtype: {}", array.dtype()) - })?; + let sum_dtype = Sum + .return_dtype(&SumAggregateOpts::default(), array.dtype()) + .ok_or_else(|| { + vortex_error::vortex_err!("Sum not supported for dtype: {}", array.dtype()) + })?; // For non-float types, try statistics short-circuit with accumulator. if !matches!(&sum_dtype, DType::Primitive(p, _) if p.is_float()) @@ -836,10 +808,10 @@ mod arithmetic_tests { let mut state = Sum.empty_partial(&options, &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); - Sum.combine_partials(&mut state, scalar1)?; + Sum.combine_partials(&mut state, partial_with_value(scalar1)?)?; let scalar2 = Scalar::primitive(50i64, Nullable); - Sum.combine_partials(&mut state, scalar2)?; + Sum.combine_partials(&mut state, partial_with_value(scalar2)?)?; let result = Sum.to_scalar(&state)?; Sum.reset(&mut state); diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index bc6d9b5cba7..cb021023e15 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -7,11 +7,11 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; +use vortex_mask::Mask; use super::SumState; use super::checked_add_i64; use super::checked_add_u64; -use crate::ExecutionCtx; use crate::arrays::PrimitiveArray; use crate::dtype::NativePType; use crate::dtype::PType; @@ -24,10 +24,9 @@ const SUM_CHUNK: usize = 1 << 16; pub(super) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, - ctx: &mut ExecutionCtx, + mask: &Mask, skip_nans: bool, ) -> VortexResult { - let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), AllOr::All => accumulate_primitive_all(inner, p, skip_nans), diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs index 5c2fd65dd71..09abe564c4c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -12,6 +12,7 @@ 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; @@ -39,6 +40,7 @@ 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; @@ -51,6 +53,11 @@ fn sum_with_options(arr: &ArrayRef, options: SumAggregateOpts) -> VortexResult 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(); @@ -76,7 +83,7 @@ fn sum_uses_new_partial_shape_by_default() { } #[test] -fn legacy_options_use_scalar_partial_and_zero_on_empty() -> VortexResult<()> { +fn legacy_options_only_describe_the_stored_partial() -> VortexResult<()> { let options = SumAggregateOpts::deserialize(&NumericalAggregateOpts::skip_nans().serialize())?; assert_eq!( options, @@ -87,17 +94,20 @@ fn legacy_options_use_scalar_partial_and_zero_on_empty() -> VortexResult<()> { ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - assert_eq!( + assert!(matches!( Sum.partial_dtype(&options, &input_dtype), - Some(DType::Primitive(PType::I64, Nullable)) - ); + Some(DType::Struct(..)) + )); let mut acc = Accumulator::try_new(Sum, options, input_dtype)?; assert_eq!( - acc.partial_scalar()?.as_primitive().typed_value::(), - Some(0) + acc.partial_scalar()? + .as_struct() + .field("is_empty") + .and_then(|is_empty| is_empty.as_bool().value()), + Some(true) ); - assert_eq!(acc.finish()?.as_primitive().typed_value::(), Some(0)); + assert!(acc.finish()?.is_null()); Ok(()) } @@ -146,7 +156,10 @@ 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, Scalar::primitive(100i64, Nullable))?; + 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)?; @@ -167,8 +180,14 @@ 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, Scalar::primitive(i64::MAX, Nullable))?; - Sum.combine_partials(&mut overflowed, Scalar::primitive(1i64, Nullable))?; + 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!( @@ -191,9 +210,15 @@ fn sum_state_overflow_sets_flag_and_poisons() -> VortexResult<()> { ); let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; - Sum.combine_partials(&mut state, Scalar::primitive(5i64, Nullable))?; + Sum.combine_partials( + &mut state, + partial_with_value(Scalar::primitive(5i64, Nullable))?, + )?; Sum.combine_partials(&mut state, overflowed)?; - Sum.combine_partials(&mut state, Scalar::primitive(7i64, Nullable))?; + Sum.combine_partials( + &mut state, + partial_with_value(Scalar::primitive(7i64, Nullable))?, + )?; let partial = Sum.to_scalar(&state)?; assert_eq!( @@ -247,13 +272,12 @@ fn sum_all_nan_is_zero_not_null() -> VortexResult<()> { } #[test] -fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { +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()); - // A scalar `Stat::Sum` is an old partial. Its zero identity cannot encode emptiness, so its - // historical zero-on-empty result is preserved when it is encountered. arr.statistics() .set(Stat::Sum, Precision::Exact(ScalarValue::from(0i64))); assert_eq!( @@ -263,40 +287,10 @@ fn legacy_scalar_partial_preserves_zero_on_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut state = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; - Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable))?; - let partial = Sum.to_scalar(&state)?; - assert_eq!( - partial - .as_struct() - .field("sum") - .and_then(|sum| sum.as_primitive().typed_value::()), - Some(0) - ); - - let mut overflowed = Sum.empty_partial(&SumAggregateOpts::default(), &dtype)?; - Sum.combine_partials( - &mut overflowed, - Scalar::null(DType::Primitive(PType::I64, 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(0) - ); - 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!( + Sum.combine_partials(&mut state, Scalar::primitive(0i64, Nullable)) + .is_err(), + "live Sum only accepts canonical struct partials" ); Ok(()) } @@ -514,13 +508,10 @@ fn sum_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> #[test] fn sum_uses_cached_stat_sum() -> VortexResult<()> { - // A planted exact `Stat::Sum` with a known null count is consumed instead of a scan - // (the planted value differs from the actual data to prove it). + // 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))); - arr.statistics() - .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) @@ -534,8 +525,6 @@ fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); arr.statistics() .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); - arr.statistics() - .set(Stat::NullCount, Precision::Exact(ScalarValue::from(0u64))); let result = sum_with_options(&arr, SumAggregateOpts::include_nans())?; assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); Ok(()) @@ -585,6 +574,22 @@ fn sum_checked_overflow_is_null_and_saturates() -> VortexResult<()> { 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()); @@ -626,11 +631,11 @@ fn sum_decimal_near_precision_boundary() -> VortexResult<()> { DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Add a small value that keeps us just under 10^14. let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable); - Sum.combine_partials(&mut state, small)?; + Sum.combine_partials(&mut state, partial_with_value(small)?)?; let result = Sum.to_scalar(&state)?; let fields = result.as_struct(); @@ -668,14 +673,14 @@ fn sum_decimal_precision_overflow_within_i256( DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; let one_more = Scalar::decimal( DecimalValue::from(one_more), DecimalDType::new(14, 0), Nullable, ); - Sum.combine_partials(&mut state, one_more)?; + Sum.combine_partials(&mut state, partial_with_value(one_more)?)?; let result = Sum.to_scalar(&state)?; assert_eq!( @@ -701,7 +706,7 @@ fn sum_decimal_accumulate_precision_overflow() -> VortexResult<()> { // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; let near_limit = Scalar::decimal(DecimalValue::from(near_limit_val), return_dtype, Nullable); - Sum.combine_partials(&mut state, near_limit)?; + Sum.combine_partials(&mut state, partial_with_value(near_limit)?)?; // Now accumulate a real i128 array with a single element = 1 to overflow precision. let decimal = DecimalArray::new(buffer![1i128], DecimalDType::new(27, 0), Validity::AllValid); diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index a3e72ba8e0c..9580474d048 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -195,7 +195,7 @@ mod tests { } #[test] - fn legacy_sum_options_select_scalar_partial() -> VortexResult<()> { + fn legacy_sum_options_mark_a_stored_scalar_partial() -> VortexResult<()> { let session = crate::array_session(); let proto = pb::AggregateFn { id: Sum.id().to_string(), diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 5ffa94e2eb6..629387bea72 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -90,8 +91,21 @@ pub(crate) fn aggregate_stats_table_dtype( ) -> DType { DType::Struct( StructFields::from_iter(aggregate_fns.iter().filter_map(|aggregate_fn| { - aggregate_state_dtype(column_dtype, aggregate_fn) - .map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable())) + let dtype = if aggregate_fn + .as_opt::() + .is_some_and(|options| !options.struct_partial) + { + aggregate_fn.return_dtype(column_dtype).or_else(|| { + if let DType::Extension(ext) = column_dtype { + aggregate_fn.return_dtype(ext.storage_dtype()) + } else { + None + } + }) + } else { + aggregate_state_dtype(column_dtype, aggregate_fn) + }; + dtype.map(|dtype| (aggregate_fn.to_string(), dtype.as_nullable())) })), Nullability::NonNullable, ) @@ -276,7 +290,7 @@ mod tests { } #[test] - fn sum_stats_table_dtype_uses_option_partial_shape() -> VortexResult<()> { + 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())?; @@ -292,6 +306,7 @@ mod tests { current.as_struct_fields().field("vortex.sum()"), Some(DType::Struct(..)) )); + Ok(()) } } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 629148bde68..2248aaed97c 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -17,7 +17,7 @@ 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_partial_array; +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; @@ -187,7 +187,9 @@ pub(super) fn normalize_sum_partial_fields( let Some(index) = names.find(aggregate_fn.to_string()) else { continue; }; - fields[index] = normalize_partial_array(fields[index].clone())?; + 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()) @@ -448,7 +450,7 @@ mod tests { } #[test] - fn legacy_scalar_sum_field_is_normalized_once() -> VortexResult<()> { + fn legacy_scalar_sum_field_is_normalized_on_read() -> VortexResult<()> { let options = SumAggregateOpts::deserialize(&NumericalAggregateOpts::default().serialize())?; let sum = Sum.bind(options); @@ -462,7 +464,6 @@ mod tests { 1, 3, )?; - let result_expr = zone_map .aggregate_field_expr(&sum) .expect("normalized Sum field"); From 0bc5ac52dc7060120bf385acb90a82305da4d784 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 3 Aug 2026 16:46:34 -0400 Subject: [PATCH 13/18] Optimize grouped sum finalization Signed-off-by: Matt Katz --- .../src/aggregate_fn/accumulator_grouped.rs | 5 ++++- .../src/aggregate_fn/fns/sum/grouped.rs | 21 ++++++++++++------- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 286f4194cd5..8bd26dad961 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -273,7 +273,10 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn flush(&mut self) -> VortexResult { - let states = std::mem::take(&mut self.partials); + let mut states = std::mem::take(&mut self.partials); + if states.len() == 1 { + return Ok(states.pop().vortex_expect("checked one partial")); + } Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index af655df0db9..65a92fff0e9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -4,7 +4,6 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; -use vortex_mask::AllOr; use vortex_mask::Mask; use super::IS_EMPTY_FIELD; @@ -89,7 +88,7 @@ fn grouped_sum( .as_ref() .validity()? .execute_mask(elements.as_ref().len(), ctx)?; - let all_valid = matches!(elem_mask.slices(), AllOr::All); + let all_valid = elem_mask.all_true(); let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), unsigned: |T| { @@ -167,16 +166,22 @@ fn sum_masked_group( elem_mask: &Mask, sum_run: &impl Fn(&mut A, &[T]) -> bool, ) -> (bool, bool) { - match elem_mask.slice(offset..offset + size).slices() { - AllOr::All => (sum_run(acc, &values[offset..offset + size]), size > 0), - AllOr::None => (false, false), - AllOr::Some(runs) => { - for &(start, end) in runs { + match elem_mask { + Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), + Mask::AllFalse(_) => (false, false), + Mask::Values(mask_values) => { + let validity = mask_values + .bit_buffer() + .as_view() + .slice(offset..offset + size); + let mut any_valid = false; + for (start, end) in validity.set_slices() { + any_valid = true; if sum_run(acc, &values[offset + start..offset + end]) { return (true, true); } } - (false, !runs.is_empty()) + (false, any_valid) } } } From 25d8216c931ad7a953d7cc5ce1f350cca5476e71 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 4 Aug 2026 10:07:22 -0700 Subject: [PATCH 14/18] Update null-on-empty Sum expectations Signed-off-by: Matt Katz --- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 6 ++++++ vortex-python/src/io.rs | 4 ++-- vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt | 3 +-- vortex-sqllogictest/slt/duckdb/nan_aggregates.slt | 2 +- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 9416f5dfc4b..e9af5759e1e 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -621,6 +621,12 @@ fn try_accumulate_cached_sum( 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)?; diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index aece56aa9d6..096db5adedf 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -314,7 +314,7 @@ impl PyVortexWriteOptions { /// >>> vx.io.VortexWriteOptions.default().write(sprl, "chonky.vortex") /// >>> import os /// >>> os.path.getsize('chonky.vortex') - /// 215932 + /// 216108 /// /// Wow, Vortex manages to use about two bytes per integer! So advanced. So tiny. /// @@ -324,7 +324,7 @@ impl PyVortexWriteOptions { /// /// >>> vx.io.VortexWriteOptions.compact().write(sprl, "tiny.vortex") /// >>> os.path.getsize('tiny.vortex') - /// 55060 + /// 55232 /// /// Random numbers are not (usually) composed of random bytes! #[staticmethod] diff --git a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt index fbf9d2b74fe..c0c7f9128e9 100644 --- a/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt +++ b/vortex-sqllogictest/slt/duckdb/aggregates_edge_cases.slt @@ -13,8 +13,7 @@ statement ok COPY (SELECT CAST(x AS INTEGER) AS x FROM (VALUES (CAST(NULL AS INTEGER)),(CAST(NULL AS INTEGER))) t(x)) TO '${WORK_DIR}/i-null.vortex'; -# TODO(https://github.com/vortex-data/vortex/issues/9084): this should be fixed (expected NULL) query I SELECT sum(x) FROM '${WORK_DIR}/i-null.vortex'; ---- -0 +NULL diff --git a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt index 59931886824..b6640cf0f39 100644 --- a/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt +++ b/vortex-sqllogictest/slt/duckdb/nan_aggregates.slt @@ -123,7 +123,7 @@ COPY ( query IRR SELECT count(x), sum(x), max(x) FROM '${WORK_DIR}/all-null.vortex'; ---- -0 0 NULL +0 NULL NULL query R SELECT min(x) FROM '${WORK_DIR}/all-null.vortex'; From 86f1af8100c33671237f1c208be89fb225bc58f5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 4 Aug 2026 10:34:44 -0700 Subject: [PATCH 15/18] fix Signed-off-by: Matt Katz --- vortex-layout/src/layouts/zoned/mod.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 01de0b2b452..bc3d7d0626e 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -267,6 +267,7 @@ impl LegacyStatsLayout { usize::try_from(self.children().child_row_count(1)) .vortex_expect("Invalid number of zones, cannot handle more than usize zones") } + /// Returns display names for the zone-map aggregates stored by this layout. pub fn present_aggregates(&self) -> Arc<[String]> { present_aggregates(&self.zone_map_schema) @@ -420,7 +421,11 @@ impl DeserializeMetadata for ZonedMetadata { vortex_bail!("Zoned metadata missing protobuf version"); }; - vortex_ensure_eq!(version, ZONED_METADATA_PROTO_VERSION); + vortex_ensure!( + version == ZONED_METADATA_PROTO_VERSION, + "Unsupported zoned metadata version: {}", + version + ); vortex_ensure!(!proto_bytes.is_empty(), "Zoned metadata missing protobuf"); let proto = ZonedMetadataProto::decode(proto_bytes)?; From 1664756966a9b7e50a6ce4ab1505a1eae52873cc Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 4 Aug 2026 10:36:07 -0700 Subject: [PATCH 16/18] fix Signed-off-by: Matt Katz --- vortex-layout/src/layouts/zoned/writer.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index c9b563900a1..92c4b4c9870 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -209,9 +209,11 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A }; let mut aggregate_fns = vec![max, min]; - let sum_options = SumAggregateOpts::skip_nans(); - if Sum.return_dtype(&sum_options, dtype).is_some() { - aggregate_fns.push(Sum.bind(sum_options)); + if Sum + .return_dtype(&SumAggregateOpts::skip_nans(), dtype) + .is_some() + { + aggregate_fns.push(Sum.bind(SumAggregateOpts::skip_nans())); } aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); From bf3c8f52203c190692dbef7c3685501e8578daf2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 4 Aug 2026 11:21:53 -0700 Subject: [PATCH 17/18] perf: avoid grouped sum bitmap appends Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/sum/grouped.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 65a92fff0e9..164bfea23b2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -135,12 +135,11 @@ fn collect_sums( all_valid: bool, sum_run: impl Fn(&mut A, &[T]) -> bool, ) -> (PrimitiveArray, BitBuffer, BitBuffer) { - let mut is_overflow = BitBufferMut::with_capacity(group_ranges.len()); - let mut is_empty = BitBufferMut::with_capacity(group_ranges.len()); + let group_count = group_ranges.len(); + let mut is_overflow = BitBufferMut::new_unset(group_count); + let mut is_empty = BitBufferMut::new_unset(group_count); let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { if !group_validity.value(i) { - is_overflow.append(false); - is_empty.append(true); return A::default(); } let mut acc = A::default(); @@ -149,8 +148,16 @@ fn collect_sums( } else { sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) }; - is_overflow.append(overflow); - is_empty.append(!any_valid); + if overflow { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_overflow.set_unchecked(i) }; + } + if !any_valid { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_empty.set_unchecked(i) }; + } acc }); let sums = PrimitiveArray::from_iter(sums); From e2b15cbb249c06703face67204ecdede8841e25a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 4 Aug 2026 12:48:56 -0700 Subject: [PATCH 18/18] perf: streamline sum partial handling Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/sum/grouped.rs | 36 ++-- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 190 ++++++++---------- .../src/aggregate_fn/fns/sum/tests.rs | 23 +++ 3 files changed, 125 insertions(+), 124 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 164bfea23b2..64b0ce81995 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -6,13 +6,11 @@ use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; -use super::IS_EMPTY_FIELD; -use super::IS_OVERFLOW_FIELD; -use super::SUM_FIELD; 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; @@ -24,8 +22,6 @@ use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::match_each_native_ptype; @@ -109,20 +105,22 @@ fn grouped_sum( } ); - Ok(StructArray::try_new( - FieldNames::from_iter([ - FieldName::from(SUM_FIELD), - FieldName::from(IS_OVERFLOW_FIELD), - FieldName::from(IS_EMPTY_FIELD), - ]), - vec![ - sums.into_array(), - BoolArray::new(is_overflow, Validity::NonNullable).into_array(), - BoolArray::new(is_empty, Validity::NonNullable).into_array(), - ], - group_validity.len(), - Validity::from_mask(group_validity.clone(), Nullability::Nullable), - )? + 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()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index e9af5759e1e..89b19b5575f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -12,9 +12,9 @@ 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; @@ -250,77 +250,23 @@ impl AggregateFnVTable for Sum { } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - let other = other.cast(&sum_partial_dtype(partial.return_dtype.clone()))?; - if other.is_null() { - partial.is_empty = false; - partial.is_overflow = true; + let (other_sum, other_is_overflow, other_is_empty) = decode_sum_partial_scalar(other)?; + validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; + + if partial.is_overflow { return Ok(()); } - - let fields = other.as_struct(); - let other = fields - .field(SUM_FIELD) - .ok_or_else(|| vortex_err!("Sum partial is missing the `{SUM_FIELD}` field"))?; - let other_is_overflow = fields - .field(IS_OVERFLOW_FIELD) - .and_then(|is_overflow| is_overflow.as_bool().value()) - .ok_or_else(|| vortex_err!("Sum partial has an invalid `{IS_OVERFLOW_FIELD}` field"))?; - let other_is_empty = fields - .field(IS_EMPTY_FIELD) - .and_then(|is_empty| is_empty.as_bool().value()) - .ok_or_else(|| vortex_err!("Sum partial has an invalid `{IS_EMPTY_FIELD}` field"))?; - - partial.is_empty &= other_is_empty; - if partial.is_overflow || other_is_overflow { + if other_is_overflow { partial.is_overflow = true; + partial.is_empty = false; return Ok(()); } if other_is_empty { return Ok(()); } - let saturated = match &mut partial.sum { - 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) if r.fits_in_precision(*dtype) => { - *value = r; - false - } - Some(_) | None => true, - } - } - }; - if saturated { - partial.is_overflow = true; - } else { - partial.is_empty = false; - } + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; + partial.is_empty = false; Ok(()) } @@ -328,7 +274,7 @@ impl AggregateFnVTable for Sum { Ok(Scalar::struct_( sum_partial_dtype(partial.return_dtype.clone()), vec![ - sum_state_scalar(partial), + sum_state_scalar(partial, Nullability::NonNullable), Scalar::bool(partial.is_overflow, Nullability::NonNullable), Scalar::bool(partial.is_empty, Nullability::NonNullable), ], @@ -481,6 +427,60 @@ pub enum SumState { }, } +fn decode_sum_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { + vortex_ensure!(!scalar.is_null(), "Sum partial must not be null"); + + let Some(fields) = scalar.as_struct_opt() else { + vortex_bail!("Sum partial must be a struct, got {}", scalar.dtype()); + }; + let sum = fields + .field(SUM_FIELD) + .ok_or_else(|| vortex_err!("Sum partial is missing the `{SUM_FIELD}` field"))?; + let is_overflow = + bool::try_from(&fields.field(IS_OVERFLOW_FIELD).ok_or_else(|| { + vortex_err!("Sum partial is missing the `{IS_OVERFLOW_FIELD}` field") + })?)?; + let is_empty = bool::try_from( + &fields + .field(IS_EMPTY_FIELD) + .ok_or_else(|| vortex_err!("Sum partial is missing the `{IS_EMPTY_FIELD}` field"))?, + )?; + + Ok((sum, is_overflow, is_empty)) +} + +fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + sum.dtype().nullability() == Nullability::NonNullable + && sum.dtype().eq_ignore_nullability(return_dtype), + "Sum partial value has dtype {}, expected {}", + sum.dtype(), + return_dtype.as_nonnullable(), + ); + Ok(()) +} + +fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { + Ok(match state { + SumState::Unsigned(acc) => checked_add_u64(acc, u64::try_from(other)?), + SumState::Signed(acc) => checked_add_i64(acc, i64::try_from(other)?), + SumState::Float(acc) => { + *acc += f64::try_from(other)?; + false + } + SumState::Decimal { value, dtype } => { + let other = DecimalValue::try_from(other)?; + match value.checked_add(&other) { + Some(result) if result.fits_in_precision(*dtype) => { + *value = result; + false + } + Some(_) | None => true, + } + } + }) +} + fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { @@ -497,35 +497,30 @@ fn make_zero_state(return_dtype: &DType) -> SumState { } fn sum_partial_dtype(sum_dtype: DType) -> DType { - DType::Struct( - 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), - ], - ), - Nullability::Nullable, + 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) -> Scalar { +fn sum_state_scalar(partial: &SumPartial, nullability: Nullability) -> Scalar { match &partial.sum { - SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::NonNullable), - SumState::Signed(v) => Scalar::primitive(*v, Nullability::NonNullable), - SumState::Float(v) => Scalar::primitive(*v, Nullability::NonNullable), - 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::NonNullable) - } + 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), } } @@ -534,22 +529,7 @@ fn sum_value_scalar(partial: &SumPartial) -> Scalar { return Scalar::null(partial.return_dtype.as_nullable()); } - nullable_sum_state_scalar(partial) -} - -fn nullable_sum_state_scalar(partial: &SumPartial) -> Scalar { - match &partial.sum { - SumState::Unsigned(v) => Scalar::primitive(*v, Nullability::Nullable), - SumState::Signed(v) => Scalar::primitive(*v, Nullability::Nullable), - SumState::Float(v) => Scalar::primitive(*v, Nullability::Nullable), - 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) - } - } + sum_state_scalar(partial, Nullability::Nullable) } /// Convert scalar legacy Sum partials read from storage (a single nullable primitive) to the struct partial shape. diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs index 09abe564c4c..c168b39e757 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -113,6 +113,29 @@ fn legacy_options_only_describe_the_stored_partial() -> VortexResult<()> { // 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