This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23685-a10d19374ffcb0f24253f819a9e254336561e0bb in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 67502a364c283144c085988b93f014232eca5768 Author: sakshichitnis27 <[email protected]> AuthorDate: Sat Aug 15 11:08:30 2026 +0000 Document decimal AVG wrapping arithmetic (#23685) ## Which issue does this PR close? - Closes #23665 ## Rationale for this change Decimal AVG intentionally uses wrapping arithmetic with a widened intermediate sum type. The existing direct `add_wrapping` and `sub_wrapping` calls made that invariant difficult to audit. ## What changes are included in this PR? - Added small private helpers documenting the widened/headroom contract. - Updated scalar, grouped, merge, retract, and `AVG(DISTINCT)` decimal paths to use the helpers. - Kept the existing state types, return types, and arithmetic behavior unchanged. ## Are these changes tested? - `cargo fmt --all -- --check` - `cargo test -p datafusion-functions-aggregate avg --lib` - `cargo test -p datafusion-functions-aggregate-common avg_distinct --lib` - `cargo clippy -p datafusion-functions-aggregate-common --all-targets -- -D warnings` ## Are there any user-facing changes? No. Co-authored-by: Adam Gutglick <[email protected]> --- .../src/aggregate/avg_distinct/decimal.rs | 15 ++++++++- datafusion/functions-aggregate/src/average.rs | 37 ++++++++++++++++++---- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs b/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs index 1b2fb01189..d9e06d3952 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs @@ -65,6 +65,19 @@ impl<I: DecimalType + Debug, S: DecimalType + Debug> DecimalDistinctAvgAccumulat } } +/// Adds a distinct input value to AVG's widened intermediate sum. +/// Wrapping is intentional because the caller selects `S` with the same +/// `avg_sum_data_type` headroom contract as the non-distinct AVG path. +#[inline] +fn add_avg_distinct_sum<I, S>(sum: S::Native, value: I::Native) -> S::Native +where + I: ArrowNumericType, + S: ArrowNumericType, + I::Native: Into<S::Native>, +{ + sum.add_wrapping(value.into()) +} + impl<I, S> Accumulator for DecimalDistinctAvgAccumulator<I, S> where I: DecimalType + ArrowNumericType + Debug, @@ -95,7 +108,7 @@ where // overflow the input's native width (mirrors the non-distinct path). let mut sum = S::Native::usize_as(0); for value in self.sum_accumulator.distinct_values() { - sum = sum.add_wrapping(value.into()); + sum = add_avg_distinct_sum::<I, S>(sum, value); } let Some(count) = S::Native::from_usize(count) else { diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index e5030bf39e..89b6330852 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -700,6 +700,29 @@ where /// Wraps on overflow, matching the `sum` aggregate and [`arrow::compute::sum`]. /// [`avg_sum_data_type`] gives `S` enough headroom that this is unreachable for /// any realistic row count. +#[inline] +fn add_avg_sum<I, S>(sum: S::Native, value: I::Native) -> S::Native +where + I: ArrowNumericType, + S: ArrowNumericType, + I::Native: Into<S::Native>, +{ + sum.add_wrapping(value.into()) +} + +/// Subtracts a value already represented by AVG's widened state. The wrapping +/// operation is intentional because `avg_sum_data_type` reserves the +/// documented headroom for the intermediate sum. +#[inline] +fn sub_avg_sum<I, S>(sum: S::Native, value: I::Native) -> S::Native +where + I: ArrowNumericType, + S: ArrowNumericType, + I::Native: Into<S::Native>, +{ + sum.sub_wrapping(value.into()) +} + fn decimal_sum_as<I, S>(values: &PrimitiveArray<I>) -> Option<S::Native> where I: DecimalType + ArrowNumericType, @@ -714,11 +737,11 @@ where let mut sum = S::Native::default(); if values.null_count() == 0 { for value in values.values() { - sum = sum.add_wrapping((*value).into()); + sum = add_avg_sum::<I, S>(sum, *value); } } else { for value in values.iter().flatten() { - sum = sum.add_wrapping(value.into()); + sum = add_avg_sum::<I, S>(sum, value); } } @@ -738,7 +761,7 @@ where if let Some(x) = decimal_sum_as::<I, S>(values) { let v = self.sum.unwrap_or_default(); - self.sum = Some(v.add_wrapping(x)); + self.sum = Some(add_avg_sum::<S, S>(v, x)); } Ok(()) } @@ -774,7 +797,7 @@ where // sums are summed if let Some(x) = sum(states[1].as_primitive::<S>()) { let v = self.sum.unwrap_or_default(); - self.sum = Some(v.add_wrapping(x)); + self.sum = Some(add_avg_sum::<S, S>(v, x)); } Ok(()) } @@ -783,7 +806,7 @@ where self.count -= (values.len() - values.null_count()) as u64; if let Some(x) = decimal_sum_as::<I, S>(values) { let v = self.sum.unwrap_or_default(); - self.sum = Some(v.sub_wrapping(x)); + self.sum = Some(sub_avg_sum::<S, S>(v, x)); } Ok(()) } @@ -990,7 +1013,7 @@ where |group_index, new_value| { // SAFETY: group_index is guaranteed to be in bounds let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; - *sum = sum.add_wrapping(new_value.into()); + *sum = add_avg_sum::<I, S>(*sum, new_value); self.counts[group_index] += 1; }, @@ -1090,7 +1113,7 @@ where |group_index, new_value: <S as ArrowPrimitiveType>::Native| { // SAFETY: group_index is guaranteed to be in bounds let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; - *sum = sum.add_wrapping(new_value); + *sum = add_avg_sum::<S, S>(*sum, new_value); }, ); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
