Shuffle performance audit: candidate optimizations needing investigation

[andygrove]

Summary

I did a read-through audit of the entire shuffle surface looking for performance
opportunities, and want to record the candidates somewhere before they're lost.

Nothing here has been benchmarked. These are code-reading findings: places where the
implementation does per-row, per-batch, or per-partition work that looks avoidable. Each one
needs to be measured before it's worth acting on, and some may turn out to be noise next to
the surrounding I/O and encode costs. Treat this as a work list for investigation, not a list
of agreed changes.

Scope covered:

• native/shuffle/ — partitioners, writers, IPC block encoding, row-to-Arrow conversion
• native/core/src/execution/operators/shuffle_scan.rs — native shuffle read
• spark/src/main/java/org/apache/spark/shuffle/ and
  spark/src/main/{java,scala}/org/apache/spark/sql/comet/execution/shuffle/ — JVM shuffle
  read and write paths

The native multi-partition writer has had recent attention (#5006,
#5003, #5004), so most of what's left is
concentrated in the JVM columnar path and the read path.

Candidates

1. process_sorted_row_partition rebuilds the schema and IPC writer once per batch

native/shuffle/src/spark_unsafe/row.rs:1422

let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec.clone())?;
written += block_writer.write_batch(&batch, &mut cursor, &ipc_time)?;

This sits inside the while current_row < row_num loop, so every batch pays for:

• make_batch (row.rs:1490): one format!("c{i}") String allocation per column, plus a
  fresh Schema
• ShuffleBlockWriter::try_new: a flattened_fields() walk and a full IPC schema flatbuffer
  encode

This is the same per-block cost that #5006 removed from the native
writer, still present on the JVM-columnar encode path.

Complication: builder_to_array (row.rs:1448-1485) decides per batch whether a Utf8/Binary
column is emitted as Dictionary or plain, so the schema genuinely can vary between batches
when preferDictionaryRatio > 1.0. A plain hoist is only safe when no column is
dictionary-capable; otherwise the writer would need caching keyed on the produced schema.

Measurable with native/shuffle/benches/row_columnar.rs.

2. The shuffle reader discards its thread-local direct buffer once per block

NativeBatchDecoderIterator.scala:193-197, called from :178

close() calls resetDataBuf(), which throws away the thread-local ByteBuffer whenever it
has grown past 128 KB. CometBlockStoreShuffleReader.read() (:109-118) constructs and
closes one NativeBatchDecoderIterator per shuffle block, so for any shuffle whose
compressed blocks exceed 128 KB we appear to get two allocateDirect calls per block: the
discard, then the regrow at :143. allocateDirect zeroes the region and goes through
Bits.reserveMemory, which can force a System.gc() under direct-memory pressure.

If that reading is right, the thread-local isn't buying anything on this path. Worth noting
that the sibling class on the ShuffleScan path,
CometShuffleBlockIterator.java:103-106, keeps its buffer across blocks and never resets, so
the two paths are inconsistent.

Needs an end-to-end shuffle-read benchmark; the existing benches don't cover this.

3. RowPartition boxes two objects per shuffled row

spark/src/main/scala/org/apache/spark/shuffle/sort/RowPartition.scala:25-31

private var rowAddresses: ArrayBuffer[Long] = new ArrayBuffer[Long](initialSize)
private var rowSizes: ArrayBuffer[Int] = new ArrayBuffer[Int](initialSize)

Scala's ArrayBuffer is not specialized — it's backed by Array[AnyRef]. So each addRow
should box a java.lang.Long (addresses are always outside the Long cache range, so always
an allocation) and a java.lang.Integer (row sizes are usually >127, so usually an
allocation). getRowAddresses/getRowSizes then unbox element-by-element into primitive
arrays for the JNI call.

That's on the order of two allocations plus an unboxing pass per row, across the whole JVM
columnar shuffle write path. Manually grown Array[Long]/Array[Int] with a count would
remove the allocations and reduce toArray to an Arrays.copyOf memcpy — or eliminate it, if
the JNI signature took an explicit length.

Of everything here this looks like the largest steady-state win, but it touches a JNI
boundary, so it needs the most care.

4. Per-partition setup cost in the bypass-merge writer

CometBypassMergeSortShuffleWriter.java:180-199 allocates one CometDiskBlockWriter per
partition, and each one's ArrowIPCWriter constructor calls serializeSchema(schema)
(CometDiskBlockWriter.java:298 → SpillWriter.java:88-104). That's
numPartitions × numColumns protobuf serializations per task, with numPartitions up to
spark.shuffle.sort.bypassMergeThreshold (200 by default), all producing the identical
result. Same shape for the per-partition new Native() (:135) and the four CometConf
lookups in the constructor (:144-152).

Separately, partitioner.getPartition(key) is called twice per record on that path:

int partition_id = partitioner.getPartition(key);
partitionWriters[partitioner.getPartition(key)].insertRow((UnsafeRow) record._2(), partition_id);

CometBypassMergeSortShuffleWriter.java:214-216 — partition_id is already in hand.

5. Index-file parsing allocates two objects per partition

CometNativeShuffleWriter.scala:158-171

Files.readAllBytes(tempIndexFilePath).grouped(OFFSET_LENGTH).drop(1).map(indexBytes =>
  ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong ...)

grouped allocates an intermediate array per partition and the map allocates a ByteBuffer
per partition. Wrapping the whole byte array in a single little-endian ByteBuffer and
looping getLong into a pre-sized Array[Long] would drop 2 × numPartitions allocations
per map task. Only likely to matter at high partition counts.

6. PartitionedBatchIterator widens the whole partition's index array up front

native/shuffle/src/partitioners/partitioned_batch_iterator.rs:85-89

let record_batches = buffered_batches.iter().collect::<Vec<_>>();
let current_indices = indices.iter().map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize)).collect::<Vec<_>>();

interleave_record_batch takes &[(usize, usize)] (arrow-select 58.4, interleave.rs:562),
so the u32 → usize widening is unavoidable — but it doesn't have to happen for the whole
partition at once. As written, every spill and every shuffle_write allocates
rows_in_partition × 16 bytes summed across all partitions, roughly a second full copy of the
index array, and record_batches is rebuilt per partition.

Building record_batches once in PartitionedBatchesProducer and widening indices
batch_size at a time into a single reusable scratch buffer would bound this to one
allocation per spill.

Related, in the caller: multi_partition.rs:497-500 replaces partition_indices with
vec![vec![]; num_output_partitions] on every spill, discarding all per-partition Vec
capacity so each inner Vec re-grows from zero after each spill. Recycling the drained Vecs
would avoid that.

Measurable with native/shuffle/benches/shuffle_writer.rs / src/bin/shuffle_bench.rs.

7. Config lookup inside the spill loop

SpillWriter.java:199

int batchSize = (int) CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE().get();

doSpilling runs once per partition per spill, so SpillSorter.writeSortedFileNative
(:247, :283) hits this numPartitions times per spill file. CometDiskBlockWriter
already caches the identical value in columnarBatchSize (:144) and then doesn't use it.

8. Smaller items

• ArrowIPCWriter.spill comparator recomputes memory usage.
  CometDiskBlockWriter.java:382-391 sorts the static currentWriters list with a comparator
  calling getActiveMemoryUsage() twice per comparison; each call is synchronized and walks
  the allocatedPages linked list (SpillWriter.java:249-258). Precompute the sort keys.
• Two JNI round-trips per block on the ShuffleScan path. shuffle_scan.rs:151-171 calls
  has_next() then get_buffer(). The returned DirectByteBuffer is the same object unless
  it was regrown (CometShuffleBlockIterator.java:105), so caching a global ref and
  re-fetching only on capacity change would halve the JNI calls per block. Probably small next
  to decode cost.
• count_new_buffers calls to_data() per column per batch. multi_partition.rs:166
  allocates a Vec<Buffer> per array on the insert hot path just to read buffer pointers and
  capacities.
• Dead initial fetch in the decoder. NativeBatchDecoderIterator.scala:48 runs
  fetchNext() before channel is assigned at :52, so it always returns None via the
  channel == null guard at :93. Harmless today because hasNext re-fetches, but it
  silently depends on field declaration order.

Suggested next steps

Roughly in order of effort-to-impact:

1. Benchmark and fix item 1 and item 2 — both are localized, and both re-do per-item work the
   codebase already avoids elsewhere.
2. Items 4 through 7 are mostly hoisting and should be quick to confirm.
3. Item 3 is likely the biggest win but needs care at the JNI boundary; worth its own issue and
   its own before/after numbers.

Benchmark coverage: native/shuffle/benches/shuffle_writer.rs,
native/shuffle/benches/row_columnar.rs, and native/shuffle/src/bin/shuffle_bench.rs cover
items 1, 6, and indirectly 3. Item 2 has no coverage and would need an end-to-end shuffle-read
benchmark.

Happy for anyone to pick individual items off into their own issues/PRs. Please post numbers
either way, including for the ones that turn out not to matter — a measured "no effect" is
useful information to record here.