dwsmith1983 commented on code in PR #6336:
URL: https://github.com/apache/datafusion-comet/pull/6336#discussion_r4380808224
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -372,14 +452,45 @@ impl PartitionWriter for LocalPartitionWriter {
// The shuffle output is complete; nothing else encodes through this
context.
self.codec_context.release_zstd();
+ self.sync_buffer_reservation();
Ok(())
}
fn write_burst_complete(&mut self) {
- // A spill burst just ended and the next encode may be a long time
coming; the zstd
- // workspace is native memory no reservation tracks, so don't sit on
it.
+ // A spill burst just ended and the next encode may be a long time
coming, so don't
+ // sit on the zstd workspace.
self.codec_context.release_zstd();
+ self.sync_buffer_reservation();
+ // A pressure spill opened its file while the batches still held the
pool; now they don't.
+ if let (Some(buffers), DataOutput::Multi { spill, .. }) =
+ (&self.buffers, &mut self.data_output)
+ {
+ spill.restore_buffer(buffers);
+ }
+ }
+
+ fn attach_buffer_reservation(
+ &mut self,
+ reservation: MemoryReservation,
+ ) -> datafusion::common::Result<()> {
+ let DataOutput::Multi { output_writer, .. } = &mut self.data_output
else {
+ return Ok(());
+ };
+ // Attached once, before any write: the rebuild below must not discard
buffered bytes.
+ debug_assert!(self.buffers.is_none() &&
output_writer.buffer().is_empty());
+ // The data file buffer was allocated before the consumer existed;
shrink it if refused.
+ let capacity = super::reserve_buffer(&reservation,
self.write_buffer_size);
Review Comment:
The data file writer now starts unbuffered and nothing is charged at attach.
The first `finish_partition` sizes it, so an operator feeding the shuffle gets
the pool until the output is written.
`writer_charges_nothing_before_the_first_row` builds the repartitioner with a 1
MiB pool and the default 1 MiB write buffer, checks that nothing is reserved,
has a second consumer take 128 bytes, and compares the output with an unbounded
pool.
--
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]
