andygrove opened a new issue, #5905: URL: https://github.com/apache/datafusion-comet/issues/5905
## 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 53d87c86e (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. ## 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<int64>, 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<struct>, 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_exact`s 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`. ```java 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<primitive> and four Map<int,int> combinations. List<Struct>, 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<struct> 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<primitive> 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<Struct> and Map; validated vs unvalidated; a cached-schema variant. - Row-to-Arrow: strings, dictionary encoding, List<Struct>, Map<Utf8,_>, nullable data. - Hashing: struct, list<struct> and map keys through the native murmur3 kernel against a flat-key baseline. -- 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]
