dwsmith1983 commented on code in PR #5568:
URL: https://github.com/apache/datafusion-comet/pull/5568#discussion_r3900480201
##########
native/shuffle/src/writers/buf_batch_writer.rs:
##########
@@ -47,23 +47,34 @@ pub(crate) struct BufBatchWriter<S:
Borrow<ShuffleBlockWriter>, W: Write> {
}
impl<S: Borrow<ShuffleBlockWriter>, W: Write> BufBatchWriter<S, W> {
+ /// `buffer` is the caller-owned byte buffer to serialize into. Passing a
buffer recovered
+ /// from a previous writer's [`Self::into_buffer`] reuses its capacity
instead of regrowing
+ /// a fresh allocation toward `buffer_max_size` for every writer.
pub(crate) fn new(
shuffle_block_writer: S,
writer: W,
buffer_max_size: usize,
batch_size: usize,
+ mut buffer: Vec<u8>,
Review Comment:
Done, write and flush borrow the scratch per call now, same shape as the
codec context, and the constructor param and into_buffer are gone. You are
right that it also fixes the take/put-back wart; an error no longer ends
recycling. The one thing ownership gave for free was buffer identity, so the
writer now records the scratch address on first use and debug-asserts every
later call passes the same one, a swapped buffer would silently drop un-flushed
bytes otherwise.
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -53,6 +53,11 @@ enum DataOutput {
spill_writers: Vec<SpillWriter>,
/// Runtime used to allocate the temporary spill files.
runtime: Arc<RuntimeEnv>,
+ /// Byte buffer recycled through the short-lived per-partition
`BufBatchWriter`s.
+ /// Partitions are written strictly one at a time, so a single buffer
keeps its
+ /// grown capacity across the whole task instead of every partition
regrowing a
+ /// fresh allocation toward the write buffer size.
+ recycled_buffer: Vec<u8>,
Review Comment:
Ran it with your shape: an 8,192-row hot batch of 8 KiB Binary payloads
followed by ~500k rows of 16-byte payloads, plus reversed-order and narrow-only
controls; codec none, 2,000 partitions, max-buffer-bytes 4 MiB (about 10 spills
per run), 3 iterations after warmup, base vs head on the same build, machine,
and input seed:
| input | base avg | head avg |
|---|---|---|
| wide-then-narrow | 0.401s | 0.384s |
| narrow-then-wide | 0.421s | 0.413s |
| narrow-only | 0.398s | 0.397s |
Head is slightly faster on all three, most on the wide-then-narrow case, so
no time regression from retention. On the memory side the answer changed while
you were writing the comment: flush now shrinks the scratch back to the
configured write-buffer size, so retained capacity after the wide batch is
bounded by construction and pinned by a test rather than instrumented at
runtime — which I think lands on the "cap" option among cap/release/account.
Spilled bytes and outputs match within the run-to-run variation of the spill
trigger (Arrow's memory estimates shift block boundaries by a spill or so even
between runs of the same binary, so byte-exact cross-tree diffs are not
meaningful under forced spilling; row-level equality is pinned by the unit
tests and the digest matrix in the review above). I did not instrument peak RSS
or live reservation traces.
--
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]
