dwsmith1983 opened a new pull request, #5565: URL: https://github.com/apache/datafusion-comet/pull/5565
## Which issue does this PR close? Part of #5002 (the compression-context reuse item; the issue stays open for its remaining items). ## Rationale for this change Every shuffle block currently creates and destroys its own zstd context: a fresh `CCtx` per encoded block in `ShuffleBlockWriter`, and a fresh `DCtx` per decoded frame in `read_ipc_compressed`. Context setup is pure overhead that scales with block count, so high-partition shuffles with small blocks pay the most. ## What changes are included in this PR? - New `ShuffleCodecContext` / `ShuffleDecodeContext` (`native/shuffle/src/codec_context.rs`) wrapping a lazily-created `zstd_safe::CCtx`/`DCtx`, reused via `Encoder::with_context` / `Decoder::with_context`. The session is reset and the level re-applied per frame, so writers with different levels can share a context and a failed encode/decode can't poison the next one. - Contexts are task-scoped, never per-output-partition (a shuffle can have thousands of partitions): `LocalPartitionWriter` owns the encode context and the per-partition `BufBatchWriter`s / `SpillWriter` borrow it; `RssPartitionWriter` is already one per task. Decode reuses a per-thread context behind the existing `read_ipc_compressed` entry points, with `_with` variants for caller-owned contexts. - The RSS path frees the zstd workspace at the end of each admitted encode (success or error): its memory accounting charges the workspace per admitted invocation and releases it afterward, so the context must not outlive that window. - Wire format is unchanged. lz4 and snappy keep per-block encoders — no reset API in the pinned crates, and their setup cost is far smaller than zstd's ~1MB workspace. Benchmarks (M-series macOS, 4M-row hash shuffle via `shuffle_bench`, 3 iterations after warmup; Linux numbers may differ — happy to see re-runs): | Shape | main | this PR | Δ | |---|---|---|---| | 2,000 partitions, zstd level 1 | 0.329s | 0.333s | within noise | | 2,000 partitions, zstd level 6 | 0.767s (encode 0.624s) | 0.757s (encode 0.611s) | −1.3% wall | | 10,000 partitions, zstd level 3 | 0.592s (encode 0.330s) | 0.552s (encode 0.295s) | **−6.8% wall, −10.6% encode** | The saving is per block, so it grows with partition count / shrinking block size; large-block shapes are compression-bound and unchanged. Criterion `shuffle_writer` zstd cases (16 partitions, 8192-row blocks) are within noise, as expected. ## How are these changes tested? Six new tests alongside the existing suites (98 total in the shuffle crate, all passing, plus the core crate's 201): - reuse across blocks, writers, and codecs round-trips and every block decodes independently - two writers with different zstd levels sharing one context each keep their own level - a mid-frame write failure doesn't poison the context for the next block (and the decode-side mirror with a truncated frame) - the RSS release-vs-local-retain contract is pinned via a test accessor - decode-context reuse across mixed-codec frames matches fresh per-frame decoders `cargo clippy --all-targets -- -D warnings` and `cargo fmt` clean. -- 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]
