perf: shuffle reader and writer performance review (native, JVM columnar, reader, Celeborn)

[andygrove]

What

This is a fresh, code-level performance review of every shuffle reader and writer in Comet, produced with Claude Fable 5.1 against main at 53d87c8 (2026-09-13). It covers the native local writer, the JVM columnar (row-based) writer, the reader/decoder, and the Celeborn remote writer, for both flat primitive schemas and nested List/Struct/Map/dictionary columns.

The review was done without cross-referencing the issue tracker, so several findings likely overlap with existing issues and open PRs. Where that is the case, please link them here rather than treating this as a duplicate report; the value of this issue is the ranked, cross-path view and the measurements.

The per-row and per-value loops are already in good shape: partition id hashing, arrow's nested interleave, IPC child-buffer slicing and the primitive row-to-Arrow loop are vectorized or hoisted. The large remaining costs are structural: fixed work that scales with partition count instead of bytes, payloads copied twice, and plan shapes that fall into a slower path because native shuffle rejects a partitioning it could handle.

Related issues and PRs

An overlap check against the tracker after filing. Several findings below were already recorded, and two have open PRs from other contributors. Those PRs have priority; the PRs opened from this issue were converted to draft pending coordination.

Finding here	Already tracked in	Open PR
W1 fresh compressor per block	#5002 (compression context item)	#5565 reuses zstd contexts (approved). lz4 and snappy remain per block because the pinned crates expose no reset.
W2 tail chunks copied twice through BatchCoalescer	#5002 (coalescer copy item; the full-batch case was fixed earlier, the tail case remained)	#5908 (draft)
W3 per-partition IPC context, flush and lseek	#5002 (per-partition BufBatchWriter churn item)	#5908 (draft). Conflicts textually with #5565 and with #5807.
W3 offsets, and the index file itself	#5790	#5807 returns partition offsets over JNI instead of the index file
W6 round-robin hashes every column	#5002 and #5397	#5449 (draft) assigns whole batches round-robin for nested schemas
W8 pmod divisions	#5002 (micro item)	none
R2 schema re-parsed per block	#5792	#5809 caches the schema (opened first); #5909 (draft) is an alternative streaming design, see the comparison on #5809
R5 decoder buffer reset per block	#5198 item 2	#5913 (draft)
J2 pointer array growth	not previously tracked	#5907 (draft)
J3 one JNI call per (spill, partition)	#5198 item 4 (bypass writer per-partition setup)	none
J4 RowPartition boxing, per-row lock and log	#5198 item 3	none
per-batch schema and ShuffleBlockWriter rebuild in row.rs (not numbered above; sits next to J6)	#5198 item 1	none

Earlier audits of the same surface: #5002 (native writer, July) and #5198 (whole shuffle surface, August). The findings not listed above (R1, R3, R4, R6, J1, J5, J7, C1 to C4, and the nested-type items) do not appear to have an issue or PR yet.

Measurements

From the existing criterion benches in native/shuffle/benches, run locally on an Apple Silicon machine with 3 s measurement windows. Sizing only, not a regression baseline.

Encoding one 8192-row batch (ShuffleBlockWriter::write_batch)

Batch	None	LZ4	Snappy	Zstd(1)
mixed: Int32, Utf8, Date32, Decimal128, nullable	13 µs	234 µs	231 µs	333 µs
flat: 50 x Int32 (1.6 MB)	78 µs	250 µs	516 µs	213 µs
nested: 4 x struct depth 6 over (Int64, Utf8, Float64)	28 µs	85 µs	232 µs	107 µs

Compression is 75 to 95 percent of encode time at full batch size. Reusing the IpcWriteContext saves at most 13 percent with codec None and nothing measurable with a codec.

Fixed cost per block, 128-row mixed batch

Codec	Per block	Over codec None
None	1.5 µs
LZ4	5.9 µs	+4.4 µs
Snappy	7.7 µs	+6.2 µs
Zstd(1)	14.5 µs	+13 µs

This is encoder construction and teardown that every block pays. Blocks per task is roughly partitions x (spills + 1).

End to end, 81,920 rows of the mixed schema, codec None, hash partitioning

Partitions	Rows per partition	Time
200	410	3.6 ms
2000	41	11.9 ms
8000	10	26.7 ms

Same bytes, 7.4x the time. Everything between 200 and 8000 partitions is per-partition fixed cost.

Reader: decode one uncompressed block vs parsing its schema alone

Block	Full decode	Schema parse only	Share
5 columns, 64 rows	1.6 µs	0.8 µs	50%
5 columns, 8192 rows	8.5 µs	0.8 µs	9%
50 columns, 64 rows	11.8 µs	6.0 µs	51%
50 columns, 512 rows	15.5 µs	5.9 µs	38%
50 columns, 8192 rows	83.5 µs	5.9 µs	7%

JVM columnar shuffle, UnsafeRow to Arrow (row.rs), 10,000 rows

Shape	Time	Per row
100 primitive columns	3.09 ms	309 ns
struct, 20 flat fields	720 µs	72 ns
struct, 20 fields nested 3 deep	799 µs	80 ns
list, 100 elements	2.21 ms	221 ns
map<int64,int64>, 10 entries	640 µs	64 ns

These are the shapes that have fast paths. The bench does not cover strings, dictionary encoding, list, map<utf8, _> or nullable data, which are the shapes that take the per-row, per-element path.

Findings

Impact ratings are relative to the path they sit in.

Native local writer

W1. Fresh compressor per block (high). native/shuffle/src/writers/shuffle_block_writer.rs:276-301. zstd::Encoder::new creates a CCtx (several hundred KiB to over 1 MiB at level 1) and a 32 KiB output Vec per block; lz4_flex::FrameEncoder::with_frame_info zeroes a 16 KiB hash table and allocates 64 KiB+ src/dst per block; snappy is similar. Proposal: hold codec state per task next to the recycled scratch buffer. For zstd use a persistent zstd::bulk::Compressor over the materialized IPC bytes (output is still one standard zstd frame, which the reader already accepts). lz4_flex has no reusable frame encoder, so that needs an upstream change or a thin frame writer over block::compress_into. Also pin the LZ4 block size: the default resolves from the first write, which is the tiny schema message, so local LZ4 always uses 64 KiB blocks.

W2. Tail chunks are materialized twice through BatchCoalescer (high). buf_batch_writer.rs:120-141, partitioners/partitioned_batch_iterator.rs:117-132. The iterator emits full batch_size chunks plus one tail; full chunks bypass the coalescer but the tail is copy_rows'd and re-emitted. In the multi-partition path each partition gets a fresh BufBatchWriter per cycle and there is never anything to coalesce with. With 200 partitions and 8192-row inputs a partition sees about 41 rows per input batch, so the entire payload is gathered by interleave and then copied again. Proposal: passthrough mode for multi-partition writers; keep coalescing only for the single-partition output.

W3. Per-partition IPC context, flush and lseek (medium, high at large partition counts). buf_batch_writer.rs:71,202, local_partition_writer.rs:238,256, spill.rs:68. Each partition writer owns an IpcWriteContext::default(), so the first block per partition per cycle regrows the body scratch from empty. stream_position() on a BufWriter flushes and seeks, and BufBatchWriter::flush flushes again, so every partition leaves as its own write(2) plus an lseek. Proposal: thread a task-scoped &mut IpcWriteContext into BufBatchWriter; track offsets arithmetically and flush once in finish_all.

W4. One spill file and one file handle per partition (medium when spilling). local_partition_writer.rs:119-127,246-248, spill.rs:89,111-124. P spill writers per task, each with its own open temp file. Proposal: one spill file per spill event with (offset, len) per partition, segment copies at finish.

W5. Fixed framing overhead per block on wide and nested schemas (medium). shuffle_block_writer.rs:159,173-182. Every block carries the full schema message, a RecordBatch flatbuffer with 16 bytes per node and per buffer, and 64-byte alignment padding per buffer. A 100-column Int64 block of 40 rows is about 32 KiB of data and 15 to 25 KiB of overhead. Cheap step: drop IPC alignment to 8 or 16. Larger step: schema once per partition segment with a reader-side cache (pairs with R2, needs a format bump).

W6. Native round-robin hashes every column by default (medium on wide tables). multi_partition.rs:369-383, maxHashColumns default 0. Proposal: choose a cheap subset by type rather than position.

W7. Dictionary columns take the StreamWriter fallback and a concatenating interleave (low today, latent). shuffle_block_writer.rs:166-171,200-209. ScanExec unpacks dictionaries so scan-fed shuffles never see this. Proposal: unpack at the writer input (SchemaAlignExec already sits there) and delete the fallback.

W8. Extra body memcpy in the precoded-schema fast path; two divisions per row in pmod (low). shuffle_block_writer.rs:214-226, comet_partitioning.rs:51-57.

Reader and decoder

R1. Fresh decoder context per block and the decompressed body written two to three times (high). native/shuffle/src/ipc.rs:39-48,86-95. Each block constructs a new streaming decoder (zstd DCtx plus window buffers, about 1 MB of malloc/memset/free per block at level 1; lz4_flex src/dst up to twice the block size; snappy zero-fills about 140 KB). Then arrow's MessageReader memsets a body buffer and read_exacts into it, so the decompressor writes into its window and the bytes are copied again. The writer never pledges a content size. Proposal, step one: one decoder context per ShuffleScanExec and per remote decoder handle. Step two: record the uncompressed IPC length in the block header, decompress once into an exact-size MutableBuffer, decode with RecordBatchDecoder over that buffer (format change; Celeborn frames need versioning).

R2. Schema re-verified and re-parsed for every block (high for wide schemas and small blocks). ipc.rs:87, shuffle_scan.rs:170-228, jni_api.rs:1329-1352. The writer emits a byte-identical pre-encoded schema message per block; the reader verifies the flatbuffer and allocates a Schema per block. Proposal: cache (raw schema bytes, parsed SchemaRef) on the exec and on the remote decoder handle, memcmp per block, and go straight to RecordBatchDecoder on a hit. This also lets cast_and_stamp_schema decide once per schema instead of a recursive DataType compare per column per batch.

R3. Shuffle read feeding ColumnarToRow pays two full FFI round trips per batch (high where a stage ends in C2R). NativeBatchDecoderIterator.scala:124-143, jni_api.rs:1403-1421,851-878, NativeUtil.scala:187-214,259-292, NativeColumnarToRowConverter.scala:80-98. Decode natively, export N columns to Java, import in Arrow Java, export N columns back, re-import natively to make rows; recursive per child for nested types. Proposal: a fused decode-plus-C2R JNI entry; fallback improvement is exporting one struct-typed array per batch into a reused VectorSchemaRoot. Likely overlaps with #5112 / #5119.

R4. Decode is serialized with downstream compute and with the next block's fetch (medium to high, higher risk). jni_api.rs:889-899,967-1004, shuffle_scan.rs:128-150. Proposal: bounded decode-ahead queue inside native, with errors routed back to the task thread so on_decode_failure still has a JNIEnv.

R5. NativeBatchDecoderIterator reallocates its direct buffer for every map-output block (medium, JVM-consumer and non-direct paths). NativeBatchDecoderIterator.scala:207-214,253,264-272, CometBlockStoreShuffleReader.scala:106-110. The reader closes the decoder iterator per fetched block and close resets the thread-local buffer to 128 KB, so every block over 128 KB compressed triggers two allocateDirect calls (zero-fill plus Bits.reserveMemory, which can trigger System.gc()). Proposal: release in the task completion listener already registered in read().

R6. Smaller items (low to medium). Channels.newChannel(InputStream) stages compressed bytes through the JDK's 8 KB heap array (CometShuffleBlockIterator.java:61,122-129); dictionary strings from the JVM columnar shuffle are materialized per batch (shuffle_scan.rs:257-263); cast_and_stamp_schema does a recursive compare per batch (shuffle_scan.rs:385-390); four JNI calls per block where one would do (shuffle_scan.rs:172-226).

JVM columnar (row-based) writer

J1. Native child, JVM row shuffle: an avoidable Arrow to row to Arrow round trip (high for affected plans). CometShuffleExchangeExec.scala:136-139,335-404,413-474, operators.scala:393-394, CometConf.scala:434,458. Native shuffle is rejected when the child is a CometNativeExec but the partitioning is round-robin (off by default), hash on struct/array/map keys (off by default pending measurement of the nested hasher), an unsupported hash expression, or range partitioning on an unsupported type. repartition(n) over a native plan is the common case. Proposal: the native writer already accepts JVM Arrow batches through the synthetic ShuffleWriterInput scan; add a CometPartitioning::PrecomputedIds(column) variant and append an Int32 partition-id column computed by a codegen projection on the JVM, so Spark computes the id and semantics are unchanged.

J2. Pointer array growth is sized from data-page bytes, not pointer-array size (high on memory, one-line fix). CometShuffleExternalSorter.java:254-286, SpillSorter.java:168-181.

long used = activeSpillSorter.getMemoryUsage();   // pages + inMemSorter
array = allocator.allocateArray(used / 8 * 2);

Spark's ShuffleExternalSorter uses inMemSorter.getMemoryUsage() (the array alone). Comet's override adds all data pages, so the first growth after 2048 records requests 2 x (pageSize + 32 KB) / 8 entries: with a 64 MB page that is a 128 MB pointer array for 2048 rows, and later growths compound. Verified against current source.

J3. One JNI call per (spill, partition), each with heavy fixed cost (high at large partition counts). SpillSorter.java:248-317, CometDiskBlockWriter.java:322-349, jni_api.rs:1162-1179,1193-1272, row.rs:1274-1283,1371-1391. Per call: decode every column's protobuf DataType, open the spill file, allocate full-capacity builders including a 4096-slot dictionary table per string column, a fresh IpcWriteContext, and copy the address/size arrays. Proposal: one JNI call per spill that takes the sorted packed-pointer array and page table and returns partition lengths and checksums.

J4. Per-row serializer round trip, double copy into pages, boxing and a lock per row (medium to high on the insert loop). CometUnsafeShuffleWriter.java:332-346, CometDiskBlockWriter.java:208-240, SpillSorter.java:285-286,358-366, RowPartition.scala:24-49. UnsafeRow to DataOutputStream to heap array to Platform.copyMemory; native then skips the 4-byte length it never uses; record._1() boxes the partition id; the bypass writer computes getPartition twice, takes a synchronized per row and logs at INFO every 8192 rows; RowPartition boxes every address and size.

J5. Nested-type fallback in row.rs dispatches and allocates per element (medium to high for nested schemas). row.rs:211-366,646-689,748-797, list.rs:343-482, map.rs:70-122. Batch dispatch exists only for List and four Map<int,int> combinations. List, List, Struct containing List/Map, Map<Utf8,> and Map<,Struct> go per row with a type match and downcast per row (and per field per element for structs); append_map_elements allocates three Strings per row for MapFieldNames and discards them. Proposal: typed appender tree built once per column per batch; short term, hoist get_map_key_value_fields and add Utf8 key/value fast paths.

J6. Strings always go through a dictionary builder then are un-dictionaried; builder capacity is lost per batch; UTF-8 validated per value (medium on string-heavy data). row.rs:1141-1172,1274-1295,1371-1378,1412-1418,1443-1491, unsafe_object.rs:93-103. Proposal: decide dictionary vs plain per column adaptively from the first batch; reserve builders per batch; validate the whole values buffer once per batch with a per-string lossy fallback.

J7. Native radix sort does six fixed passes with 8x temporary allocation; the JVM reserves scratch it never uses (medium on large spills). native/core/src/execution/sort.rs:60-160, SpillSorter.java:127,219-225. Proposal: use Spark's ShuffleInMemorySorter.getSortedIterator() (skips constant-byte passes, uses the reserved half, removes a JNI call) or port that scheme.

Remote shuffle (Celeborn) writer

C1. Reservation is about 7x the uncompressed batch and held through the network round trip (high on multi-task executors). rss_partition_writer.rs:202-214,244-253, CelebornShufflePartitionPusher.java:695-715,815-830. After the push returns, the native Vec is dropped and the JNI byte[] is unreachable, yet releasableBytes keeps the full reservation until transport completion. With the 512 MiB default that is about 9 frames in flight per executor for 8 MiB batches. Proposal: shrink to frame length plus header on native release; charge encode scratch to the DataFusion pool.

C2. One frame per partition chunk per spill, no coalescing, doPush=true bypasses Celeborn's merge (high at large partition counts). rss_partition_writer.rs:137-141,346-347, CelebornShufflePartitionPusher.java:570-584. Proposal: stage small per-partition batches natively across write calls; pack several complete blocks for the same partition into one push.

C3. Five JNI crossings and two admission round trips per frame; schema re-serialized per block; fresh buffer and encoder per frame; four to_data() estimator passes; unconditional dictionary GC (medium, dominant for small frames). rss_partition_writer.rs:184-194,235-238,253-304,347-352,594-607,904-938, shuffle_block_writer.rs:125-129,164-165,200-209. A fixed 4 KiB metadata slack also causes the first encode of a wide-schema frame to overflow with codec None or incompressible data, dropping and re-encoding the batch.

C4. Completion polled every 10 ms with O(pending) reflective scans; java.lang.reflect.Proxy on the transport client's channel for every tenant of the client (medium). CelebornShufflePartitionPusher.java:781-788,801,851-917, CelebornTransportCallbackTracker.java:341-348,350-354,391-405,685-726,847-890.

Nested and complex types, consolidated

Already right: arrow-select's interleave recurses per struct field with the same indices and copies only referenced child ranges for lists and maps; IPC writing re-encodes offsets and slices children for List, Map, Utf8 and Binary. Neither needs a Comet-side fast path.

• Hash partitioning on nested keys is off by default, which pushes any such shuffle to the JVM row path (J1). The blocker named in the config doc is unmeasured performance of the nested murmur3 kernel. Its map path slices one entry at a time with per-entry dispatch (hash_funcs/utils.rs:1193-1215); list uses a gather per element position with an eligibility check on retained child size (utils.rs:516-761). Benchmarking those two shapes is cheap and would either flip the default or point at the map path as the one to vectorize.
• Row-to-Arrow for nested types is per-row and per-element for every shape except List and Map<int,int> (J5), and the bench does not cover the slow shapes.
• Per-block framing grows with nesting: schema message size and one node plus buffers per level, each buffer padded to 64 bytes (W5); the reader re-parses that schema per block (R2) and does a recursive DataType compare per batch (R6).
• Nested FFI export and import in the JVM-consumer reader path recurses per child, twice (R3).
• Dictionary columns take a re-encoding fallback on the writer (W7) and are always materialized on the reader (R6).
• RSS size estimation walks every nested child four times per frame and rebases List and Map offsets on every split (C3).

Suggested order

Ordered by expected win per unit of risk.

1. J2 pointer-array growth (one line, memory).
2. R1 step one and W1: reuse codec contexts on both sides. Add a compressed small-block case to the writer bench and codec cases to the reader bench first.
3. W2 and W3: passthrough for multi-partition writers, task-scoped IPC context, arithmetic offsets. The high-partition bench should drop toward the 200-partition figure.
4. R2 cached schema on the reader, folding in R6's schema reconciliation.
5. R5 stop resetting the decoder buffer per block.
6. J1 precomputed partition ids so native children never take the row path, plus a nested-hash benchmark to decide the default.
7. R3 fused decode-plus-C2R for stages that end in rows.
8. C1 to C3 for Celeborn users, starting with the reservation shrink.
9. J3, J5, W4, R4: larger restructures once the above have landed.

Benchmark coverage to add

• Writer: end-to-end wide schema (100+ columns), List/Struct/Map through interleave, a spill run with a memory limit, a compressed high-partition case.
• Reader: zstd/lz4/snappy decode; dictionary-encoded Utf8; List and Map; validated vs unvalidated; a cached-schema variant.
• Row-to-Arrow: strings, dictionary encoding, List, Map<Utf8,_>, nullable data.
• Hashing: struct, list and map keys through the native murmur3 kernel against a flat-key baseline.