From 60067efea72a3e469f77925526661a27437cb091 Mon Sep 17 00:00:00 2001 From: breken-ai <312387581+breken-ai@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:07:42 -0700 Subject: [PATCH 1/2] fix: avoid divide-by-zero in avg for negative-scale decimals DecimalAverager computed 10^sum_scale and 10^target_scale with pow_wrapping(scale as u32). A negative scale wraps to a zero factor and avg() then divided by it. Store the single 10^(target_scale - sum_scale) multiplier instead, computed with checked exponentiation. Closes #24898. --- .../functions-aggregate-common/src/utils.rs | 54 ++++++++----------- .../sqllogictest/test_files/aggregate.slt | 26 +++++++++ 2 files changed, 49 insertions(+), 31 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/utils.rs b/datafusion/functions-aggregate-common/src/utils.rs index abe2c9a79c78d..4cb35086f3dc0 100644 --- a/datafusion/functions-aggregate-common/src/utils.rs +++ b/datafusion/functions-aggregate-common/src/utils.rs @@ -26,9 +26,7 @@ use datafusion_common::cast::{as_list_array, as_primitive_array}; use datafusion_common::hash_utils::RandomState; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::utils::memory::estimate_memory_size; -use datafusion_common::{ - HashSet, Result, ScalarValue, exec_err, internal_datafusion_err, -}; +use datafusion_common::{HashSet, Result, ScalarValue, exec_err}; use datafusion_expr_common::accumulator::Accumulator; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use std::sync::Arc; @@ -97,10 +95,9 @@ impl Eq for Hashable {} /// For example, the precision is 3, the max of value is `999` and the min /// value is `-999` pub struct DecimalAverager { - /// scale factor for sum values (10^sum_scale) - sum_mul: T::Native, - /// scale factor for target (10^target_scale) - target_mul: T::Native, + /// factor that rescales sum values to the target scale + /// (`10^(target_scale - sum_scale)`) + scale_mul: T::Native, /// the output precision target_precision: u8, /// the output scale @@ -120,31 +117,26 @@ impl DecimalAverager { target_precision: u8, target_scale: i8, ) -> Result { - let sum_mul = T::Native::from_usize(10_usize) - .map(|b| b.pow_wrapping(sum_scale as u32)) - .ok_or_else(|| { - internal_datafusion_err!("Failed to compute sum_mul in DecimalAverager") - })?; - - let target_mul = T::Native::from_usize(10_usize) - .map(|b| b.pow_wrapping(target_scale as u32)) - .ok_or_else(|| { - internal_datafusion_err!( - "Failed to compute target_mul in DecimalAverager" - ) - })?; - - if target_mul >= sum_mul { - Ok(Self { - sum_mul, - target_mul, - target_precision, - target_scale, - }) - } else { + // Only the ratio `10^target_scale / 10^sum_scale` is needed, and the + // scale difference is non-negative even when both scales are negative + // (e.g. `Decimal128(10, -2)`). + let scale_diff = i16::from(target_scale) - i16::from(sum_scale); + if scale_diff < 0 { // can't convert the lit decimal to the returned data type - exec_err!("Arithmetic Overflow in AvgAccumulator") + return exec_err!("Arithmetic Overflow in AvgAccumulator"); } + + let Some(scale_mul) = T::Native::from_usize(10_usize) + .and_then(|b| b.pow_checked(scale_diff as u32).ok()) + else { + return exec_err!("Arithmetic Overflow in AvgAccumulator"); + }; + + Ok(Self { + scale_mul, + target_precision, + target_scale, + }) } /// Returns the `sum`/`count` as a i128/i256 Decimal128/Decimal256 with @@ -155,7 +147,7 @@ impl DecimalAverager { /// * count: total count, stored as a i128/i256 (*NOT* a Decimal128/Decimal256 value) #[inline(always)] pub fn avg(&self, sum: T::Native, count: T::Native) -> Result { - if let Ok(value) = sum.mul_checked(self.target_mul.div_wrapping(self.sum_mul)) { + if let Ok(value) = sum.mul_checked(self.scale_mul) { let new_value = value.div_wrapping(count); let validate = T::validate_decimal_precision( diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index b3f8a5515d7fe..e15a91744a7b5 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3479,6 +3479,32 @@ select avg(c1), arrow_typeof(avg(c1)) from t; statement ok drop table t; +# avg_decimal_negative_scale +statement ok +create table t (k int, c1 decimal(10, -2)) as values + (1, arrow_cast(10000, 'Decimal128(10, -2)')), + (1, arrow_cast(20000, 'Decimal128(10, -2)')), + (2, arrow_cast(20000, 'Decimal128(10, -2)')); + +query RT +select avg(c1), arrow_typeof(avg(c1)) from t where k = 1; +---- +15000.00 Decimal128(14, 2) + +query IR rowsort +select k, avg(c1) from t group by k; +---- +1 15000.00 +2 20000.00 + +query R +select avg(distinct c1) from t; +---- +15000.00 + +statement ok +drop table t; + # avg_decimal_with_nulls statement ok create table t (c1 decimal(10, 0)) as values (1), (NULL), (3), (4), (5); From ac0e79125a90a74a4f1c93bbe50a5aa8ff95161d Mon Sep 17 00:00:00 2001 From: breken-ai <312387581+breken-ai@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:29:32 -0700 Subject: [PATCH 2/2] keep the overflow cause when the avg rescale multiplier does not fit Review follow-up: `pow_checked(..).ok()` dropped the ArrowError. Propagate it, together with the source and target scales, in the existing "Arithmetic Overflow in AvgAccumulator" execution error, and report the impossible `from_usize(10)` failure as an internal error as before. Add unit tests for the negative-scale multiplier and for the error detail. Co-Authored-By: Claude Opus 5.5 --- .../functions-aggregate-common/src/utils.rs | 51 ++++++++++++++++--- 1 file changed, 45 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/utils.rs b/datafusion/functions-aggregate-common/src/utils.rs index 4cb35086f3dc0..043ce215c0052 100644 --- a/datafusion/functions-aggregate-common/src/utils.rs +++ b/datafusion/functions-aggregate-common/src/utils.rs @@ -26,7 +26,9 @@ use datafusion_common::cast::{as_list_array, as_primitive_array}; use datafusion_common::hash_utils::RandomState; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::utils::memory::estimate_memory_size; -use datafusion_common::{HashSet, Result, ScalarValue, exec_err}; +use datafusion_common::{ + HashSet, Result, ScalarValue, exec_datafusion_err, exec_err, internal_datafusion_err, +}; use datafusion_expr_common::accumulator::Accumulator; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use std::sync::Arc; @@ -126,11 +128,14 @@ impl DecimalAverager { return exec_err!("Arithmetic Overflow in AvgAccumulator"); } - let Some(scale_mul) = T::Native::from_usize(10_usize) - .and_then(|b| b.pow_checked(scale_diff as u32).ok()) - else { - return exec_err!("Arithmetic Overflow in AvgAccumulator"); - }; + let ten = T::Native::from_usize(10_usize).ok_or_else(|| { + internal_datafusion_err!("Failed to compute scale_mul in DecimalAverager") + })?; + let scale_mul = ten.pow_checked(scale_diff as u32).map_err(|e| { + exec_datafusion_err!( + "Arithmetic Overflow in AvgAccumulator: cannot rescale from scale {sum_scale} to {target_scale}: {e}" + ) + })?; Ok(Self { scale_mul, @@ -256,3 +261,37 @@ impl GenericDistinctBuffer { estimate_memory_size::(num_elements, fixed_size).unwrap() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Decimal128Type; + + #[test] + fn decimal_averager_negative_scale() { + // Decimal128(10, -2) averaged into Decimal128(14, 2): sums are + // rescaled by 10^(2 - -2) + let averager = DecimalAverager::::try_new(-2, 14, 2).unwrap(); + // 100 + 200 at scale -2 is 30000; their average 15000 at scale 2 is 1500000 + assert_eq!(averager.avg(300, 2).unwrap(), 1_500_000); + } + + #[test] + fn decimal_averager_unrepresentable_multiplier_keeps_cause() { + // 10^255 does not fit in an i128; the error keeps the arrow overflow + // detail instead of dropping it + let err = DecimalAverager::::try_new(-128, 38, 127) + .err() + .expect("10^255 overflows i128"); + let msg = err.to_string(); + assert!( + msg.contains("Arithmetic Overflow in AvgAccumulator"), + "unexpected error: {msg}" + ); + assert!( + msg.contains("cannot rescale from scale -128 to 127"), + "unexpected error: {msg}" + ); + assert!(msg.contains("10 ^ 255"), "unexpected error: {msg}"); + } +}