viirya opened a new pull request, #26726: URL: https://github.com/apache/datafusion/pull/26726
## Which issue does this PR close? - Closes #26725. ## 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 distinct `Int64` keys 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? - On a refused reservation, `ClusteredPartialAggregateStream` now calls `take_state_batch` under 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 in `GroupedHashAggregateStream`. A side effect: when the only group in the table is the active one, the stream now emits it instead of returning the error. - Updated doc comments. ## What is the testing strategy for this PR? - `clustered_partial_stream::tests::memory_pressure_releases_retained_capacity` (new): one and two `Int64` keys, 8192 distinct keys per batch, 100 KiB `GreedyMemoryPool`. It fails on `main` with `ResourcesExhausted`. - `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 sorted `GROUP BY` from SQL with 4 partitions. It asserts that the plan uses a `Partial` aggregate with `group_clustering_mode=Full`, and compares the result with an unlimited run. The pool limits only `ClusteredPartialAggregateStream` reservations. 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 on `main` with `ResourcesExhausted`. ## 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. -- 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]
