Skip to content

perf: validate shuffle IPC context reuse savings - #5727

Open
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:verify-shuffle-ipc-context-reuse
Open

perf: validate shuffle IPC context reuse savings#5727
peterxcli wants to merge 2 commits into
apache:mainfrom
peterxcli:verify-shuffle-ipc-context-reuse

Conversation

@peterxcli

@peterxcli peterxcli commented Sep 5, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5446.

Rationale for this change

Validate the IPC metadata allocation savings from #5038 after the Arrow upgrade in #5262.

What changes are included in this PR?

  • Add fresh-versus-reused context benchmarks across schemas, batch sizes, and codecs.
  • Add allocation regression coverage with byte equality and IPC round-trips. Production behavior is unchanged.

How are these changes tested?

119 native shuffle tests passed, along with the full native/JVM build and formatting checks.

Three benchmark passes on Apple M4 with Arrow 59.3.0 show:

Columns Allocation calls saved per block Requested bytes saved per block
4 9 1,064
50 12 8,232

128-row batches encode 2.0–16.8% faster; full-size compressed batches are approximately unchanged. These are block-encoding measurements, not end-to-end Spark results. RSS and dictionary fallback paths show no allocation savings.

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

Correctness

This validates context reuse that is already in the production shuffle writer after the Arrow upgrade. The PR adds a paired Criterion benchmark, extends the existing test-only allocator observer, and records the measurement methodology and results. It changes no production writer logic, compression settings, dependencies or context lifetime.

The allocation regression writes through the real local/RSS block encoder. Each arm warms its destination buffer, writes 100 identical batches, then compares the final frames byte-for-byte and decodes the reused frame back to the input batch. Those checks run outside allocation measurement. The observer counts successful allocation/reallocation requests on the synchronous encoding thread, and its reset guard clears observation on exit. I found no P1/P2 correctness issue.

Maintained Spark branch-3.5-openai at 5947fd6e74a1b2b04e4f83b7a659b02a9a2bac8b and branch-4.0-openai at 03f28fc4318024830a2ee8da7e83c42e0994d37a serialize records into partition streams and deserialize those streams back to records. This PR leaves Comet's production representation and behavior unchanged. Its IPC round-trips are native component checks, with no new Spark expression, ANSI-mode or null-semantics implementation. The allocation fixture is non-null Int32 data; the timing fixture also includes nullable strings, dates, decimals and nested structs. Maintained Spark 3.4/4.1 sources were unavailable for independent qualification.

Validation

Reviewed HEAD d16818ef6d87b653d948c9efa1eb49f4c02fb17a against authoritative BASE 0fc677962ab64f3884b95db8fb4b9723ef5870bf. The topic branches from 75fdddc9285ec61c0cd326977c61dd41fca39a8b. The Rust CI job checked out merge 936d42fde9d9823a62176e923ace01619919fc10, whose parents are the exact base/head. All three PR files and the lockfile match HEAD; the merge also contains the newer base fixes.

That job passed the new allocation regression and all 122 shuffle tests, with 1,168 tests passing overall and four skipped. It also passed all-target lint and benchmark compilation. The final current-head CI refresh at 2026-09-05 23:41:07 UTC reported 55 successful, ten running and seven skipped checks, with no failures. The separate benchmark check was skipped, and benchmark compilation does not execute Criterion. The author's 119-test release run and M4 measurements remain author-reported. I performed no local Rust/JNI execution or benchmark run.

Performance

The paired benchmark changes only context lifetime while retaining the same writer and warmed destination capacity. It covers 128/8,192 rows, mixed/wide/nested schemas and four codecs. Fixture construction and writer setup are outside timing. Fresh-context replacement and destruction are included in the fresh arm, which is appropriate for the per-block counterfactual.

The source supports the claimed mechanism: locked Arrow 59.3 copies completed metadata before resetting its retained FlatBufferBuilder. Scratch reservation is off by default and IPC buffer compression is disabled. Outer compressors are recreated for every block. RSS and dictionary paths construct their own StreamWriter context, so the equal-allocation negative controls exercise the expected unaffected paths. The observer is absent from the Criterion executable and adds no production overhead.

The report attributes 9 calls/1,064 requested bytes saved for four primitive columns and 12 calls/8,232 bytes for 50 columns. These are gross Rust allocation requests, excluding Zstd C allocations, rather than peak memory. The quieter-host timing report limits its conclusion to repeated warmed blocks: small-block gains, approximately neutral large compressed blocks, and no end-to-end Spark speedup. I checked the table arithmetic and methodology but did not independently reproduce those numeric results.

Design

Using the production encoder keeps both benchmark arms aligned with the code being evaluated. BufBatchWriter retains its context across emitted batches, and JVM sorted-row serialization retains one across its batch loop. The latter rebuilds its block writer per batch, a cost this benchmark explicitly excludes. First-use allocation, coalescing, partitioning and filesystem work remain outside these measurements, so they do not predict gains for a writer that emits only one block.

Abstraction & complexity

The change reuses the existing allocator observer and adds one focused benchmark group. Thread-local counters remain confined to tests, and correctness checks stay separate from allocation and timing measurements. No new production abstraction or shared mutable lifetime is introduced. The boolean fresh/reused comparison is small enough to follow directly without additional machinery.

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

Correctness

Rechecked 66226b03 after the PR description was shortened. Source is unchanged from the reviewed head. Relative to approved d16818ef, only the report file was deleted. The benchmark, allocation regression, production encoder and dependencies remain unchanged. I found no new P1/P2 issue.

The maintained Spark 3.5/4.0 serialization contract remains unaffected. No expression, ANSI-mode, type or null behavior changed. Maintained Spark 3.4/4.1 sources remain unavailable for independent qualification.

At the 2026-09-06 04:31:09 UTC CI check, there were 33 successful, 34 running and seven skipped checks, with no failures. The fetched merge ref has a different base parent from the assigned review base. I have no completed current-pair runtime result to credit and reran no local Rust/JNI tests or benchmarks for this description update.

Performance

The current description contains a compact summary and two allocation rows. It no longer contains the full methodology, commands or timing tables. The complete report at the previously approved commit remains publicly accessible, and the benchmark/test code is unchanged. That retained evidence is sufficient to inspect the reported comparison, so I do not see a new reproducibility issue requiring a change before merge.

The figures remain author-reported. They describe repeated warmed block encoding, excluding first-use context allocation and end-to-end Spark work. Both arms retain the writer and output capacity. Fresh-context replacement remains inside timing, while allocation correctness checks remain outside measurement.

Design

The description edit does not change context ownership, reset boundaries, dictionary fallback or allocator lifetime. The historical report linked above preserves the detailed measurement context despite its removal from the current tree and description.

Abstraction & complexity

No code or abstraction was added by this update. I found no new complexity concern.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf: verify shuffle IPC context-reuse allocation savings with Arrow 59.1+

2 participants