andygrove opened a new pull request, #26429: URL: https://github.com/apache/datafusion/pull/26429
## Which issue does this PR close? - Closes #25606. ## 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 interning `num_groups` new 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`, `GroupValuesColumn` and `GroupValuesRows`. The estimate comes from the new `datafusion_common::utils::memory::estimate_hash_table_growth`, which builds on `estimate_memory_size`. `ArrowBytesMap` and `ArrowBytesViewMap` get a `hash_table_growth_size` method for the byte maps. - Before aggregating a batch, the final, single, partial and partial reduce hash aggregate streams grow their reservation by the table that their group values would grow to if every row of the batch, in every grouping set, were a new group. If the pool cannot provide it while the table holds groups, they spill (final and single) or emit (partial and partial reduce) those groups first, so that the table grows from empty. Either way, the reservation is resized to the table's actual size after the batch, as before. - The spill and early emit code that each stream runs on memory pressure moves into a function that both paths use. 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_growth` checks 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_table` checks 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_grows` and `partial_hash_emits_before_the_hash_table_grows` aggregate 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 a `GROUP BY` over 5M string keys that spills, the query fails with `ResourcesExhausted` less often under tight limits (failures in 10 runs each): | Limit | main | This PR | | ----- | ---- | ------- | | 64 MB | 3 | 3 | | 48 MB | 7 | 3 | | 32 MB | 10 | 6 | 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 of `main` over 5 rounds, with a median of 1.00x. The cases without `GROUP 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 spilling `GROUP BY` over 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. `GroupValues` has a new provided method, and `datafusion_common::utils::memory`, `ArrowBytesMap` and `ArrowBytesViewMap` have new public functions. No existing API changes. -- 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]
