geoffreyclaude opened a new pull request, #26366: URL: https://github.com/apache/datafusion/pull/26366
## Which issue does this PR close? Follow-up to [#25372](https://github.com/apache/datafusion/pull/25372), related to [#23174](https://github.com/apache/datafusion/issues/23174). This can merge independently of the sort-task concurrency change in [#26242](https://github.com/apache/datafusion/pull/26242). ## Rationale for this change Multi-key sorts over many small batches spend quadratic time accounting for retained row buffers, delaying their first merged output. The merge caches one encoded row buffer per input stream. Each new batch currently scans every cache slot to recalculate the reservation. Reading the first batch from each of N streams therefore requires N full scans. A large merge can spend seconds doing this bookkeeping in a single Tokio poll. ## What changes are included in this PR? - Track each stream's reserved bytes and update the total by that stream's change, making each accounting update constant-time. - Keep the existing charge while a buffer is reused. Failed buffer growth and stream completion release only that stream's charge. - Add a SQL benchmark that sorts two million rows loaded in 32-row batches. Two sort keys exercise the affected row-format merge. Data generation happens before timing. ## What is the testing strategy for this PR? A focused regression test fills a bounded memory pool, then rejects growth of a reused row buffer. It checks that the failed stream releases its previous reservation while another stream remains accounted for. The SQL benchmark runs a single query, with no concurrent workload: ```bash cargo run --release --bin benchmark_runner -- sort_small_batches --iterations 5 ``` With four Tokio workers and one partition, nine executions per version in interleaved groups gave: | Single-query wall time | `upstream/main` | This PR | |---|---:|---:| | Median | 4.82 s | 1.30 s | | Range | 4.52–5.64 s | 1.23–1.46 s | That is about **3.7× faster**. Both optimized builds use the same upstream base and return all two million rows. ## Are there any user-facing changes? Large row-format merges perform less bookkeeping while preserving their memory reservations. The SQL benchmark is listed in the benchmark guide, and the internal reservation invariant is documented alongside the per-stream sizes. Developed with AI assistance. -- 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]
