perf: decode shuffle blocks against a cached schema instead of re-parsing per block

[peterxcli]

Which issue does this PR close?

Closes #5792. Builds on #5805, which added the benchmark.

Rationale for this change

Every shuffle block is a self-contained Arrow IPC stream, so read_single_batch built a fresh StreamReader per block and parsed the schema flatbuffer once per block, even though every block in a shuffle carries the same schema and the reducer already knows it from the plan protobuf.

The win is not the parse, and not where #5792 predicted. The parse is worth about 2 us; the real cost is that StreamReader allocates a MutableBuffer::from_len_zeroed(bodyLength) per message and copies the body into it (reader.rs:1845), so it zero-fills and copies every block body. Decoding in place against a known schema skips that, and the saving scales with body size.

Measured on an idle 16-core x86_64 Linux host, alternating branches base/cached/base/cached so drift moves both together, criterion defaults (3s warmup, 5s measurement, 100 samples). Both runs of each branch shown, and parse_schema_only, which this change does not touch, held within 2 percent across all four rounds:

shape	base	cached	change
5 col x 64 row	4.31 us / 4.25 us	5.07 us / 4.76 us	+14.9%
5 col x 512 row	5.26 us / 5.17 us	5.44 us / 5.39 us	+3.8%
5 col x 8192 row	24.92 us / 25.88 us	18.17 us / 18.01 us	-28.8%
50 col x 64 row	38.78 us / 38.14 us	37.66 us / 37.53 us	-2.2%
50 col x 512 row	50.33 us / 50.48 us	46.44 us / 45.22 us	-9.1%
50 col x 8192 row	443.9 us / 405.0 us	302.6 us / 287.9 us	-30.4%

This is a trade, not a free win. Narrow schemas regress until blocks get large: 15 percent slower at 64 rows and still 4 percent slower at 512, crossing over to a 29 percent gain by 8192. Materializing the block and walking its messages is not repaid when the body is small. Wide schemas are neutral to positive throughout, since their bodies are already large enough.

The 8192 row rows are the block size a default 200 partition shuffle produces, which is the case this is aimed at. But a shuffle with many partitions and few rows each lands in the regressing region, and that is also the case #5792 originally expected to benefit.

If reviewers would rather not take the regression, the fast path can be gated on body size, which would keep the large-block gain and leave small blocks on the existing path. I have not implemented that here because it adds a threshold to tune; happy to if preferred.

An earlier revision of this description carried numbers from a loaded laptop that overstated the gains (up to -51.7%) and understated the small-block regression (+6.7%), and had the wrong sign on the 5 col x 512 row row. The table above replaces them.

What changes are included in this PR?

• A per-thread cache keyed on the raw schema message, so a hit costs one memcmp. It holds four schemas: a reduce task can interleave blocks from several shuffles, and a single entry would thrash.
• On a hit, the block is decoded in place with RecordBatchDecoder over the already-decompressed buffer, so arrays borrow it instead of being copied into a fresh zeroed buffer.
• On a miss, the original StreamReader path runs unchanged and its parsed schema is cached for later blocks.
• arrow-data becomes an explicit dependency for UnsafeFlag, which the trusted-local path needs to keep skipping validation. Already in the tree via arrow, so no build cost.

The fast path never reports an error of its own. A cache miss, a dictionary message, more than one record batch, trailing bytes after the end-of-stream marker, or a block that fails to decode all return None and fall back to the general decoder. Validation behaviour and every error message are unchanged, and the fast path is always safe to skip.

One subtlety: read_message reports both an explicit end-of-stream marker and a clean message boundary as "no more messages". Relying on that alone would have let trailing bytes after the marker pass, which the general decoder rejects. expect_end_of_stream checks for that specifically.

How are these changes tested?

Four new tests in ipc.rs, all exercising the warm-cache path the existing tests never reached:

• cached_schema_decode_matches_the_first_decode decodes each block twice across all four codecs and both entry points, asserting the warm decode equals the cold one and the original batch.
• dictionary_blocks_keep_decoding_with_a_warm_cache covers a schema that never takes the fast path.
• trailing_data_still_fails_with_a_warm_cache covers the end-of-stream gap above.
• truncated_block_fails_with_a_warm_cache checks a body-truncated block fails cold and warm, and that a stream ending on a message boundary without the marker stays valid, as before.

datafusion-comet-shuffle 129 passed, datafusion-comet --lib 333 passed, CometNativeShuffleSuite 53, CometShuffleSuite 44, clippy clean.

The benchmark gains a decode_block_uncached arm that clears the cache each iteration. It is not a stand-in for the base branch, since it also pays the cache insert every time, but it bounds what the cache is worth within a single run.

🤖 Generated with Claude Code

[andygrove]

The description on this one is unusually careful and I appreciate the retraction of the earlier numbers. I traced the fast path against StreamReader::next_ipc_message in arrow-ipc 59.2.0 and could not find an input that decodes differently. The cache key is the raw schema message bytes compared byte for byte and only Schema messages are inserted, so a hit cannot reuse a schema across blocks that do not share one. Dictionary blocks, multi-batch blocks, trailing bytes, truncation and a second schema message all fall back correctly. My concerns are all on the measurement and cost side.

First, PR Build (Linux) / ubuntu-latest/rust-test is red and is taking Required Checks down with it. The only annotation is exit code 100, which is what cargo nextest returns when a test fails rather than when the build breaks. That job was green on cc6b3d805 before the merge from main, and it is green on main itself, so the red appeared with the merge commit. Could you confirm whether that is the schema cache or an unrelated flake?

The benchmark only encodes with CompressionCodec::None, so every row of the results table uses a codec that is off by default. spark.comet.shuffle.compression.codec defaults to lz4. That matters more than usual because the change has a different shape per codec. On NONE the new code pays a fresh aligned allocation and a full block memcpy in Buffer::from(encoded) that the old path did not pay, which is probably a good part of the small-block regression you measured. On the compressed codecs the block is instead built by read_to_end into a growing Vec, so it reallocates and recopies as it doubles, where the old path had the decompressor write straight into the destination buffers. Could you add Lz4Frame to the matrix and re-run? Whether the small-block regression is worth taking really depends on what the default codec does.

None of the four new tests would fail if the fast path were removed. Changing try_decode_with_cached_schema to return None unconditionally leaves all of them passing, because each only checks that the batch decodes correctly or that the error cases still error, and the general decoder does both. Nothing proves the fast path ever runs, which is the entire change. Would a #[cfg(test)] counter bumped when the fast path returns a batch work? That would also let dictionary_blocks_keep_decoding_with_a_warm_cache assert the opposite, that the fast path was declined, and let the tests call reset_schema_cache so the cold and warm phases are explicit rather than depending on nextest giving each test its own process.

The fast path also calls root_as_message four times per block. read_message parses and verifies the schema message, then try_decode_with_cached_schema parses the same bytes again for header_type, and the record batch message gets the same treatment in read_message and decode_with_known_schema. root_as_message runs the verifier over the whole message, so on a wide schema this repeats most of the parse the cache exists to avoid. Given your numbers that looks like a plausible source of the regression. Could IpcMessage carry the Message it already parsed? On a cache hit the schema message does not need parsing at all, since a hit on the exact bytes already proves it is a schema message. I would rather close the regression that way than gate the fast path on a body-size threshold that then needs tuning.

Dictionary blocks can never take the fast path, because the message after the schema is a DictionaryBatch and expect_end_of_stream then rejects the record batch that follows. Comet's native shuffle does dictionary encode string and binary columns, which is why ShuffleBlockWriter has the SchemaEncoding::Fallback arm and ShuffleScanExec needs unpack_dictionary. Those blocks now pay full block materialization and a failed probe every time with none of the gain, and the benchmark builds only plain Int64 and Utf8 columns so it is not measured. Could you add a dictionary-encoded string column to the matrix? If it is a real regression, recording in the cache that a schema never takes the fast path would let the probe be skipped from the second block onward.

Last one is memory. Buffer::from_vec keeps the vector's allocation as it is, and read_to_end grows geometrically, so the buffer handed to the decoder can be close to twice the decompressed block. Every array in the batch is a slice of that buffer, so the whole allocation including spare capacity stays alive as long as any downstream operator holds the batch. The previous path allocated an exactly sized body buffer per message. There is a knock-on too, since Buffer::capacity returns the layout size and get_array_memory_size sums that per buffer, so what DataFusion reserves for a shuffle-read batch in a sort or join now tracks vector capacity rather than body length. Have you measured resident footprint across a shuffle read before and after? A shrink_to_fit costs a copy so it may not be the answer, but I would like to see the number before this lands.