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-23399-e271f65c9f8188cb6b1876da099646f396e34e00 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 32f5b4772255ab4926b63210b4d12f05b9fe3eba Author: Oleks V <[email protected]> AuthorDate: Wed Jul 22 10:21:05 2026 -0700 chore: fix `SlidingDistinctCountAccumulator::size()` to include budget for distinct values (#23399) ## Rationale for this change Related to https://github.com/apache/datafusion/issues/23393 ## Rationale for this change `SlidingDistinctCountAccumulator::size()` returned a flat `size_of_val(self)` (64 B) regardless of how many distinct values were held. Since `MemoryPool` reservations trust this value, memory limits were silently bypassed for `COUNT(DISTINCT …)` window aggregates — real heap could reach GBs while the pool reported ~64 B. ## What changes are included in this PR? Fix `SlidingDistinctCountAccumulator::size()` in `datafusion/functions-aggregate/src/count.rs` to mirror the sibling `DistinctCountAccumulator::full_size` pattern: ```rust size_of_val(self) + (size_of::<ScalarValue>() + size_of::<usize>()) * self.counts.capacity() + self.counts.keys().map(|k| k.size() - size_of_val(k)).sum::<usize>() + self.data_type.size() - size_of_val(&self.data_type) ``` Accounts for every persistent field of the struct: `self` bytes + HashMap bucket capacity + per-key inner heap + `DataType` inner heap. Uses `.capacity()` (not `.len()`) so post-`retract_batch` state still reports reserved memory correctly. --- datafusion/expr-common/src/accumulator.rs | 2 ++ datafusion/expr-common/src/groups_accumulator.rs | 2 ++ datafusion/functions-aggregate/src/count.rs | 10 ++++++++++ datafusion/physical-plan/src/aggregates/group_values/mod.rs | 4 +++- 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 59fb6a5952..7e9a4ae525 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -92,6 +92,8 @@ pub trait Accumulator: Send + Sync + Debug + std::any::Any { /// /// "Allocated" means that for internal containers such as `Vec`, /// the `capacity` should be used not the `len`. + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; /// Returns the intermediate state of the accumulator, consuming the diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index b021674cbe..13b2f853c9 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -248,6 +248,8 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// This function is called once per batch, so it should be `O(n)` to /// compute, not `O(num_groups)` + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; } diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index f0ce8c82a1..983828ea90 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -555,7 +555,17 @@ impl Accumulator for SlidingDistinctCountAccumulator { } fn size(&self) -> usize { + // Mirrors `DistinctCountAccumulator::full_size`: self + HashMap + // bucket array + per-key inner heap + DataType inner heap. size_of_val(self) + + (size_of::<ScalarValue>() + size_of::<usize>()) * self.counts.capacity() + + self + .counts + .keys() + .map(|k| k.size() - size_of_val(k)) + .sum::<usize>() + + self.data_type.size() + - size_of_val(&self.data_type) } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7a..1101d53531 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -99,7 +99,9 @@ pub trait GroupValues: Send { /// assigned. fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>; - /// Returns the number of bytes of memory used by this [`GroupValues`] + /// Returns the number of bytes of memory used by this [`GroupValues`]. + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; /// Returns true if this [`GroupValues`] is empty --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
