You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/11 08:25:00 UTC
(datafusion) branch gh-readonly-queue/main/pr-24099-a1f64e94b4d4d2eec6d4a8af448cedb86a075f75 created (now 8cb7c84d6c)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24099-a1f64e94b4d4d2eec6d4a8af448cedb86a075f75
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 8cb7c84d6c Fix aggregate accumulator capacity accounting (#24099)
This branch includes the following new commits:
new 8cb7c84d6c Fix aggregate accumulator capacity accounting (#24099)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Fix aggregate accumulator capacity accounting (#24099)
Posted by gi...@apache.org.
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-24099-a1f64e94b4d4d2eec6d4a8af448cedb86a075f75
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 8cb7c84d6c3e4b0157985ef8d0b076cde7bf7e59
Author: kosiew <ko...@gmail.com>
AuthorDate: Tue Aug 11 16:24:37 2026 +0800
Fix aggregate accumulator capacity accounting (#24099)
## Which issue does this PR close?
* Part of #23393
## Rationale for this change
Two aggregate accumulators underreport owned memory in their `size()`
implementations:
* `CountGroupsAccumulator` calculates vector backing storage using the
wrong element type.
* `SlidingDistinctSumAccumulator` reports only the accumulator struct
size and omits its owned hash map allocation.
This change corrects those capacity-based estimates while preserving the
different `Accumulator::size()` and `GroupsAccumulator::size()`
self-inclusion contracts.
## What changes are included in this PR?
* Add `vec_capacity_bytes` to calculate a vector’s backing allocation
using its actual element type.
* Update `CountGroupsAccumulator::size()` to report the capacity-based
allocation of its counts vector while remaining self-excluded.
* Update `SlidingDistinctSumAccumulator::size()` to include
`size_of_val(self)` and an estimate of its hash map’s allocated entry
capacity.
* Document that the hash map estimate excludes implementation-specific
control bytes.
## Are these changes tested?
Yes. This PR adds the following unit tests:
* `count_groups_size_includes_vec_capacity`
* `vec_capacity_bytes_uses_element_type`
* `sliding_distinct_sum_size_includes_hash_map_capacity`
The tests verify that:
* grouped count size grows when vector backing storage is allocated;
* vector capacity accounting uses the vector’s element type;
* sliding distinct-sum size includes hash map capacity, grows after the
map expands, and continues to reflect retained capacity after entries
are retracted.
## Are there any user-facing changes?
No user-facing API or aggregate result changes are intended. This PR
only corrects internal memory-size estimates used for aggregate
accounting.
## LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated
content has been manually reviewed.
---
datafusion/functions-aggregate/src/count.rs | 20 ++++++++++++-
datafusion/functions-aggregate/src/sum.rs | 44 +++++++++++++++++++++++++++--
2 files changed, 60 insertions(+), 4 deletions(-)
diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
index f0de9d9848..1e72d8ac3d 100644
--- a/datafusion/functions-aggregate/src/count.rs
+++ b/datafusion/functions-aggregate/src/count.rs
@@ -29,6 +29,7 @@ use arrow::{
},
};
use datafusion_common::hash_utils::RandomState;
+use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx};
use datafusion_common::{
HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err,
stats::Precision, utils::expr::COUNT_STAR_EXPANSION,
@@ -774,7 +775,7 @@ impl GroupsAccumulator for CountGroupsAccumulator {
Ok(vec![state_array])
}
fn size(&self) -> usize {
- self.counts.capacity() * size_of::<usize>()
+ self.counts.heap_size(&mut DFHeapSizeCtx::default())
}
}
@@ -932,6 +933,23 @@ mod tests {
)?)
}
+ #[test]
+ fn count_groups_size_includes_vec_capacity() -> Result<()> {
+ let mut acc = CountGroupsAccumulator::new();
+ let empty_size = acc.size();
+ assert_eq!(empty_size, 0);
+ let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
+ acc.update_batch(&[values], &[0, 1, 2], None, 3)?;
+
+ assert!(acc.counts.capacity() > 0);
+ let allocated_size = acc.counts.heap_size(&mut DFHeapSizeCtx::default());
+ assert_eq!(allocated_size, acc.counts.capacity() * size_of::<i64>());
+ assert_eq!(acc.size(), allocated_size);
+ assert!(acc.size() > empty_size);
+
+ Ok(())
+ }
+
#[test]
fn count_accumulator_nulls() -> Result<()> {
let mut accumulator = CountAccumulator::new();
diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs
index c124c7a1a0..71932c5f0b 100644
--- a/datafusion/functions-aggregate/src/sum.rs
+++ b/datafusion/functions-aggregate/src/sum.rs
@@ -48,7 +48,7 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_o
use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator;
use datafusion_macros::user_doc;
use datafusion_physical_expr::expressions::{CastExpr, Column};
-use std::mem::size_of_val;
+use std::mem::{size_of, size_of_val};
make_udaf_expr_and_func!(
Sum,
@@ -671,7 +671,8 @@ impl Accumulator for SlidingDistinctSumAccumulator {
}
fn size(&self) -> usize {
- size_of_val(self)
+ // Estimate the owned map buckets; implementation-specific control bytes are excluded.
+ size_of_val(self) + self.counts.capacity() * size_of::<(i64, usize)>()
}
fn state(&mut self) -> Result<Vec<ScalarValue>> {
@@ -722,7 +723,10 @@ mod tests {
array::{Decimal128Array, Int64Array},
buffer::{NullBuffer, ScalarBuffer},
};
- use std::sync::Arc;
+ use std::{
+ mem::{size_of, size_of_val},
+ sync::Arc,
+ };
#[test]
fn sliding_distinct_sum_ignores_null_slots() -> Result<()> {
@@ -750,6 +754,40 @@ mod tests {
Ok(())
}
+ fn expected_sliding_distinct_sum_size(acc: &SlidingDistinctSumAccumulator) -> usize {
+ size_of_val(acc) + acc.counts.capacity() * size_of::<(i64, usize)>()
+ }
+
+ #[test]
+ fn sliding_distinct_sum_size_includes_hash_map_capacity() -> Result<()> {
+ let mut acc = SlidingDistinctSumAccumulator::try_new(&DataType::Int64)?;
+ let empty_size = acc.size();
+ let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
+ acc.update_batch(&[Arc::clone(&values)])?;
+
+ let expected = expected_sliding_distinct_sum_size(&acc);
+ assert!(acc.counts.capacity() > 0);
+ assert_eq!(acc.size(), expected);
+ assert!(acc.size() > empty_size);
+
+ let initial_capacity = acc.counts.capacity();
+ let additional_values: ArrayRef =
+ Arc::new(Int64Array::from_iter(4..4 + initial_capacity as i64 + 1));
+ acc.update_batch(&[Arc::clone(&additional_values)])?;
+
+ let grown_size = expected_sliding_distinct_sum_size(&acc);
+ assert!(acc.counts.capacity() > initial_capacity);
+ assert_eq!(acc.size(), grown_size);
+ assert!(acc.size() > expected);
+
+ acc.retract_batch(&[values])?;
+ acc.retract_batch(&[additional_values])?;
+ assert!(acc.counts.is_empty());
+ assert_eq!(acc.size(), grown_size);
+
+ Ok(())
+ }
+
#[test]
fn sliding_distinct_sum_returns_null_for_all_null_frame() -> Result<()> {
let mut acc = SlidingDistinctSumAccumulator::try_new(&DataType::Int64)?;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org