andygrove opened a new pull request, #2809:
URL: https://github.com/apache/datafusion-ballista/pull/2809
> **Experimental / WIP.** Opened as a draft to get early feedback on the
format and the approach. Not ready to merge.
>
> **Stacked on #2805** (which is stacked on #2802). Review only the last
commit, `perf: store shuffle string offsets as lengths before zstd compression`.
# Which issue does this PR close?
Closes #2804.
# Rationale for this change
Shuffle data is Arrow IPC with arrow-ipc's per-buffer compression. The
offsets buffers of `Utf8`, `Binary`, list and map columns hold running sums.
Those never repeat, so zstd compresses them poorly even when every value has
the same length, and they often cost about as much as the string bytes. Stored
as value lengths split into byte planes, they compress to almost nothing.
DataFusion Comet made the same change in apache/datafusion-comet#6943.
`shuffle_bench`, one task, 16 hash partitions, 8,192-row batches, zstd,
TPC-H SF100 tables with all columns, two interleaved runs each:
| Table (rows) | Offsets as is | Offsets as lengths | Smaller by |
|---|---:|---:|---:|
| customer (15M) | 977.2 MB | 803.1 MB | 17.8% |
| part (20M) | 791.7 MB | 530.1 MB | 33.0% |
| orders (first 30M) | 1,106.3 MB | 845.7 MB | 23.6% |
| lineitem (first 30M) | 1,240.1 MB | 943.6 MB | 23.9% |
Write time and the time to read every partition back stayed within
run-to-run noise (for example lineitem write 12.2 s / 12.0 s off vs 12.4 s /
12.2 s on, read back 2.63 s / 2.63 s off vs 2.51 s / 2.57 s on). These are
single-machine numbers. Cluster runs are still to do.
# What changes are included in this PR?
- `ballista_core::shuffle_ipc` (new):
- `offsets.rs` ports Comet's buffer-role mapping and the in-place rewrite
between offsets and byte-plane lengths.
- `ShuffleStreamWriter` encodes each batch uncompressed with
`IpcDataGenerator`, rewrites the offsets, compresses the body as one zstd frame
and rebuilds the message. `ShuffleEncoding` decides once per schema whether to
do that: only with zstd, only for schemas with an offsets buffer whose types
all have a mapped layout, and only when arrow-ipc lays out an empty batch of
the schema as the mapping expects. Everything else gets arrow-ipc's
`StreamWriter` and identical bytes.
- `ShuffleIpcReader` (a `Read` adapter for `StreamReader`) and
`ShuffleBytesTranslator` (push-based, for block transport) turn rewritten
messages back into standard uncompressed IPC messages and pass all other
messages through byte for byte.
- Format: the stream keeps Arrow IPC framing. A rewritten record-batch
message has a `BodyCompression` codec Arrow does not define (`0x40`), buffer
offsets that describe the uncompressed body, and a body holding the
uncompressed and compressed lengths followed by the zstd frame. Readers that do
not know it fail with arrow-ipc's "compression type not supported" error
instead of misreading the data. Each stream stays self-contained, so the sort
shuffle's byte-copied spill streams and its index are unchanged.
- Writers: the passthrough `ShuffleWriterExec` (`write_stream_to_disk`),
sort shuffle spills and the sort shuffle's in-memory partition encoding.
- Readers: local passthrough reads, `MultiStreamPartitionStream`, the Flight
`do_get` serve path (which still re-encodes with LZ4, so clients are
unaffected) and `BlockDataStream`.
- Not covered: range shuffle (IPC file format with per-message slicing),
pipelined in-memory shuffle, and LZ4 (lz4_flex decodes long runs slowly,
PSeitz/lz4_flex#227).
- `ballista.shuffle.compression.offsets_as_lengths` (default `true`, only
has an effect with zstd) as an off switch.
- `BALLISTA_PROTOCOL_VERSION` 3 → 4, because block transport ships raw
shuffle bytes to the consuming executor.
- `shuffle_bench` gains `--codec` and `--offsets-as-lengths`, and reports
shuffle bytes and read-back time.
- Tests: Comet's role and round-trip unit tests, plus round trips of string,
binary, list, large list, fixed-size list, map and struct columns with nulls,
empty values, sliced inputs and empty batches under every codec and setting.
They cover both readers, with block-transport bytes fed in blocks of 1 to N
bytes, and check that:
- plain streams match arrow-ipc's byte for byte,
- unmapped schemas (views, dictionaries) and schemas without offsets fall
back,
- concatenated streams are crossed correctly,
- an unaware arrow reader rejects the format,
- truncated or unmapped rewritten batches are errors.
The existing `sort_shuffle` end-to-end matrix (local, Flight, block
transport) passes with the setting on by default.
# Are there any user-facing changes?
Yes:
- Shuffle files written with zstd are smaller for schemas with string,
binary, list or map columns.
- The new `ballista.shuffle.compression.offsets_as_lengths` setting.
- `BALLISTA_PROTOCOL_VERSION` goes to 4, so the scheduler and all executors
must be upgraded together.
- `write_stream_to_disk` and `SpillManager::new` take a `ShuffleCompression`
instead of an `Option<CompressionType>`.
The upgrade guide has entries for these.
--
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]
