andygrove opened a new pull request, #6943:
URL: https://github.com/apache/datafusion-comet/pull/6943
## 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::with_offsets_as_lengths` (off by default, zstd only)
rewrites the encoded body on the dictionary-free local path and sets the high
bit of the 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, schemas
without offsets, and types whose layout isn't mapped (views, unions, run-end
encoding).
- `shuffle_bench --offsets-as-lengths`, and a paragraph in the shuffle
tuning guide.
## 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.
- 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`,
the latter run through `CometShuffleSuite` and `DisableAQECometShuffleSuite`.
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. 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.
- 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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
