andygrove opened a new pull request, #5005: URL: https://github.com/apache/datafusion-comet/pull/5005
> **Experimental / not ready for review.** Opening as a draft to run larger-scale benchmarks before this is ready. Please do not review yet. ## Which issue does this PR close? Addresses two of the "high impact" items in #5002: the per-block zstd compression-context allocation and the per-block IPC schema re-encoding. That issue is an umbrella tracking several shuffle-writer optimizations, so this PR does not close it. ## Rationale for this change `ShuffleBlockWriter::write_batch` allocated a fresh zstd context (a `ZSTD_CCtx` plus its workspace) and re-serialized the schema flatbuffer (via `StreamWriter::try_new`) for every block. `read_ipc_compressed` likewise allocated a fresh `ZSTD_DCtx` per block. With small blocks (high partition counts), per-block context setup is a meaningful fraction of compression time. ## What changes are included in this PR? - **Reuse the zstd compression/decompression context across blocks** via a per-thread context. Compression/decompression is streamed (not buffered whole), so large blocks are not copied an extra time. An earlier attempt using zstd's one-shot bulk API regressed large blocks by ~3% because it had to buffer the entire uncompressed stream; the streaming approach avoids that. - **Pre-encode the IPC schema message once** in `try_new` and write it verbatim per block, instead of re-serializing it per block. Schemas that contain dictionary types keep using `StreamWriter`, whose dictionary-id bookkeeping ties schema and batch encoding together. Output framing is unchanged: blocks remain standalone Arrow IPC streams in the same length-prefixed format, so the read side and the wire format are unaffected. ### Preliminary benchmark (local, zstd, 5M-row input; `shuffle_bench`) Higher partition count means smaller blocks, where per-block context setup dominates. | Partitions | ~rows/block | Encode time (main → this PR) | | ---: | ---: | :--- | | 2,000 | ~5,000 | 0.301s → 0.303s (flat) | | 20,000 | ~500 | 0.492s → 0.449s (~9% faster) | | 50,000 | ~200 | 0.752s → 0.669s (~11% faster) | No regression at large block sizes; the win grows as blocks shrink. Larger-scale benchmarks are still to be run, which is why this is a draft. ## How are these changes tested? Existing `datafusion-comet-shuffle` and `datafusion-comet` shuffle tests pass. New tests were added for a dictionary-typed column roundtrip (exercising the `StreamWriter` fallback and dictionary-id matching) and for reading multiple distinct zstd blocks on one thread (exercising cross-frame reuse of the thread-local compression/decompression contexts, which a dirty context would corrupt). -- 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]
