Repository navigation
fix: release per-batch accumulator scratch after a clustered partial emit under memory pressure - #26782
Draft
viirya wants to merge 7 commits into
Draft
fix: release per-batch accumulator scratch after a clustered partial emit under memory pressure#26782viirya wants to merge 7 commits into
viirya wants to merge 7 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>
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>
…rtial-oom Co-authored-by: Isaac <no-reply@databricks.com>
…table Resetting the table under full group clustering split the active group into two state rows within one partition. The aggregate still declares the group keys unique and ordered, so a sort on the keys and another column above it could pass rows through out of order. Emit only the completed groups again, as before, and then release the capacity the table kept for them: a new GroupValues::shrink_to_fit, implemented for GroupValuesPrimitive, GroupValuesClustered, GroupValuesColumn and GroupValuesRows, and the group_indices scratch. Merge the two unit tests into one that overlaps batch boundaries and covers all four implementations, restore the refusal snapshot, and drop the SQL test, which duplicated the unit test. Co-authored-by: Isaac <no-reply@databricks.com>
…emit Add GroupsAccumulator::shrink_to_fit, a no-op by default, and call it for every accumulator from ClusteredAggregateTable::shrink_to_fit. HllGroupsAccumulator drops its per-batch hashes, and GroupsAccumulatorAdapter drops its per-batch scratch, which it otherwise keeps while the active group remains. Closes apache#26781. Co-authored-by: Isaac <no-reply@databricks.com>
github-actions
Bot
added
logical-expr
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.
(Stacked on #26726. Until #26726 merges, this branch's diff also contains #26726's commits. Only the last commit belongs to this PR.)
Which issue does this PR close?
Rationale for this change
#26726 lets a fully clustered partial aggregate continue under memory pressure: it emits the completed groups, then releases the capacity that the group values and
group_indicesgrew for the batch, so that reserving the table fits again. Accumulators can also keep per-batch scratch thatsize()counts, andGroupsAccumulatorhad no way to release it. Such a query still fails withResourcesExhausted, although the table holds only the active group.For example, with 8192-row batches and a 100 KiB pool:
APPROX_DISTINCT(k), APPROX_DISTINCT(v) GROUP BY k: eachHllGroupsAccumulatorkeeps a 64 KiBhashesbuffer, and the stream fails to reserve ~123 KB.ANY_VALUE(v) GROUP BY k, or any aggregate without a nativeGroupsAccumulator:GroupsAccumulatorAdapterkeeps its per-batchScratch(groups_with_rows,offsets,batch_indices, ~160 KiB). It releases that only when every group is emitted, but the clustered emit keeps the active group. The stream fails to reserve ~283 KB.What changes are included in this PR?
GroupsAccumulator::shrink_to_fit, which does nothing by default, documented likeGroupValues::shrink_to_fit.ClusteredAggregateTable::shrink_to_fitcalls it for every accumulator, throughHashAggregateAccumulator::shrink_to_fit.HllGroupsAccumulatordropshashesand shrinksgroups.GroupsAccumulatorAdapterdrops itsScratchand shrinksstates, and updatesallocation_bytes.The other
GroupsAccumulatorimplementations were checked and keep only per-group state (or no counted per-batch scratch), so they keep the default.ForeignGroupsAccumulator(FFI) also keeps the default, because forwarding the call would need a new function pointer inFFI_GroupsAccumulator. That can be a follow-up if a foreign accumulator needs it.What is the testing strategy for this PR?
New test
accumulator_scratch_under_memory_pressureinclustered_partial_stream.rs, next tocompleted_groups_under_memory_pressureand using the same pattern: a sortedInt64key, batches whose boundaries share a key, a 100 KiBGreedyMemoryPool, and an exact comparison with an unbounded run. It coversAPPROX_DISTINCT(k0), APPROX_DISTINCT(v)(HllGroupsAccumulator) andANY_VALUE(v)(GroupsAccumulatorAdapter). Before this change, each case fails withResourcesExhausted.Are there any user-facing changes?
A new provided method,
GroupsAccumulator::shrink_to_fit, with a no-op default. Existing implementations don't need to change.This pull request and its description were written by Isaac.