Skip to content

bench: add a shuffle read benchmark covering the per-block schema parse - #5805

Open
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:perf/shuffle-read-schema-parse-bench
Open

bench: add a shuffle read benchmark covering the per-block schema parse#5805
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:perf/shuffle-read-schema-parse-bench

Conversation

@peterxcli

@peterxcli peterxcli commented Sep 9, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #5792. This is the measurement that issue asks for; it does not change the decode path. The implementation is #5809.

Rationale for this change

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, encoding the schema once in ShuffleBlockWriter::try_new and writing the pre-encoded bytes verbatim.

#5198 recorded that this read-side item had no benchmark coverage, and #5792 was filed as unmeasured. Removing the parse means replacing StreamReader with RecordBatchDecoder, which has to reproduce the existing guards for truncated LZ4, trailing data, multi-batch frames, and the skip_validation split. Worth doing against a number rather than a hunch.

What changes are included in this PR?

native/shuffle/benches/shuffle_reader.rs, parameterized by column count and rows per block, measuring the schema parse separately from the full block decode. No production code changes.

Run on an idle 16-core x86_64 Linux host with criterion's default window (3s warmup, 5s measurement, 100 samples). Medians, with the confidence interval for each decode row:

shape decode schema parse share
5 col x 64 row 4.13 us [4.08, 4.18] 1.99 us 48%
5 col x 512 row 5.05 us [5.00, 5.11] 1.98 us 39%
5 col x 8192 row 24.30 us [24.19, 24.44] 1.97 us 8%
50 col x 64 row 39.07 us [38.76, 39.39] 19.62 us 50%
50 col x 512 row 51.11 us [50.81, 51.41] 19.49 us 38%
50 col x 8192 row 429.9 us [423.5, 437.2] 18.72 us 4%

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. At the 8192 row blocks a default 200 partition shuffle produces it is 4 to 8 percent.

How are these changes tested?

Benchmark-only, so there is no behaviour to test:

cargo bench --bench shuffle_reader -p datafusion-comet-shuffle

It encodes through ShuffleBlockWriter and decodes through read_ipc_compressed, the same paths the write and read sides use, so it fails if either changes incompatibly.

🤖 Generated with Claude Code

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>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@peterxcli
peterxcli marked this pull request as ready for review September 10, 2026 15:10

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed bf16d59703e3 against 5627ab8c0147. No verified P1/P2 findings.

Correctness

The existing reader constructs a new Arrow stream reader for each shuffle block. This PR adds a Criterion baseline for that repeated work: six combinations of 5/50 columns and 64/512/8192 rows, with a full-block decode and a schema-construction measurement for each. Production decoding, the wire format, and Spark expression/operator behavior are unchanged.

The fixture uses the production ShuffleBlockWriter. Removing the 16-byte outer header leaves the codec tag expected by read_ipc_compressed. The schema case then skips the four-byte NONE tag. In the pinned Arrow 59.3.0 implementation, StreamReader::try_new consumes the schema message without reading a record batch. Both measurements therefore start from the same encoded schema, and the full decode exercises the trusted local reader path, including its end-of-stream checks. The selected arrays are non-null Int64/Utf8. This benchmark does not add coverage for null, nested, dictionary, overflow, or error semantics.

For the Spark comparison, the maintained 3.5 and 4.0 branches both default shuffle compression to enabled and SQL shuffle partitions to 200. Comet defaults its codec to LZ4 and defines 8192 as a maximum batch row count. Those defaults do not establish the block-size distribution of a workload. Maintained 3.4/4.1 source was unavailable, so I do not claim source coverage for those versions.

The successful Rust CI job ran all-target Clippy, cargo check --benches, and nextest (1263 passed, 5 skipped). It tested merge commit a6f9a08f1641; I verified that its benchmark, shuffle reader/writer, dependency lock, and Rust action blobs match this head. The benchmark-specific job was skipped. I did not run a local benchmark or Spark/JVM test.

Performance

Fixture construction and encoding are outside the timed loops, and inputs/results are passed through black_box. Varying schema width and rows separately is useful for observing fixed schema work versus payload-dependent decode work.

The measurement includes more than the flatbuffer conversion: reader construction, schema ownership operations, and teardown are included. Criterion 0.7.0 iter also times destruction of the returned schema or record batch. These are useful lifecycle measurements, but their ratio is not a measured saving from a cached-schema implementation.

The reported medians remain author-provided. They describe repeatedly decoding an in-memory, uncompressed fixture. They do not measure compressed shuffle, remote validation, mapper fan-in, I/O, or end-to-end query improvement. Before using the ratios to justify a production optimization, a follow-up comparison should measure the proposed reader with representative codecs and observed block sizes.

Design

Using the existing writer and reader keeps the baseline tied to the actual format and avoids a second decoder implementation. Keeping the schema-caching implementation separate is appropriate: this change supplies a reproducible starting point without coupling a performance experiment to a wire-format or compatibility change. The public behavior added here is the new shuffle_reader benchmark target.

Abstraction & complexity

The schema, batch, and encoding helpers each serve a concrete fixture-building purpose, and the two timed closures remain easy to inspect. The explicit 16-byte and four-byte offsets match the current format and are documented at their use sites. No new production abstraction, dependency, or configuration is introduced; I found no actionable complexity issue in this scope.

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)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants