[andygrove]

Experimental / WIP. Opened as a draft to get early feedback on the format and the approach. Not ready to merge.

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 at the shuffle's level 1 in both columns, TPC-H SF100 tables with all columns, two interleaved runs each:

Table (rows)	Offsets as is	Offsets as lengths	Smaller by
customer (15M)	980.0 MB	823.5 MB	16.0%
part (20M)	814.6 MB	559.1 MB	31.4%
orders (first 30M)	1,127.4 MB	873.9 MB	22.5%
lineitem (first 30M)	1,262.8 MB	982.2 MB	22.2%

Writing and reading back are no slower, and slightly faster with smaller output: lineitem write 11.23 s / 11.25 s off vs 10.81 s / 10.90 s on, read back 2.70 s / 2.78 s off vs 2.57 s / 2.60 s on; part write 7.50 s / 7.47 s off vs 7.09 s / 7.05 s on. At level 3 the savings were 1-2 points larger (customer 17.8%, part 33.0%, orders 23.6%, lineitem 23.9%) for about 1 s more write time on lineitem. 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 (at SHUFFLE_ZSTD_COMPRESSION_LEVEL, level 1, like the plain path) 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, perf: faster overlapping copy in the decoder 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.