comphead commented on code in PR #26429:
URL: https://github.com/apache/datafusion/pull/26429#discussion_r4470270091
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -119,6 +120,40 @@ pub fn estimate_memory_size<T>(num_elements: usize,
fixed_size: usize) -> Result
})
}
+/// Estimates the size in bytes of the table that a hashbrown [`HashTable`]
+/// allocates to make room for `additional` more entries, or returns 0 if they
+/// fit in its capacity.
+///
+/// A table that grows allocates the larger table before it moves the entries
+/// and frees the current one, so for a moment it holds both. Memory accounting
+/// that measures the table before or after inserting misses the larger table,
+/// which reserving this size before inserting covers.
+///
+/// Returns `usize::MAX` if the estimate overflows.
+///
+/// # Example
+///
+/// ```rust
+/// # use datafusion_common::utils::memory::estimate_hash_table_growth;
+/// # use hashbrown::hash_table::HashTable;
+/// let table: HashTable<u64> = HashTable::with_capacity(100);
+/// let capacity = table.capacity();
+/// // The table holds `capacity` entries without growing
+/// assert_eq!(estimate_hash_table_growth(&table, capacity), 0);
+/// // One more entry makes it allocate a larger table
+/// assert!(estimate_hash_table_growth(&table, capacity + 1) >
table.allocation_size());
+/// ```
+pub fn estimate_hash_table_growth<T>(table: &HashTable<T>, additional: usize)
-> usize {
+ let required = table.len().saturating_add(additional);
+ if required <= table.capacity() {
+ return 0;
+ }
+ // A full table grows to hold at least one more entry than its capacity,
+ // which doubles its buckets
+ let capacity = required.max(table.capacity().saturating_add(1));
+ estimate_memory_size::<T>(capacity, 0).unwrap_or(usize::MAX)
Review Comment:
`intern` inserts one row at a time, so a batch that needs more than one
doubling grows the table several times. Each growth frees the previous table,
so the peak is at the last growth: the final table plus the table half its
size. This returns only the final table, so the reservation covers the current
table plus the final one, not the peak.
A standalone probe on hashbrown 0.17.1 (`(u64, usize)` entries, every row
new) shows the gap:
| table before | new rows | reserved | peak |
| --- | --- | --- | --- |
| empty | 8192 | 278,528 | 417,808 |
| 1000 groups | 8192 | 313,352 | 417,808 |
| 3000 groups | 60,000 | 2,297,864 | 3,342,352 |
Adding the half-size table when it is larger than the current one matched
the peak within the 16-byte trailing control group in every case. Also,
`required > capacity` always holds after the early return, so the `max` has no
effect:
```suggestion
let grown = estimate_memory_size::<T>(required, 0).unwrap_or(usize::MAX);
// Growing more than once frees the current table before the last growth,
// which holds the final table and the one half its size
grown.saturating_add((grown / 2).saturating_sub(table.allocation_size()))
```
The "grow it several times" case in `test_estimate_hash_table_growth` only
compares against the final allocation, so it doesn't catch this.
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -365,6 +366,42 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
self.state.building().group_values.len()
}
+ /// Returns the size in bytes of the larger hash table that aggregating
+ /// `num_rows` input rows may allocate while the current one is still
+ /// allocated. See [`GroupValues::hash_table_growth_size`].
+ pub(in crate::aggregates) fn hash_table_growth_size(&self, num_rows:
usize) -> usize {
+ let state = self.state.building();
+ // Each grouping set interns every row, so each can add a group per row
+ let num_groups =
num_rows.saturating_mul(state.group_by.groups().len().max(1));
Review Comment:
This sizes the table as if every row were a new group, so it asks for memory
whenever `len + num_rows > capacity`, even if every row matches an existing
group. In the same probe, a 100-group table (capacity 112, 2,184 bytes) asks
for 278,528 bytes for an 8192-row batch, and for 2,228,224 bytes with 8
grouping sets.
A spill or emit resets the table to capacity 0, so the next batch asks
again. When the pool can't provide that, Final and Single write one tiny run
per batch, and Partial emits after every batch. On `main` a low-cardinality
aggregate under the same pool never spills, because its table never grows. I
haven't run this end to end. Did the 32 MB and 48 MB runs record spill counts?
An option that avoids the worst-case guess is to reserve at the point of
growth. `intern` would stop before the insert that grows the table and report
how many rows it consumed. The stream then reserves one doubling, or spills,
and continues with the rest of the batch. That is exact, covers the
multi-growth peak above, and only adds a check on the growth path. It also
overlaps with EmilyMatt's slicing suggestion.
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -757,6 +778,28 @@ impl FinalHashAggregateStream {
Ok(())
}
+ /// Sorts and spills one complete in-memory state run: the groups
+ /// aggregated so far.
+ fn spill(
Review Comment:
`Aggregating::spill` in `single_stream.rs` does the same take, spill and
resize. `take_state_batch` is generic over the marker, so a single method on
`AggregateSpill` that takes `&mut AggregateHashTable<M>` and the reservation
can serve both streams.
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -429,6 +442,22 @@ impl PartialHashAggregateStream {
}
}
+ /// Takes the partial states of the groups aggregated so far, to emit them
+ /// on memory pressure.
+ fn take_early_emit_states(
Review Comment:
`take_early_emit_states` only feeds `emit_on_memory_pressure`, and both
callers pass `hash_table.memory_size()`. If `emit_on_memory_pressure` took
`&mut hash_table` and did the take and the `early_emit_count` increment itself,
this helper would go away, and both the new pre-batch block and the OOM branch
would become a single call. The same applies to `partial_reduce_stream.rs`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
