Re: [PR] perf: reuse per-partition scratch in the shuffle write path [datafusion-comet]

["dwsmith1983 (via GitHub)" <gi...@apache.org>]

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 unflushed bytes otherwise.



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org