Repository navigation
Conversation
A hash table that grows allocates its larger table before it frees the current one. The hash aggregate streams aggregate a batch first and only then grow their reservation to the table's new size, so while the table grows they hold both tables, up to three times the old one, without having reserved the larger one. Before aggregating a batch, the final, single, partial and partial reduce hash aggregate streams now reserve the table that their group values would grow to if every row were a new group. If the pool cannot provide it, they spill (final and single) or emit (partial and partial reduce) the groups aggregated so far first, so that the table grows from empty. Closes apache#25606
github-actions
Bot
added
physical-expr
|
@EmilyMatt could you help with reviews? |
Codecov Report❌ Patch coverage is @@ Coverage Diff @@
## main #26429 +/- ##
========================================
Coverage 82.89% 82.89%
========================================
Files 1153 1153
Lines 457487 457805 +318
Branches 457487 457805 +318
========================================
+ Hits 379213 379490 +277
- Misses 55264 55280 +16
- Partials 23010 23035 +25
... and 5 files with indirect coverage changes 🚀 New features to boost your workflow:
|
|
Accidentally commented this in the wrong tab 🤦♀️ What I see:
|
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Reserving before intern is the right direction. Two points below are about how big the reservation is: it is too small when one batch grows the table more than once, and too large when the rows hit existing groups. The other two are small reuse cleanups.
common.rs:398: importDataFusionErrorinstead of spelling outdatafusion_common::DataFusionError::ResourcesExhausted.final_hash_spills_before_the_hash_table_growsandsingle_hash_spills_before_the_hash_table_growsdiffer only in the mode. One test that loops over[AggregateMode::Final, AggregateMode::Single]covers both.
| // 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) |
There was a problem hiding this 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:
| // 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) | |
| 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.
| 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)); |
There was a problem hiding this 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.
|
|
||
| /// Takes the partial states of the groups aggregated so far, to emit them | ||
| /// on memory pressure. | ||
| fn take_early_emit_states( |
There was a problem hiding this 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.
|
|
||
| /// Sorts and spills one complete in-memory state run: the groups | ||
| /// aggregated so far. | ||
| fn spill( |
There was a problem hiding this 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.
# Conflicts: # datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs # datafusion/physical-plan/src/aggregates/group_values/row.rs # datafusion/physical-plan/src/aggregates/hash_stream.rs
Which issue does this PR close?
Rationale for this change
A hash aggregate can use more memory than its memory pool grants it while its hash table grows. A growing hash table allocates its larger table, at least twice the size, before it frees the current one. The hash aggregate streams aggregate a batch first and only then resize their reservation to the table's new size, so while the table grows they hold up to three times the old table, and the larger table is never reserved. As @EmilyMatt suggested in #25606, the memory for a batch should be reserved before it is inserted, spilling first if it does not fit.
What changes are included in this PR?
GroupValues::hash_table_growth_size(num_groups)returns the size of the larger hash table that interningnum_groupsnew groups would allocate, or 0 if they fit. It defaults to 0, and the group values that use a hash table implement it:GroupValuesPrimitive,GroupValuesBytes,GroupValuesBytesView,GroupValuesColumnandGroupValuesRows. The estimate comes from the newdatafusion_common::utils::memory::estimate_hash_table_growth, which builds onestimate_memory_size.ArrowBytesMapandArrowBytesViewMapget ahash_table_growth_sizemethod for the byte maps.The clustered aggregate streams, and memory that accumulators allocate as groups are added, are unchanged.
What is the testing strategy for this PR?
test_estimate_hash_table_growthchecks the estimate against the allocation of a hashbrown table after it grows once and after it grows several times.group_values_report_the_growth_of_their_hash_tablechecks that each group values implementation with a hash table reports its growth, and that boolean group values, which have none, report 0.final_hash_spills_before_the_hash_table_grows,single_hash_spills_before_the_hash_table_growsandpartial_hash_emits_before_the_hash_table_growsaggregate a batch that grows the hash table, under a memory limit that holds the table after either batch but not the larger table while the current one is held. The final and single aggregates now spill before the batch, and the partial aggregate emits the earlier groups before aggregating it. All three fail on main, where nothing spills or emits.With
datafusion-cli, 16 partitions sharing a greedy pool, and aGROUP BYover 5M string keys that spills, the query fails withResourcesExhaustedless often under tight limits (failures in 10 runs each):On both, the remaining failures come from reservations made while replaying the spilled runs.
Performance compared with
main, alternating runs on the same machine:aggregate_query_sql: every case is within 0.97x to 1.03x ofmainover 5 rounds, with a median of 1.00x. The cases withoutGROUP BY, which this PR does not change, are within 1.01x.datafusion-cli, median of 7 runs: 500M rows into 100 groups take 0.70 s on both, 100M rows into 20M integer groups take 1.39 s against 1.41 s, 50M rows into 10M string groups take 0.81 s against 0.83 s, and the spillingGROUP BYover 5M string keys at 256 MB takes 0.16 s against 0.17 s.Are there any user-facing changes?
Under a memory limit, hash aggregates spill or emit earlier than before when their hash table would grow past the limit.
GroupValueshas a new provided method, anddatafusion_common::utils::memory,ArrowBytesMapandArrowBytesViewMaphave new public functions. No existing API changes.