dwsmith1983 commented on code in PR #6336:
URL: https://github.com/apache/datafusion-comet/pull/6336#discussion_r4380930872
##########
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);
+ if capacity != output_writer.capacity() {
+ let output_file = output_writer
+ .get_ref()
+ .try_clone()
Review Comment:
The `try_clone()` is gone. The data file writer starts unbuffered around the
original `File`, and the first `finish_partition` moves that `File` into a
writer at the granted capacity, so no second descriptor is opened.
--
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]
