perf: the sort-based JVM columnar shuffle is 2x slower than Spark at 2,000 partitions because every partition write pays builder, file and codec setup #7061

What is the problem the feature request solves?

With many partitions, the sort-based JVM columnar shuffle is about twice as slow as Spark's. Each map task writes every partition with its own writeSortedFileNative call, and each call has a large fixed cost.

SELECT /*+ REPARTITION(n) */ * FROM orders (TPC-H SF2, 3M rows, noop sink), local[8], 6 GB off-heap. spark.executor.cores=8 keeps these shuffles on the sort-based writer. The task times are the stage totals of the warm runs:

Going from 400 to 2,000 partitions adds about 12,800 (map task, partition) writes. Comet's map stage grows by about 5.3 s, around 410 µs per write, and Spark's by about 0.65 s, around 50 µs. With 1 GB off-heap the query is 2.7 s against 1.1 s.

A CPU profile of the map tasks inside writeSortedFileNative at 2,000 partitions:

Each call, in process_sorted_row_partition (native/shuffle/src/spark_unsafe/row.rs), does the following for about 190 rows here (3M rows / 8 map tasks / 2,000 partitions):

• builds every column's builder with capacity for batch_size rows (8,192), including a dictionary builder per string column, whatever the row count;
• opens the output file, appends, and closes it;
• creates a new codec context and block writer, which encodes the schema again.

The measurements are from macOS on an Apple M3 Max with a release build. macOS inflates two of the costs: its allocator returns freed blocks to the OS with madvise, and close after a write is slow there. So the gap is probably smaller on Linux, but the per-call work is the same.

Describe the potential solution

This is item J3 in #5905 (and item 4 of #5198), whose proposal is one JNI call per spill that writes every partition. A first step that is a one-line change is to size the builders by min(batch_size, row_num) rather than batch_size. Keeping the file open across a spill's partitions would remove the open and close per call.

Additional context

The reduce stage is also about 0.7 ms slower per task than Spark's at 2,000 partitions. I didn't profile that side. The bypass writer, used for up to 100 / cores partitions, makes the same per-partition calls; #7058 covers its separate one-row-block problem.