dwsmith1983 commented on code in PR #6336:
URL: https://github.com/apache/datafusion-comet/pull/6336#discussion_r4380899805
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -245,88 +313,100 @@ impl PartitionWriter for LocalPartitionWriter {
let write_buffer_size = self.write_buffer_size;
let batch_size = self.batch_size;
- match &mut self.data_output {
- DataOutput::Single { writer, scratch } => {
- // Single-partition data was already streamed via `write`,
starting at
- // offset 0 (already recorded in `self.offsets[0]`). Stream
any trailing
- // batches (normally none) without flushing; the long-lived
writer is
- // flushed once in `finish_all`.
- for batch in iter.by_ref() {
- let batch = batch?;
- writer.write(
- &batch,
- scratch,
- &mut self.codec_context,
- &metrics.encode_time,
- &metrics.write_time,
- )?;
- }
- }
- DataOutput::Multi {
- output_writer,
- shuffle_block_writer,
- spill,
- spill_reader,
- recycled_buffer,
- ..
- } => {
- self.offsets[pid] = output_writer.stream_position()?;
-
- let mut flush_timer = metrics.write_time.timer();
- spill.flush()?;
- flush_timer.stop();
- let ranges = spill.ranges(pid)?;
- if !ranges.is_empty() {
- if spill_reader.is_none() {
- let path = spill.path()?.ok_or_else(|| {
- DataFusionError::Internal(
- "shuffle spill ranges recorded without a spill
file".to_string(),
- )
- })?;
- *spill_reader = Some((File::open(path)?, vec![0;
write_buffer_size]));
- }
- let (spill_file, buffer) = spill_reader.as_mut().unwrap();
- let mut write_timer = metrics.write_time.timer();
- for range in ranges {
- copy_spill_range(spill_file, buffer, range,
output_writer)?;
- }
- write_timer.stop();
- }
-
- // Write in memory batches to output data file. Each partition
uses its
- // own writer so coalescing does not cross partition
boundaries, but the
- // scratch buffer is shared so its capacity carries over to
the next one.
- let mut buf_batch_writer = BufBatchWriter::new(
- shuffle_block_writer,
- output_writer,
- write_buffer_size,
- batch_size,
- );
- let codec_context = &mut self.codec_context;
- let result: datafusion::common::Result<()> = (|| {
+ // Synced on every exit, errors included, so it matches what a failed
partition left.
+ let result: datafusion::common::Result<()> = (|| {
+ match &mut self.data_output {
+ DataOutput::Single { writer, scratch } => {
+ // Single-partition data was already streamed via `write`,
starting at
+ // offset 0 (already recorded in `self.offsets[0]`).
Stream any trailing
+ // batches (normally none) without flushing; the
long-lived writer is
+ // flushed once in `finish_all`.
for batch in iter.by_ref() {
let batch = batch?;
- buf_batch_writer.write(
+ writer.write(
&batch,
- recycled_buffer,
- codec_context,
+ scratch,
+ &mut self.codec_context,
&metrics.encode_time,
&metrics.write_time,
)?;
}
- buf_batch_writer.flush(
- recycled_buffer,
- codec_context,
- &metrics.encode_time,
- &metrics.write_time,
- )
- })();
- // An errored partition must hand back a drained buffer, or
its bytes
- // leak into the next partition's block.
- result.inspect_err(|_| recycled_buffer.clear())?;
+ }
+ DataOutput::Multi {
+ output_writer,
+ shuffle_block_writer,
+ spill,
+ spill_reader,
+ recycled_buffer,
+ ..
+ } => {
+ self.offsets[pid] = output_writer.stream_position()?;
+
+ let mut flush_timer = metrics.write_time.timer();
+ spill.flush()?;
+ flush_timer.stop();
+ let ranges = spill.ranges(pid)?;
+ if !ranges.is_empty() {
+ if spill_reader.is_none() {
+ let path = spill.path()?.ok_or_else(|| {
+ DataFusionError::Internal(
+ "shuffle spill ranges recorded without a
spill file"
+ .to_string(),
+ )
+ })?;
+ // Open first, so a failed open leaves nothing
charged.
+ let spill_file = File::open(path)?;
+ let buffer_size =
+
self.buffers.as_ref().map_or(write_buffer_size, |buffers| {
+ super::reserve_buffer(buffers,
write_buffer_size)
Review Comment:
The first `finish_partition` now flushes the spill file and releases its
write buffer before anything else, then sizes the copy buffer and only then the
data file buffer. The copy buffer serves every spilled range; the data file
buffer matters less, because blocks of the write buffer size bypass it, so it
is the one that falls back to 8 KiB when the pool is short.
`output_buffers_are_sized_after_the_spill_buffer_is_released` checks the
capacities at pools that fit two full buffers, one, and none, with identical
output, and the last case still copies a large range through `io::copy`.
`failed_first_partition_keeps_the_reservation_in_sync` fails a partition after
the release and checks that the reservation still equals what the writer holds.
--
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]
