andygrove opened a new issue, #5008:
URL: https://github.com/apache/datafusion-comet/issues/5008
## 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:
```rust
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.
--
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]
