bench: add a shuffle read benchmark covering the per-block schema parse - #5805
bench: add a shuffle read benchmark covering the per-block schema parse#5805peterxcli wants to merge 2 commits into
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>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
sunchao
left a comment
There was a problem hiding this comment.
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.
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_batchbuilds a freshStreamReaderper 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 inShuffleBlockWriter::try_newand 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
StreamReaderwithRecordBatchDecoder, which has to reproduce the existing guards for truncated LZ4, trailing data, multi-batch frames, and theskip_validationsplit. 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:
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:
It encodes through
ShuffleBlockWriterand decodes throughread_ipc_compressed, the same paths the write and read sides use, so it fails if either changes incompatibly.🤖 Generated with Claude Code