[andygrove]

Which issue does this PR close?

Part of #6941.

Rationale for this change

Arrow's offsets buffers are running sums, so they barely compress: seven short string columns carry 28 bytes per row of offsets, and LZ4 keeps all of them while zstd keeps about 16 (#5678 has the investigation). Lengths compress far better. Stored as lengths in byte planes, the offsets of a block mostly turn into a few long runs.

With zstd that makes shuffle files with string, binary, array and map columns 9% to 40% smaller, and zstd gets no slower, because it has less to do. LZ4 is left out for now: lz4_flex decodes the long runs that lengths compress into two bytes at a time, which makes LZ4 decoding much slower (PSeitz/lz4_flex#227 would fix that). Snappy also spends more time than it saves.

shuffle_bench, one map task, 16 partitions, 8192-row batches. CPU is per block on real shuffle blocks, the minimum of 21 interleaved runs on an Apple M3 Max.

Data	zstd bytes per row, before → after	Encode	Decode
Short IDs, fixed length	31.2 → 25.3 (−19%)	−2%	+1%
Short IDs, variable length	40.5 → 26.4 (−35%)
Vocabulary strings	48.1 → 31.0 (−36%)	−13%	−15%
Random strings	68.8 → 49.8 (−28%)
TPC-H SF1 lineitem	39.9 → 31.0 (−22%)	−4%	−3%
Structs nested 6 deep	31.5 → 20.8 (−34%)	−3%	−10%
One Int64 column	4.6 → 4.6 (0%)	0%	0%

From 16 to 4,000 partitions the savings range from 9% to 40%. Spark 4.1 with zstd writes 28.1, 26.4, 39.9, 69.1, 47.5 and 37.6 bytes per row for the first six, so at 16 partitions Comet's zstd shuffle is now at or below Spark's on all of them. At 4,000 partitions the two short-ID datasets stay 6% to 9% above Spark.

What changes are included in this PR?

• native/shuffle/src/offsets_as_lengths.rs (new): which IPC body buffer of each type holds offsets, mirroring arrow-ipc's write_array_data, and the in-place rewrite between offsets and lengths in byte planes. The first entry is kept, so any offsets round-trip, including ones that don't start at zero.
• ShuffleBlockWriter::try_new_with_offsets_as_lengths builds a writer that, with zstd, rewrites each encoded body on the dictionary-free local path. The writer decides when it is built: it checks once, on an empty batch, that arrow-ipc writes the buffers the mapped roles describe (a type's buffer count doesn't depend on its data), keeps plain offsets if not, and otherwise sets the high bit of its codec tag's first byte. The existing tags are ASCII, and the JVM never reads the tag.
• read_ipc_compressed_impl, which every read path goes through, rebuilds the offsets of flagged blocks under any codec, so turning lengths on for LZ4 later only changes the writer. It bounds-checks every buffer it touches, and rejects a flagged block whose schema has no mapped layout.
• spark.comet.shuffle.compression.offsetsAsLengths.enabled (default true, effective only with zstd), passed through the ShuffleWriter protobuf to native shuffle and through writeSortedFileNative to the JVM columnar shuffle.
• Unchanged byte for byte: Celeborn writers, dictionary schemas (including JVM shuffle blocks in which spark.comet.shuffle.jvm.preferDictionary.ratio dictionary-encodes a column), schemas without offsets, and types whose layout isn't mapped (views, unions, run-end encoding).
• shuffle_bench --offsets-as-lengths, a paragraph in the shuffle tuning guide, the codec-tag flag and the new file in native_shuffle.md, and the config in the native and JVM shuffle config tables.

How are these changes tested?

• Native unit tests:
  • Buffer roles match arrow-ipc's layout for every mapped type.
  • The rewrite round-trips any bit pattern, leaves other buffers and trailing bytes alone, and rejects out-of-range buffers.
  • Blocks of every mapped type round-trip under all four codecs, on both the trusted and the validating decode. The data includes nulls, empty values, and sliced inputs whose offsets don't start at zero.
  • Lengths are used only when asked for, and only with zstd.
  • The once-per-writer layout check agrees with arrow-ipc for every mapped type and rejects a role list that doesn't match.
  • Blocks without offsets, dictionary blocks and Celeborn blocks keep the plain tag.
  • A flagged block with an unmapped layout is an error.
• New tests in CometNativeShuffleSuite and CometColumnarShuffleSuite (sharing cometShuffleMetric and rowsWithOffsets helpers in CometTestBase), the latter run through CometShuffleSuite and DisableAQECometShuffleSuite. Their rows have offsets in string, binary, array, map and struct columns, with nulls, empty values and null array elements. They check answers against Spark, check that zstd writes fewer shuffle bytes with the setting on, and check that LZ4 writes identical bytes either way. They also check the answer with zstd while the writer spills, under a 64k spark.comet.shuffle.native.maxBufferBytes for native shuffle and at the suite's 10-row spill threshold for the columnar one. The columnar test covers both the bypass-merge and the sort-based writer.
• Run locally, 210 tests in total, all passing: CometNativeShuffleSuite, CometShuffleSuite, DisableAQECometShuffleSuite, CometShuffleEncryptionSuite, CometShuffleInputConversionSuite, SpillSorterSuite, CometDiskBlockWriterSuite, plus the shuffle crate's tests and core's shuffle tests. The 210 also all pass with zstd as the default codec, and a reader that ignores the flag fails 107 of them there.
• Before narrowing this to zstd, I checked the same rewrite with every codec on 7 datasets. For 112 combinations of dataset, codec and partition count, files written with and without it decoded to identical batches, 496M rows in all. With the rewrite enabled for all codecs, a reader that ignored the flag failed 26 of 73 CometNativeShuffleSuite tests, so the suites do exercise the new blocks.

[viirya]

Thanks for this, and for the careful write-up in #6941. The design holds up well. The rewrite is a lossless in-place transform on byte ranges, and the writer and reader share buffer_roles, so a layout mismatch after a future arrow upgrade can only cost compression, never correctness. A buffer-count mismatch falls back to the plain format. I checked push_roles against write_array_data in arrow-ipc 59.3 and the order and type predicates match. Keeping the reader codec-agnostic, so enabling LZ4 later is a writer-only change, is a nice touch.

Two documentation points, one inline and one here because the files aren't in the diff:

This changes the on-disk block format, but docs/source/contributor-guide/native_shuffle.md still describes the old one. The Encoding step of the write path and the block layout description don't say that the high bit of the codec tag's first byte now marks blocks whose offsets buffers hold lengths. Could you add a sentence there? Could you also add spark.comet.shuffle.compression.offsetsAsLengths.enabled to its Configuration table next to compression.zstd.level, and to the config table in jvm_shuffle.md? Both tables are hand-maintained, unlike configs.md.

The Lint Java (Spark 4.2) failure is a ClassNotFoundException for scala.tools.nsc.Main during the Maven bootstrap, so a re-run should clear it. I'd like to see the Spark 4.1 SQL label run finish before this is queued.

[andygrove]

Thanks for the review. I added the codec-tag flag to the Encoding step and to the block header description in native_shuffle.md, and the config to the tables in native_shuffle.md and jvm_shuffle.md (1e08791).

I also pushed a simplification of the writer (3194f1a): it now decides when it is built instead of per block. The fallback you mentioned still holds, just per writer. It encodes an empty batch once and keeps plain offsets if arrow-ipc's buffers don't match the mapped roles, which is enough because a type's buffer count doesn't depend on its data.

The push restarts the Lint Java job and the Spark 4.1 SQL run. Agreed on waiting for that run before this is queued.

[kazuyukitanimura]

One thing I noticed is that every other Comet suite and the Spark SQL run use the default lz4, so these two tests are the only place a flagged block goes end to end through the writer, the spill, the fetch and the decode. Could the data include nulls, empty strings, a binary column and a struct with a string field? Could the answer check also run once with spilling on, for example the columnar suite's default spill threshold and a small spark.comet.shuffle.native.maxBufferBytes for native? The byte-count comparison can keep the higher threshold. Alternatively, a one-off local run of CometNativeShuffleSuite and CometShuffleSuite with spark.comet.shuffle.compression.codec=zstd, with the result posted here, would give similar confidence. The PR description mentions you ran something like that before narrowing this to zstd.

[andygrove]

Done in 057f7a5, and I also did the one-off run.

Both tests now take their rows from a rowsWithOffsets helper in CometTestBase, which has offsets in a string, a binary, an array, a map and a struct column, with nulls, empty strings, empty binary values and null array elements among them. Each test also checks the answer once more with spilling. The native test sets spark.comet.shuffle.native.maxBufferBytes to 64k and asserts that the writer spilled, and the columnar test runs at the suite's threshold of 10 rows through both the bypass-merge and the sort-based writer. The byte-count comparisons keep the higher threshold. I also ran just the spill checks against a reader that ignores the lengths flag. Each of them fails with First offset ... of Utf8 is larger than values length, so they do read blocks that hold lengths.

For the one-off run I made zstd the default codec in CometTestBase. All 210 tests in CometNativeShuffleSuite, CometShuffleSuite, DisableAQECometShuffleSuite, CometShuffleEncryptionSuite, CometShuffleInputConversionSuite, SpillSorterSuite and CometDiskBlockWriterSuite pass. With the reader that ignores the flag, 107 of them fail, so that run went through flagged blocks well beyond these two tests.

[kazuyukitanimura]

Could you add offsets_as_lengths.rs to this table? It now defines part of the block format (buffer roles, the byte-plane layout and LENGTHS_TAG_FLAG), so someone reading the "Rust Side" table to find the format code should see it next to shuffle_block_writer.rs and ipc.rs.

[andygrove]

Added in b41b08d, next to ipc.rs.

[sunchao]

Reviewed all 26 changed files and all three commits from base 55ba50d45137781f50ec1b7323da7109e66f36f9 through head 1e0879139b7e27bc5bd6b9bcb48f7ca3dd69d49e. The PR is not a draft. No introduced P1/P2 issues found within this review. Existing non-Copilot discussion contains no substantiated unresolved P1/P2 defect.

Summary

• Prior state and problem: Shuffle IPC stored cumulative offsets directly. These compress poorly for short variable-width values, increasing shuffle I/O even when the values themselves compress well.
• Design approach: Eligible writers replace offsets with lengths arranged in byte planes before zstd compression. A codec-tag flag tells the shared reader to restore the original offsets before Arrow decoding.
• Correctness: Checked buffer_roles against Arrow IPC 59.3.0’s write_array_data, including nested layouts and 32/64-bit offsets. Wrapping subtraction and addition preserve the original bits, including nonzero starting offsets. Buffer-range checks, restoration before Arrow validation, and unchanged checksum/partition-length accounting are consistent. The native round-trip tests cover nulls, empty values and sliced inputs.
• Compatibility analysis: Compared Spark’s ShuffleExchangeExec and IndexShuffleBlockResolver sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Partition assignment and Spark’s index-file contracts are unchanged. JNI arguments and protobuf plumbing agree across the changed call sites. New readers accept existing blocks. Older readers cannot read newly flagged blocks, so the format extension is compatible in one direction only.
• Key design decisions: Restricting writes to zstd follows the reported codec measurements. Dictionary-containing schemas, Celeborn, unmapped layouts and schemas without offsets retain their existing encoding. Eligibility is established once per writer, and the supported configuration provides an opt-out.
• Implementation sketch: The native writer receives the setting through ShuffleWriter; JVM bypass and sort writers pass it through writeSortedFileNative. Both reach ShuffleBlockWriter. Both read paths reach the shared IPC decoder, which restores offsets in its mutable body buffer before exposing arrays.
• Performance: The author reports 9–40% smaller shuffle files across the measured workloads and partition counts. CI passes the native compression-size assertion and JVM shuffle-byte comparisons. The transform adds linear passes and temporary scratch buffers, reused across buffers of each width within a block. It does not add another complete IPC-body copy. I found no evidence-backed P1/P2 performance regression, but did not independently reproduce CPU benchmarks.
• Design: Applying the rewrite at the shared block boundary keeps partitioning, spill-file assembly and fetch behavior straightforward. The reader derives behavior from each block’s flag rather than the current configuration, allowing flagged and plain blocks to coexist.
• Abstraction & complexity: The shared BufferRole traversal and width-parameterized transform keep encoding and decoding consistent. Immutable writer metadata remains shared through Arc. These abstractions fit the problem, with no complexity issue meeting the reporting bar identified.
• Behavioral changes worth calling out: Compared the touched encoding paths with branch-1.1 at 8ac6fccde53d96fec22184be287747eeae481cac. The intended change is the flagged representation for eligible zstd blocks, enabled by default through the new setting. The default codec remains lz4, and excluded blocks keep their prior representation. No unintended result or fallback change was identified.
• Suggested improvements: No additional change meets the P1/P2 reporting bar. The existing request for broader zstd end-to-end spill coverage remains open and is not duplicated as a finding.

Routed skills: review-comet-pr, review-comet-shuffle-pr, review-comet-ffi-pr and review-comet-memory-pr, with their relevant contributor documentation.

Exact-head CI: 36 checks passed and 18 were skipped, with no failed or pending checks. The Comet CI run tested a merge of the specified head and base whose tree exactly matches the reviewed checkout. Logs confirm 2,435 passing Rust tests, 573 passing JVM shuffle tests, the new offsets-as-lengths tests, and the Spark 4.1.3 SQL matrix. Other Spark SQL versions, macOS, Iceberg and the benchmark check were skipped.

Validation limits: A standalone harness using the exact transform source passed 16,384 deterministic cases against an independent byte-plane reference, plus range/count checks. The focused cargo test --locked -p datafusion-comet-shuffle offsets_as_lengths run could not resolve pinned datafusion 55.2.0 from the configured registry. Full local JVM tests, additional forced-spill zstd runs and performance benchmarks were therefore not executed. Integration evidence comes from the verified CI logs. The project checkout remains unchanged.

[viirya]

Thanks for the updates. The docs now describe the format and its scope accurately. Deciding once per writer is a nice simplification. A type's buffer count doesn't depend on its data in arrow-ipc's write_array_data, so the empty-batch check settles the layout for the writer's lifetime, and a mismatch that somehow got through now fails the task instead of writing a bad block. LGTM.