{{ message }}
Conversation
ShuffleBlockWriter pre-encodes the IPC schema message once and writes it verbatim into every block, but the reader parsed it again per block through StreamReader::try_new, which for small blocks is about half of the decode. Add ShuffleBlockDecoder, which reads a block message by message with the same framing and checks as before but keeps the raw bytes of the last schema message next to the parsed schema. A byte-identical message reuses the schema; anything else parses and replaces the cache, so a block is always decoded against the schema it carries. ShuffleScanExec, the remote decoder handle and the handle-less local JNI entry point each hold a decoder. shuffle_reader bench, cached vs fresh: 5 and 50 columns at 64 rows -54%, 50 columns at 512 rows -39%, 50 columns at 8192 rows -7%.
andygrove
marked this pull request as draft
September 13, 2026 23:03
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

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.
ShuffleBlockWriterpre-encodes that message once and writes it verbatim into each block, but the reader went throughStreamReader::try_newper block, which verifies the flatbuffer and allocates aSchemawith oneArc<Field>andStringper column every time. For the small blocks that high partition counts produce, that fixed cost is about half of the decode:What changes are included in this PR?
ShuffleBlockDecoderindatafusion-comet-shufflethat decodes a block message by message (the same framingStreamReaderuses: schema, any dictionary batches, one record batch, end-of-stream), but keeps the raw bytes of the last schema message together with the parsedSchemaRef. 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 versusskip_validationfor local ones.ShuffleScanExecholds a decoder for the life of the scan (the native-consumer path). The remote decoder JNI handle holds one, and the handle-less localdecodeShuffleBlockentry point uses a thread-local decoder; both are safe because the byte comparison makes a stale cache miss rather than mis-decode.read_ipc_compressedandread_ipc_compressed_validatedremain as thin wrappers over a throwaway decoder for callers that decode a single block.arrow-datais added as a direct dependency forUnsafeFlag, whichRecordBatchDecoder::with_skip_validationtakes; it was already in the dependency tree.Benchmark
shuffle_readerbench, codec None, Apple Silicon.decode_blockuses 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_schemaholds the decoder across blocks.How are these changes tested?
schema_cache_hits_identical_messages_and_misses_different_onesdecodes 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.dictionary_blocks_decode_with_cached_schemacovers the dictionary-batch-before-record-batch layout the JVM columnar shuffle produces for strings, with dictionaries scoped per block.ipc.rsmalformed-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 coreshuffle_scantests and clippy pass.CometNativeShuffleSuiteandCometShuffleSuite(101 tests) pass against the rebuilt native library, covering theShuffleScanExecand thread-local decode paths end to end.