peterxcli opened a new pull request, #5809: URL: https://github.com/apache/datafusion-comet/pull/5809
## Which issue does this PR close? Closes #5792. Builds on #5805, which added the benchmark; this PR is the implementation. ## 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 write side already avoids the mirror image of this: `ShuffleBlockWriter` encodes the schema once in `try_new` and writes the pre-encoded bytes verbatim. **The measured win is not where the issue predicted, and that is worth stating up front.** #5792 expected the gain at small blocks, where the constant per-block parse is the largest share of decode. It is the other way round. The parse is worth under a microsecond; 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 entirely, and the saving scales with body size. Comparing this commit against its parent back to back on the same machine, with `parse_schema_only` as a control arm that this change does not touch (it drifted within 5% between the two runs, so they are comparable): | shape | before | after | change | | --- | --- | --- | --- | | 5 col x 64 row | 1.663 us | 1.775 us | +6.7% | | 5 col x 512 row | 2.120 us | 1.913 us | -9.8% | | 5 col x 8192 row | 11.098 us | 7.841 us | -29.3% | | 50 col x 64 row | 13.479 us | 12.849 us | -4.7% | | 50 col x 512 row | 18.606 us | 16.090 us | -13.5% | | 50 col x 8192 row | 159.49 us | 77.03 us | -51.7% | The 8192 row rows are the ones that matter: that is the block size a default 200 partition shuffle produces. Small blocks are marginally slower, because materializing the block and walking its messages is not repaid when the body is tiny. If reviewers would rather not take that trade, the fast path could be gated on body size, though it adds a knob for a tenth of a microsecond. ## What changes are included in this PR? - A per-thread cache keyed on the raw schema message, so a hit costs one memcmp rather than a flatbuffer parse. It holds four schemas: a reduce task can interleave blocks from more than one shuffle (a join reading both sides), 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` is added as an explicit dependency for `UnsafeFlag`, which the trusted-local path needs to keep skipping validation. It is already in the tree via `arrow`, so this adds 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 simply fails to decode all return `None` and fall back to the general decoder. Validation behaviour and every error message are therefore unchanged, and the fast path is always safe to skip. This is the property that makes the change reviewable: the new code can only be faster or bypassed, never differently correct. One subtlety worth a reviewer's eye: `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 unnoticed, which the general decoder rejects. `expect_end_of_stream` checks for that specifically, and `trailing_data_still_fails_with_a_warm_cache` covers it. ## How are these changes tested? Four new tests in `ipc.rs`, all exercising the warm-cache path that 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. Runs on this branch: - `cargo test -p datafusion-comet-shuffle` — 129 passed - `cargo test -p datafusion-comet --lib` — 333 passed - `CometNativeShuffleSuite` — 53 passed - `CometShuffleSuite` — 44 passed - `cargo clippy --all-targets -D warnings` clean The benchmark gains a `decode_block_uncached` arm that clears the cache each iteration, so the cached and uncached paths can be compared inside one run where machine drift moves both together. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
