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