perf: reuse Arrow IPC CompressionContext across shuffle blocks instead of per-batch allocation

[andygrove]

Background

Follow-on from review of #5006, which encodes the shuffle IPC schema once per writer instead of per block.

In ShuffleBlockWriter::encode_ipc_stream (native/shuffle/src/writers/shuffle_block_writer.rs), the fast path constructs a fresh CompressionContext (an alias of arrow-ipc's IpcWriteContext) for every batch:

let mut compression_context = CompressionContext::default();
let (encoded_dictionaries, encoded_batch) = data_gen.encode(
    batch,
    &mut dictionary_tracker,
    &options,
    &mut compression_context,
)?;

Unlike the other per-batch objects on this path (IpcDataGenerator is zero-sized; DictionaryTracker does no heap allocation until used and must be fresh per standalone stream), CompressionContext holds reusable buffers: a FlatBufferBuilder and, for zstd, a Compressor. arrow-ipc's own documentation describes it as intended to be reused across writes to avoid re-initializing the context (and re-allocating the flatbuffer builder / zstd compressor) for every block.

Proposal

Reuse a single CompressionContext across the blocks written by a given partition writer.

The wiring is not trivial: CompressionContext is not Clone (the zstd Compressor isn't), while ShuffleBlockWriter is Clone and shared across partitions via Borrow, and write_batch takes &self. So the context can't simply become a field on ShuffleBlockWriter. The natural owner is the per-partition BufBatchWriter (the actually-mutable, per-partition writer); reuse means threading a &mut CompressionContext down through write_batch and its call sites.

Related

Part of the shuffle-writer optimization umbrella #5002.

[peterxcli]

take