From 08e60a40ea956a6899927dbc5308148868ed7f09 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 8 Sep 2026 15:46:26 +0100 Subject: [PATCH 1/5] Add sum oracle to the fuzzer Signed-off-by: Robert Kruszewski --- Cargo.lock | 2 + fuzz/Cargo.toml | 4 + fuzz/src/array/mod.rs | 12 +- fuzz/src/array/sum.rs | 15 -- fuzz/src/array/sum/mod.rs | 144 ++++++++++++ fuzz/src/array/sum/tests.rs | 424 ++++++++++++++++++++++++++++++++++++ 6 files changed, 579 insertions(+), 22 deletions(-) delete mode 100644 fuzz/src/array/sum.rs create mode 100644 fuzz/src/array/sum/mod.rs create mode 100644 fuzz/src/array/sum/tests.rs diff --git a/Cargo.lock b/Cargo.lock index f7d3103c9d5..bcd26545527 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11117,6 +11117,8 @@ dependencies = [ "arbitrary", "itertools 0.14.0", "libfuzzer-sys", + "num-traits", + "rstest", "strum 0.28.0", "tokio", "tracing", diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 13f307e8b14..0f31d29c29c 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -28,6 +28,7 @@ cuda = ["vortex-cuda", "tokio"] # Always needed - arbitrary is used for input generation arbitrary = { workspace = true } itertools = { workspace = true } +num-traits = { workspace = true } strum = { workspace = true, features = ["derive"] } # Vortex core - no default features for WASM compatibility (files feature pulls in tokio) @@ -55,6 +56,9 @@ vortex-cuda = { workspace = true, optional = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } +[dev-dependencies] +rstest = { workspace = true } + [lints] workspace = true diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index e513c5daf81..f2528042d72 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -341,13 +341,11 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { return Err(EmptyChoose); } - // Sum - returns a scalar, does NOT update current_array (terminal operation) - let current_array_canonical = current_array - .clone() - .execute::(&mut ctx) - .vortex_expect("execute canonical should succeed in fuzz test"); - let sum_result = sum_canonical_array(current_array_canonical, &mut ctx) - .vortex_expect("sum_canonical_array should succeed in fuzz test"); + let Some(sum_result) = sum_canonical_array(¤t_array, &mut ctx) + .vortex_expect("sum_canonical_array should succeed in fuzz test") + else { + return Err(EmptyChoose); + }; (Action::Sum, ExpectedValue::Scalar(sum_result)) } ActionType::MinMax => { diff --git a/fuzz/src/array/sum.rs b/fuzz/src/array/sum.rs deleted file mode 100644 index 9cb2eb9bb90..00000000000 --- a/fuzz/src/array/sum.rs +++ /dev/null @@ -1,15 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray as _; -use vortex_array::aggregate_fn::fns::sum::sum; -use vortex_array::scalar::Scalar; -use vortex_error::VortexResult; - -/// Compute sum on the canonical form of the array to get a consistent baseline. -pub fn sum_canonical_array(canonical: Canonical, ctx: &mut ExecutionCtx) -> VortexResult { - // TODO(joe): replace with baseline not using canonical - sum(&canonical.into_array(), ctx) -} diff --git a/fuzz/src/array/sum/mod.rs b/fuzz/src/array/sum/mod.rs new file mode 100644 index 00000000000..07e617743d3 --- /dev/null +++ b/fuzz/src/array/sum/mod.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::CheckedAdd; +use num_traits::Zero; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::dtype::BigCast; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::MAX_PRECISION; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::dtype::i256; +use vortex_array::match_each_decimal_value_type; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +#[cfg(test)] +mod tests; + +/// Reference sum of chunked or canonical arrays using native arithmetic. +/// Checks overflow after each group, preserving chunk boundaries. +/// Returns `None` only for floats, whose rounding depends on addition order. +pub fn sum_canonical_array( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + Ok(Some(match array.dtype() { + DType::Bool(_) => Scalar::from(accumulate(array, Some(0u64), &|_| true, ctx)?), + DType::Primitive(ptype, _) if ptype.is_unsigned_int() => { + Scalar::from(accumulate(array, Some(0u64), &|_| true, ctx)?) + } + DType::Primitive(ptype, _) if ptype.is_signed_int() => { + Scalar::from(accumulate(array, Some(0i64), &|_| true, ctx)?) + } + DType::Primitive(..) => return Ok(None), + DType::Decimal(input_dtype, _) => { + let output_dtype = DecimalDType::new( + (input_dtype.precision() + 10).min(MAX_PRECISION), + input_dtype.scale(), + ); + let limit = i256::from_i128(10) + .checked_pow(output_dtype.precision().into()) + .vortex_expect("10^76 fits in i256"); + let values_type = DecimalType::smallest_decimal_value_type(&output_dtype); + match_each_decimal_value_type!(values_type, |I| { + let limit = ::from(limit) + .vortex_expect("precision limit fits native accumulator"); + match accumulate( + array, + Some(I::zero()), + &|&value| -limit < value && value < limit, + ctx, + )? { + Some(value) => { + Scalar::decimal(DecimalValue::from(value), output_dtype, Nullable) + } + None => Scalar::null(DType::Decimal(output_dtype, Nullable)), + } + }) + } + _ => vortex_bail!("Unsupported sum dtype: {}", array.dtype()), + })) +} + +fn accumulate( + array: &ArrayRef, + initial: Option, + fits: &impl Fn(&T) -> bool, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + T: BigCast + Copy + Zero + CheckedAdd, +{ + let Some(initial) = initial else { + return Ok(None); + }; + if array.is_empty() { + return Ok(Some(initial)); + } + + if let Some(chunked) = array.as_opt::() { + // A nested chunked array has its own partial; canonical chunks share its running total. + let mut partial = Some(T::zero()); + for chunk in chunked.non_empty_chunks() { + partial = accumulate(chunk, partial, fits, ctx)?; + } + Ok(partial + .and_then(|partial| initial.checked_add(&partial)) + .filter(fits)) + } else { + // Canonical batches share the parent's running total. Native overflow is checked on + // each addition, while decimal precision is checked only at the end of this group. + Ok(native_values::(array, ctx)? + .into_iter() + .try_fold(initial, |sum, value| sum.checked_add(&value)) + .filter(fits)) + } +} + +fn native_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let canonical = array.clone().execute::(ctx)?; + let valid = canonical + .clone() + .into_array() + .validity()? + .execute_mask(canonical.len(), ctx)? + .to_bit_buffer(); + Ok(match canonical { + Canonical::Bool(array) => vec![ + T::from((array.to_bit_buffer() & valid).true_count() as u64) + .vortex_expect("boolean count fits accumulator"), + ], + Canonical::Primitive(array) => match_each_integer_ptype!(array.ptype(), |P| { + array + .as_slice::

() + .iter() + .zip(valid.iter()) + .filter_map(|(&value, valid)| valid.then_some(value)) + .map(|value| T::from(value).vortex_expect("integer value fits accumulator")) + .collect() + }), + Canonical::Decimal(array) => match_each_decimal_value_type!(array.values_type(), |D| { + array + .buffer::() + .iter() + .zip(valid.iter()) + .filter_map(|(&value, valid)| valid.then_some(value)) + .map(|value| T::from(value).vortex_expect("decimal value fits accumulator")) + .collect() + }), + _ => vortex_bail!("Unsupported sum dtype: {}", array.dtype()), + }) +} diff --git a/fuzz/src/array/sum/tests.rs b/fuzz/src/array/sum/tests.rs new file mode 100644 index 00000000000..8de7601e095 --- /dev/null +++ b/fuzz/src/array/sum/tests.rs @@ -0,0 +1,424 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::i256; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::match_each_decimal_value_type; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::sum_canonical_array; +use crate::Action; +use crate::ExpectedValue; +use crate::FuzzArrayAction; +use crate::SESSION; +use crate::run_fuzz_action; + +fn reference(array: impl IntoArray, ctx: &mut ExecutionCtx) -> VortexResult> { + sum_canonical_array(&array.into_array(), ctx) +} + +#[test] +fn test_sum_ignores_cached_statistics() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([1i64, 2, 3]); + array + .as_ref() + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(999i64))); + assert_eq!( + reference( + Canonical::Primitive(array), + &mut SESSION.create_execution_ctx() + )?, + Some(Scalar::from(6i64)) + ); + Ok(()) +} + +#[rstest] +#[case::i8(buffer![i8::MAX, i8::MAX, -1].into_array(), 253i64.into())] +#[case::i16(buffer![i16::MAX, i16::MAX, -1].into_array(), 65_533i64.into())] +#[case::i32(buffer![i32::MAX, i32::MAX, -1].into_array(), 4_294_967_293i64.into())] +#[case::i64(buffer![1i64, 2, 3].into_array(), 6i64.into())] +#[case::u8(buffer![u8::MAX, u8::MAX].into_array(), 510u64.into())] +#[case::u16(buffer![u16::MAX, u16::MAX].into_array(), 131_070u64.into())] +#[case::u32(buffer![u32::MAX, u32::MAX].into_array(), 8_589_934_590u64.into())] +#[case::u64(buffer![u64::MAX, 0].into_array(), u64::MAX.into())] +#[case::unsigned_overflow(buffer![u64::MAX, 1].into_array(), Scalar::from(None::))] +#[case::unsigned_nulls( + PrimitiveArray::new(buffer![u64::MAX, u64::MAX], Validity::from_iter([true, false])).into_array(), + u64::MAX.into() +)] +#[case::unsigned_all_null(PrimitiveArray::from_option_iter([None::, None]).into_array(), 0u64.into())] +#[case::unsigned_empty(Buffer::::empty().into_array(), 0u64.into())] +#[case::signed_empty(Buffer::::empty().into_array(), 0i64.into())] +#[case::bool(BoolArray::from_iter([true, false, true]).into_array(), 2u64.into())] +#[case::bool_nulls(BoolArray::from_iter([Some(true), None, Some(false)]).into_array(), 1u64.into())] +#[case::bool_all_null(BoolArray::from_iter([None::, None]).into_array(), 0u64.into())] +#[case::bool_empty(BoolArray::from_iter([] as [bool; 0]).into_array(), 0u64.into())] +fn test_sum_integer_and_bool_values( + #[case] array: ArrayRef, + #[case] expected: Scalar, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let result = reference(array.execute::(&mut ctx)?, &mut ctx)? + .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; + assert_eq!(result.dtype(), &expected.dtype().as_nullable()); + assert_eq!(result, expected); + Ok(()) +} + +#[rstest] +#[case(DecimalType::I8)] +#[case(DecimalType::I16)] +#[case(DecimalType::I32)] +#[case(DecimalType::I64)] +#[case(DecimalType::I128)] +#[case(DecimalType::I256)] +fn test_sum_decimal_storage(#[case] values_type: DecimalType) -> VortexResult<()> { + let array = match_each_decimal_value_type!(values_type, |D| { + let value = DecimalValue::I8(99) + .cast::() + .ok_or_else(|| vortex_err!("99 fits in every decimal storage type"))?; + DecimalArray::new( + buffer![value, value, -value], + DecimalDType::new(2, 0), + Validity::NonNullable, + ) + }); + let result = reference( + Canonical::Decimal(array), + &mut SESSION.create_execution_ctx(), + )? + .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(12, 0), Nullability::Nullable) + ); + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I64(99)) + ); + Ok(()) +} + +#[test] +fn test_sum_decimal_widens_beyond_i128() -> VortexResult<()> { + let value = 10i128.pow(38) - 1; + let array = DecimalArray::new( + buffer![value, value], + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + assert_eq!( + reference( + Canonical::Decimal(array), + &mut SESSION.create_execution_ctx() + )?, + Some(Scalar::decimal( + DecimalValue::I256(i256::from_i128(value) * i256::from_i128(2)), + DecimalDType::new(48, 2), + Nullability::Nullable, + )) + ); + Ok(()) +} + +#[test] +fn test_sum_decimal_cancellation_within_one_group() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(76, -76); + let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + + for value in [six_e75, -six_e75] { + let first = DecimalArray::new(buffer![value, value], decimal_dtype, Validity::NonNullable) + .into_array(); + let last = + DecimalArray::new(buffer![-value], decimal_dtype, Validity::NonNullable).into_array(); + let dtype = first.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![first, last], dtype)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + // The first chunk exceeds precision 76, although cancellation makes the final sum fit. + assert!(sum(&chunked, &mut ctx)?.is_null()); + assert_eq!( + reference(chunked.clone(), &mut ctx)?, + Some(Scalar::null(chunked.dtype().as_nullable())) + ); + let canonical = chunked.execute::(&mut ctx)?; + assert_eq!( + reference(canonical.clone(), &mut ctx)?, + Some(Scalar::decimal( + DecimalValue::I256(value), + decimal_dtype, + Nullability::Nullable + )) + ); + assert_eq!( + sum(&canonical.clone().into_array(), &mut ctx)? + .as_decimal() + .decimal_value(), + Some(DecimalValue::I256(value)) + ); + assert_eq!( + reference(canonical, &mut ctx)?, + Some(Scalar::decimal( + DecimalValue::I256(value), + decimal_dtype, + Nullability::Nullable + )) + ); + } + Ok(()) +} + +#[test] +fn test_sum_decimal_precision_boundary_and_nulls() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let max = i256::from_i128(10).wrapping_pow(76) - i256::ONE; + let mut ctx = SESSION.create_execution_ctx(); + for validity in [ + Validity::from_iter([true, false, true]), + Validity::AllInvalid, + ] { + let array = DecimalArray::new(buffer![max, max, -max], dtype, validity); + assert_eq!( + reference(Canonical::Decimal(array), &mut ctx)?, + Some(Scalar::decimal( + DecimalValue::I256(i256::ZERO), + dtype, + Nullability::Nullable + )) + ); + } + Ok(()) +} + +#[test] +fn test_sum_decimal_keeps_definite_overflow() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + let mut ctx = SESSION.create_execution_ctx(); + for value in [six_e75, -six_e75] { + let array = DecimalArray::new(buffer![value, value], dtype, Validity::NonNullable); + assert_eq!( + reference(Canonical::Decimal(array), &mut ctx)?, + Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) + ); + } + Ok(()) +} + +#[test] +fn test_sum_decimal_native_overflow() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + let mut ctx = SESSION.create_execution_ctx(); + for value in [six_e75, -six_e75] { + let values = iter::repeat_n(value, 10) + .chain(iter::repeat_n(-value, 10)) + .collect::>(); + let array = DecimalArray::new(values, dtype, Validity::NonNullable); + assert_eq!( + reference(Canonical::Decimal(array), &mut ctx)?, + Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) + ); + } + Ok(()) +} + +#[test] +fn test_sum_decimal_uses_widened_precision() -> VortexResult<()> { + let array = DecimalArray::new( + buffer![99i8, 99, -99], + DecimalDType::new(2, 0), + Validity::NonNullable, + ); + assert_eq!( + reference( + Canonical::Decimal(array), + &mut SESSION.create_execution_ctx() + )?, + Some(Scalar::decimal( + DecimalValue::I64(99), + DecimalDType::new(12, 0), + Nullability::Nullable + )) + ); + Ok(()) +} + +#[test] +fn test_sum_signed_overflow_and_cancellation() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for (values, expected) in [ + (vec![Some(i64::MAX), Some(1), Some(-1)], Some(None)), + (vec![Some(i64::MIN), Some(-1), Some(1)], Some(None)), + ( + vec![Some(i64::MAX), Some(-1), Some(1)], + Some(Some(i64::MAX)), + ), + (vec![Some(i64::MAX), Some(i64::MIN)], Some(Some(-1i64))), + (vec![Some(i64::MAX), Some(1)], Some(None)), + (vec![Some(i64::MIN), Some(-1)], Some(None)), + (vec![Some(i64::MAX), None, Some(i64::MIN)], Some(Some(-1))), + (vec![None, None], Some(Some(0))), + (vec![], Some(Some(0))), + ] { + let array = PrimitiveArray::from_option_iter(values); + assert_eq!( + reference(Canonical::Primitive(array), &mut ctx)?, + expected.map(Scalar::from) + ); + } + Ok(()) +} + +fn decimal_chunks(groups: &[&[i8]]) -> VortexResult { + let unit = i256::from_i128(10).wrapping_pow(75); + let dtype = DecimalDType::new(76, -76); + ChunkedArray::try_new( + groups.iter().map(|values| { + DecimalArray::new( + values + .iter() + .map(|&value| unit * i256::from_i128(i128::from(value))) + .collect::>(), + dtype, + Validity::NonNullable, + ) + .into_array() + }), + DType::Decimal(dtype, Nullability::NonNullable), + ) + .map(IntoArray::into_array) +} + +#[rstest] +#[case::overflow_in_first_group(decimal_chunks(&[&[6, 6], &[-6]]), None)] +#[case::overflow_after_second_group(decimal_chunks(&[&[6], &[6], &[-6]]), None)] +#[case::negative_overflow(decimal_chunks(&[&[-6, -6], &[6]]), None)] +#[case::shared_running_total(decimal_chunks(&[&[-6], &[6, 6]]), Some(6))] +#[case::cancellation_within_group(decimal_chunks(&[&[6, 6, -6]]), Some(6))] +#[case::empty_groups(decimal_chunks(&[&[], &[6], &[], &[-6], &[]]), Some(0))] +fn test_sum_checks_decimal_group_boundaries( + #[case] array: VortexResult, + #[case] expected: Option, +) -> VortexResult<()> { + let array = array?; + let mut ctx = SESSION.create_execution_ctx(); + let expected = match expected { + Some(value) => Scalar::decimal( + DecimalValue::I256( + i256::from_i128(10).wrapping_pow(75) * i256::from_i128(i128::from(value)), + ), + DecimalDType::new(76, -76), + Nullability::Nullable, + ), + None => Scalar::null(array.dtype().as_nullable()), + }; + assert_eq!(reference(array.clone(), &mut ctx)?, Some(expected.clone())); + assert_eq!(sum(&array, &mut ctx)?, expected); + Ok(()) +} + +#[test] +fn test_sum_nested_group_has_independent_partial() -> VortexResult<()> { + let inner = decimal_chunks(&[&[6, 6], &[-6]])?; + let array = ChunkedArray::try_new( + vec![decimal_chunks(&[&[-6]])?, inner.clone()], + inner.dtype().clone(), + )? + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + reference(array.clone(), &mut ctx)?, + Some(Scalar::null(array.dtype().as_nullable())) + ); + assert!(sum(&array, &mut ctx)?.is_null()); + Ok(()) +} + +#[test] +fn test_sum_cached_statistics_do_not_change_native_overflow() -> VortexResult<()> { + let array = buffer![i64::MAX, 1, -1].into_array(); + array + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(i64::MAX))); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(reference(array, &mut ctx)?, Some(Scalar::from(None::))); + Ok(()) +} + +#[test] +fn test_sum_cached_statistics_do_not_change_groups() -> VortexResult<()> { + let array = decimal_chunks(&[&[6, 6], &[-6]])?; + array.statistics().set( + Stat::Sum, + Precision::Exact(ScalarValue::from(DecimalValue::I256( + i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6), + ))), + ); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + reference(array.clone(), &mut ctx)?, + Some(Scalar::null(array.dtype().as_nullable())) + ); + Ok(()) +} + +#[rstest] +#[case::chunked(false, None)] +#[case::canonical(true, Some(6))] +fn test_sum_action_stores_reference_scalar( + #[case] canonical: bool, + #[case] expected_sum: Option, +) -> VortexResult<()> { + let mut array = decimal_chunks(&[&[6, 6], &[-6]])?; + let mut ctx = SESSION.create_execution_ctx(); + if canonical { + array = array.execute::(&mut ctx)?.into_array(); + } + let expected_sum = match expected_sum { + Some(value) => Scalar::decimal( + DecimalValue::I256( + i256::from_i128(10).wrapping_pow(75) * i256::from_i128(i128::from(value)), + ), + DecimalDType::new(76, -76), + Nullability::Nullable, + ), + None => Scalar::null(array.dtype().as_nullable()), + }; + let reference_sum = sum_canonical_array(&array, &mut ctx)? + .ok_or_else(|| vortex_err!("expected a decimal sum"))?; + assert_eq!(reference_sum, expected_sum); + let fuzz_action = FuzzArrayAction { + array, + actions: vec![(Action::Sum, ExpectedValue::Scalar(reference_sum))], + }; + assert!(run_fuzz_action(fuzz_action).map_err(|error| vortex_err!("{error}"))?); + Ok(()) +} From 54163a3994f7395fee848679fcea681cb1d39553 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 9 Sep 2026 14:58:46 +0100 Subject: [PATCH 2/5] fixes Signed-off-by: Robert Kruszewski --- fuzz/src/array/mod.rs | 4 +- fuzz/src/array/sum/mod.rs | 73 ++++-- fuzz/src/array/sum/tests.rs | 165 ++++++------ vortex-array/benches/aggregate_sum.rs | 41 +++ .../aggregate_fn/fns/sum_v2/decimal/mod.rs | 241 ++++++++++++++++++ .../aggregate_fn/fns/sum_v2/decimal/tests.rs | 183 +++++++++++++ .../src/aggregate_fn/fns/sum_v2/mod.rs | 56 ++-- 7 files changed, 634 insertions(+), 129 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs create mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index f2528042d72..bc889fa314c 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -48,7 +48,7 @@ use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; use vortex_array::aggregate_fn::fns::min_max::min_max; -use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::arbitrary::ArbitraryArray; @@ -670,7 +670,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { current_array = cast_result; } Action::Sum => { - let sum_result = sum(¤t_array, &mut ctx) + let sum_result = sum_v2(¤t_array, &mut ctx) .vortex_expect("sum operation should succeed in fuzz test"); assert_scalar_eq(&expected.scalar(), &sum_result, i)?; } diff --git a/fuzz/src/array/sum/mod.rs b/fuzz/src/array/sum/mod.rs index 07e617743d3..92de619e223 100644 --- a/fuzz/src/array/sum/mod.rs +++ b/fuzz/src/array/sum/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::CheckedAdd; +use num_traits::WrappingAdd; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -29,20 +29,45 @@ use vortex_error::vortex_bail; mod tests; /// Reference sum of chunked or canonical arrays using native arithmetic. -/// Checks overflow after each group, preserving chunk boundaries. +/// Decimal precision is checked only after wrapping accumulation of all chunks. +/// Empty and all-null arrays return null, matching SQL SUM. /// Returns `None` only for floats, whose rounding depends on addition order. pub fn sum_canonical_array( array: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + let mut any_valid = false; Ok(Some(match array.dtype() { - DType::Bool(_) => Scalar::from(accumulate(array, Some(0u64), &|_| true, ctx)?), - DType::Primitive(ptype, _) if ptype.is_unsigned_int() => { - Scalar::from(accumulate(array, Some(0u64), &|_| true, ctx)?) - } - DType::Primitive(ptype, _) if ptype.is_signed_int() => { - Scalar::from(accumulate(array, Some(0i64), &|_| true, ctx)?) - } + DType::Bool(_) => Scalar::from( + accumulate( + array, + Some(0u64), + &|sum, value| sum.checked_add(value), + &mut any_valid, + ctx, + )? + .filter(|_| any_valid), + ), + DType::Primitive(ptype, _) if ptype.is_unsigned_int() => Scalar::from( + accumulate( + array, + Some(0u64), + &|sum, value| sum.checked_add(value), + &mut any_valid, + ctx, + )? + .filter(|_| any_valid), + ), + DType::Primitive(ptype, _) if ptype.is_signed_int() => Scalar::from( + accumulate( + array, + Some(0i64), + &|sum, value| sum.checked_add(value), + &mut any_valid, + ctx, + )? + .filter(|_| any_valid), + ), DType::Primitive(..) => return Ok(None), DType::Decimal(input_dtype, _) => { let output_dtype = DecimalDType::new( @@ -59,9 +84,12 @@ pub fn sum_canonical_array( match accumulate( array, Some(I::zero()), - &|&value| -limit < value && value < limit, + &|sum, value| Some(WrappingAdd::wrapping_add(&sum, &value)), + &mut any_valid, ctx, - )? { + )? + .filter(|&value| any_valid && -limit < value && value < limit) + { Some(value) => { Scalar::decimal(DecimalValue::from(value), output_dtype, Nullable) } @@ -76,11 +104,12 @@ pub fn sum_canonical_array( fn accumulate( array: &ArrayRef, initial: Option, - fits: &impl Fn(&T) -> bool, + add: &impl Fn(T, T) -> Option, + any_valid: &mut bool, ctx: &mut ExecutionCtx, ) -> VortexResult> where - T: BigCast + Copy + Zero + CheckedAdd, + T: BigCast + Copy + Zero, { let Some(initial) = initial else { return Ok(None); @@ -93,18 +122,13 @@ where // A nested chunked array has its own partial; canonical chunks share its running total. let mut partial = Some(T::zero()); for chunk in chunked.non_empty_chunks() { - partial = accumulate(chunk, partial, fits, ctx)?; + partial = accumulate(chunk, partial, add, any_valid, ctx)?; } - Ok(partial - .and_then(|partial| initial.checked_add(&partial)) - .filter(fits)) + Ok(partial.and_then(|partial| add(initial, partial))) } else { - // Canonical batches share the parent's running total. Native overflow is checked on - // each addition, while decimal precision is checked only at the end of this group. - Ok(native_values::(array, ctx)? - .into_iter() - .try_fold(initial, |sum, value| sum.checked_add(&value)) - .filter(fits)) + let values = native_values::(array, ctx)?; + *any_valid |= !values.is_empty(); + Ok(values.into_iter().try_fold(initial, add)) } } @@ -116,6 +140,9 @@ fn native_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vortex .validity()? .execute_mask(canonical.len(), ctx)? .to_bit_buffer(); + if valid.true_count() == 0 { + return Ok(Vec::new()); + } Ok(match canonical { Canonical::Bool(array) => vec![ T::from((array.to_bit_buffer() & valid).true_count() as u64) diff --git a/fuzz/src/array/sum/tests.rs b/fuzz/src/array/sum/tests.rs index 8de7601e095..b6a5e6ec9e6 100644 --- a/fuzz/src/array/sum/tests.rs +++ b/fuzz/src/array/sum/tests.rs @@ -6,10 +6,9 @@ use std::iter; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; -use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::DecimalArray; @@ -38,10 +37,6 @@ use crate::FuzzArrayAction; use crate::SESSION; use crate::run_fuzz_action; -fn reference(array: impl IntoArray, ctx: &mut ExecutionCtx) -> VortexResult> { - sum_canonical_array(&array.into_array(), ctx) -} - #[test] fn test_sum_ignores_cached_statistics() -> VortexResult<()> { let array = PrimitiveArray::from_iter([1i64, 2, 3]); @@ -50,10 +45,7 @@ fn test_sum_ignores_cached_statistics() -> VortexResult<()> { .statistics() .set(Stat::Sum, Precision::Exact(ScalarValue::from(999i64))); assert_eq!( - reference( - Canonical::Primitive(array), - &mut SESSION.create_execution_ctx() - )?, + sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, Some(Scalar::from(6i64)) ); Ok(()) @@ -73,20 +65,23 @@ fn test_sum_ignores_cached_statistics() -> VortexResult<()> { PrimitiveArray::new(buffer![u64::MAX, u64::MAX], Validity::from_iter([true, false])).into_array(), u64::MAX.into() )] -#[case::unsigned_all_null(PrimitiveArray::from_option_iter([None::, None]).into_array(), 0u64.into())] -#[case::unsigned_empty(Buffer::::empty().into_array(), 0u64.into())] -#[case::signed_empty(Buffer::::empty().into_array(), 0i64.into())] +#[case::unsigned_all_null(PrimitiveArray::from_option_iter([None::, None]).into_array(), Scalar::from(None::))] +#[case::unsigned_empty(Buffer::::empty().into_array(), Scalar::from(None::))] +#[case::signed_empty(Buffer::::empty().into_array(), Scalar::from(None::))] #[case::bool(BoolArray::from_iter([true, false, true]).into_array(), 2u64.into())] #[case::bool_nulls(BoolArray::from_iter([Some(true), None, Some(false)]).into_array(), 1u64.into())] -#[case::bool_all_null(BoolArray::from_iter([None::, None]).into_array(), 0u64.into())] -#[case::bool_empty(BoolArray::from_iter([] as [bool; 0]).into_array(), 0u64.into())] +#[case::bool_all_null(BoolArray::from_iter([None::, None]).into_array(), Scalar::from(None::))] +#[case::bool_empty(BoolArray::from_iter([] as [bool; 0]).into_array(), Scalar::from(None::))] fn test_sum_integer_and_bool_values( #[case] array: ArrayRef, #[case] expected: Scalar, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let result = reference(array.execute::(&mut ctx)?, &mut ctx)? - .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; + let result = sum_canonical_array( + &array.execute::(&mut ctx)?.into_array(), + &mut ctx, + )? + .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; assert_eq!(result.dtype(), &expected.dtype().as_nullable()); assert_eq!(result, expected); Ok(()) @@ -110,11 +105,8 @@ fn test_sum_decimal_storage(#[case] values_type: DecimalType) -> VortexResult<() Validity::NonNullable, ) }); - let result = reference( - Canonical::Decimal(array), - &mut SESSION.create_execution_ctx(), - )? - .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; + let result = sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())? + .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; assert_eq!( result.dtype(), &DType::Decimal(DecimalDType::new(12, 0), Nullability::Nullable) @@ -135,10 +127,7 @@ fn test_sum_decimal_widens_beyond_i128() -> VortexResult<()> { Validity::NonNullable, ); assert_eq!( - reference( - Canonical::Decimal(array), - &mut SESSION.create_execution_ctx() - )?, + sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, Some(Scalar::decimal( DecimalValue::I256(i256::from_i128(value) * i256::from_i128(2)), DecimalDType::new(48, 2), @@ -149,7 +138,7 @@ fn test_sum_decimal_widens_beyond_i128() -> VortexResult<()> { } #[test] -fn test_sum_decimal_cancellation_within_one_group() -> VortexResult<()> { +fn test_sum_decimal_cancellation_across_groups() -> VortexResult<()> { let decimal_dtype = DecimalDType::new(76, -76); let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); @@ -162,35 +151,20 @@ fn test_sum_decimal_cancellation_within_one_group() -> VortexResult<()> { let chunked = ChunkedArray::try_new(vec![first, last], dtype)?.into_array(); let mut ctx = SESSION.create_execution_ctx(); - // The first chunk exceeds precision 76, although cancellation makes the final sum fit. - assert!(sum(&chunked, &mut ctx)?.is_null()); - assert_eq!( - reference(chunked.clone(), &mut ctx)?, - Some(Scalar::null(chunked.dtype().as_nullable())) - ); - let canonical = chunked.execute::(&mut ctx)?; - assert_eq!( - reference(canonical.clone(), &mut ctx)?, - Some(Scalar::decimal( - DecimalValue::I256(value), - decimal_dtype, - Nullability::Nullable - )) - ); - assert_eq!( - sum(&canonical.clone().into_array(), &mut ctx)? - .as_decimal() - .decimal_value(), - Some(DecimalValue::I256(value)) - ); - assert_eq!( - reference(canonical, &mut ctx)?, - Some(Scalar::decimal( - DecimalValue::I256(value), - decimal_dtype, - Nullability::Nullable - )) + // The first chunk exceeds precision 76, but the final sum fits after cancellation. + let expected = Scalar::decimal( + DecimalValue::I256(value), + decimal_dtype, + Nullability::Nullable, ); + let canonical = chunked.clone().execute::(&mut ctx)?.into_array(); + for array in [chunked, canonical] { + assert_eq!(sum_v2(&array, &mut ctx)?, expected); + assert_eq!( + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) + ); + } } Ok(()) } @@ -205,13 +179,14 @@ fn test_sum_decimal_precision_boundary_and_nulls() -> VortexResult<()> { Validity::AllInvalid, ] { let array = DecimalArray::new(buffer![max, max, -max], dtype, validity); + let expected = if matches!(array.as_ref().validity()?, Validity::AllInvalid) { + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) + } else { + Scalar::decimal(DecimalValue::I256(i256::ZERO), dtype, Nullability::Nullable) + }; assert_eq!( - reference(Canonical::Decimal(array), &mut ctx)?, - Some(Scalar::decimal( - DecimalValue::I256(i256::ZERO), - dtype, - Nullability::Nullable - )) + sum_canonical_array(&array.into_array(), &mut ctx)?, + Some(expected) ); } Ok(()) @@ -225,7 +200,7 @@ fn test_sum_decimal_keeps_definite_overflow() -> VortexResult<()> { for value in [six_e75, -six_e75] { let array = DecimalArray::new(buffer![value, value], dtype, Validity::NonNullable); assert_eq!( - reference(Canonical::Decimal(array), &mut ctx)?, + sum_canonical_array(&array.into_array(), &mut ctx)?, Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) ); } @@ -233,7 +208,7 @@ fn test_sum_decimal_keeps_definite_overflow() -> VortexResult<()> { } #[test] -fn test_sum_decimal_native_overflow() -> VortexResult<()> { +fn test_sum_decimal_native_overflow_cancels() -> VortexResult<()> { let dtype = DecimalDType::new(76, 0); let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); let mut ctx = SESSION.create_execution_ctx(); @@ -242,10 +217,14 @@ fn test_sum_decimal_native_overflow() -> VortexResult<()> { .chain(iter::repeat_n(-value, 10)) .collect::>(); let array = DecimalArray::new(values, dtype, Validity::NonNullable); + let expected = + Scalar::decimal(DecimalValue::I256(i256::ZERO), dtype, Nullability::Nullable); + let array = array.into_array(); assert_eq!( - reference(Canonical::Decimal(array), &mut ctx)?, - Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) ); + assert_eq!(sum_v2(&array, &mut ctx)?, expected); } Ok(()) } @@ -258,10 +237,7 @@ fn test_sum_decimal_uses_widened_precision() -> VortexResult<()> { Validity::NonNullable, ); assert_eq!( - reference( - Canonical::Decimal(array), - &mut SESSION.create_execution_ctx() - )?, + sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, Some(Scalar::decimal( DecimalValue::I64(99), DecimalDType::new(12, 0), @@ -285,12 +261,12 @@ fn test_sum_signed_overflow_and_cancellation() -> VortexResult<()> { (vec![Some(i64::MAX), Some(1)], Some(None)), (vec![Some(i64::MIN), Some(-1)], Some(None)), (vec![Some(i64::MAX), None, Some(i64::MIN)], Some(Some(-1))), - (vec![None, None], Some(Some(0))), - (vec![], Some(Some(0))), + (vec![None, None], Some(None)), + (vec![], Some(None)), ] { let array = PrimitiveArray::from_option_iter(values); assert_eq!( - reference(Canonical::Primitive(array), &mut ctx)?, + sum_canonical_array(&array.into_array(), &mut ctx)?, expected.map(Scalar::from) ); } @@ -318,13 +294,15 @@ fn decimal_chunks(groups: &[&[i8]]) -> VortexResult { } #[rstest] -#[case::overflow_in_first_group(decimal_chunks(&[&[6, 6], &[-6]]), None)] -#[case::overflow_after_second_group(decimal_chunks(&[&[6], &[6], &[-6]]), None)] -#[case::negative_overflow(decimal_chunks(&[&[-6, -6], &[6]]), None)] +#[case::cancellation_across_groups(decimal_chunks(&[&[6, 6], &[-6]]), Some(6))] +#[case::cancellation_across_three_groups(decimal_chunks(&[&[6], &[6], &[-6]]), Some(6))] +#[case::negative_cancellation(decimal_chunks(&[&[-6, -6], &[6]]), Some(-6))] #[case::shared_running_total(decimal_chunks(&[&[-6], &[6, 6]]), Some(6))] #[case::cancellation_within_group(decimal_chunks(&[&[6, 6, -6]]), Some(6))] #[case::empty_groups(decimal_chunks(&[&[], &[6], &[], &[-6], &[]]), Some(0))] -fn test_sum_checks_decimal_group_boundaries( +#[case::final_overflow(decimal_chunks(&[&[6], &[6]]), None)] +#[case::all_empty(decimal_chunks(&[&[], &[]]), None)] +fn test_sum_decimal_ignores_group_boundaries( #[case] array: VortexResult, #[case] expected: Option, ) -> VortexResult<()> { @@ -340,13 +318,16 @@ fn test_sum_checks_decimal_group_boundaries( ), None => Scalar::null(array.dtype().as_nullable()), }; - assert_eq!(reference(array.clone(), &mut ctx)?, Some(expected.clone())); - assert_eq!(sum(&array, &mut ctx)?, expected); + assert_eq!( + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) + ); + assert_eq!(sum_v2(&array, &mut ctx)?, expected); Ok(()) } #[test] -fn test_sum_nested_group_has_independent_partial() -> VortexResult<()> { +fn test_sum_nested_group_cancellation() -> VortexResult<()> { let inner = decimal_chunks(&[&[6, 6], &[-6]])?; let array = ChunkedArray::try_new( vec![decimal_chunks(&[&[-6]])?, inner.clone()], @@ -354,11 +335,16 @@ fn test_sum_nested_group_has_independent_partial() -> VortexResult<()> { )? .into_array(); let mut ctx = SESSION.create_execution_ctx(); + let expected = Scalar::decimal( + DecimalValue::I256(i256::ZERO), + DecimalDType::new(76, -76), + Nullability::Nullable, + ); assert_eq!( - reference(array.clone(), &mut ctx)?, - Some(Scalar::null(array.dtype().as_nullable())) + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) ); - assert!(sum(&array, &mut ctx)?.is_null()); + assert_eq!(sum_v2(&array, &mut ctx)?, expected); Ok(()) } @@ -369,7 +355,10 @@ fn test_sum_cached_statistics_do_not_change_native_overflow() -> VortexResult<() .statistics() .set(Stat::Sum, Precision::Exact(ScalarValue::from(i64::MAX))); let mut ctx = SESSION.create_execution_ctx(); - assert_eq!(reference(array, &mut ctx)?, Some(Scalar::from(None::))); + assert_eq!( + sum_canonical_array(&array, &mut ctx)?, + Some(Scalar::from(None::)) + ); Ok(()) } @@ -384,14 +373,18 @@ fn test_sum_cached_statistics_do_not_change_groups() -> VortexResult<()> { ); let mut ctx = SESSION.create_execution_ctx(); assert_eq!( - reference(array.clone(), &mut ctx)?, - Some(Scalar::null(array.dtype().as_nullable())) + sum_canonical_array(&array, &mut ctx)?, + Some(Scalar::decimal( + DecimalValue::I256(i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6)), + DecimalDType::new(76, -76), + Nullability::Nullable + )) ); Ok(()) } #[rstest] -#[case::chunked(false, None)] +#[case::chunked(false, Some(6))] #[case::canonical(true, Some(6))] fn test_sum_action_stores_reference_scalar( #[case] canonical: bool, diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index 31856e22f9f..d703ce1c935 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -9,8 +9,16 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; use vortex_array::expr::stats::Stat; +use vortex_array::match_each_decimal_value_type; +use vortex_array::scalar::DecimalValue; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_session::VortexSession; fn main() { @@ -23,6 +31,39 @@ const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(array_session); +#[divan::bench(args = [8, 28, 66])] +fn sum_v2_decimal(bencher: Bencher, precision: u8) { + bench_decimal_sum(bencher, precision, false); +} + +#[divan::bench(args = [8, 28, 66])] +fn sum_v2_decimal_nulls(bencher: Bencher, precision: u8) { + bench_decimal_sum(bencher, precision, true); +} + +fn bench_decimal_sum(bencher: Bencher, precision: u8, nullable: bool) { + let dtype = DecimalDType::new(precision, 2); + let values_type = DecimalType::smallest_decimal_value_type(&dtype); + let array = match_each_decimal_value_type!(values_type, |I| { + let values = (0..N) + .map(|i| { + DecimalValue::I64(i as i64 % 1000 - 500) + .cast::() + .vortex_expect("benchmark value fits decimal storage") + }) + .collect::>(); + let validity = if nullable { + Validity::from_iter((0..N).map(|i| i % 5 != 0)) + } else { + Validity::NonNullable + }; + DecimalArray::new(values, dtype, validity).into_array() + }); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| sum_v2(array, ctx)); +} + #[divan::bench] fn sum_i32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(1); diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs new file mode 100644 index 00000000000..4ce31b04bda --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use num_traits::AsPrimitive; +use num_traits::CheckedAdd; +use num_traits::CheckedMul; +use num_traits::NumOps; +use num_traits::WrappingAdd; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_mask::Mask; + +use super::SumState; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::FieldNames; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; + +const VALUE_FIELD: &str = "value"; +const CARRY_FIELD: &str = "carry"; + +fn carry_dtype() -> DecimalDType { + DecimalDType::new(38, 0) +} + +pub(super) fn decimal_partial_dtype(dtype: DType) -> DType { + let DType::Decimal(..) = dtype else { + return dtype; + }; + DType::Struct( + StructFields::new( + FieldNames::from_iter([VALUE_FIELD, CARRY_FIELD]), + vec![ + dtype.as_nonnullable(), + DType::Decimal(carry_dtype(), Nullability::NonNullable), + ], + ), + dtype.nullability(), + ) +} + +pub(super) fn decimal_partial_scalar( + value: DecimalValue, + dtype: DecimalDType, + nullability: Nullability, +) -> Scalar { + let value = value.cast::().vortex_expect("decimal fits i256"); + let limit = i256::from_i128(10).wrapping_pow(dtype.precision().into()); + // Decimal scalars must fit their precision, including in partial states. Preserve excess + // separately so later partials can cancel it. The largest carry is at precision 39, + // where an i256 accumulator's quotient has at most 38 digits. + Scalar::struct_( + decimal_partial_dtype(DType::Decimal(dtype, nullability)), + [ + Scalar::decimal( + DecimalValue::I256(value % limit), + dtype, + Nullability::NonNullable, + ), + Scalar::decimal( + DecimalValue::I256(value / limit), + carry_dtype(), + Nullability::NonNullable, + ), + ], + ) +} + +pub(super) fn decimal_partial_value( + partial: &Scalar, + dtype: DecimalDType, +) -> VortexResult { + let fields = partial.as_struct(); + let value = fields + .field(VALUE_FIELD) + .ok_or_else(|| vortex_err!("Decimal sum partial is missing value"))?; + let carry = fields + .field(CARRY_FIELD) + .ok_or_else(|| vortex_err!("Decimal sum partial is missing carry"))?; + let value = DecimalValue::try_from(&value)? + .cast::() + .vortex_expect("decimal fits i256"); + let carry = DecimalValue::try_from(&carry)? + .cast::() + .vortex_expect("decimal fits i256"); + let limit = i256::from_i128(10).wrapping_pow(dtype.precision().into()); + let value = carry + .checked_mul(&limit) + .and_then(|carry| carry.checked_add(&value)) + .map(DecimalValue::I256) + .ok_or_else(|| vortex_err!("Decimal sum partial exceeds its native accumulator"))?; + match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { + value + .cast::() + .map(DecimalValue::from) + .ok_or_else(|| vortex_err!("Decimal sum partial exceeds its native accumulator")) + }) +} + +pub(super) fn add_decimal(value: &mut DecimalValue, other: DecimalValue, dtype: DecimalDType) { + match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { + let lhs: I = value + .cast() + .vortex_expect("partial fits native accumulator"); + let rhs: I = other + .cast() + .vortex_expect("partial fits native accumulator"); + *value = DecimalValue::from(WrappingAdd::wrapping_add(&lhs, &rhs)); + }) +} + +pub(super) fn multiply_decimal( + value: DecimalValue, + len: usize, + dtype: DecimalDType, +) -> DecimalValue { + let value = arrow_buffer::i256::from(value.cast::().vortex_expect("decimal fits i256")); + let product = i256::from(value.wrapping_mul(arrow_buffer::i256::from_i128(len as i128))); + match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { + let product: I = product.as_(); + DecimalValue::from(product) + }) +} + +pub(super) fn finalize_decimal(partials: ArrayRef) -> VortexResult { + if !matches!(partials.dtype(), DType::Struct(..)) { + return Ok(partials); + } + let value = partials.get_item(VALUE_FIELD)?; + let carry = partials.get_item(CARRY_FIELD)?; + let zero = ConstantArray::new( + Scalar::decimal( + DecimalValue::I128(0), + carry_dtype(), + Nullability::NonNullable, + ), + partials.len(), + ) + .into_array(); + value.mask(carry.binary(zero, Operator::Eq)?) +} + +/// Accumulate a decimal array into the sum state. +/// Native addition wraps; precision is checked when the aggregate is finalized. +pub(super) fn accumulate_decimal( + inner: &mut SumState, + d: &DecimalArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; + let validity = match &mask { + Mask::AllTrue(_) => None, + Mask::Values(mask_values) => Some(mask_values.bit_buffer()), + Mask::AllFalse(_) => { + return Ok(false); + } + }; + + let SumState::Decimal { value, dtype } = inner else { + vortex_panic!("expected decimal sum state for decimal input"); + }; + + let values_type = DecimalType::smallest_decimal_value_type(dtype); + match_each_decimal_value_type!(d.values_type(), |T| { + match_each_decimal_value_type!(values_type, |I| { + let initial: I = value + .cast() + .vortex_expect("cannot fail to cast initial value"); + *value = sum_decimal_value(initial, d.buffer::(), validity); + Ok(false) + }) + }) +} + +fn sum_decimal_value( + initial: I, + values: Buffer, + validity: Option<&BitBuffer>, +) -> DecimalValue +where + T: AsPrimitive, + I: NumOps + WrappingAdd + Copy + NativeDecimalType + 'static, + bool: AsPrimitive, + DecimalValue: From, +{ + let sum = match validity { + Some(v) => sum_decimal_with_validity(values, v, initial), + None => sum_decimal(values, initial), + }; + + DecimalValue::from(sum) +} + +fn sum_decimal, I: Copy + WrappingAdd + 'static>( + values: Buffer, + initial: I, +) -> I { + let mut sum = initial; + for v in values.iter() { + let v: I = v.as_(); + sum = WrappingAdd::wrapping_add(&sum, &v); + } + sum +} + +fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> I +where + T: AsPrimitive, + I: NumOps + WrappingAdd + Copy + 'static, + bool: AsPrimitive, +{ + let mut sum = initial; + for (v, valid) in values.iter().zip_eq(validity) { + let v: I = v.as_() * valid.as_(); + + sum = WrappingAdd::wrapping_add(&sum, &v); + } + sum +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs new file mode 100644 index 00000000000..30a3b63d92d --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::super::SumV2; +use super::super::sum_v2; +use super::decimal_partial_scalar; +use super::decimal_partial_value; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::Accumulator; +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::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::ListViewArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::validity::Validity; + +#[rstest] +fn partial_preserves_native_extremes( + #[values(11, 18, 19, 38, 39, 76)] precision: u8, +) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 0); + let values_type = DecimalType::smallest_decimal_value_type(&dtype); + match_each_decimal_value_type!(values_type, |I| { + for value in [I::MIN, I::MAX] { + let value = DecimalValue::from(value); + let partial = decimal_partial_scalar(value, dtype, Nullability::NonNullable); + assert_eq!(decimal_partial_value(&partial, dtype)?, value); + } + }); + Ok(()) +} + +#[rstest] +fn cancellation_across_batches( + #[values(1, 2, 7, 10, 21)] batch_size: usize, + #[values(false, true)] negative: bool, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let dtype = DecimalDType::new(76, -76); + let unit = i256::from_i128(10).wrapping_pow(75); + let value = unit * i256::from_i128(if negative { -6 } else { 6 }); + // Ten values overflow i256 before cancellation; the final total is one input value. + let values = iter::repeat_n(value, 10) + .chain(iter::repeat_n(-value, 9)) + .collect::>(); + let chunks = values + .chunks(batch_size) + .map(|values| { + let mut values = values.to_vec(); + let validity = if nullable { + values.push(value); + Validity::from_iter((0..values.len()).map(|i| i + 1 < values.len())) + } else { + Validity::NonNullable + }; + DecimalArray::new(values.into_iter().collect::>(), dtype, validity) + .into_array() + }) + .collect::>(); + let input_dtype = DType::Decimal( + dtype, + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }, + ); + let expected = Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + let mut accumulator = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + input_dtype.clone(), + )?; + let mut combined = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + input_dtype.clone(), + )?; + for chunk in &chunks { + accumulator.accumulate(chunk, &mut ctx)?; + let mut partial = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + input_dtype.clone(), + )?; + partial.accumulate(chunk, &mut ctx)?; + combined.combine_partials(partial.flush()?)?; + } + assert_eq!(accumulator.finish()?, expected); + assert_eq!(combined.finish()?, expected); + let chunked = ChunkedArray::try_new(chunks, input_dtype.clone())?.into_array(); + let nested = ChunkedArray::try_new(vec![chunked], input_dtype)?.into_array(); + assert_eq!(sum_v2(&nested, &mut ctx)?, expected); + Ok(()) +} + +#[rstest] +#[case::i64(8, DecimalValue::I64(99_999_999), 100_000_000_000)] +#[case::i128(28, DecimalValue::I128(10i128.pow(28) - 1), 100_000_000_000)] +#[case::i256(76, DecimalValue::I256(i256::from_i128(10).wrapping_pow(76) - i256::ONE), 10)] +fn constant_native_overflow_cancels( + #[case] precision: u8, + #[case] value: DecimalValue, + #[case] len: usize, +) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 0); + let scalar = Scalar::decimal(value, dtype, Nullability::NonNullable); + let negative = Scalar::decimal( + value + .checked_mul(&DecimalValue::I8(-1)) + .ok_or_else(|| vortex_err!("test value can be negated"))?, + dtype, + Nullability::NonNullable, + ); + let mut accumulator = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + scalar.dtype().clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate(&ConstantArray::new(scalar, len).into_array(), &mut ctx)?; + let partial = accumulator.flush()?; + accumulator.accumulate(&ConstantArray::new(negative, len).into_array(), &mut ctx)?; + accumulator.combine_partials(partial)?; + let expected_dtype = DecimalDType::new((precision + 10).min(76), 0); + assert_eq!( + accumulator.finish()?, + Scalar::decimal(DecimalValue::I8(0), expected_dtype, Nullability::Nullable) + ); + Ok(()) +} + +#[test] +fn grouped_final_precision_and_empty() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + let elements = + DecimalArray::new(buffer![value, value, -value], dtype, Validity::NonNullable).into_array(); + let groups = ListViewArray::new( + elements.clone(), + buffer![0u32, 0, 0].into_array(), + buffer![3u32, 2, 0].into_array(), + Validity::NonNullable, + ); + let mut accumulator = GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + elements.dtype().clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate_list(&groups.into_array(), &mut ctx)?; + let expected = DecimalArray::new( + buffer![value, i256::ZERO, i256::ZERO], + dtype, + Validity::from_iter([true, false, false]), + ) + .into_array(); + assert_arrays_eq!(accumulator.finish()?, expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index e1287448a8f..118f15c172d 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod decimal; mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; @@ -11,6 +12,13 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use self::decimal::accumulate_decimal; +use self::decimal::add_decimal; +use self::decimal::decimal_partial_dtype; +use self::decimal::decimal_partial_scalar; +use self::decimal::decimal_partial_value; +use self::decimal::finalize_decimal; +use self::decimal::multiply_decimal; use crate::ArrayRef; use crate::ArrayView; use crate::Canonical; @@ -24,7 +32,6 @@ use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumState; use crate::aggregate_fn::fns::sum::accumulate_bool; -use crate::aggregate_fn::fns::sum::accumulate_decimal; use crate::aggregate_fn::fns::sum::accumulate_primitive; use crate::aggregate_fn::fns::sum::make_zero_state; use crate::aggregate_fn::fns::sum::multiply_constant; @@ -39,7 +46,6 @@ use crate::dtype::StructFields; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; -use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; @@ -64,9 +70,12 @@ pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// Sum an array, returning null when it has no valid values or if the sum overflows. /// +/// Decimal sums use wrapping native arithmetic and check the result precision only at finalization. +/// /// This aggregate intentionally has a distinct ID and partial representation from the legacy /// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex -/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state. +/// files, while `SumV2` uses an explicit `{ sum, is_overflow, is_empty }` state. For decimals, +/// `sum` carries both a value within the result precision and any excess in a separate field. /// /// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans` /// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise, @@ -202,6 +211,16 @@ impl AggregateFnVTable for SumV2 { if !constant.scalar().is_null() && !constant.is_empty() { partial.is_empty = false; } + if let SumState::Decimal { value, dtype } = &mut partial.sum { + if let Some(constant_value) = constant.scalar().as_decimal().decimal_value() { + add_decimal( + value, + multiply_decimal(constant_value, constant.len(), *dtype), + *dtype, + ); + } + return Ok(()); + } if partial.skip_nans && constant .scalar() @@ -257,13 +276,20 @@ impl AggregateFnVTable for SumV2 { .get_item(IS_OVERFLOW_FIELD)? .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? .fill_null(true)?; - sum.mask(is_invalid.not()?) + finalize_decimal(sum)?.mask(is_invalid.not()?) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { if partial.is_overflow || partial.is_empty { return Ok(Scalar::null(partial.return_dtype.as_nullable())); } + if let SumState::Decimal { value, dtype } = &partial.sum { + return Ok(if value.fits_in_precision(*dtype) { + Scalar::decimal(*value, *dtype, Nullability::Nullable) + } else { + Scalar::null(partial.return_dtype.as_nullable()) + }); + } Ok(sum_state_scalar(partial, Nullability::Nullable)) } } @@ -284,7 +310,7 @@ fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult { } } - sum.mask(is_valid) + finalize_decimal(sum)?.mask(is_valid) } /// In-memory state for SumV2 accumulation. @@ -339,12 +365,12 @@ fn decode_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { } fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { + let partial_dtype = decimal_partial_dtype(return_dtype.as_nonnullable()); vortex_ensure!( - sum.dtype().nullability() == Nullability::NonNullable - && sum.dtype().eq_ignore_nullability(return_dtype), + sum.dtype().nullability() == Nullability::NonNullable && sum.dtype() == &partial_dtype, "SumV2 partial value has dtype {}, expected {}", sum.dtype(), - return_dtype.as_nonnullable(), + partial_dtype, ); Ok(()) } @@ -358,14 +384,8 @@ fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { - 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, - } + add_decimal(value, decimal_partial_value(other, *dtype)?, *dtype); + false } }) } @@ -382,7 +402,7 @@ fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields { FieldName::from(IS_EMPTY_FIELD), ]), vec![ - sum_dtype.as_nonnullable(), + decimal_partial_dtype(sum_dtype.as_nonnullable()), DType::Bool(Nullability::NonNullable), DType::Bool(Nullability::NonNullable), ], @@ -394,7 +414,7 @@ fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar SumState::Unsigned(value) => Scalar::primitive(*value, nullability), SumState::Signed(value) => Scalar::primitive(*value, nullability), SumState::Float(value) => Scalar::primitive(*value, nullability), - SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), + SumState::Decimal { value, dtype } => decimal_partial_scalar(*value, *dtype, nullability), } } From cea1e9d43d7ae9a4997d7d69ab861d0e5b341c27 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 9 Sep 2026 15:24:21 +0100 Subject: [PATCH 3/5] less Signed-off-by: Robert Kruszewski --- fuzz/src/array/sum/mod.rs | 64 +++------- fuzz/src/array/sum/tests.rs | 5 +- .../aggregate_fn/fns/sum_v2/decimal/mod.rs | 60 +++++---- .../aggregate_fn/fns/sum_v2/decimal/tests.rs | 116 ++++++++++++++++-- .../src/aggregate_fn/fns/sum_v2/mod.rs | 15 ++- 5 files changed, 169 insertions(+), 91 deletions(-) diff --git a/fuzz/src/array/sum/mod.rs b/fuzz/src/array/sum/mod.rs index 92de619e223..d923725be68 100644 --- a/fuzz/src/array/sum/mod.rs +++ b/fuzz/src/array/sum/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::WrappingAdd; +use num_traits::CheckedAdd; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -29,7 +29,7 @@ use vortex_error::vortex_bail; mod tests; /// Reference sum of chunked or canonical arrays using native arithmetic. -/// Decimal precision is checked only after wrapping accumulation of all chunks. +/// Native overflow is checked during accumulation; decimal precision is checked only at the end. /// Empty and all-null arrays return null, matching SQL SUM. /// Returns `None` only for floats, whose rounding depends on addition order. pub fn sum_canonical_array( @@ -38,36 +38,15 @@ pub fn sum_canonical_array( ) -> VortexResult> { let mut any_valid = false; Ok(Some(match array.dtype() { - DType::Bool(_) => Scalar::from( - accumulate( - array, - Some(0u64), - &|sum, value| sum.checked_add(value), - &mut any_valid, - ctx, - )? - .filter(|_| any_valid), - ), - DType::Primitive(ptype, _) if ptype.is_unsigned_int() => Scalar::from( - accumulate( - array, - Some(0u64), - &|sum, value| sum.checked_add(value), - &mut any_valid, - ctx, - )? - .filter(|_| any_valid), - ), - DType::Primitive(ptype, _) if ptype.is_signed_int() => Scalar::from( - accumulate( - array, - Some(0i64), - &|sum, value| sum.checked_add(value), - &mut any_valid, - ctx, - )? - .filter(|_| any_valid), - ), + DType::Bool(_) => { + Scalar::from(accumulate(array, Some(0u64), &mut any_valid, ctx)?.filter(|_| any_valid)) + } + DType::Primitive(ptype, _) if ptype.is_unsigned_int() => { + Scalar::from(accumulate(array, Some(0u64), &mut any_valid, ctx)?.filter(|_| any_valid)) + } + DType::Primitive(ptype, _) if ptype.is_signed_int() => { + Scalar::from(accumulate(array, Some(0i64), &mut any_valid, ctx)?.filter(|_| any_valid)) + } DType::Primitive(..) => return Ok(None), DType::Decimal(input_dtype, _) => { let output_dtype = DecimalDType::new( @@ -81,14 +60,8 @@ pub fn sum_canonical_array( match_each_decimal_value_type!(values_type, |I| { let limit = ::from(limit) .vortex_expect("precision limit fits native accumulator"); - match accumulate( - array, - Some(I::zero()), - &|sum, value| Some(WrappingAdd::wrapping_add(&sum, &value)), - &mut any_valid, - ctx, - )? - .filter(|&value| any_valid && -limit < value && value < limit) + match accumulate(array, Some(I::zero()), &mut any_valid, ctx)? + .filter(|&value| any_valid && -limit < value && value < limit) { Some(value) => { Scalar::decimal(DecimalValue::from(value), output_dtype, Nullable) @@ -104,12 +77,11 @@ pub fn sum_canonical_array( fn accumulate( array: &ArrayRef, initial: Option, - add: &impl Fn(T, T) -> Option, any_valid: &mut bool, ctx: &mut ExecutionCtx, ) -> VortexResult> where - T: BigCast + Copy + Zero, + T: BigCast + Copy + Zero + CheckedAdd, { let Some(initial) = initial else { return Ok(None); @@ -122,13 +94,15 @@ where // A nested chunked array has its own partial; canonical chunks share its running total. let mut partial = Some(T::zero()); for chunk in chunked.non_empty_chunks() { - partial = accumulate(chunk, partial, add, any_valid, ctx)?; + partial = accumulate(chunk, partial, any_valid, ctx)?; } - Ok(partial.and_then(|partial| add(initial, partial))) + Ok(partial.and_then(|partial| initial.checked_add(&partial))) } else { let values = native_values::(array, ctx)?; *any_valid |= !values.is_empty(); - Ok(values.into_iter().try_fold(initial, add)) + Ok(values + .into_iter() + .try_fold(initial, |sum, value| sum.checked_add(&value))) } } diff --git a/fuzz/src/array/sum/tests.rs b/fuzz/src/array/sum/tests.rs index b6a5e6ec9e6..e7cc4c43da8 100644 --- a/fuzz/src/array/sum/tests.rs +++ b/fuzz/src/array/sum/tests.rs @@ -208,7 +208,7 @@ fn test_sum_decimal_keeps_definite_overflow() -> VortexResult<()> { } #[test] -fn test_sum_decimal_native_overflow_cancels() -> VortexResult<()> { +fn test_sum_decimal_native_overflow_is_absorbing() -> VortexResult<()> { let dtype = DecimalDType::new(76, 0); let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); let mut ctx = SESSION.create_execution_ctx(); @@ -217,8 +217,7 @@ fn test_sum_decimal_native_overflow_cancels() -> VortexResult<()> { .chain(iter::repeat_n(-value, 10)) .collect::>(); let array = DecimalArray::new(values, dtype, Validity::NonNullable); - let expected = - Scalar::decimal(DecimalValue::I256(i256::ZERO), dtype, Nullability::Nullable); + let expected = Scalar::null(DType::Decimal(dtype, Nullability::Nullable)); let array = array.into_array(); assert_eq!( sum_canonical_array(&array, &mut ctx)?, diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs index 4ce31b04bda..6b5824403c3 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs @@ -6,7 +6,6 @@ use num_traits::AsPrimitive; use num_traits::CheckedAdd; use num_traits::CheckedMul; use num_traits::NumOps; -use num_traits::WrappingAdd; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -116,7 +115,11 @@ pub(super) fn decimal_partial_value( }) } -pub(super) fn add_decimal(value: &mut DecimalValue, other: DecimalValue, dtype: DecimalDType) { +pub(super) fn add_decimal( + value: &mut DecimalValue, + other: DecimalValue, + dtype: DecimalDType, +) -> bool { match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { let lhs: I = value .cast() @@ -124,7 +127,13 @@ pub(super) fn add_decimal(value: &mut DecimalValue, other: DecimalValue, dtype: let rhs: I = other .cast() .vortex_expect("partial fits native accumulator"); - *value = DecimalValue::from(WrappingAdd::wrapping_add(&lhs, &rhs)); + match CheckedAdd::checked_add(&lhs, &rhs) { + Some(sum) => { + *value = DecimalValue::from(sum); + false + } + None => true, + } }) } @@ -132,12 +141,11 @@ pub(super) fn multiply_decimal( value: DecimalValue, len: usize, dtype: DecimalDType, -) -> DecimalValue { - let value = arrow_buffer::i256::from(value.cast::().vortex_expect("decimal fits i256")); - let product = i256::from(value.wrapping_mul(arrow_buffer::i256::from_i128(len as i128))); +) -> Option { + let value = value.cast::().vortex_expect("decimal fits i256"); + let product = DecimalValue::I256(value.checked_mul(&i256::from_i128(len as i128))?); match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { - let product: I = product.as_(); - DecimalValue::from(product) + product.cast::().map(DecimalValue::from) }) } @@ -160,7 +168,7 @@ pub(super) fn finalize_decimal(partials: ArrayRef) -> VortexResult { } /// Accumulate a decimal array into the sum state. -/// Native addition wraps; precision is checked when the aggregate is finalized. +/// Native addition is checked; precision is checked when the aggregate is finalized. pub(super) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, @@ -180,14 +188,20 @@ pub(super) fn accumulate_decimal( }; let values_type = DecimalType::smallest_decimal_value_type(dtype); - match_each_decimal_value_type!(d.values_type(), |T| { + let sum = match_each_decimal_value_type!(d.values_type(), |T| { match_each_decimal_value_type!(values_type, |I| { let initial: I = value .cast() .vortex_expect("cannot fail to cast initial value"); - *value = sum_decimal_value(initial, d.buffer::(), validity); - Ok(false) + sum_decimal_value(initial, d.buffer::(), validity) }) + }); + Ok(match sum { + Some(sum) => { + *value = sum; + false + } + None => true, }) } @@ -195,10 +209,10 @@ fn sum_decimal_value( initial: I, values: Buffer, validity: Option<&BitBuffer>, -) -> DecimalValue +) -> Option where T: AsPrimitive, - I: NumOps + WrappingAdd + Copy + NativeDecimalType + 'static, + I: NumOps + CheckedAdd + Copy + NativeDecimalType + 'static, bool: AsPrimitive, DecimalValue: From, { @@ -207,34 +221,34 @@ where None => sum_decimal(values, initial), }; - DecimalValue::from(sum) + sum.map(DecimalValue::from) } -fn sum_decimal, I: Copy + WrappingAdd + 'static>( +fn sum_decimal, I: Copy + CheckedAdd + 'static>( values: Buffer, initial: I, -) -> I { +) -> Option { let mut sum = initial; for v in values.iter() { let v: I = v.as_(); - sum = WrappingAdd::wrapping_add(&sum, &v); + sum = sum.checked_add(&v)?; } - sum + Some(sum) } -fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> I +fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> Option where T: AsPrimitive, - I: NumOps + WrappingAdd + Copy + 'static, + I: NumOps + CheckedAdd + Copy + 'static, bool: AsPrimitive, { let mut sum = initial; for (v, valid) in values.iter().zip_eq(validity) { let v: I = v.as_() * valid.as_(); - sum = WrappingAdd::wrapping_add(&sum, &v); + sum = sum.checked_add(&v)?; } - sum + Some(sum) } #[cfg(test)] diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs index 30a3b63d92d..d756522c917 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs @@ -54,16 +54,16 @@ fn partial_preserves_native_extremes( #[rstest] fn cancellation_across_batches( - #[values(1, 2, 7, 10, 21)] batch_size: usize, + #[values(1, 2, 7, 9, 17)] batch_size: usize, #[values(false, true)] negative: bool, #[values(false, true)] nullable: bool, ) -> VortexResult<()> { let dtype = DecimalDType::new(76, -76); let unit = i256::from_i128(10).wrapping_pow(75); let value = unit * i256::from_i128(if negative { -6 } else { 6 }); - // Ten values overflow i256 before cancellation; the final total is one input value. - let values = iter::repeat_n(value, 10) - .chain(iter::repeat_n(-value, 9)) + // Nine values exceed precision 76 without overflowing i256. + let values = iter::repeat_n(value, 9) + .chain(iter::repeat_n(-value, 8)) .collect::>(); let chunks = values .chunks(batch_size) @@ -121,7 +121,7 @@ fn cancellation_across_batches( #[case::i64(8, DecimalValue::I64(99_999_999), 100_000_000_000)] #[case::i128(28, DecimalValue::I128(10i128.pow(28) - 1), 100_000_000_000)] #[case::i256(76, DecimalValue::I256(i256::from_i128(10).wrapping_pow(76) - i256::ONE), 10)] -fn constant_native_overflow_cancels( +fn constant_native_overflow_is_absorbing( #[case] precision: u8, #[case] value: DecimalValue, #[case] len: usize, @@ -142,13 +142,101 @@ fn constant_native_overflow_cancels( )?; let mut ctx = array_session().create_execution_ctx(); accumulator.accumulate(&ConstantArray::new(scalar, len).into_array(), &mut ctx)?; + assert!(accumulator.is_saturated()); let partial = accumulator.flush()?; accumulator.accumulate(&ConstantArray::new(negative, len).into_array(), &mut ctx)?; accumulator.combine_partials(partial)?; let expected_dtype = DecimalDType::new((precision + 10).min(76), 0); assert_eq!( accumulator.finish()?, - Scalar::decimal(DecimalValue::I8(0), expected_dtype, Nullability::Nullable) + Scalar::null(DType::Decimal(expected_dtype, Nullability::Nullable)) + ); + Ok(()) +} + +#[test] +fn constant_precision_overflow_cancels() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + let scalar = Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::NonNullable); + let mut accumulator = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + scalar.dtype().clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate(&ConstantArray::new(scalar, 2).into_array(), &mut ctx)?; + assert!(!accumulator.is_saturated()); + let partial = accumulator.flush()?; + let negative = Scalar::decimal(DecimalValue::I256(-value), dtype, Nullability::NonNullable); + accumulator.accumulate(&ConstantArray::new(negative, 1).into_array(), &mut ctx)?; + accumulator.combine_partials(partial)?; + assert_eq!( + accumulator.finish()?, + Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::Nullable) + ); + Ok(()) +} + +#[rstest] +fn native_overflow_is_absorbing( + #[values(false, true)] negative: bool, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let value = + i256::from_i128(10).wrapping_pow(75) * i256::from_i128(if negative { -6 } else { 6 }); + let values = iter::repeat_n(value, 10) + .chain(iter::repeat_n(-value, 10)) + .chain(iter::once(value)) + .collect::>(); + let validity = if nullable { + Validity::from_iter((0..values.len()).map(|i| i + 1 < values.len())) + } else { + Validity::NonNullable + }; + let array = DecimalArray::new(values, dtype, validity).into_array(); + let mut accumulator = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate(&array, &mut ctx)?; + assert!(accumulator.is_saturated()); + let partial = accumulator.flush()?; + accumulator.combine_partials(partial)?; + assert!(accumulator.is_saturated()); + assert_eq!( + accumulator.finish()?, + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) + ); + Ok(()) +} + +#[test] +fn partial_merge_native_overflow_is_absorbing() -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); + let array = DecimalArray::new(buffer![value; 5], dtype, Validity::NonNullable).into_array(); + let mut accumulator = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate(&array, &mut ctx)?; + assert!(!accumulator.is_saturated()); + let partial = accumulator.flush()?; + accumulator.combine_partials(partial.clone())?; + assert!(!accumulator.is_saturated()); + accumulator.combine_partials(partial)?; + assert!(accumulator.is_saturated()); + let negative = DecimalArray::new(buffer![-value; 5], dtype, Validity::NonNullable).into_array(); + accumulator.accumulate(&negative, &mut ctx)?; + assert_eq!( + accumulator.finish()?, + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) ); Ok(()) } @@ -157,12 +245,16 @@ fn constant_native_overflow_cancels( fn grouped_final_precision_and_empty() -> VortexResult<()> { let dtype = DecimalDType::new(76, 0); let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let elements = - DecimalArray::new(buffer![value, value, -value], dtype, Validity::NonNullable).into_array(); + let values = [value, value, -value] + .into_iter() + .chain(iter::repeat_n(value, 10)) + .chain(iter::repeat_n(-value, 10)) + .collect::>(); + let elements = DecimalArray::new(values, dtype, Validity::NonNullable).into_array(); let groups = ListViewArray::new( elements.clone(), - buffer![0u32, 0, 0].into_array(), - buffer![3u32, 2, 0].into_array(), + buffer![0u32, 0, 0, 3].into_array(), + buffer![3u32, 2, 0, 20].into_array(), Validity::NonNullable, ); let mut accumulator = GroupedAccumulator::try_new( @@ -173,9 +265,9 @@ fn grouped_final_precision_and_empty() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); accumulator.accumulate_list(&groups.into_array(), &mut ctx)?; let expected = DecimalArray::new( - buffer![value, i256::ZERO, i256::ZERO], + buffer![value, i256::ZERO, i256::ZERO, i256::ZERO], dtype, - Validity::from_iter([true, false, false]), + Validity::from_iter([true, false, false, false]), ) .into_array(); assert_arrays_eq!(accumulator.finish()?, expected, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index 118f15c172d..e5dc16d1237 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -70,7 +70,7 @@ pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// Sum an array, returning null when it has no valid values or if the sum overflows. /// -/// Decimal sums use wrapping native arithmetic and check the result precision only at finalization. +/// Decimal sums use checked native arithmetic and check the result precision only at finalization. /// /// This aggregate intentionally has a distinct ID and partial representation from the legacy /// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex @@ -213,11 +213,11 @@ impl AggregateFnVTable for SumV2 { } if let SumState::Decimal { value, dtype } = &mut partial.sum { if let Some(constant_value) = constant.scalar().as_decimal().decimal_value() { - add_decimal( - value, - multiply_decimal(constant_value, constant.len(), *dtype), - *dtype, - ); + partial.is_overflow = + match multiply_decimal(constant_value, constant.len(), *dtype) { + Some(product) => add_decimal(value, product, *dtype), + None => true, + }; } return Ok(()); } @@ -384,8 +384,7 @@ fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { - add_decimal(value, decimal_partial_value(other, *dtype)?, *dtype); - false + add_decimal(value, decimal_partial_value(other, *dtype)?, *dtype) } }) } From 69ecab681edf1291a918fdace1df76382d7f4767 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 9 Sep 2026 16:17:56 +0100 Subject: [PATCH 4/5] moore Signed-off-by: Robert Kruszewski --- fuzz/src/array/mod.rs | 5 +- fuzz/src/array/sum/mod.rs | 91 ++-- fuzz/src/array/sum/tests.rs | 423 +++++++----------- vortex-array/benches/aggregate_sum.rs | 41 -- .../aggregate_fn/fns/sum_v2/decimal/mod.rs | 255 ----------- .../aggregate_fn/fns/sum_v2/decimal/tests.rs | 275 ------------ .../src/aggregate_fn/fns/sum_v2/mod.rs | 55 +-- 7 files changed, 236 insertions(+), 909 deletions(-) delete mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs delete mode 100644 vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index bc889fa314c..d568dbe3f64 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -48,7 +48,7 @@ use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; use vortex_array::aggregate_fn::fns::min_max::min_max; -use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; +use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::arbitrary::ArbitraryArray; @@ -344,6 +344,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { let Some(sum_result) = sum_canonical_array(¤t_array, &mut ctx) .vortex_expect("sum_canonical_array should succeed in fuzz test") else { + // Reject sums whose overflow can depend on grouping or addition order. return Err(EmptyChoose); }; (Action::Sum, ExpectedValue::Scalar(sum_result)) @@ -670,7 +671,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { current_array = cast_result; } Action::Sum => { - let sum_result = sum_v2(¤t_array, &mut ctx) + let sum_result = sum(¤t_array, &mut ctx) .vortex_expect("sum operation should succeed in fuzz test"); assert_scalar_eq(&expected.scalar(), &sum_result, i)?; } diff --git a/fuzz/src/array/sum/mod.rs b/fuzz/src/array/sum/mod.rs index d923725be68..f4ce549eedf 100644 --- a/fuzz/src/array/sum/mod.rs +++ b/fuzz/src/array/sum/mod.rs @@ -7,9 +7,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arrays::Chunked; use vortex_array::arrays::bool::BoolArrayExt; -use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::dtype::BigCast; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -28,26 +26,23 @@ use vortex_error::vortex_bail; #[cfg(test)] mod tests; -/// Reference sum of chunked or canonical arrays using native arithmetic. -/// Native overflow is checked during accumulation; decimal precision is checked only at the end. -/// Empty and all-null arrays return null, matching SQL SUM. -/// Returns `None` only for floats, whose rounding depends on addition order. +/// Reference sum independent of grouping and addition order. +/// Returns `None` for floats or mixed-sign inputs whose positive or negative subtotal overflows +/// the native accumulator or decimal precision. Single-sign overflow returns a null scalar; +/// empty and all-null inputs return zero, matching `sum`. pub fn sum_canonical_array( array: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let mut any_valid = false; - Ok(Some(match array.dtype() { - DType::Bool(_) => { - Scalar::from(accumulate(array, Some(0u64), &mut any_valid, ctx)?.filter(|_| any_valid)) - } + match array.dtype() { + DType::Bool(_) => accumulate::(array, |_| true, Scalar::from, ctx), DType::Primitive(ptype, _) if ptype.is_unsigned_int() => { - Scalar::from(accumulate(array, Some(0u64), &mut any_valid, ctx)?.filter(|_| any_valid)) + accumulate::(array, |_| true, Scalar::from, ctx) } DType::Primitive(ptype, _) if ptype.is_signed_int() => { - Scalar::from(accumulate(array, Some(0i64), &mut any_valid, ctx)?.filter(|_| any_valid)) + accumulate::(array, |_| true, Scalar::from, ctx) } - DType::Primitive(..) => return Ok(None), + DType::Primitive(..) => Ok(None), DType::Decimal(input_dtype, _) => { let output_dtype = DecimalDType::new( (input_dtype.precision() + 10).min(MAX_PRECISION), @@ -60,50 +55,54 @@ pub fn sum_canonical_array( match_each_decimal_value_type!(values_type, |I| { let limit = ::from(limit) .vortex_expect("precision limit fits native accumulator"); - match accumulate(array, Some(I::zero()), &mut any_valid, ctx)? - .filter(|&value| any_valid && -limit < value && value < limit) - { - Some(value) => { - Scalar::decimal(DecimalValue::from(value), output_dtype, Nullable) - } - None => Scalar::null(DType::Decimal(output_dtype, Nullable)), - } + accumulate::( + array, + |&value| -limit < value && value < limit, + |value| match value { + Some(value) => { + Scalar::decimal(DecimalValue::from(value), output_dtype, Nullable) + } + None => Scalar::null(DType::Decimal(output_dtype, Nullable)), + }, + ctx, + ) }) } _ => vortex_bail!("Unsupported sum dtype: {}", array.dtype()), - })) + } } fn accumulate( array: &ArrayRef, - initial: Option, - any_valid: &mut bool, + fits: impl Fn(&T) -> bool, + scalar: impl Fn(Option) -> Scalar, ctx: &mut ExecutionCtx, -) -> VortexResult> +) -> VortexResult> where - T: BigCast + Copy + Zero + CheckedAdd, + T: BigCast + Copy + Zero + CheckedAdd + PartialOrd, { - let Some(initial) = initial else { - return Ok(None); - }; - if array.is_empty() { - return Ok(Some(initial)); + let mut positive = Some(T::zero()); + let mut negative = Some(T::zero()); + for value in native_values::(array, ctx)? { + let subtotal = if value < T::zero() { + &mut negative + } else { + &mut positive + }; + *subtotal = subtotal + .and_then(|subtotal| subtotal.checked_add(&value)) + .filter(&fits); } - if let Some(chunked) = array.as_opt::() { - // A nested chunked array has its own partial; canonical chunks share its running total. - let mut partial = Some(T::zero()); - for chunk in chunked.non_empty_chunks() { - partial = accumulate(chunk, partial, any_valid, ctx)?; - } - Ok(partial.and_then(|partial| initial.checked_add(&partial))) - } else { - let values = native_values::(array, ctx)?; - *any_valid |= !values.is_empty(); - Ok(values - .into_iter() - .try_fold(initial, |sum, value| sum.checked_add(&value))) - } + // Every partial sum lies between these sign-separated bounds. If either bound overflows, + // an opposite-sign value could cancel it before another grouping detects the overflow. + let value = match (positive, negative) { + (Some(positive), Some(negative)) => positive.checked_add(&negative), + (None, Some(negative)) if negative.is_zero() => None, + (Some(positive), None) if positive.is_zero() => None, + _ => return Ok(None), + }; + Ok(Some(scalar(value))) } fn native_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { diff --git a/fuzz/src/array/sum/tests.rs b/fuzz/src/array/sum/tests.rs index e7cc4c43da8..bd3c3377f39 100644 --- a/fuzz/src/array/sum/tests.rs +++ b/fuzz/src/array/sum/tests.rs @@ -8,15 +8,17 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; +use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; use vortex_array::dtype::i256; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; @@ -32,6 +34,7 @@ use vortex_error::vortex_err; use super::sum_canonical_array; use crate::Action; +use crate::CompressorStrategy; use crate::ExpectedValue; use crate::FuzzArrayAction; use crate::SESSION; @@ -39,15 +42,17 @@ use crate::run_fuzz_action; #[test] fn test_sum_ignores_cached_statistics() -> VortexResult<()> { - let array = PrimitiveArray::from_iter([1i64, 2, 3]); - array - .as_ref() - .statistics() - .set(Stat::Sum, Precision::Exact(ScalarValue::from(999i64))); - assert_eq!( - sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, - Some(Scalar::from(6i64)) - ); + let mut ctx = SESSION.create_execution_ctx(); + for (values, expected) in [ + ([1i64, 2, 3], Some(Scalar::from(6i64))), + ([i64::MAX, -1, 1], None), + ] { + let array = PrimitiveArray::from_iter(values).into_array(); + array + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(999i64))); + assert_eq!(sum_canonical_array(&array, &mut ctx)?, expected); + } Ok(()) } @@ -65,25 +70,66 @@ fn test_sum_ignores_cached_statistics() -> VortexResult<()> { PrimitiveArray::new(buffer![u64::MAX, u64::MAX], Validity::from_iter([true, false])).into_array(), u64::MAX.into() )] -#[case::unsigned_all_null(PrimitiveArray::from_option_iter([None::, None]).into_array(), Scalar::from(None::))] -#[case::unsigned_empty(Buffer::::empty().into_array(), Scalar::from(None::))] -#[case::signed_empty(Buffer::::empty().into_array(), Scalar::from(None::))] +#[case::unsigned_all_null(PrimitiveArray::from_option_iter([None::, None]).into_array(), 0u64.into())] +#[case::unsigned_empty(Buffer::::empty().into_array(), 0u64.into())] +#[case::signed_empty(Buffer::::empty().into_array(), 0i64.into())] #[case::bool(BoolArray::from_iter([true, false, true]).into_array(), 2u64.into())] #[case::bool_nulls(BoolArray::from_iter([Some(true), None, Some(false)]).into_array(), 1u64.into())] -#[case::bool_all_null(BoolArray::from_iter([None::, None]).into_array(), Scalar::from(None::))] -#[case::bool_empty(BoolArray::from_iter([] as [bool; 0]).into_array(), Scalar::from(None::))] +#[case::bool_all_null(BoolArray::from_iter([None::, None]).into_array(), 0u64.into())] +#[case::bool_empty(BoolArray::from_iter([] as [bool; 0]).into_array(), 0u64.into())] fn test_sum_integer_and_bool_values( #[case] array: ArrayRef, #[case] expected: Scalar, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let result = sum_canonical_array( - &array.execute::(&mut ctx)?.into_array(), - &mut ctx, - )? - .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; - assert_eq!(result.dtype(), &expected.dtype().as_nullable()); - assert_eq!(result, expected); + assert_eq!( + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) + ); + assert_eq!(sum(&array, &mut ctx)?, expected); + Ok(()) +} + +#[rstest] +#[case::overflow_first([i64::MAX, 1, -1])] +#[case::cancellation_first([i64::MAX, -1, 1])] +#[case::negative_overflow([i64::MIN, -1, 1])] +#[case::negative_cancellation_first([i64::MIN, 1, -1])] +fn test_sum_rejects_ambiguous_native_overflow(#[case] values: [i64; 3]) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for batch_size in 1..=values.len() { + let array = ChunkedArray::try_new( + values + .chunks(batch_size) + .map(|values| PrimitiveArray::from_iter(values.iter().copied()).into_array()), + DType::Primitive(PType::I64, Nullability::NonNullable), + )? + .into_array(); + assert_eq!(sum_canonical_array(&array, &mut ctx)?, None); + let canonical = array.execute::(&mut ctx)?.into_array(); + assert_eq!(sum_canonical_array(&canonical, &mut ctx)?, None); + } + Ok(()) +} + +#[rstest] +#[case::positive_overflow(vec![Some(i64::MAX), Some(1), Some(0)], None)] +#[case::negative_overflow(vec![Some(i64::MIN), Some(-1), Some(0)], None)] +#[case::native_extremes(vec![Some(i64::MAX), Some(i64::MIN)], Some(-1))] +#[case::nullable(vec![Some(i64::MAX), None, Some(-1)], Some(i64::MAX - 1))] +#[case::all_null(vec![None, None], Some(0))] +fn test_sum_signed_unambiguous( + #[case] values: Vec>, + #[case] expected: Option, +) -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter(values).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let expected = Scalar::from(expected); + assert_eq!( + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) + ); + assert_eq!(sum(&array, &mut ctx)?, expected); Ok(()) } @@ -104,17 +150,19 @@ fn test_sum_decimal_storage(#[case] values_type: DecimalType) -> VortexResult<() DecimalDType::new(2, 0), Validity::NonNullable, ) + .into_array() }); - let result = sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())? - .ok_or_else(|| vortex_err!("expected an unambiguous sum"))?; - assert_eq!( - result.dtype(), - &DType::Decimal(DecimalDType::new(12, 0), Nullability::Nullable) + let expected = Scalar::decimal( + DecimalValue::I64(99), + DecimalDType::new(12, 0), + Nullability::Nullable, ); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!( - result.as_decimal().decimal_value(), - Some(DecimalValue::I64(99)) + sum_canonical_array(&array, &mut ctx)?, + Some(expected.clone()) ); + assert_eq!(sum(&array, &mut ctx)?, expected); Ok(()) } @@ -125,9 +173,10 @@ fn test_sum_decimal_widens_beyond_i128() -> VortexResult<()> { buffer![value, value], DecimalDType::new(38, 2), Validity::NonNullable, - ); + ) + .into_array(); assert_eq!( - sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, + sum_canonical_array(&array, &mut SESSION.create_execution_ctx())?, Some(Scalar::decimal( DecimalValue::I256(i256::from_i128(value) * i256::from_i128(2)), DecimalDType::new(48, 2), @@ -137,141 +186,6 @@ fn test_sum_decimal_widens_beyond_i128() -> VortexResult<()> { Ok(()) } -#[test] -fn test_sum_decimal_cancellation_across_groups() -> VortexResult<()> { - let decimal_dtype = DecimalDType::new(76, -76); - let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - - for value in [six_e75, -six_e75] { - let first = DecimalArray::new(buffer![value, value], decimal_dtype, Validity::NonNullable) - .into_array(); - let last = - DecimalArray::new(buffer![-value], decimal_dtype, Validity::NonNullable).into_array(); - let dtype = first.dtype().clone(); - let chunked = ChunkedArray::try_new(vec![first, last], dtype)?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - // The first chunk exceeds precision 76, but the final sum fits after cancellation. - let expected = Scalar::decimal( - DecimalValue::I256(value), - decimal_dtype, - Nullability::Nullable, - ); - let canonical = chunked.clone().execute::(&mut ctx)?.into_array(); - for array in [chunked, canonical] { - assert_eq!(sum_v2(&array, &mut ctx)?, expected); - assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(expected.clone()) - ); - } - } - Ok(()) -} - -#[test] -fn test_sum_decimal_precision_boundary_and_nulls() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let max = i256::from_i128(10).wrapping_pow(76) - i256::ONE; - let mut ctx = SESSION.create_execution_ctx(); - for validity in [ - Validity::from_iter([true, false, true]), - Validity::AllInvalid, - ] { - let array = DecimalArray::new(buffer![max, max, -max], dtype, validity); - let expected = if matches!(array.as_ref().validity()?, Validity::AllInvalid) { - Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) - } else { - Scalar::decimal(DecimalValue::I256(i256::ZERO), dtype, Nullability::Nullable) - }; - assert_eq!( - sum_canonical_array(&array.into_array(), &mut ctx)?, - Some(expected) - ); - } - Ok(()) -} - -#[test] -fn test_sum_decimal_keeps_definite_overflow() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let mut ctx = SESSION.create_execution_ctx(); - for value in [six_e75, -six_e75] { - let array = DecimalArray::new(buffer![value, value], dtype, Validity::NonNullable); - assert_eq!( - sum_canonical_array(&array.into_array(), &mut ctx)?, - Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) - ); - } - Ok(()) -} - -#[test] -fn test_sum_decimal_native_overflow_is_absorbing() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let six_e75 = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let mut ctx = SESSION.create_execution_ctx(); - for value in [six_e75, -six_e75] { - let values = iter::repeat_n(value, 10) - .chain(iter::repeat_n(-value, 10)) - .collect::>(); - let array = DecimalArray::new(values, dtype, Validity::NonNullable); - let expected = Scalar::null(DType::Decimal(dtype, Nullability::Nullable)); - let array = array.into_array(); - assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(expected.clone()) - ); - assert_eq!(sum_v2(&array, &mut ctx)?, expected); - } - Ok(()) -} - -#[test] -fn test_sum_decimal_uses_widened_precision() -> VortexResult<()> { - let array = DecimalArray::new( - buffer![99i8, 99, -99], - DecimalDType::new(2, 0), - Validity::NonNullable, - ); - assert_eq!( - sum_canonical_array(&array.into_array(), &mut SESSION.create_execution_ctx())?, - Some(Scalar::decimal( - DecimalValue::I64(99), - DecimalDType::new(12, 0), - Nullability::Nullable - )) - ); - Ok(()) -} - -#[test] -fn test_sum_signed_overflow_and_cancellation() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - for (values, expected) in [ - (vec![Some(i64::MAX), Some(1), Some(-1)], Some(None)), - (vec![Some(i64::MIN), Some(-1), Some(1)], Some(None)), - ( - vec![Some(i64::MAX), Some(-1), Some(1)], - Some(Some(i64::MAX)), - ), - (vec![Some(i64::MAX), Some(i64::MIN)], Some(Some(-1i64))), - (vec![Some(i64::MAX), Some(1)], Some(None)), - (vec![Some(i64::MIN), Some(-1)], Some(None)), - (vec![Some(i64::MAX), None, Some(i64::MIN)], Some(Some(-1))), - (vec![None, None], Some(None)), - (vec![], Some(None)), - ] { - let array = PrimitiveArray::from_option_iter(values); - assert_eq!( - sum_canonical_array(&array.into_array(), &mut ctx)?, - expected.map(Scalar::from) - ); - } - Ok(()) -} - fn decimal_chunks(groups: &[&[i8]]) -> VortexResult { let unit = i256::from_i128(10).wrapping_pow(75); let dtype = DecimalDType::new(76, -76); @@ -293,108 +207,93 @@ fn decimal_chunks(groups: &[&[i8]]) -> VortexResult { } #[rstest] -#[case::cancellation_across_groups(decimal_chunks(&[&[6, 6], &[-6]]), Some(6))] -#[case::cancellation_across_three_groups(decimal_chunks(&[&[6], &[6], &[-6]]), Some(6))] -#[case::negative_cancellation(decimal_chunks(&[&[-6, -6], &[6]]), Some(-6))] -#[case::shared_running_total(decimal_chunks(&[&[-6], &[6, 6]]), Some(6))] -#[case::cancellation_within_group(decimal_chunks(&[&[6, 6, -6]]), Some(6))] -#[case::empty_groups(decimal_chunks(&[&[], &[6], &[], &[-6], &[]]), Some(0))] -#[case::final_overflow(decimal_chunks(&[&[6], &[6]]), None)] -#[case::all_empty(decimal_chunks(&[&[], &[]]), None)] -fn test_sum_decimal_ignores_group_boundaries( +#[case::overflow_first(decimal_chunks(&[&[6, 6], &[-6]]))] +#[case::cancellation_first(decimal_chunks(&[&[6, -6], &[6]]))] +#[case::negative_overflow(decimal_chunks(&[&[-6, -6], &[6]]))] +#[case::negative_cancellation_first(decimal_chunks(&[&[-6, 6], &[-6]]))] +#[case::zero_total(decimal_chunks(&[&[6, 6], &[-6, -6]]))] +fn test_sum_rejects_ambiguous_decimal_precision( #[case] array: VortexResult, - #[case] expected: Option, ) -> VortexResult<()> { let array = array?; + let nested = ChunkedArray::try_new(vec![array.clone()], array.dtype().clone())?.into_array(); let mut ctx = SESSION.create_execution_ctx(); - let expected = match expected { - Some(value) => Scalar::decimal( - DecimalValue::I256( - i256::from_i128(10).wrapping_pow(75) * i256::from_i128(i128::from(value)), - ), - DecimalDType::new(76, -76), - Nullability::Nullable, - ), - None => Scalar::null(array.dtype().as_nullable()), - }; - assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(expected.clone()) - ); - assert_eq!(sum_v2(&array, &mut ctx)?, expected); + let canonical = array.clone().execute::(&mut ctx)?.into_array(); + for array in [array, nested, canonical] { + assert_eq!(sum_canonical_array(&array, &mut ctx)?, None); + } Ok(()) } -#[test] -fn test_sum_nested_group_cancellation() -> VortexResult<()> { - let inner = decimal_chunks(&[&[6, 6], &[-6]])?; - let array = ChunkedArray::try_new( - vec![decimal_chunks(&[&[-6]])?, inner.clone()], - inner.dtype().clone(), - )? - .into_array(); +#[rstest] +fn test_sum_decimal_native_overflow(#[values(false, true)] negative: bool) -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let value = + i256::from_i128(10).wrapping_pow(75) * i256::from_i128(if negative { -6 } else { 6 }); let mut ctx = SESSION.create_execution_ctx(); - let expected = Scalar::decimal( - DecimalValue::I256(i256::ZERO), - DecimalDType::new(76, -76), - Nullability::Nullable, - ); - assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(expected.clone()) - ); - assert_eq!(sum_v2(&array, &mut ctx)?, expected); + for mixed in [false, true] { + let values = iter::repeat_n(value, 10) + .chain(iter::repeat_n(-value, if mixed { 10 } else { 0 })) + .collect::>(); + let array = DecimalArray::new(values, dtype, Validity::NonNullable).into_array(); + let expected = (!mixed).then(|| Scalar::null(DType::Decimal(dtype, Nullability::Nullable))); + assert_eq!(sum_canonical_array(&array, &mut ctx)?, expected); + } Ok(()) } -#[test] -fn test_sum_cached_statistics_do_not_change_native_overflow() -> VortexResult<()> { - let array = buffer![i64::MAX, 1, -1].into_array(); - array - .statistics() - .set(Stat::Sum, Precision::Exact(ScalarValue::from(i64::MAX))); +#[rstest] +#[case::boundary(Validity::from_iter([true, false, false]), true)] +#[case::cancellation(Validity::from_iter([true, false, true]), true)] +#[case::ambiguous(Validity::AllValid, false)] +#[case::all_null(Validity::AllInvalid, true)] +fn test_sum_decimal_precision_boundary_and_nulls( + #[case] validity: Validity, + #[case] accepted: bool, +) -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let max = i256::from_i128(10).wrapping_pow(76) - i256::ONE; + let array = DecimalArray::new(buffer![max, i256::ONE, -max], dtype, validity).into_array(); let mut ctx = SESSION.create_execution_ctx(); - assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(Scalar::from(None::)) - ); + let result = sum_canonical_array(&array, &mut ctx)?; + assert_eq!(result.is_some(), accepted); + if let Some(expected) = result { + assert_eq!(sum(&array, &mut ctx)?, expected); + } Ok(()) } -#[test] -fn test_sum_cached_statistics_do_not_change_groups() -> VortexResult<()> { - let array = decimal_chunks(&[&[6, 6], &[-6]])?; - array.statistics().set( - Stat::Sum, - Precision::Exact(ScalarValue::from(DecimalValue::I256( - i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6), - ))), +#[rstest] +#[case::positive(6)] +#[case::negative(-6)] +fn test_sum_constant_decimal_definite_overflow(#[case] value: i128) -> VortexResult<()> { + let dtype = DecimalDType::new(76, 0); + let scalar = Scalar::decimal( + DecimalValue::I256(i256::from_i128(10).wrapping_pow(75) * i256::from_i128(value)), + dtype, + Nullability::NonNullable, ); - let mut ctx = SESSION.create_execution_ctx(); + let array = ConstantArray::new(scalar, 2).into_array(); assert_eq!( - sum_canonical_array(&array, &mut ctx)?, - Some(Scalar::decimal( - DecimalValue::I256(i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6)), - DecimalDType::new(76, -76), - Nullability::Nullable - )) + sum_canonical_array(&array, &mut SESSION.create_execution_ctx())?, + Some(Scalar::null(DType::Decimal(dtype, Nullability::Nullable))) ); Ok(()) } #[rstest] -#[case::chunked(false, Some(6))] -#[case::canonical(true, Some(6))] -fn test_sum_action_stores_reference_scalar( - #[case] canonical: bool, - #[case] expected_sum: Option, +#[case::safe_cancellation(decimal_chunks(&[&[4, 4], &[-4]]), Some(4))] +#[case::negative_cancellation(decimal_chunks(&[&[-4, -4], &[4]]), Some(-4))] +#[case::positive_overflow(decimal_chunks(&[&[6, 5]]), None)] +#[case::negative_overflow(decimal_chunks(&[&[-6, -5]]), None)] +#[case::empty(decimal_chunks(&[&[], &[]]), Some(0))] +fn test_sum_action_accepts_unambiguous_inputs( + #[case] array: VortexResult, + #[case] expected: Option, + #[values(false, true)] compress: bool, ) -> VortexResult<()> { - let mut array = decimal_chunks(&[&[6, 6], &[-6]])?; - let mut ctx = SESSION.create_execution_ctx(); - if canonical { - array = array.execute::(&mut ctx)?.into_array(); - } - let expected_sum = match expected_sum { + let array = array?; + let expected = match expected { Some(value) => Scalar::decimal( DecimalValue::I256( i256::from_i128(10).wrapping_pow(75) * i256::from_i128(i128::from(value)), @@ -404,13 +303,31 @@ fn test_sum_action_stores_reference_scalar( ), None => Scalar::null(array.dtype().as_nullable()), }; - let reference_sum = sum_canonical_array(&array, &mut ctx)? - .ok_or_else(|| vortex_err!("expected a decimal sum"))?; - assert_eq!(reference_sum, expected_sum); - let fuzz_action = FuzzArrayAction { - array, - actions: vec![(Action::Sum, ExpectedValue::Scalar(reference_sum))], - }; - assert!(run_fuzz_action(fuzz_action).map_err(|error| vortex_err!("{error}"))?); + assert_eq!( + sum_canonical_array(&array, &mut SESSION.create_execution_ctx())?, + Some(expected.clone()) + ); + let mut actions = Vec::new(); + if compress { + actions.push(( + Action::Compress(CompressorStrategy::Default), + ExpectedValue::Array(array.clone()), + )); + } + actions.push((Action::Sum, ExpectedValue::Scalar(expected))); + assert!( + run_fuzz_action(FuzzArrayAction { array, actions }) + .map_err(|error| vortex_err!("{error}"))? + ); + Ok(()) +} + +#[test] +fn test_sum_rejects_floats() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(); + assert_eq!( + sum_canonical_array(&array, &mut SESSION.create_execution_ctx())?, + None + ); Ok(()) } diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index d703ce1c935..31856e22f9f 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -9,16 +9,8 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::array_session; -use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::DecimalType; use vortex_array::expr::stats::Stat; -use vortex_array::match_each_decimal_value_type; -use vortex_array::scalar::DecimalValue; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_session::VortexSession; fn main() { @@ -31,39 +23,6 @@ const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(array_session); -#[divan::bench(args = [8, 28, 66])] -fn sum_v2_decimal(bencher: Bencher, precision: u8) { - bench_decimal_sum(bencher, precision, false); -} - -#[divan::bench(args = [8, 28, 66])] -fn sum_v2_decimal_nulls(bencher: Bencher, precision: u8) { - bench_decimal_sum(bencher, precision, true); -} - -fn bench_decimal_sum(bencher: Bencher, precision: u8, nullable: bool) { - let dtype = DecimalDType::new(precision, 2); - let values_type = DecimalType::smallest_decimal_value_type(&dtype); - let array = match_each_decimal_value_type!(values_type, |I| { - let values = (0..N) - .map(|i| { - DecimalValue::I64(i as i64 % 1000 - 500) - .cast::() - .vortex_expect("benchmark value fits decimal storage") - }) - .collect::>(); - let validity = if nullable { - Validity::from_iter((0..N).map(|i| i % 5 != 0)) - } else { - Validity::NonNullable - }; - DecimalArray::new(values, dtype, validity).into_array() - }); - bencher - .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) - .bench_refs(|(array, ctx)| sum_v2(array, ctx)); -} - #[divan::bench] fn sum_i32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(1); diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs deleted file mode 100644 index 6b5824403c3..00000000000 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/mod.rs +++ /dev/null @@ -1,255 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use itertools::Itertools; -use num_traits::AsPrimitive; -use num_traits::CheckedAdd; -use num_traits::CheckedMul; -use num_traits::NumOps; -use vortex_buffer::BitBuffer; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_error::vortex_panic; -use vortex_mask::Mask; - -use super::SumState; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::DecimalArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::FieldNames; -use crate::dtype::NativeDecimalType; -use crate::dtype::Nullability; -use crate::dtype::StructFields; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::operators::Operator; - -const VALUE_FIELD: &str = "value"; -const CARRY_FIELD: &str = "carry"; - -fn carry_dtype() -> DecimalDType { - DecimalDType::new(38, 0) -} - -pub(super) fn decimal_partial_dtype(dtype: DType) -> DType { - let DType::Decimal(..) = dtype else { - return dtype; - }; - DType::Struct( - StructFields::new( - FieldNames::from_iter([VALUE_FIELD, CARRY_FIELD]), - vec![ - dtype.as_nonnullable(), - DType::Decimal(carry_dtype(), Nullability::NonNullable), - ], - ), - dtype.nullability(), - ) -} - -pub(super) fn decimal_partial_scalar( - value: DecimalValue, - dtype: DecimalDType, - nullability: Nullability, -) -> Scalar { - let value = value.cast::().vortex_expect("decimal fits i256"); - let limit = i256::from_i128(10).wrapping_pow(dtype.precision().into()); - // Decimal scalars must fit their precision, including in partial states. Preserve excess - // separately so later partials can cancel it. The largest carry is at precision 39, - // where an i256 accumulator's quotient has at most 38 digits. - Scalar::struct_( - decimal_partial_dtype(DType::Decimal(dtype, nullability)), - [ - Scalar::decimal( - DecimalValue::I256(value % limit), - dtype, - Nullability::NonNullable, - ), - Scalar::decimal( - DecimalValue::I256(value / limit), - carry_dtype(), - Nullability::NonNullable, - ), - ], - ) -} - -pub(super) fn decimal_partial_value( - partial: &Scalar, - dtype: DecimalDType, -) -> VortexResult { - let fields = partial.as_struct(); - let value = fields - .field(VALUE_FIELD) - .ok_or_else(|| vortex_err!("Decimal sum partial is missing value"))?; - let carry = fields - .field(CARRY_FIELD) - .ok_or_else(|| vortex_err!("Decimal sum partial is missing carry"))?; - let value = DecimalValue::try_from(&value)? - .cast::() - .vortex_expect("decimal fits i256"); - let carry = DecimalValue::try_from(&carry)? - .cast::() - .vortex_expect("decimal fits i256"); - let limit = i256::from_i128(10).wrapping_pow(dtype.precision().into()); - let value = carry - .checked_mul(&limit) - .and_then(|carry| carry.checked_add(&value)) - .map(DecimalValue::I256) - .ok_or_else(|| vortex_err!("Decimal sum partial exceeds its native accumulator"))?; - match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { - value - .cast::() - .map(DecimalValue::from) - .ok_or_else(|| vortex_err!("Decimal sum partial exceeds its native accumulator")) - }) -} - -pub(super) fn add_decimal( - value: &mut DecimalValue, - other: DecimalValue, - dtype: DecimalDType, -) -> bool { - match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { - let lhs: I = value - .cast() - .vortex_expect("partial fits native accumulator"); - let rhs: I = other - .cast() - .vortex_expect("partial fits native accumulator"); - match CheckedAdd::checked_add(&lhs, &rhs) { - Some(sum) => { - *value = DecimalValue::from(sum); - false - } - None => true, - } - }) -} - -pub(super) fn multiply_decimal( - value: DecimalValue, - len: usize, - dtype: DecimalDType, -) -> Option { - let value = value.cast::().vortex_expect("decimal fits i256"); - let product = DecimalValue::I256(value.checked_mul(&i256::from_i128(len as i128))?); - match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&dtype), |I| { - product.cast::().map(DecimalValue::from) - }) -} - -pub(super) fn finalize_decimal(partials: ArrayRef) -> VortexResult { - if !matches!(partials.dtype(), DType::Struct(..)) { - return Ok(partials); - } - let value = partials.get_item(VALUE_FIELD)?; - let carry = partials.get_item(CARRY_FIELD)?; - let zero = ConstantArray::new( - Scalar::decimal( - DecimalValue::I128(0), - carry_dtype(), - Nullability::NonNullable, - ), - partials.len(), - ) - .into_array(); - value.mask(carry.binary(zero, Operator::Eq)?) -} - -/// Accumulate a decimal array into the sum state. -/// Native addition is checked; precision is checked when the aggregate is finalized. -pub(super) fn accumulate_decimal( - inner: &mut SumState, - d: &DecimalArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let mask = d.as_ref().validity()?.execute_mask(d.as_ref().len(), ctx)?; - let validity = match &mask { - Mask::AllTrue(_) => None, - Mask::Values(mask_values) => Some(mask_values.bit_buffer()), - Mask::AllFalse(_) => { - return Ok(false); - } - }; - - let SumState::Decimal { value, dtype } = inner else { - vortex_panic!("expected decimal sum state for decimal input"); - }; - - let values_type = DecimalType::smallest_decimal_value_type(dtype); - let sum = match_each_decimal_value_type!(d.values_type(), |T| { - match_each_decimal_value_type!(values_type, |I| { - let initial: I = value - .cast() - .vortex_expect("cannot fail to cast initial value"); - sum_decimal_value(initial, d.buffer::(), validity) - }) - }); - Ok(match sum { - Some(sum) => { - *value = sum; - false - } - None => true, - }) -} - -fn sum_decimal_value( - initial: I, - values: Buffer, - validity: Option<&BitBuffer>, -) -> Option -where - T: AsPrimitive, - I: NumOps + CheckedAdd + Copy + NativeDecimalType + 'static, - bool: AsPrimitive, - DecimalValue: From, -{ - let sum = match validity { - Some(v) => sum_decimal_with_validity(values, v, initial), - None => sum_decimal(values, initial), - }; - - sum.map(DecimalValue::from) -} - -fn sum_decimal, I: Copy + CheckedAdd + 'static>( - values: Buffer, - initial: I, -) -> Option { - let mut sum = initial; - for v in values.iter() { - let v: I = v.as_(); - sum = sum.checked_add(&v)?; - } - Some(sum) -} - -fn sum_decimal_with_validity(values: Buffer, validity: &BitBuffer, initial: I) -> Option -where - T: AsPrimitive, - I: NumOps + CheckedAdd + Copy + 'static, - bool: AsPrimitive, -{ - let mut sum = initial; - for (v, valid) in values.iter().zip_eq(validity) { - let v: I = v.as_() * valid.as_(); - - sum = sum.checked_add(&v)?; - } - Some(sum) -} - -#[cfg(test)] -mod tests; diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs deleted file mode 100644 index d756522c917..00000000000 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/decimal/tests.rs +++ /dev/null @@ -1,275 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::iter; - -use rstest::rstest; -use vortex_buffer::Buffer; -use vortex_buffer::buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use super::super::SumV2; -use super::super::sum_v2; -use super::decimal_partial_scalar; -use super::decimal_partial_value; -use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::aggregate_fn::Accumulator; -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::ChunkedArray; -use crate::arrays::ConstantArray; -use crate::arrays::DecimalArray; -use crate::arrays::ListViewArray; -use crate::assert_arrays_eq; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::Nullability; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; -use crate::validity::Validity; - -#[rstest] -fn partial_preserves_native_extremes( - #[values(11, 18, 19, 38, 39, 76)] precision: u8, -) -> VortexResult<()> { - let dtype = DecimalDType::new(precision, 0); - let values_type = DecimalType::smallest_decimal_value_type(&dtype); - match_each_decimal_value_type!(values_type, |I| { - for value in [I::MIN, I::MAX] { - let value = DecimalValue::from(value); - let partial = decimal_partial_scalar(value, dtype, Nullability::NonNullable); - assert_eq!(decimal_partial_value(&partial, dtype)?, value); - } - }); - Ok(()) -} - -#[rstest] -fn cancellation_across_batches( - #[values(1, 2, 7, 9, 17)] batch_size: usize, - #[values(false, true)] negative: bool, - #[values(false, true)] nullable: bool, -) -> VortexResult<()> { - let dtype = DecimalDType::new(76, -76); - let unit = i256::from_i128(10).wrapping_pow(75); - let value = unit * i256::from_i128(if negative { -6 } else { 6 }); - // Nine values exceed precision 76 without overflowing i256. - let values = iter::repeat_n(value, 9) - .chain(iter::repeat_n(-value, 8)) - .collect::>(); - let chunks = values - .chunks(batch_size) - .map(|values| { - let mut values = values.to_vec(); - let validity = if nullable { - values.push(value); - Validity::from_iter((0..values.len()).map(|i| i + 1 < values.len())) - } else { - Validity::NonNullable - }; - DecimalArray::new(values.into_iter().collect::>(), dtype, validity) - .into_array() - }) - .collect::>(); - let input_dtype = DType::Decimal( - dtype, - if nullable { - Nullability::Nullable - } else { - Nullability::NonNullable - }, - ); - let expected = Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::Nullable); - let mut ctx = array_session().create_execution_ctx(); - let mut accumulator = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - input_dtype.clone(), - )?; - let mut combined = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - input_dtype.clone(), - )?; - for chunk in &chunks { - accumulator.accumulate(chunk, &mut ctx)?; - let mut partial = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - input_dtype.clone(), - )?; - partial.accumulate(chunk, &mut ctx)?; - combined.combine_partials(partial.flush()?)?; - } - assert_eq!(accumulator.finish()?, expected); - assert_eq!(combined.finish()?, expected); - let chunked = ChunkedArray::try_new(chunks, input_dtype.clone())?.into_array(); - let nested = ChunkedArray::try_new(vec![chunked], input_dtype)?.into_array(); - assert_eq!(sum_v2(&nested, &mut ctx)?, expected); - Ok(()) -} - -#[rstest] -#[case::i64(8, DecimalValue::I64(99_999_999), 100_000_000_000)] -#[case::i128(28, DecimalValue::I128(10i128.pow(28) - 1), 100_000_000_000)] -#[case::i256(76, DecimalValue::I256(i256::from_i128(10).wrapping_pow(76) - i256::ONE), 10)] -fn constant_native_overflow_is_absorbing( - #[case] precision: u8, - #[case] value: DecimalValue, - #[case] len: usize, -) -> VortexResult<()> { - let dtype = DecimalDType::new(precision, 0); - let scalar = Scalar::decimal(value, dtype, Nullability::NonNullable); - let negative = Scalar::decimal( - value - .checked_mul(&DecimalValue::I8(-1)) - .ok_or_else(|| vortex_err!("test value can be negated"))?, - dtype, - Nullability::NonNullable, - ); - let mut accumulator = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - scalar.dtype().clone(), - )?; - let mut ctx = array_session().create_execution_ctx(); - accumulator.accumulate(&ConstantArray::new(scalar, len).into_array(), &mut ctx)?; - assert!(accumulator.is_saturated()); - let partial = accumulator.flush()?; - accumulator.accumulate(&ConstantArray::new(negative, len).into_array(), &mut ctx)?; - accumulator.combine_partials(partial)?; - let expected_dtype = DecimalDType::new((precision + 10).min(76), 0); - assert_eq!( - accumulator.finish()?, - Scalar::null(DType::Decimal(expected_dtype, Nullability::Nullable)) - ); - Ok(()) -} - -#[test] -fn constant_precision_overflow_cancels() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let scalar = Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::NonNullable); - let mut accumulator = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - scalar.dtype().clone(), - )?; - let mut ctx = array_session().create_execution_ctx(); - accumulator.accumulate(&ConstantArray::new(scalar, 2).into_array(), &mut ctx)?; - assert!(!accumulator.is_saturated()); - let partial = accumulator.flush()?; - let negative = Scalar::decimal(DecimalValue::I256(-value), dtype, Nullability::NonNullable); - accumulator.accumulate(&ConstantArray::new(negative, 1).into_array(), &mut ctx)?; - accumulator.combine_partials(partial)?; - assert_eq!( - accumulator.finish()?, - Scalar::decimal(DecimalValue::I256(value), dtype, Nullability::Nullable) - ); - Ok(()) -} - -#[rstest] -fn native_overflow_is_absorbing( - #[values(false, true)] negative: bool, - #[values(false, true)] nullable: bool, -) -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let value = - i256::from_i128(10).wrapping_pow(75) * i256::from_i128(if negative { -6 } else { 6 }); - let values = iter::repeat_n(value, 10) - .chain(iter::repeat_n(-value, 10)) - .chain(iter::once(value)) - .collect::>(); - let validity = if nullable { - Validity::from_iter((0..values.len()).map(|i| i + 1 < values.len())) - } else { - Validity::NonNullable - }; - let array = DecimalArray::new(values, dtype, validity).into_array(); - let mut accumulator = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - array.dtype().clone(), - )?; - let mut ctx = array_session().create_execution_ctx(); - accumulator.accumulate(&array, &mut ctx)?; - assert!(accumulator.is_saturated()); - let partial = accumulator.flush()?; - accumulator.combine_partials(partial)?; - assert!(accumulator.is_saturated()); - assert_eq!( - accumulator.finish()?, - Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) - ); - Ok(()) -} - -#[test] -fn partial_merge_native_overflow_is_absorbing() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let array = DecimalArray::new(buffer![value; 5], dtype, Validity::NonNullable).into_array(); - let mut accumulator = Accumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - array.dtype().clone(), - )?; - let mut ctx = array_session().create_execution_ctx(); - accumulator.accumulate(&array, &mut ctx)?; - assert!(!accumulator.is_saturated()); - let partial = accumulator.flush()?; - accumulator.combine_partials(partial.clone())?; - assert!(!accumulator.is_saturated()); - accumulator.combine_partials(partial)?; - assert!(accumulator.is_saturated()); - let negative = DecimalArray::new(buffer![-value; 5], dtype, Validity::NonNullable).into_array(); - accumulator.accumulate(&negative, &mut ctx)?; - assert_eq!( - accumulator.finish()?, - Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) - ); - Ok(()) -} - -#[test] -fn grouped_final_precision_and_empty() -> VortexResult<()> { - let dtype = DecimalDType::new(76, 0); - let value = i256::from_i128(10).wrapping_pow(75) * i256::from_i128(6); - let values = [value, value, -value] - .into_iter() - .chain(iter::repeat_n(value, 10)) - .chain(iter::repeat_n(-value, 10)) - .collect::>(); - let elements = DecimalArray::new(values, dtype, Validity::NonNullable).into_array(); - let groups = ListViewArray::new( - elements.clone(), - buffer![0u32, 0, 0, 3].into_array(), - buffer![3u32, 2, 0, 20].into_array(), - Validity::NonNullable, - ); - let mut accumulator = GroupedAccumulator::try_new( - SumV2, - NumericalAggregateOpts::default(), - elements.dtype().clone(), - )?; - let mut ctx = array_session().create_execution_ctx(); - accumulator.accumulate_list(&groups.into_array(), &mut ctx)?; - let expected = DecimalArray::new( - buffer![value, i256::ZERO, i256::ZERO, i256::ZERO], - dtype, - Validity::from_iter([true, false, false, false]), - ) - .into_array(); - assert_arrays_eq!(accumulator.finish()?, expected, &mut ctx); - Ok(()) -} diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index e5dc16d1237..e1287448a8f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod decimal; mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; @@ -12,13 +11,6 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use self::decimal::accumulate_decimal; -use self::decimal::add_decimal; -use self::decimal::decimal_partial_dtype; -use self::decimal::decimal_partial_scalar; -use self::decimal::decimal_partial_value; -use self::decimal::finalize_decimal; -use self::decimal::multiply_decimal; use crate::ArrayRef; use crate::ArrayView; use crate::Canonical; @@ -32,6 +24,7 @@ use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumState; use crate::aggregate_fn::fns::sum::accumulate_bool; +use crate::aggregate_fn::fns::sum::accumulate_decimal; use crate::aggregate_fn::fns::sum::accumulate_primitive; use crate::aggregate_fn::fns::sum::make_zero_state; use crate::aggregate_fn::fns::sum::multiply_constant; @@ -46,6 +39,7 @@ use crate::dtype::StructFields; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; +use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; @@ -70,12 +64,9 @@ pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// Sum an array, returning null when it has no valid values or if the sum overflows. /// -/// Decimal sums use checked native arithmetic and check the result precision only at finalization. -/// /// This aggregate intentionally has a distinct ID and partial representation from the legacy /// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex -/// files, while `SumV2` uses an explicit `{ sum, is_overflow, is_empty }` state. For decimals, -/// `sum` carries both a value within the result precision and any excess in a separate field. +/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state. /// /// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans` /// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise, @@ -211,16 +202,6 @@ impl AggregateFnVTable for SumV2 { if !constant.scalar().is_null() && !constant.is_empty() { partial.is_empty = false; } - if let SumState::Decimal { value, dtype } = &mut partial.sum { - if let Some(constant_value) = constant.scalar().as_decimal().decimal_value() { - partial.is_overflow = - match multiply_decimal(constant_value, constant.len(), *dtype) { - Some(product) => add_decimal(value, product, *dtype), - None => true, - }; - } - return Ok(()); - } if partial.skip_nans && constant .scalar() @@ -276,20 +257,13 @@ impl AggregateFnVTable for SumV2 { .get_item(IS_OVERFLOW_FIELD)? .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? .fill_null(true)?; - finalize_decimal(sum)?.mask(is_invalid.not()?) + sum.mask(is_invalid.not()?) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { if partial.is_overflow || partial.is_empty { return Ok(Scalar::null(partial.return_dtype.as_nullable())); } - if let SumState::Decimal { value, dtype } = &partial.sum { - return Ok(if value.fits_in_precision(*dtype) { - Scalar::decimal(*value, *dtype, Nullability::Nullable) - } else { - Scalar::null(partial.return_dtype.as_nullable()) - }); - } Ok(sum_state_scalar(partial, Nullability::Nullable)) } } @@ -310,7 +284,7 @@ fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult { } } - finalize_decimal(sum)?.mask(is_valid) + sum.mask(is_valid) } /// In-memory state for SumV2 accumulation. @@ -365,12 +339,12 @@ fn decode_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { } fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { - let partial_dtype = decimal_partial_dtype(return_dtype.as_nonnullable()); vortex_ensure!( - sum.dtype().nullability() == Nullability::NonNullable && sum.dtype() == &partial_dtype, + sum.dtype().nullability() == Nullability::NonNullable + && sum.dtype().eq_ignore_nullability(return_dtype), "SumV2 partial value has dtype {}, expected {}", sum.dtype(), - partial_dtype, + return_dtype.as_nonnullable(), ); Ok(()) } @@ -384,7 +358,14 @@ fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { - add_decimal(value, decimal_partial_value(other, *dtype)?, *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, + } } }) } @@ -401,7 +382,7 @@ fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields { FieldName::from(IS_EMPTY_FIELD), ]), vec![ - decimal_partial_dtype(sum_dtype.as_nonnullable()), + sum_dtype.as_nonnullable(), DType::Bool(Nullability::NonNullable), DType::Bool(Nullability::NonNullable), ], @@ -413,7 +394,7 @@ fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar SumState::Unsigned(value) => Scalar::primitive(*value, nullability), SumState::Signed(value) => Scalar::primitive(*value, nullability), SumState::Float(value) => Scalar::primitive(*value, nullability), - SumState::Decimal { value, dtype } => decimal_partial_scalar(*value, *dtype, nullability), + SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), } } From eff1c5b1c326230549c83e97538bf6a68dc44718 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 9 Sep 2026 18:12:14 +0100 Subject: [PATCH 5/5] simpler Signed-off-by: Robert Kruszewski --- fuzz/src/array/sum/mod.rs | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/fuzz/src/array/sum/mod.rs b/fuzz/src/array/sum/mod.rs index f4ce549eedf..98b3f69ce2c 100644 --- a/fuzz/src/array/sum/mod.rs +++ b/fuzz/src/array/sum/mod.rs @@ -72,6 +72,9 @@ pub fn sum_canonical_array( } } +/// Sum positive and negative values separately to bound every possible partial sum. +/// Reject overflowing bounds with mixed signs, since cancellation can make overflow depend on +/// grouping or addition order. fn accumulate( array: &ArrayRef, fits: impl Fn(&T) -> bool, @@ -81,21 +84,20 @@ fn accumulate( where T: BigCast + Copy + Zero + CheckedAdd + PartialOrd, { - let mut positive = Some(T::zero()); - let mut negative = Some(T::zero()); - for value in native_values::(array, ctx)? { - let subtotal = if value < T::zero() { - &mut negative - } else { - &mut positive - }; - *subtotal = subtotal - .and_then(|subtotal| subtotal.checked_add(&value)) - .filter(&fits); - } + let values = native_values::(array, ctx)?; + let positive = values + .iter() + .copied() + .filter(|&value| value > T::zero()) + .try_fold(T::zero(), |sum, value| sum.checked_add(&value)) + .filter(&fits); + let negative = values + .iter() + .copied() + .filter(|&value| value < T::zero()) + .try_fold(T::zero(), |sum, value| sum.checked_add(&value)) + .filter(&fits); - // Every partial sum lies between these sign-separated bounds. If either bound overflows, - // an opposite-sign value could cancel it before another grouping detects the overflow. let value = match (positive, negative) { (Some(positive), Some(negative)) => positive.checked_add(&negative), (None, Some(negative)) if negative.is_zero() => None,