andygrove commented on code in PR #5568:
URL: https://github.com/apache/datafusion-comet/pull/5568#discussion_r3896060699
##########
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:
One thing I want to make sure we've thought about. `write_batch_to_buffer`
appends a whole block and only then checks `pos >= buffer_max_size`, so the
buffer's high-water mark is the write buffer size plus one full block, not the
write buffer size. Previously the writer was dropped at the end of each
partition and that peak went back to the allocator. Now it lives in
`recycled_buffer` until the `LocalPartitionWriter` drops, and nothing charges
it to the shuffle reservation. The spill path is where I'd worry, since spills
happen under memory pressure and we'd now be holding an extra untracked
allocation between them. Would a `shrink_to(write_buffer_size)` when you store
it back be enough to bound that? #5565 bounds retention for its codec context,
so it would be good for these two to agree.
Related, and worth fixing since we squash-merge: the description says this
is "a strict improvement for memory at high partition counts". Partitions are
already written one at a time, so at most one buffer was ever live and peak is
unchanged. The real change is that the peak is retained rather than released.
Could you reword that so the tradeoff is on the record?
Separately, `into_buffer()` will happily hand back a buffer with bytes still
in it, and `new()` clears whatever it's given. Both callers flush first so this
is fine today, but if someone later reuses a buffer from an unflushed writer
the encoded bytes just disappear and it surfaces as missing rows two stages
downstream. Could `into_buffer()` carry a
`debug_assert!(self.buffer.is_empty())`? The doc comment already states the
invariant.
##########
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:
How do you feel about matching the shape #5565 uses? It threads its
task-scoped state as a `&mut ShuffleCodecContext` parameter on `write` and
`flush`, where this moves a `Vec<u8>` in through the constructor and back out
through `into_buffer()`. If both land we end up with two conventions for the
same idea in the same struct, and whichever rebases second is unlikely to
revisit it.
Borrowing would also clean up a small wart. The `mem::take` empties the
caller's buffer up front and the write-back only runs on success, so any error
in between quietly ends recycling for the rest of the task. Not a correctness
issue, but it goes away on its own if the buffer is borrowed.
##########
native/shuffle/src/partitioners/partitioned_batch_iterator.rs:
##########
@@ -42,14 +42,22 @@ impl PartitionedBatchesProducer {
}
}
+ /// References to all buffered batches. Build this once per write cycle
and share it
+ /// across every partition's [`Self::produce`] call instead of rebuilding
a fresh
+ /// `Vec<&RecordBatch>` over all buffered batches for each partition.
+ pub(super) fn batch_refs(&self) -> Vec<&RecordBatch> {
+ self.buffered_batches.iter().collect()
+ }
+
pub(super) fn produce<'a>(
- &'a mut self,
+ &'a self,
+ refs: &'a [&'a RecordBatch],
Review Comment:
The indices are positions into `self.buffered_batches`, but `refs` here is
unconstrained, so a slice built from the wrong producer or a filtered one
compiles fine and shuffles the wrong rows instead of failing. Would you mind
adding a `debug_assert_eq!(refs.len(), self.buffered_batches.len())` at the
top? Cheap, and it catches the realistic version of that mistake.
--
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]
