viirya opened a new pull request, #26782: URL: https://github.com/apache/datafusion/pull/26782
(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? - Closes #26781. ## 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_indices` grew for the batch, so that reserving the table fits again. Accumulators can also keep per-batch scratch that `size()` counts, and `GroupsAccumulator` had no way to release it. Such a query still fails with `ResourcesExhausted`, 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`: each `HllGroupsAccumulator` keeps a 64 KiB `hashes` buffer, and the stream fails to reserve ~123 KB. - `ANY_VALUE(v) GROUP BY k`, or any aggregate without a native `GroupsAccumulator`: `GroupsAccumulatorAdapter` keeps its per-batch `Scratch` (`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? - Add `GroupsAccumulator::shrink_to_fit`, which does nothing by default, documented like `GroupValues::shrink_to_fit`. - `ClusteredAggregateTable::shrink_to_fit` calls it for every accumulator, through `HashAggregateAccumulator::shrink_to_fit`. - `HllGroupsAccumulator` drops `hashes` and shrinks `groups`. - `GroupsAccumulatorAdapter` drops its `Scratch` and shrinks `states`, and updates `allocation_bytes`. The other `GroupsAccumulator` implementations 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 in `FFI_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_pressure` in `clustered_partial_stream.rs`, next to `completed_groups_under_memory_pressure` and using the same pattern: a sorted `Int64` key, batches whose boundaries share a key, a 100 KiB `GreedyMemoryPool`, and an exact comparison with an unbounded run. It covers `APPROX_DISTINCT(k0), APPROX_DISTINCT(v)` (`HllGroupsAccumulator`) and `ANY_VALUE(v)` (`GroupsAccumulatorAdapter`). Before this change, each case fails with `ResourcesExhausted`. ## 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. -- 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]
