dwsmith1983 opened a new pull request, #5568: URL: https://github.com/apache/datafusion-comet/pull/5568
## Which issue does this PR close? Part of #5002 (the two per-partition allocation items in the medium tier; the issue stays open for its remaining items). Independent of #5565 — whichever merges second gets a small mechanical rebase over the shared `BufBatchWriter` constructor. ## Rationale for this change The multi-partition write path pays two allocation costs per partition per event: - `BufBatchWriter` starts from an empty byte buffer and regrows it toward the write-buffer size (1MB default). `finish_partition` and the spill path construct one per partition per event, so a 2,000-partition pass runs 2,000 growth cycles — roughly an extra copy of the payload through doubling-growth memcpy, plus the transient allocations. - `PartitionedBatchIterator` materializes the whole partition's `(u32,u32)` index list as a fresh `(usize,usize)` vector (16 bytes per row, re-materialized every write cycle) and rebuilds a batch-ref vector over all buffered batches for every non-empty partition. ## What changes are included in this PR? - `BufBatchWriter::new` takes a caller-owned buffer and `into_buffer()` hands it back drained (capacity kept). The task-level `LocalPartitionWriter` owns one buffer and recycles it through both sequential partition loops (spill and finish), so an event allocates it once and holds exactly one regardless of partition count — a strict improvement for memory at high partition counts. The long-lived single-partition writer keeps owning its buffer as before. - `PartitionedBatchesProducer::produce` takes `&self` plus a caller-built batch-ref slice (`batch_refs()`, built once per write cycle); the iterator borrows the raw `(u32,u32)` indices and converts per output chunk into a small reusable scratch (capacity ≤ batch size) instead of widening the full list up front. The empty-partition path stays allocation-free. Wire format and produced batches are unchanged, pinned by byte-level tests (recycled vs fresh buffers produce identical bytes; chunked conversion interleaves identically to full up-front conversion, tail chunk included). Benchmarks (M-series macOS, 4M-row hash shuffle via `shuffle_bench`, 3 iterations after warmup): write time drops ~27% at 2,000 partitions (0.015s → 0.011s) with the other phases flat, and the criterion end-to-end suite shows no regressions (one config improved). Honest framing: a single-task bench understates this change — the removed churn is per-partition allocator traffic, which matters most under concurrent tasks sharing the allocator, and the bounded buffer count is the memory-profile win. ## How are these changes tested? Two new tests alongside the existing suites (94 in the shuffle crate, all passing, plus the core crate's 201): - a recycled buffer produces byte-identical output to fresh per-partition buffers, comes back drained, and keeps its grown capacity across partitions - chunked index conversion interleaves exactly like converting the whole partition up front, including the short tail chunk, sharing one batch-ref slice across partitions `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]
