This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6263df2ca2 Fix AvgGroupsAccumulator::size() self-overcount (#25117)
6263df2ca2 is described below
commit 6263df2ca2a7dac79124814288880ff76499c5ad
Author: kosiew <[email protected]>
AuthorDate: Fri Sep 11 11:20:34 2026 +0000
Fix AvgGroupsAccumulator::size() self-overcount (#25117)
## Which issue does this PR close?
* Part of #23393
## Rationale for this change
`GroupsAccumulator::size()` is expected to report retained state owned
by the accumulator without counting the accumulator object itself.
`AvgGroupsAccumulator::size()` was also charging for the inline `Vec`
descriptors for `counts` and `sums`, even though those descriptors are
part of the accumulator object rather than separately retained heap
state. This caused AVG group accumulators to over-report their memory
usage.
## What changes are included in this PR?
This PR:
* Removes the `size_of::<Vec<_>>()` charges for the inline `counts` and
`sums` vector descriptors from `AvgGroupsAccumulator::size()`.
* Keeps accounting for the heap capacities of `counts` and `sums`.
* Keeps accounting for retained null-tracking state through
`null_state.size()`.
* Adds coverage verifying that the sum buffer is accounted for using the
actual sum native type, including the `Decimal128` input / `Decimal256`
sum case.
## Are these changes tested?
Yes. This PR adds the following focused unit tests in
`datafusion/functions-aggregate/src/average.rs`:
* `avg_groups_size_excludes_accumulator_storage`
* `avg_groups_size_uses_sum_native_type`
The first verifies that an empty accumulator reports zero retained state
and that, after group state is allocated, `size()` equals the capacities
of `counts` and `sums` plus `null_state.size()`.
The second verifies that memory accounting uses the sum accumulator's
native type by checking a `Decimal128` AVG whose sum state uses
`Decimal256`.
No test execution results are shown in the patch.
## Are there any user-facing changes?
No direct user-facing API or query-result changes are included.
This changes the memory-size accounting reported by
`AvgGroupsAccumulator` so that it follows the
`GroupsAccumulator::size()` contract and no longer includes inline
vector descriptors.
## LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated
content has been manually reviewed.
---
datafusion/functions-aggregate/src/average.rs | 59 ++++++++++++++++++++++++---
1 file changed, 53 insertions(+), 6 deletions(-)
diff --git a/datafusion/functions-aggregate/src/average.rs
b/datafusion/functions-aggregate/src/average.rs
index 8ce893c17d..f59ac816a9 100644
--- a/datafusion/functions-aggregate/src/average.rs
+++ b/datafusion/functions-aggregate/src/average.rs
@@ -1246,12 +1246,9 @@ where
fn size(&self) -> usize {
// Heap buffers
self.counts.capacity() * size_of::<u64>()
- + self.sums.capacity() * size_of::<S::Native>()
- // Vec struct overhead (ptr, len, cap) for each field
- + size_of::<Vec<u64>>()
- + size_of::<Vec<S::Native>>()
- // Null tracking buffers
- + self.null_state.size()
+ + self.sums.capacity() * size_of::<S::Native>()
+ // Null tracking buffers
+ + self.null_state.size()
}
}
@@ -1643,6 +1640,56 @@ mod tests {
Ok(())
}
+ #[test]
+ fn avg_groups_size_excludes_accumulator_storage() -> Result<()> {
+ let mut accumulator = AvgGroupsAccumulator::<Float64Type, _>::new(
+ &DataType::Float64,
+ &DataType::Float64,
+ |sum, count| Ok(sum / count as f64),
+ );
+
+ assert_eq!(accumulator.size(), 0);
+
+ let values = Arc::new(Float64Array::from(vec![Some(2.0), None,
Some(4.0)]));
+ accumulator.update_batch(&[values], &[0, 1, 2], None, 3)?;
+
+ let expected_size = accumulator.counts.capacity() * size_of::<u64>()
+ + accumulator.sums.capacity() * size_of::<f64>()
+ + accumulator.null_state.size();
+ assert_eq!(accumulator.size(), expected_size);
+ assert!(accumulator.size() > 0);
+
+ Ok(())
+ }
+
+ #[test]
+ fn avg_groups_size_uses_sum_native_type() -> Result<()> {
+ let input_type = DataType::Decimal128(26, 0);
+ let sum_type = avg_sum_data_type(&input_type);
+ let return_type =
Avg::new().return_type(std::slice::from_ref(&input_type))?;
+ assert_eq!(sum_type, DataType::Decimal256(76, 0));
+ assert_eq!(return_type, DataType::Decimal128(30, 4));
+ let mut accumulator = AvgGroupsAccumulator::<
+ Decimal128Type,
+ _,
+ Decimal256Type,
+ Decimal128Type,
+ >::new(&sum_type, &return_type, |_, _| Ok(0_i128));
+
+ let values = Arc::new(
+ Decimal128Array::from(vec![Some(2), None, Some(4)])
+ .with_precision_and_scale(26, 0)?,
+ );
+ accumulator.update_batch(&[values], &[0, 1, 2], None, 3)?;
+
+ let expected_size = accumulator.counts.capacity() * size_of::<u64>()
+ + accumulator.sums.capacity() * size_of::<i256>()
+ + accumulator.null_state.size();
+ assert_eq!(accumulator.size(), expected_size);
+
+ Ok(())
+ }
+
#[test]
fn average_groups_preserving_reads() -> Result<()> {
let mut accumulator = AvgGroupsAccumulator::<Float64Type, _>::new(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
