Repository navigation
fix: release the clustered partial aggregate table on memory pressure under full group clustering - #26726
Open
viirya wants to merge 4 commits into
Open
fix: release the clustered partial aggregate table on memory pressure under full group clustering#26726viirya wants to merge 4 commits into
viirya wants to merge 4 commits into
Conversation
…acity After the completed groups are emitted under memory pressure with full group clustering, only one group remains, but the table keeps the capacity it grew for the whole batch (the single-key hash table, the clustered new_groups scratch and group_indices). Re-reserving that retained size fails in a 100 KiB pool for 8192 distinct keys per batch. This test fails on purpose; it demonstrates the issue and is not a fix. Co-authored-by: Isaac <no-reply@databricks.com>
Run a sorted GROUP BY end to end with a pool that limits only the clustered partial aggregate streams, since the downstream merge and final aggregate cannot spill and a shared limit cannot single out the partial stream. With 8192 distinct keys per batch the retained table needs ~344 KB (one key) or ~131 KB (two keys) after the completed groups are emitted, and a 100 KiB limit fails. This test fails on purpose; it demonstrates the issue and is not a fix. Co-authored-by: Isaac <no-reply@databricks.com>
… under full group clustering On a refused reservation under full group clustering, the stream emitted only the completed groups. The remaining table holds one group but keeps the capacity it grew for the batch (the single-key hash table, the clustered new_groups scratch and group_indices), so re-reserving it could still fail with ResourcesExhausted. Use take_state_batch as partial group clustering already does: emit every state, including the active group, and reset the table. The final stage merges the split active group, as it merges states of the same key from different partitions. Co-authored-by: Isaac <no-reply@databricks.com>
github-actions
Bot
added
core
Under full group clustering, the partial stream now emits the state of the one incomplete group on a refusal and finishes, instead of returning the refusal. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
#26576 made a fully sorted partial aggregate emit its completed groups when a reservation is refused. The table that remains holds one group but keeps the capacity it grew for the batch, so re-reserving it can still fail with
ResourcesExhausted. With 8192 distinctInt64keys per batch, the table still needs 344 KB, while a table holding only the active group needs 5.6 KB. The issue has the details.What changes are included in this PR?
ClusteredPartialAggregateStreamnow callstake_state_batchunder full group clustering too, as it already does under partial group clustering. It emits every state, including the active group, and resets the table, which releases the capacity. The final stage merges the active group's two state rows, as it merges states of the same key from different partitions or from early emission inGroupedHashAggregateStream. A side effect: when the only group in the table is the active one, the stream now emits it instead of returning the error.What is the testing strategy for this PR?
clustered_partial_stream::tests::memory_pressure_releases_retained_capacity(new): one and twoInt64keys, 8192 distinct keys per batch, 100 KiBGreedyMemoryPool. It fails onmainwithResourcesExhausted.clustered_partial_stream::tests::completed_groups_under_memory_pressure(updated): key 256 spans both batches, so its state is now split into two rows. The test checks for the extra row and compares the states merged by key with the unlimited run.memory_limit::clustered_partial_oom::sorted_partial_aggregate_completes_under_memory_pressure(new,core_integration): runs a sortedGROUP BYfrom SQL with 4 partitions. It asserts that the plan uses aPartialaggregate withgroup_clustering_mode=Full, and compares the result with an unlimited run. The pool limits onlyClusteredPartialAggregateStreamreservations. Under a shared limit, the non-spilling merge and final aggregate need about the same memory and fail at the same limits, so a shared limit cannot single out this stream. It fails onmainwithResourcesExhausted.aggregate_memory_call_sequence::full_clustering_refusal_is_an_error(updated snapshot): with a limit too small for even one group, the partial stream now emits that group's state and finishes instead of returning the refusal. The final and single cases are unchanged.Are there any user-facing changes?
No API changes. Under memory pressure, a fully sorted partial aggregate can now emit two state rows for the active group, which the final stage merges.
This pull request and its description were written by Isaac.