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: [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]
