Conversation
Every shuffle block is a self-contained Arrow IPC stream, so read_single_batch builds a fresh StreamReader per block and parses the schema flatbuffer once per block, even though every block in a shuffle carries the same schema. The write side already avoids the mirror image of this, encoding the schema once in ShuffleBlockWriter::try_new and writing the pre-encoded bytes verbatim, but there was no read-side benchmark to say whether the reader's half is worth removing. This adds one, parameterized by column count and rows per block, measuring the schema parse separately from the full block decode. On an M-series laptop: shape decode schema parse share 5 col x 64 row 1.93 us 1.14 us 59% 5 col x 512 row 2.38 us 0.91 us 38% 5 col x 8192 row 10.99 us 0.86 us 8% 50 col x 64 row 12.77 us 6.03 us 47% 50 col x 512 row 17.89 us 6.05 us 34% 50 col x 8192 row 218 us 6.05 us 3% The parse cost is constant per block and independent of row count, so its share is set by how many rows land in a block. That is largest exactly where the issue predicted: wide shuffles, where rows per partition are few, and repeated spilling, where each spill round emits its own block per partition. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…sing per block 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. The write side already avoids the mirror image of this, encoding the schema once in ShuffleBlockWriter::try_new and writing the pre-encoded bytes verbatim. Blocks are now decoded against a per-thread cache keyed on the raw schema message, so a hit costs one memcmp. On a hit the block is decoded in place with RecordBatchDecoder; on a miss the original StreamReader path runs unchanged and its parsed schema is cached for later blocks. The cache holds four schemas, since a reduce task can interleave blocks from more than one shuffle and a single entry would thrash. 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 fall back to the general decoder, so validation behaviour and every error message are unchanged and the fast path is always safe to skip. The measured win is not where apache#5792 predicted. Comparing this commit against its parent back to back, with the parse_schema_only arm as a control that this change does not touch (it drifted within 5% between the runs): 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 issue 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, while decoding in place avoids the per-body MutableBuffer that StreamReader allocates and zero-fills before copying into it, and that cost scales with body size. Small blocks are marginally slower, since materializing the block and walking its messages is not repaid when the body is tiny. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
enhancement
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…chema-cache Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
apache#5805 landed on main as a squash, so the benchmark it added conflicted with the original commits on this branch. The conflict was one-sided: resolved to main's file plus this branch's reset_schema_cache import and decode_block_uncached arm, with no main-only content dropped. The auto-merged manifests keep both main's DataFusion 55.1.0 and this branch's arrow-data dependency. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
The description on this one is unusually careful and I appreciate the retraction of the earlier numbers. I traced the fast path against StreamReader::next_ipc_message in arrow-ipc 59.2.0 and could not find an input that decodes differently. The cache key is the raw schema message bytes compared byte for byte and only Schema messages are inserted, so a hit cannot reuse a schema across blocks that do not share one. Dictionary blocks, multi-batch blocks, trailing bytes, truncation and a second schema message all fall back correctly. My concerns are all on the measurement and cost side.
First, PR Build (Linux) / ubuntu-latest/rust-test is red and is taking Required Checks down with it. The only annotation is exit code 100, which is what cargo nextest returns when a test fails rather than when the build breaks. That job was green on cc6b3d805 before the merge from main, and it is green on main itself, so the red appeared with the merge commit. Could you confirm whether that is the schema cache or an unrelated flake?
The benchmark only encodes with CompressionCodec::None, so every row of the results table uses a codec that is off by default. spark.comet.shuffle.compression.codec defaults to lz4. That matters more than usual because the change has a different shape per codec. On NONE the new code pays a fresh aligned allocation and a full block memcpy in Buffer::from(encoded) that the old path did not pay, which is probably a good part of the small-block regression you measured. On the compressed codecs the block is instead built by read_to_end into a growing Vec, so it reallocates and recopies as it doubles, where the old path had the decompressor write straight into the destination buffers. Could you add Lz4Frame to the matrix and re-run? Whether the small-block regression is worth taking really depends on what the default codec does.
None of the four new tests would fail if the fast path were removed. Changing try_decode_with_cached_schema to return None unconditionally leaves all of them passing, because each only checks that the batch decodes correctly or that the error cases still error, and the general decoder does both. Nothing proves the fast path ever runs, which is the entire change. Would a #[cfg(test)] counter bumped when the fast path returns a batch work? That would also let dictionary_blocks_keep_decoding_with_a_warm_cache assert the opposite, that the fast path was declined, and let the tests call reset_schema_cache so the cold and warm phases are explicit rather than depending on nextest giving each test its own process.
The fast path also calls root_as_message four times per block. read_message parses and verifies the schema message, then try_decode_with_cached_schema parses the same bytes again for header_type, and the record batch message gets the same treatment in read_message and decode_with_known_schema. root_as_message runs the verifier over the whole message, so on a wide schema this repeats most of the parse the cache exists to avoid. Given your numbers that looks like a plausible source of the regression. Could IpcMessage carry the Message it already parsed? On a cache hit the schema message does not need parsing at all, since a hit on the exact bytes already proves it is a schema message. I would rather close the regression that way than gate the fast path on a body-size threshold that then needs tuning.
Dictionary blocks can never take the fast path, because the message after the schema is a DictionaryBatch and expect_end_of_stream then rejects the record batch that follows. Comet's native shuffle does dictionary encode string and binary columns, which is why ShuffleBlockWriter has the SchemaEncoding::Fallback arm and ShuffleScanExec needs unpack_dictionary. Those blocks now pay full block materialization and a failed probe every time with none of the gain, and the benchmark builds only plain Int64 and Utf8 columns so it is not measured. Could you add a dictionary-encoded string column to the matrix? If it is a real regression, recording in the cache that a schema never takes the fast path would let the probe be skipped from the second block onward.
Last one is memory. Buffer::from_vec keeps the vector's allocation as it is, and read_to_end grows geometrically, so the buffer handed to the decoder can be close to twice the decompressed block. Every array in the batch is a slice of that buffer, so the whole allocation including spare capacity stays alive as long as any downstream operator holds the batch. The previous path allocated an exactly sized body buffer per message. There is a knock-on too, since Buffer::capacity returns the layout size and get_array_memory_size sums that per buffer, so what DataFusion reserves for a shuffle-read batch in a sort or join now tracks vector capacity rather than body length. Have you measured resident footprint across a shuffle read before and after? A shrink_to_fit costs a copy so it may not be the answer, but I would like to see the number before this lands.
|
I opened #5909 for the same problem before finding this PR, sorry for the overlap. It is now a draft, and this PR has priority since it came first. Posting the comparison here because the two take different routes to the same cache and the numbers suggest they are complementary rather than competing. Where the cost is. Both PRs agree the schema parse is a fixed cost per block. On my machine (Apple Silicon, codec None, the same What #5909 does differently. It keeps the streaming Relative change of a decoder held across blocks versus a fresh decoder per block, same branch, same machine:
The fresh-decoder numbers in #5909 are unchanged from Suggestion. Both PRs key the cache on the raw schema message bytes, so the cache itself is the same idea. The difference is what happens to the body: this PR materializes the block and decodes in place, #5909 streams it. The two stack naturally: the streaming loop for blocks under a size threshold, where materialization does not pay for itself, and your in-place decode above it. That would give the small-block gain without the regression you flagged and keep the large-block gain, which is the size gate you offered in the description. Happy to close #5909 once that lands, or to rebase #5909 on top of this if you would rather keep the in-place decode as the base. Whichever you prefer. One more thing worth a test if it is not already covered: the JVM columnar shuffle writes dictionary-encoded strings, so blocks on that path carry a dictionary batch before the record batch, and the dictionary must be scoped to its own block rather than to the cached schema. |
…ema without parsing it Replaces the materialize-then-probe fast path with one message loop that mirrors StreamReader. A cached schema is matched on its raw bytes and never verified or parsed again; the record batch and any dictionary batches are parsed once each. Bodies are read into exactly sized buffers, so a decoded batch reports the same memory as before, and dictionary blocks decode from the cache with dictionaries scoped to their own block. A #[cfg(test)] hit/miss counter proves which path each decode took; the tests reset the cache so cold and warm phases are explicit. The benchmark adds Lz4Frame, the default codec, and a dictionary-encoded string column. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Summary
Reviewed 0f61f3510f588c1a80969239b1e81571d1f269b0 against f69c4c81b9429e327ea95658530ae4ed4ed19635, including the existing discussion.
Prior state and problem
Each shuffle block contains a complete Arrow IPC stream. The previous reader built a new StreamReader per block, reparsed its schema, and allocated a zero-filled body buffer. Repeating that work can matter when a reduce task reads many blocks with the same schema.
Design approach and implementation
This change fully materializes the IPC payload, keeps up to four raw-schema-message/SchemaRef pairs per thread, and uses RecordBatchDecoder for a cache hit with exactly one record batch. Unsupported layouts and failed probes return to StreamReader. The wire format and writer are unchanged; the behavioral change is in reader allocation, parsing, and ownership.
I found no additional correctness defect in the paths inspected. The cache compares the complete schema-message bytes, so a different field type, nullability, nested layout, dictionary identity, or metadata cannot silently reuse a different schema. The cached state contains no batch data or dictionary values. Dictionary messages use a fresh general reader for each block, preserving block-scoped dictionaries.
The fast decoder uses the same batch metadata version and validation choice as Arrow's general reader. Remote reads still call read_ipc_compressed_validated, then validate logical types before normalization; the JNI failure callback remains in place. A schema previously cached by a trusted local read does not disable validation on a later remote read. Owned, reference-counted buffers also survive reuse of the JVM input buffer and schema-cache eviction.
The one-batch/end checks reject the fast path for additional messages, trailing data, or truncated bodies; errors then come from the general decoder. Legacy clean message-boundary termination remains accepted. This is a source-level assessment, not a new malformed-input execution run.
The existing request to prove fast-path execution is still relevant: output equality also passes if the optimized path always declines. Add explicit cold reset/hit assertions, dictionary-decline assertions, alternating schemas and eviction, and a warm-cache validated corrupt-array case before relying on these tests as optimization coverage.
Validation and disposition
The lockfile and native-build log resolve Arrow 59.3.0, despite the manifest's 59.2.0 lower bound. I checked the relevant 59.3.0 reader/buffer sources against the tagged upstream files.
The Rust job failed at cancelling_abort_keeps_the_guard_armed, whose source is unchanged by this PR. It ran 143 passing tests and one failure, then skipped 1,300 remaining tests because of fail-fast. That does not identify a shuffle failure, and it does not verify the new Rust shuffle tests. Resolve the failure and obtain a completed relevant run.
The Spark 3.5 shuffle job passed 442 tests overall: CometShuffleSuite lists 44 passes; CometNativeShuffleSuite lists 51 passes and six cancellations. This includes dictionary, direct-read and multiple-shuffle cases. Its checkout and the native-build checkout were 4e2296af, whose entire source tree equals the reviewed head.
COMMENT: the existing performance concern and validation requests remain; no duplicate inline finding. No local native build, benchmark, or new runtime probe was run for this review. Maintained Spark 3.4 and 4.1 source branches were unavailable, so no source-compatibility claim is made for those versions.
Performance
P2 — the existing whole-block materialization/fallback concern remains. I independently confirmed the mechanism described in the earlier review: NONE now copies the complete payload into a new Buffer; compressed blocks first grow a full decoded Vec. Dictionary blocks then always run StreamReader, which still allocates and copies its message bodies. They pay for the extra full-block buffer without taking the optimized decode path. Cold misses have the same additional materialization cost. Please address that fallback cost and qualify the resulting implementation before merging.
The checked-in benchmark only covers uncompressed, non-dictionary Int64/Utf8 blocks. LZ4 is the default. The PR reports small-block regressions of 14.9% and 3.8%, alongside larger-block gains; I have not reproduced those timings or established that their binaries came from this exact head. Clearing the cache on every iteration measures the new cold path, including insertion, rather than the previous implementation.
A concrete comparison should run the actual base and revised head alternately on the same idle host, with matching release settings and dependency locks: 5/50 columns × 64/512/8192 rows, at least NONE and LZ4, both plain and dictionary string/binary data, and both validation entry points. Assert decoded values and schema equality; record cache hits, allocation/copy volume, peak live memory, and retained capacity while batches remain live. Include interleaved schemas exceeding the four-entry cache. Report commit/binary identity with the results.
Design
Raw-byte schema identity and retaining Arrow as the fallback keep schema reuse easy to audit. Keeping dictionary values out of this cache is the right lifetime boundary. The full-block allocation strategy, however, couples schema caching to a separate memory tradeoff. A cached schema alone does not establish that this allocation strategy benefits ordinary compressed or dictionary shuffles.
On successful in-place decoding, array slices share the full block allocation, including spare Vec capacity; that allocation survives while a referencing array is retained. The previous reader also shared body buffers, so this is not a new per-column duplication of the entire batch. The additional capacity and metadata retention still need the memory measurement already requested in the existing review.
Abstraction & complexity
The helpers are localized, but the successful fast path currently verifies schema metadata twice and record-batch metadata twice. Carry the already parsed message through the helpers, as requested in the earlier review, to remove that repeated work before adding a size threshold. Keep one explicit fallback boundary and make the hit/decline behavior observable in tests; a second decoder path needs evidence that it is actually exercised.
Thanks for the careful read. Pushed a5290ca, which restructures the decoder rather than patching the probe, and reworked the description to match. Point by point: CI. The red Lz4Frame in the matrix. Added, alongside Tests that would survive removing the fast path. There is no separate fast path any more, but the equivalent gap is covered: under Four Dictionary blocks. They take the same loop: dictionary batches decode with Memory. The whole-block |
Thanks for the comparison, and for #5909. You were right that the streaming loop is the better base, so this revision (a5290ca) is built on it rather than stacked beside it: each block is walked message by message, the schema message is served from the cache on a byte match without being parsed, and bodies are read into exactly sized buffers. There is no whole-block materialization left to gate behind a threshold, so the "streaming below, in-place above" split turned out not to be needed. What remains of the earlier design: the cache is per thread and holds four schemas, so a reduce task that interleaves blocks from two shuffles does not thrash and no decoder handle has to be threaded through Re-measured against The dictionary case you flagged is covered by |
…ys on a warm cache, bench the validated entry point The RSS tests' allocation observer is shared with the reader tests, which compare one warm decode against the StreamReader path this change replaced: no more allocations, bytes or peak live memory on any codec. The corrupt offsets test now also fails validation with the schema served from the cache, and the benchmark times read_ipc_compressed_validated as well. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Thanks for the thorough pass. That review is against 0f61f35; a5290ca (pushed shortly before it landed) replaced the design it describes, and fe93ff2 adds what was still missing. Mapping the findings to the current head: P2, whole-block materialization and the dictionary fallback. Gone. The reader walks each block message by message the way Schema and record-batch metadata verified twice. Gone. A hit does not parse the schema message at all (length prefix, memcmp, move on); every other message is parsed once and the parsed Spare Tests that prove the optimized path runs. Under Allocation and copy volume, peak live memory, retained capacity. fe93ff2 shares the RSS tests' allocation observer with the reader tests. One warm decode, counted with the crate's test allocation observer (Rust allocations only; zstd's C-side workspace is outside it, which is why ZSTD reads like NONE). Old is the
The difference is the schema parse and the metadata buffer, about ten allocations and 1.2 KiB per block; the body allocation is identical. On LZ4 and Snappy the decompressors' own working buffers dominate both columns. Benchmark. Re-run on an idle 16-core x86_64 host against Validated decode (
CI. The failure was Not covered: dictionary binary columns are the same code path as dictionary strings and are not benchmarked separately, and interleaved schemas beyond the cache capacity are tested for correctness but not timed. |
3 participants
Which issue does this PR close?
Closes #5792. Builds on #5805, which added the benchmark. Takes the same route as #5909, which @andygrove opened for the same problem: this revision replaces the earlier materialize-and-probe fast path with a message loop of the kind #5909 proposed, keyed on the raw schema bytes, so the two are no longer competing designs.
Rationale for this change
Every shuffle block is a self-contained Arrow IPC stream, so the reader built a fresh
StreamReaderper block.StreamReader::try_newverifies the schema flatbuffer and allocates aSchemawith oneArc<Field>andStringper column every time, even thoughShuffleBlockWriterwrites the same pre-encoded schema message into every block and the reducer sees the same schema thousands of times. For the small blocks that high partition counts produce, that fixed cost is about half of the decode.The earlier revision of this PR materialized the whole block and probed it against the cache. Review found it re-parsed the schema message twice on a hit, built the block through a growing
Vecthat kept its spare capacity alive under every array, and made dictionary blocks pay for a probe that could never succeed. All three are gone.What changes are included in this PR?
read_ipc_compressedwalks each block message by message, the wayStreamReaderdoes: schema, any dictionary batches, one record batch, end of stream.StreamReaderis no longer used on this path.SchemaRef, four entries deep so a reduce task interleaving blocks from several shuffles does not thrash. A hit is one memcmp; the schema message is neither verified nor parsed. A miss parses it, exactly asStreamReaderdid, and caches the result.MutableBuffer::from_len_zeroed(bodyLength)StreamReaderused, for uncompressed blocks a copy of the body slice with no zero fill first. No growth slack, so a decoded batch reports the sameget_array_memory_sizeas before, and the whole-block materialization is gone.read_dictionary_impl, scoped to their own block, so the JVM columnar shuffle's dictionary-encoded strings decode from the cached schema too.skip_validationfor local ones. One malformed case is stricter than before on every codec: one to three stray bytes where a message length should be were previously swallowed as an end of stream, and are now an error.Benchmark
shuffle_readeron an idle 16-core x86_64 Linux host,mainat 481aefe versus this branch (the decoder as of a5290ca; fe93ff2 changes only tests and the benchmark source), alternating base/new/base/new with the same benchmark source on both sides (maingets a no-opreset_schema_cache). Criterion defaults.decode_blockholds the cache across iterations;decode_block_uncachedclears it every iteration, so onmainit equalsdecode_block.parse_schema_onlyis untouched by this change and serves as a control.Cached decode (
decode_block, the steady state of a shuffle read):Cache miss (
decode_block_uncached, cleared every iteration; onmainthis is the same code asdecode_block). This is the cost of the new reader with no help from the cache:Validated decode (
decode_block_validated, the remote entry point, schema from the cache):parse_schema_only, the control, moved -1.5% to +3.1% across the six shapes.No shape regresses. The narrow small blocks that the previous revision made 4 to 15 percent slower are now the largest wins, because a hit no longer verifies the schema flatbuffer at all. Under LZ4, the default, the gain shrinks with block size because decompression dominates: a 50 column, 8192 row block takes 3 ms to decompress and decode, of which the schema parse was 20 us. The large-block gain on uncompressed blocks, 40 percent at 50 col x 8192 rows, holds even on a cache miss: it comes from copying the body without zero-filling it first, which the compressed path cannot do because
Read::read_exactneeds an initialized destination.Memory
Bodies are exactly sized, so nothing changes in what a batch holds:
decoded_arrays_report_the_same_memory_size_as_stream_readerdecodes a 100,000-row block on every codec, cold and warm, and assertsget_array_memory_sizeequals a plainStreamReaderdecode of the same bytes.warm_decode_allocates_no_more_than_stream_readercompares one warm decode against theStreamReaderpath this PR replaced, on the same bytes in the same test, using the allocation observer the RSS writer tests already had, and asserts the new path never allocates more in count, bytes, or peak live memory. Rust allocations only; zstd's C-side workspace is outside the observer, which is why ZSTD reads like NONE:The difference is the schema parse and the metadata buffer, about ten allocations and 1.2 KiB per block; the body allocation is identical. On LZ4 and Snappy the decompressors' own working buffers dominate both columns.
I have not measured executor resident size across a Spark shuffle read; with per-body allocations identical to the previous path there is nothing left for it to show, but I am happy to run one if wanted.
How are these changes tested?
Under
cfg(test)the cache counts hits and misses, andreset_schema_cacheclears both, so each test states its cold and warm phases explicitly rather than relying onnextestprocess isolation:warm_decodes_hit_the_cache_and_match_the_cold_one: on every codec and both entry points, the second decode is a hit and equals the first and the original batch, for a plain schema and a dictionary schema.dictionaries_are_scoped_to_their_block_under_a_cached_schema: two blocks with the same schema but different dictionaries decode to their own values with the schema served from the cache, validated and not, on every codec.distinct_schemas_miss_once_and_recent_ones_stay_cached: alternating schemas stay resident; one more than the capacity evicts only the least recently used.decoded_arrays_report_the_same_memory_size_as_stream_readerandwarm_decode_allocates_no_more_than_stream_reader, as above.invalid_array_offsets_fail_validation_cold_and_warm: the corrupt-offsets block fails validation both when its schema is parsed and when it is served from the cache by an earlier valid block.oversized_metadata_length_is_an_error_and_releases_the_scratch: a forged 2 MiB metadata length fails cleanly on both the streamed and in-place readers and leaves no oversized scratch behind.partial_length_prefix_is_an_error: the stricter malformed case above, on every codec.trailing_data_still_fails_with_a_warm_cacheandtruncated_block_fails_with_a_warm_cachenow assert which path they failed on.datafusion-comet-shuffle134 passed,datafusion-comet --lib414 passed, clippy clean with-D warningson all targets.CometNativeShuffleSuiteandCometShuffleSuite(101 tests) pass on Spark 4.1 against the rebuilt native library.The red
rust-testjob on the earlier push failed iniceberg_write::tests::cancelling_abort_keeps_the_guard_armed, which this PR does not touch; #5919 makes that test deterministic. On a5290ca the fullrust-testjob completed green.🤖 Generated with Claude Code