adriangb opened a new pull request, #25605: URL: https://github.com/apache/datafusion/pull/25605
> [!NOTE] > This PR is stacked on https://github.com/apache/datafusion/pull/25123. Only the last commit, `perf: reuse GroupsAccumulatorAdapter scratch buffers across batches`, belongs to this PR. I will rebase it onto `main` after #25123 merges. ## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/25604. ## Rationale for this change `GroupsAccumulatorAdapter` runs every aggregate that has no `GroupsAccumulator` of its own. For each input batch, `invoke_per_accumulator` allocated two new buffers, `groups_with_rows` and `offsets`, and dropped them at the end of the batch. @alamb suggested to reuse them in https://github.com/apache/datafusion/pull/25123#discussion_r4066870146. ## What changes are included in this PR? - A new `Scratch` struct on the adapter holds `groups_with_rows` and `offsets`. Each batch clears them and fills them again, so they allocate only when a batch needs more capacity than an earlier batch. - `invoke_per_accumulator` takes the buffers out of `self` for the batch and puts them back on every return path, errors included. The body moves unchanged into `invoke_per_accumulator_with_scratch`. - The adapter charges the change in the capacity of the buffers to `allocation_bytes`, thus `GroupsAccumulator::size` includes them. It releases them when an emit leaves no group, thus an adapter that has emitted all groups still reports 0 bytes. `batch_indices` stays a new `Vec` for each batch. `batch_indices.into()` moves it into the `UInt32Array` that `take_arrays` and `get_filter_at_indices` read, so to reuse it would add a copy of the same size as the allocation that it saves. I have not measured a speed change. This removes two allocations for each batch, while the same batch calls `Accumulator::update_batch` once for each group that it touches. Thus I expect the change to be small. ## What is the testing strategy for this PR? The three existing adapter tests that check the exact memory charge now include the scratch buffers: - `adapter_charges_retained_indices_once_and_releases_them` also checks that the adapter keeps some scratch capacity after a batch, and that `size()` is 0 again after `evaluate(EmitTo::All)`. - `adapter_clears_successful_group_indices_after_later_error` and `adapter_reconciles_allocation_after_later_error_with_grouped_metric` check the charge after an update error. To make sure that these tests catch an error, I removed the charge for the scratch buffers. All three tests failed. Without the release on emit, `adapter_charges_retained_indices_once_and_releases_them` failed. The unit tests of `datafusion-functions-aggregate-common`, the `memory_limit` tests in `core_integration` and the sqllogictest suite pass. ## Are there any user-facing changes? No. There are no changes to any public API or to any result. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
