perf: cache the IPC schema across shuffle blocks on the reader by andygrove · Pull Request #5909 · apache/datafusion-comet · GitHub
Skip to content

perf: cache the IPC schema across shuffle blocks on the reader - #5909

Draft
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf/shuffle-reader-schema-cache
Draft

andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf/shuffle-reader-schema-cache

Conversation

@andygrove

Copy link
Copy Markdown
Member

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.

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%.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant