neilconway opened a new pull request, #24034: URL: https://github.com/apache/datafusion/pull/24034
## Which issue does this PR close? - Part of #23982 ## Rationale for this change `LinearSearch::evaluate_partition_batches` issued one `take_record_batch` call per partition present in the input batch. For batches with many partitions, this is inefficient. Instead, we can build a Vec of the batch's row indices that groups the rows by partition, gather all rows with a single `take_record_batch` call, and hand each partition a slice of the result. Batches that contain a single partition skip the gather entirely. Memory caveat: the emitted slices share the gathered batch's buffers. A partition that never receives rows again retains its slice and therefore pins the gathered batch's buffers (up to one input batch worth of memory per input batch in the worst case). This could be addressed, e.g., with a compaction pass to copy long-lived slices into owned buffers, but I have omitted that for now. Benchmarks: (using #24032) - linear 100 partitions: 44.3 ms -> 44.1 ms (within noise) - linear 10000 partitions: 199.8 ms -> 170.1 ms (-14.9%) - linear sparse 32768 partitions: 224.6 ms -> 205.3 ms (-8.6%) - linear rows 10000 partitions: 169.0 ms -> 142.0 ms (-15.9%) - linear multi 10000 partitions: 295.9 ms -> 268.2 ms (-9.4%) - sorted 10000 partitions: 34.0 ms -> 34.5 ms (+1.3%; unchanged code path) ## What changes are included in this PR? - Rewrite `get_per_partition_indices` and rename to `compute_partition_permutation` - Add focused unit test for computing partition permutations correctly ## Are these changes tested? Yes, new test added. ## Are there any user-facing changes? No. -- 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]
