You are viewing a plain text version of this content. The canonical link for it is here.

Posted to github@datafusion.apache.org by "peterxcli (via GitHub)" <gi...@apache.org> on 2026/09/09 17:16:25 UTC

[PR] perf: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

peterxcli opened a new pull request, #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809

## Which issue does this PR close?

Closes #5792. Builds on #5805, which added the benchmark; this PR is the implementation.

## Rationale for this change

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 and the reducer already knows it from the plan protobuf. The write side already avoids the mirror image of this: `ShuffleBlockWriter` encodes the schema once in `try_new` and writes the pre-encoded bytes verbatim.

**The measured win is not where the issue predicted, and that is worth stating up front.** #5792 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; the real cost is that `StreamReader` allocates a `MutableBuffer::from_len_zeroed(bodyLength)` per message and copies the body into it (`reader.rs:1845`), so it zero-fills and copies every block body. Decoding in place against a known schema skips that entirely, and the saving scales with body size.

Comparing this commit against its parent back to back on the same machine, with `parse_schema_only` as a control arm that this change does not touch (it drifted within 5% between the two runs, so they are comparable):

| 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 8192 row rows are the ones that matter: that is the block size a default 200 partition shuffle produces. Small blocks are marginally slower, because materializing the block and walking its messages is not repaid when the body is tiny. If reviewers would rather not take that trade, the fast path could be gated on body size, though it adds a knob for a tenth of a microsecond.

## What changes are included in this PR?

- A per-thread cache keyed on the raw schema message, so a hit costs one memcmp rather than a flatbuffer parse. It holds four schemas: a reduce task can interleave blocks from more than one shuffle (a join reading both sides), and a single entry would thrash.
- On a hit, the block is decoded in place with `RecordBatchDecoder` over the already-decompressed buffer, so arrays borrow it instead of being copied into a fresh zeroed buffer.
- On a miss, the original `StreamReader` path runs unchanged and its parsed schema is cached for later blocks.
- `arrow-data` is added as an explicit dependency for `UnsafeFlag`, which the trusted-local path needs to keep skipping validation. It is already in the tree via `arrow`, so this adds no build cost.

**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 return `None` and fall back to the general decoder. Validation behaviour and every error message are therefore unchanged, and the fast path is always safe to skip. This is the property that makes the change reviewable: the new code can only be faster or bypassed, never differently correct.

One subtlety worth a reviewer's eye: `read_message` reports both an explicit end-of-stream marker and a clean message boundary as "no more messages". Relying on that alone would have let trailing bytes after the marker pass unnoticed, which the general decoder rejects. `expect_end_of_stream` checks for that specifically, and `trailing_data_still_fails_with_a_warm_cache` covers it.

## How are these changes tested?

Four new tests in `ipc.rs`, all exercising the warm-cache path that the existing tests never reached:

- `cached_schema_decode_matches_the_first_decode` decodes each block twice across all four codecs and both entry points, asserting the warm decode equals the cold one and the original batch.
- `dictionary_blocks_keep_decoding_with_a_warm_cache` covers a schema that never takes the fast path.
- `trailing_data_still_fails_with_a_warm_cache` covers the end-of-stream gap above.
- `truncated_block_fails_with_a_warm_cache` checks a body-truncated block fails cold and warm, and that a stream ending on a message boundary without the marker stays valid, as before.

Runs on this branch:

- `cargo test -p datafusion-comet-shuffle` — 129 passed
- `cargo test -p datafusion-comet --lib` — 333 passed
- `CometNativeShuffleSuite` — 53 passed
- `CometShuffleSuite` — 44 passed
- `cargo clippy --all-targets -D warnings` clean

The benchmark gains a `decode_block_uncached` arm that clears the cache each iteration, so the cached and uncached paths can be compared inside one run where machine drift moves both together.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

--
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: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.

andygrove commented on PR #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809#issuecomment-5656824227

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 `shuffle_reader` bench shapes) `parse_schema_only` is 0.8 µs at 5 columns and 6 µs at 50 columns, which is about half of a 64-row block's decode and under 10 percent of an 8192-row block's. So the parse matters most for small blocks, and the body zero-fill plus copy that this PR removes matters most for large ones.

**What #5909 does differently.** It keeps the streaming `Read`-based path and replaces `StreamReader` with a message-level loop over the same framing (schema, dictionary batches, one record batch, end of stream, plus the existing trailing-data and single-batch checks). The decoder keeps the raw bytes of the last schema message next to the parsed `SchemaRef`; each block's schema message is compared with `==` on the bytes and reused on a match, parsed and cached otherwise. It does not materialize the block, so it never pays for a block the old path did not pay for. The body is still read into a `MutableBuffer` as before, so it does not get the large-block win this PR measures.

Relative change of a decoder held across blocks versus a fresh decoder per block, same branch, same machine:

| shape | #5909 cached vs fresh | this PR (from the description) |
| --- | --- | --- |
| 5 col x 64 row | -54% | +15% |
| 5 col x 512 row | -33% | +4% |
| 5 col x 8192 row | -8% | -29% |
| 50 col x 64 row | -54% | -2% |
| 50 col x 512 row | -39% | -9% |
| 50 col x 8192 row | -7% | -30% |

The fresh-decoder numbers in #5909 are unchanged from `main`, so the message loop is free on a cache miss.

**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.

--
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

Re: [PR] perf: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

Posted by "peterxcli (via GitHub)" <gi...@apache.org>.

peterxcli commented on PR #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809#issuecomment-5677467957

   > Replying to https://github.com/apache/datafusion-comet/pull/5809#pullrequestreview-5204781228
   
   Thanks for the thorough pass. That review is against 0f61f3510; a5290ca76 (pushed shortly before it landed) replaced the design it describes, and fe93ff220 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 `StreamReader` does; there is no full-block `Buffer`, no growing `Vec`, and no fallback decoder. Dictionary batches decode in the same loop with `read_dictionary_impl`, scoped to their block, from the cached schema. Cold misses parse the schema exactly as `StreamReader` did and read the same exactly sized bodies.
   
   **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 `Message` is carried through the decode.
   
   **Spare `Vec` capacity under live arrays.** Gone with the materialization. Compressed bodies use `MutableBuffer::from_len_zeroed(bodyLength)` as before; uncompressed bodies are an exactly sized copy. `decoded_arrays_report_the_same_memory_size_as_stream_reader` asserts `get_array_memory_size` equals a plain `StreamReader` decode on every codec.
   
   **Tests that prove the optimized path runs.** Under `cfg(test)` the cache counts hits and misses and `reset_schema_cache` clears both, so every cache test states its cold and warm phases and asserts exact counts: cold reset then hit, dictionary blocks served from the cache with per-block dictionaries, alternating schemas staying resident, eviction past the four-entry capacity, and (fe93ff220) the corrupt-offsets block still failing validation when the schema is served from the cache by an earlier valid block.
   
   **Allocation and copy volume, peak live memory, retained capacity.** fe93ff220 shares the RSS tests' allocation observer with the reader tests. `warm_decode_allocates_no_more_than_stream_reader` runs the old `StreamReader` path and the cached decode on the same bytes and asserts the new one never allocates more, in count, bytes, or peak:
   
   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 `StreamReader` path this PR replaced, run in the same test on the same bytes:
   
   | block | codec | old: allocations / bytes / peak live | this PR: allocations / bytes / peak live |
   | --- | --- | --- | --- |
   | 3 rows, 3 columns | none | 22 / 2,451 / 1,675 | 12 / 1,248 / 920 |
   | 3 rows, 3 columns | lz4 | 24 / 133,523 / 132,747 | 14 / 132,320 / 131,992 |
   | 3 rows, 3 columns | zstd | 22 / 2,451 / 1,675 | 12 / 1,248 / 920 |
   | 3 rows, 3 columns | snappy | 24 / 144,477 / 143,701 | 14 / 143,274 / 142,946 |
   | 8192 rows, 2 columns | none | 18 / 150,402 / 149,554 | 9 / 149,032 / 148,816 |
   | 8192 rows, 2 columns | lz4 | 20 / 674,690 / 673,842 | 11 / 673,320 / 673,104 |
   | 8192 rows, 2 columns | zstd | 18 / 150,402 / 149,554 | 9 / 149,032 / 148,816 |
   | 8192 rows, 2 columns | snappy | 20 / 292,428 / 291,580 | 11 / 291,058 / 290,842 |
   
   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 `main` at 481aefea9, base and head alternated over two rounds, same benchmark source on both sides, commit identity in the description. `None` and `Lz4Frame`, 5 and 50 columns by 64, 512 and 8192 rows, plain and dictionary strings, cached and cache-cleared arms, and the untouched `parse_schema_only` control (within 3%). No shape regresses; the tables are in the description. fe93ff220 adds `decode_block_validated` for the remote entry point:
   
   Validated decode (`decode_block_validated`, the remote entry point, schema from the cache):
   
   | shape | codec | base (two rounds) | this PR (two rounds) | change |
   | --- | --- | --- | --- | --- |
   | 5 col x 64 row | lz4 | 9.71 us / 9.15 us | 6.56 us / 6.72 us | -29.6% |
   | 5 col x 512 row | lz4 | 25.31 us / 26.36 us | 21.18 us / 20.84 us | -18.7% |
   | 5 col x 8192 row | lz4 | 341.45 us / 341.73 us | 335.25 us / 335.59 us | -1.8% |
   | 50 col x 64 row | lz4 | 72.82 us / 72.26 us | 45.98 us / 47.36 us | -35.7% |
   | 50 col x 512 row | lz4 | 244.81 us / 247.56 us | 191.44 us / 198.29 us | -20.8% |
   | 50 col x 8192 row | lz4 | 3.76 ms / 3.83 ms | 3.73 ms / 3.68 ms | -2.3% |
   | 5 col x 64 row (dictionary strings) | lz4 | 14.60 us / 14.51 us | 10.04 us / 9.77 us | -31.9% |
   | 5 col x 8192 row (dictionary strings) | lz4 | 283.60 us / 281.65 us | 265.88 us / 269.28 us | -5.3% |
   | 5 col x 64 row | none | 5.09 us / 5.03 us | 2.98 us / 2.90 us | -41.9% |
   | 5 col x 512 row | none | 8.44 us / 8.46 us | 5.90 us / 5.93 us | -30.0% |
   | 5 col x 8192 row | none | 67.79 us / 66.41 us | 57.49 us / 59.16 us | -13.1% |
   | 50 col x 64 row | none | 48.92 us / 47.37 us | 26.70 us / 26.73 us | -44.5% |
   | 50 col x 512 row | none | 89.48 us / 88.57 us | 64.11 us / 66.84 us | -26.5% |
   | 50 col x 8192 row | none | 1.08 ms / 1.11 ms | 905.54 us / 884.60 us | -18.3% |
   | 5 col x 64 row (dictionary strings) | none | 9.24 us / 9.49 us | 6.65 us / 5.97 us | -32.6% |
   | 5 col x 8192 row (dictionary strings) | none | 40.48 us / 38.39 us | 30.70 us / 30.27 us | -22.7% |
   
   **CI.** The failure was `iceberg_write::tests::cancelling_abort_keeps_the_guard_armed`, untouched by this PR and fixed by #5919. On a5290ca76 the full `rust-test` job completed green with the new shuffle tests included.
   
   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.
   


-- 
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: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

Posted by "peterxcli (via GitHub)" <gi...@apache.org>.

peterxcli commented on PR #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809#issuecomment-5676938081

> Replying to https://github.com/apache/datafusion-comet/pull/5809#pullrequestreview-5191026120

Thanks for the careful read. Pushed a5290ca76, which restructures the decoder rather than patching the probe, and reworked the description to match. Point by point:

**CI.** The red `rust-test` is `execution::operators::iceberg_write::tests::cancelling_abort_keeps_the_guard_armed` panicking on "expected the deletes to yield so the abort can be cancelled mid-flight". It is scheduling-sensitive and nothing in this PR touches `iceberg_write`; it went red on the merge commit because that is the run it happened to land on. Your #5919 is the fix. The new push re-runs the job.

**Lz4Frame in the matrix.** Added, alongside `None`, and re-run on the same idle 16-core host with base/new alternated. Table in the description. Under LZ4 the cached decode is 30 to 40 percent faster at 64 rows, 20 to 24 percent at 512, and 2 to 5 percent at 8192 rows where decompression dominates. The miss path is flat to slightly faster on every LZ4 shape, so the cache costs nothing on the default codec. The `read_to_end` growth you described went with the materialization: compressed bodies are now read straight into an exactly sized buffer, as before this PR.

**Tests that would survive removing the fast path.** There is no separate fast path any more, but the equivalent gap is covered: under `cfg(test)` the cache counts hits and misses, `reset_schema_cache` clears both, and every cache test now asserts the exact counts it expects, so a decode that silently re-parsed the schema would fail `warm_decodes_hit_the_cache_and_match_the_cold_one`. The dictionary test asserts a hit rather than a declined probe, because dictionary blocks now decode from the cache too (below), and it checks that each block's record batch is decoded against that block's dictionary and not the previous block's.

**Four `root_as_message` calls per block.** On a hit the schema message is now not parsed at all: the reader locates it by its length prefix, memcmps it against the cache, and moves on. The record batch message is parsed once, dictionary messages once each. `IpcMessage` is gone; the loop keeps the parsed `Message` for the duration of the message it describes. This is what turned the small-block regression into a gain, so no body-size threshold.

**Dictionary blocks.** They take the same loop: dictionary batches decode with `read_dictionary_impl` into a map scoped to the block, then the record batch decodes against it. No failed probe and no materialization. The benchmark gained `5col_64row_dict` and `5col_8192row_dict` arms on both codecs. Cached dictionary blocks are 38 percent faster uncompressed and 28 percent under LZ4 at 64 rows, 39 and 4.5 percent at 8192 rows. On a miss the 64-row shapes are within noise, so no probe cost remains.

**Memory.** The whole-block `read_to_end` is gone. Compressed bodies use the same `MutableBuffer::from_len_zeroed(bodyLength)` `StreamReader` used; uncompressed bodies are an exactly sized copy of the body slice without the zero fill. `decoded_arrays_report_the_same_memory_size_as_stream_reader` asserts `get_array_memory_size` equals a plain `StreamReader` decode on every codec, cold and warm, so what DataFusion reserves for a shuffle-read batch is unchanged. I also bounded the per-thread metadata scratch: a corrupt length that grows it past 1 MiB releases it after the block. I have not measured executor RSS across a Spark shuffle read; with per-body allocations identical to the previous path I do not expect it to show anything, but I can run one if you would like the number.

--
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

Re: [PR] perf: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

Posted by "peterxcli (via GitHub)" <gi...@apache.org>.

peterxcli commented on PR #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809#issuecomment-5676938442

   > Replying to https://github.com/apache/datafusion-comet/pull/5809#issuecomment-5656824227
   
   Thanks for the comparison, and for #5909. You were right that the streaming loop is the better base, so this revision (a5290ca76) 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 `ShuffleScanExec` or the JNI entry points; uncompressed blocks are walked in place, so their bodies are copied once without a zero fill; and the metadata scratch is reused across blocks.
   
   Re-measured against `main` on the idle host, base/new alternated, `None` and `Lz4Frame`, plain and dictionary-encoded strings. Table in the description. Against `main`, 5 col x 64 row is now -53% uncompressed and -30% under LZ4, and 50 col x 64 row -55% and -40%, in line with the -54% you measured for #5909's cached decoder on the uncompressed shape. The large-block gain from the earlier revision survives on uncompressed blocks (-40% at 50 col x 8192 rows, even on a cache miss, from copying the body without a zero fill) and is small under LZ4 (-2% to -5%) where decompression dominates.
   
   The dictionary case you flagged is covered by `dictionaries_are_scoped_to_their_block_under_a_cached_schema`: two blocks that share a schema but carry different dictionaries decode to their own values with the schema served from the cache, on every codec, validated and not.
   


-- 
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: decode shuffle blocks against a cached schema instead of re-parsing per block [datafusion-comet]

Posted by "sunchao (via GitHub)" <gi...@apache.org>.

sunchao merged PR #5809:
URL: https://github.com/apache/datafusion-comet/pull/5809


-- 
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