Skip to content

perf: the JVM shuffle's native radix sort always runs six passes, and the JVM reserves scratch it never uses #7075

Open

Description

@andygrove

What is the problem the feature request solves?

The sort-based JVM columnar shuffle sorts each spill's packed pointers in native code with a radix sort that always runs six passes, and the JVM reserves scratch memory the native sort never uses.

  • native/core/src/execution/sort.rs (RdxSort) sorts on the 24-bit partition id with 16 buckets and cfg_nrounds() = 6 (two rounds per byte). It always does all six rounds, and it moves every element through two sets of 16 Vec buckets, so it needs about 8x the input in temporary memory. Spark's ShuffleInMemorySorter.getSortedIterator() skips the passes for bytes that are constant across the input (common when the partition count is below 256) and sorts in the reserved half of the pointer array.
  • SpillSorter constructs new ShuffleInMemorySorter(allocator, 1, true) with radix sort enabled, so the pointer array reserves a second half for Spark's radix sort. Comet sorts through Native.sortRowPartitionsNative instead, so that half is never used.

Describe the potential solution

Either use Spark's ShuffleInMemorySorter.getSortedIterator() (removes a JNI call and the 8x temporary allocation, skips constant-byte passes, uses the reserved half), or port its scheme to Rust (sort in place against a same-size scratch array, skip passes whose bucket histogram has one entry). If the native sort stays, construct the in-memory sorter with useRadixSort=false so the unused half is not reserved.

Benchmark with a native sorted_spill-style bench at 100, 1000 and 10000 partitions and 100K rows per spill, and check peak memory.

Additional context

Finding J7 in #5905. Related: #7065 (one native call per spill, which also touches SpillSorter.writeSortedFileNative).

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions