viirya opened a new pull request, #7125: URL: https://github.com/apache/datafusion-comet/pull/7125
## Which issue does this PR close? Closes #7075. ## Rationale for this change When a JVM shuffle spills, `SpillSorter` sorts the spill's packed record pointers by partition with `Native.sortRowPartitionsNative`, before writing every partition in one native call. That sort was a copy of the rdxsort-rs crate. It always ran six 4-bit passes over the 24-bit partition id, whatever the partition count, and allocated two sets of 16 bucket vectors, each sized n/4. That is about 8n longs of virtual memory per call (roughly 2n resident), none of it reserved from any memory pool. And it is allocated exactly when the task is spilling because it is short of memory. That scratch space is not needed. `SpillSorter` builds its `ShuffleInMemorySorter` with radix sort enabled, so the sorter fills only half of its pointer array (`getUsableCapacity()` is `array.size() / 2`). The other half stays free as scratch space for Spark's `RadixSort`, and the task memory manager already accounts for it. This follows the issue's second option: port Spark's scheme to Rust and keep the native call, rather than switching to `ShuffleInMemorySorter.getSortedIterator()`. The sort takes only a small share of a spill's CPU next to the row-to-Arrow conversion. The main gain here is dropping the untracked allocation, not an end-to-end speedup. ## What changes are included in this PR? - `execution/row_partition_sort.rs` replaces `execution/sort.rs`. It is an LSD counting sort over the partition id's three bytes, as in Spark's `RadixSort.sort`. One pass counts every byte. A byte that has the same value in every pointer is skipped, so a spill with up to 256 partitions needs one scatter pass. Each remaining byte scatters between the records and the scratch half of the same array, and the result is copied back if it ends in the scratch half. The sort allocates nothing and is stable, so its output is the same as before: pointers sorted by partition, in insertion order within a partition. - `sortRowPartitionsNative` takes the array's capacity in longs. It returns an error if the capacity is less than `2 * size`. `SpillSorter.writeSortedPartitions`, its only caller, passes `sorterArray.size()`. - A new `sort_row_partitions` criterion bench. Bench results on 1M records (Apple M4 Max; `cargo bench --bench sort_row_partitions`, baseline taken on `main`): | partitions | before | after | speedup | |---|---|---|---| | 1 | 6.52 ms | 0.54 ms | 12x | | 200 | 6.69 ms | 1.17 ms | 5.7x | | 4000 | 6.94 ms | 1.87 ms | 3.7x | | 100,000 | 7.42 ms | 4.77 ms | 1.6x | | 200, 90% in one partition | 9.08 ms | 3.73 ms | 2.4x | ## How are these changes tested? - The existing `sort.rs` and `jni_api` sort tests are ported to the new module. New tests cover no records, one partition, partitions that fit in one byte and in two bytes, more than 65536 partitions (three passes, ending in the scratch half), partitions that differ only in the upper bytes, skewed partitions, an array longer than `2 * size`, and an array without room for scratch space. Each test checks the result against a stable sort by partition id, and the scratch space starts out filled with garbage. - The `write_sorted_partitions` tests in `jni_api` use the new sort. - `SpillSorterSuite`, `CometShuffleSuite` and `DisableAQECometShuffleSuite` pass. 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]
