You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/13 21:02:28 UTC
[PR] perf: cache the IPC schema across shuffle blocks on the reader [datafusion-comet]
andygrove opened a new pull request, #5909:
URL: https://github.com/apache/datafusion-comet/pull/5909
## Which issue does this PR close?
Part of #5905 (finding R2). Does not close it.
## Rationale for this change
Every Comet shuffle block is a complete Arrow IPC stream, so every block starts with a schema message. `ShuffleBlockWriter` pre-encodes that message once and writes it verbatim into each block, but the reader went through `StreamReader::try_new` per block, which verifies the flatbuffer and allocates a `Schema` with one `Arc<Field>` and `String` per column every time. For the small blocks that high partition counts produce, that fixed cost is about half of the decode:
| Block (codec None) | Full decode | Schema parse alone |
|---|---:|---:|
| 5 columns, 64 rows | 1.6 µs | 0.8 µs |
| 50 columns, 64 rows | 11.8 µs | 6.0 µs |
| 50 columns, 512 rows | 15.5 µs | 5.9 µs |
| 50 columns, 8192 rows | 83.5 µs | 5.9 µs |
## What changes are included in this PR?
- A `ShuffleBlockDecoder` in `datafusion-comet-shuffle` that decodes a block message by message (the same framing `StreamReader` uses: schema, any dictionary batches, one record batch, end-of-stream), but keeps the raw bytes of the last schema message together with the parsed `SchemaRef`. Each block's schema message is compared against the cached bytes and served from the cache on a match; a mismatch parses and replaces the cache, so a block is always decoded against the schema it actually carries. All the existing checks are preserved: exactly one record batch per frame, no trailing bytes after the IPC stream or after the compressed stream, LZ4 end-mark enforcement, and full validation for remote blocks versus `skip_validation` for local ones.
- `ShuffleScanExec` holds a decoder for the life of the scan (the native-consumer path). The remote decoder JNI handle holds one, and the handle-less local `decodeShuffleBlock` entry point uses a thread-local decoder; both are safe because the byte comparison makes a stale cache miss rather than mis-decode.
- `read_ipc_compressed` and `read_ipc_compressed_validated` remain as thin wrappers over a throwaway decoder for callers that decode a single block.
- `arrow-data` is added as a direct dependency for `UnsafeFlag`, which `RecordBatchDecoder::with_skip_validation` takes; it was already in the dependency tree.
## Benchmark
`shuffle_reader` bench, codec None, Apple Silicon. `decode_block` uses a fresh decoder per block (what the old path did; the numbers are unchanged from before this PR, so the new message-level reader costs nothing on a miss). `decode_block_cached_schema` holds the decoder across blocks.
| Block | Fresh decoder | Cached schema | Change |
|---|---:|---:|---:|
| 5 columns, 64 rows | 1.62 µs | 0.74 µs | -54% |
| 5 columns, 512 rows | 2.43 µs | 1.62 µs | -33% |
| 5 columns, 8192 rows | 8.55 µs | 7.90 µs | -8% |
| 50 columns, 64 rows | 11.6 µs | 5.35 µs | -54% |
| 50 columns, 512 rows | 15.5 µs | 9.48 µs | -39% |
| 50 columns, 8192 rows | 83.5 µs | 77.6 µs | -7% |
## How are these changes tested?
- New unit test `schema_cache_hits_identical_messages_and_misses_different_ones` decodes repeated blocks, then a block with a different schema, then the first schema again, under every codec and with validation on and off, checking both the decoded batches and the parse counter.
- New unit test `dictionary_blocks_decode_with_cached_schema` covers the dictionary-batch-before-record-batch layout the JVM columnar shuffle produces for strings, with dictionaries scoped per block.
- All existing `ipc.rs` malformed-input tests (truncated codec tag, empty stream, multiple batches, trailing data, truncated LZ4 end mark, invalid offsets under validation) pass unchanged against the new reader.
- `datafusion-comet-shuffle` (127 tests), the core `shuffle_scan` tests and clippy pass.
- `CometNativeShuffleSuite` and `CometShuffleSuite` (101 tests) pass against the rebuilt native library, covering the `ShuffleScanExec` and thread-local decode paths end to end.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: cache the IPC schema across shuffle blocks on the reader [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #5909:
URL: https://github.com/apache/datafusion-comet/pull/5909#issuecomment-5832262263
Closing as superseded by #5809, which landed the same fix on main: blocks are decoded message by message against a schema cache keyed on the raw schema message bytes, so a hit is one memcmp and the schema is not re-verified or re-parsed. The cache lives per thread inside `read_ipc_compressed` and `read_ipc_compressed_validated`, so `ShuffleScanExec`, the local `decodeShuffleBlock` entry point and the remote decoder handle all use it without a decoder being passed through them, which was the rest of this PR. #6098 has since capped how much the cache holds per thread.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: cache the IPC schema across shuffle blocks on the reader [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove closed pull request #5909: perf: cache the IPC schema across shuffle blocks on the reader
URL: https://github.com/apache/datafusion-comet/pull/5909
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org