andygrove commented on PR #5565: URL: https://github.com/apache/datafusion-comet/pull/5565#issuecomment-5552528334
Here's my AI review: > This seems lot a lot of additional complexity for little performance gain? Digging into the numbers, I think that instinct is right, and I'd narrow the scope rather than land this as is. **The measured gain is about 3 µs per shuffle block, and the defaults never reach it.** Two things from `CometConf.scala` matter here: the default shuffle codec is `lz4` (line 510) and the default zstd level is `1` (line 518). This PR does nothing for lz4 or snappy, so out of the box it's inert — you need a non-default codec *and* a non-default level *and* a high partition count before it does anything. Working backwards from the benchmark table gives a consistent mechanism: | shape | encode base → PR | blocks | implied per-block | |---|---|---|---| | 10,000 parts, zstd-3 | 0.326 → 0.297s | ~10,000 | **~2.9 µs** | | 2,000 parts, zstd-1 | 0.188 → 0.187s | ~2,000 | predicted ~6 ms, below noise | | 2,000 parts, zstd-6 | 0.606 → 0.610s | ~2,000 | predicted ~6 ms, below noise | (4M rows / 10k partitions is ~400 rows each, under the 8192 batch size, so blocks ≈ partitions.) Worth noting that level 6 has a *larger* workspace than level 3 and still showed nothing, which confirms the variable is block count, not level. So the honest claim is ~3 µs saved per block: 29 ms at 10,000 blocks, 0.6 ms at Spark's default 200 partitions. Given how sensitive these shuffle benchmarks are to machine load, I'd also want the 10k result reproduced on Linux with more than 3 iterations before treating the 4% as real. **The decode half has no measured benefit, and it introduces a retention hazard main doesn't have.** The comment above puts it at 56 ns/frame of setup against ~114 µs/frame of work. The stated justification — "the decode-side value is the bounded retained workspace, not speed" — is circular: on main, `zstd::Decoder::with_buffer` builds a fresh DCtx per frame, so retention between frames is already zero. The PR introduces retention and then spends the 8 MiB cap, the `sizeof()` checks, the error-path releases, and the EOF hook bounding retention it created. The thread-local in `ipc.rs` also backs the live production path for the JVM shuffle reader (`NativeBatchDecoderIterator` → `Native.decodeShuffleBlock`), so every Spark task thread now holds roughly 1-2.6 MB of native zstd workspace that Comet's memory accounting doesn't track. That's exactly the property used to justify releasing per-invocation on the RSS encode path — same hazard, opposite decision, for a 0.05% saving. **Where the complexity actually lands.** It's ~290 net production lines and ~1,080 test/bench lines, but the part that costs us long term is four things: 1. `MAX_RETAINED_ZSTD_CONTEXT_BYTES` is a magic constant backed by a hand-measured table of zstd workspace sizes, with a test whose only job is to fail when a zstd bump silently disables the optimization (level 8 sits 3% under the cap). That's a tax on every dependency bump. 2. Two divergent lifetime contracts for one object — RSS releases, local retains — enforced by a bool threaded through `write_batch_with_codec_limits`. 3. A new `write_burst_complete()` on the `PartitionWriter` trait that exists only to drop one codec's allocator, i.e. a lifecycle hook in the partitioner API serving zstd specifically. 4. The thread-local above. **What I'd suggest instead:** keep the encode-side reuse for local shuffle, and drop the decode half entirely — `ShuffleDecodeContext`, the four `_with` entry points, `decode_remote_shuffle_batch_with`, the thread-local, the `ScanInputState` restructure, `benches/ipc_decode.rs`, and the decode tests. That's roughly half the production diff, removed for something already measured as neutral. Then, could the `sizeof()`-based cap be replaced with a retain/don't-retain decision made once in `ShuffleBlockWriter::try_new` from the configured level? Codec and level are uniform per query, so the alternating-level test is defending a case that can't occur in production, and `reset(SessionOnly)` preserves parameters anyway, which makes the per-block `set_parameter` redundant. That would drop the constant, the measured table, and both `release_zstd_if_oversized` sites, leaving the mechanical `&mut` threading as the only real cost. One hypothesis I checked and discarded, in case anyone else wonders the same: hoisting `CompressionContext` out of `BufBatchWriter` is *not* a hidden memory win at high partition counts. `IpcWriteOptions::try_new(64, false, V5)` leaves Arrow's inner buffer compression off, so those per-partition contexts were an empty `Vec` plus a `None`, about 40 bytes each. Good hygiene, not a saving. -- 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]
