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-25117-1fb898093211dc544adb58ae517f32b333273efd in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
