sunchao commented on code in PR #6336:
URL: https://github.com/apache/datafusion-comet/pull/6336#discussion_r4368030211
##########
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:
[P2] The empty output writer can take the entire available pool before its
input produces a row. A hash join feeding the shuffle then fails its first
allocation, even though this buffer receives its first bytes during
finalization. Defer allocating and charging it until the first
`finish_partition`.
Evidence: With two output partitions, a 16-row `HashJoinExec`, and the
default 1 MiB write buffer, the base wrote all rows with a 1 MiB pool. The head
denied `HashJoinInput[0]` its first 128 bytes. Reducing the pool by one byte
made the head succeed because the writer took the fallback. With eight
partitions and 256 batches of 1,024 integers, the same eager allocation caused
256 spills versus four on the base. Deferring it restored four spills and
preserved all 262,144 rows.
##########
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();
Review Comment:
[P2] After a pressure spill, `recycled_buffer` is empty but its full
capacity stays charged. One wide row can leave that charge larger than the
whole pool, forcing every subsequent narrow batch to spill. Shrink or drop
drained scratch that the pool cannot fund, then synchronize the reservation
before restoring the spill buffer.
Evidence: With LZ4, a 512 KiB pool, and one 768 KiB incompressible value,
the head held 1,062,608 bytes after the first spill, including 1,046,224 bytes
of empty scratch. A separate 16 KiB request failed. Adding 128 one-row batches
produced 129 spills versus one on the base. Releasing only the empty scratch
restored the grant and the one-spill result. All 129 rows decoded.
##########
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:
[P2] During final replay, the flushed spill writer still holds its charged
buffer when the reader asks for another. A refused read buffer sends
medium-sized ranges through `seek` and `io::copy` even though a full-size
allocation is sitting idle. Transfer the drained spill buffer to the reader, or
release its capacity and charge before choosing the read-buffer size.
Evidence: On Linux, replaying 2,048 identical 12,828-byte ranges with a 2
MiB + 128 KiB pool took a median 31.8 ms on the base and 64.3 ms on the head
over seven interleaved rounds. The head held a 1 MiB spill-write buffer and
received an 8 KiB read buffer. Reusing the existing allocation took 35.8 ms at
the same budget. A full repartitioner run produced 256 spills, 1,024 ranges,
and output CRC `01ffa3c0` on both revisions.
##########
native/shuffle/src/writers/local/spill.rs:
##########
@@ -214,17 +217,63 @@ impl PartitionedSpill {
fn ensure_spill_file_created(
&mut self,
runtime: &RuntimeEnv,
+ buffers: Option<&MemoryReservation>,
) -> datafusion::common::Result<()> {
if self.spill_file.is_none() {
let temp_file = runtime
.disk_manager
.create_tmp_file("shuffle writer spill")?;
- let writer = BufWriter::with_capacity(self.write_buffer_size,
temp_file.open_writer()?);
+ // Open before charging, so a failed open leaves nothing on the
reservation.
+ let file = temp_file.open_writer()?;
+ let capacity = buffers.map_or(self.write_buffer_size,
|reservation| {
+ super::reserve_buffer(reservation, self.write_buffer_size)
+ });
+ let writer = BufWriter::with_capacity(capacity, file);
self.spill_file = Some(ActiveSpillFile { temp_file, writer });
}
Ok(())
}
+ /// Grows a fallback-sized buffer back to the configured size if `buffers`
grants it.
+ /// Buffered bytes move to the new buffer; nothing is written to the file.
+ pub(crate) fn restore_buffer(&mut self, buffers: &MemoryReservation) {
+ if self.failed {
+ return;
+ }
+ let Some(spill_file) = self.spill_file.take() else {
+ return;
+ };
+ let capacity = spill_file.writer.capacity();
+ if capacity >= self.write_buffer_size
+ || buffers.try_grow(self.write_buffer_size - capacity).is_err()
Review Comment:
[P2] A successful full-size restoration can leave too little memory for even
the next input batch. The writer then spills on every subsequent batch while
retaining the enlarged buffer. Account for the input headroom needed after a
pressure spill when choosing the restore size. An intermediate capacity can
retain buffered I/O without forcing immediate re-spills.
Evidence: With eight hash partitions, LZ4, `maxBufferBytes=0`, and 256
batches of 1,024 integers, raising the head's pool from 2,096 to 2,112 KiB
increased spills from four to 198. The restored 1 MiB spill buffer left just
4,066 bytes free. The base spilled twice at 2,112 KiB. Capping only restoration
at 256 KiB produced six spills and reduced median time from 284 to 177 ms over
five interleaved runs. All rows decoded. The same 198-to-six difference
persisted with the data-output buffer deferred at a 1,088 KiB pool.
##########
docs/source/user-guide/latest/tuning/memory.md:
##########
@@ -182,11 +187,16 @@ alongside the JVM's own non-heap memory. To size the
overhead from it:
non-heap memory, and add the most untracked memory seen on any executor.
3. Add a margin on top. The log can miss the true peak between samples, and
none of the figures
includes the allocator's fragmentation and retained pages, or memory
allocated by native C
- libraries such as zstd. With native Arrow UDFs, budget the embedded Python
interpreter and
- PyArrow allocations in addition to the untracked memory calculated from the
log. A PyArrow
- result buffer later imported into the JVM can appear in the `imported from
native` figure that
- the estimate subtracts, even though Rust's `allocated` does not count it.
Use the executor's
- peak resident memory to size this additional overhead.
+ libraries such as zstd. The zstd context of the native shuffle writer is
the exception: when
+ `spark.comet.shuffle.compression.codec` is `zstd`, a shuffle task writing
to more than one local
+ partition reserves its context while spilling or writing its output, so
that memory is in
Review Comment:
[P2] The sizing guide says the local zstd workspace appears in `reserved`
while writing. When a context exceeds the 8 MiB retention cap,
`release_zstd_if_oversized()` drops it before the writer samples it. That peak
appears in neither `reserved` nor Rust's `allocated` figure and can reach tens
of MiB per concurrent task. Distinguish retained contexts from transient
compression workspace here and in the contributor guide.
Evidence: The same batch round-tripped at zstd levels 1, 9, and 19. At level
1 the context measured 1,369,625 bytes both after compression and at the
accounting boundary. Levels 9 and 19 measured 15,459,865 and 93,848,215 bytes
after compression, respectively, and zero at the accounting boundary.
##########
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:
[P2] `try_clone()` requires a spare file descriptor solely to resize an
empty `BufWriter`. If opening the data file used the executor's last descriptor
and the pool refuses the configured buffer, attachment now fails with `EMFILE`
before processing a row. Move the original `File` into the resized writer or
wrap it once at the granted capacity.
Evidence: A child process with `RLIMIT_NOFILE=32` left one descriptor for
the output file. With a 64 KiB pool, a 1 MiB write buffer, and a two-partition,
16-row input requiring no spill, the base published `[0, 262, 536]`. The head
opened the file, failed in `try_clone()` with OS error 24, and published no
offsets.
##########
native/shuffle/src/writers/local/mod.rs:
##########
@@ -17,3 +17,18 @@
pub(crate) mod local_partition_writer;
mod spill;
+
+use datafusion::execution::memory_pool::MemoryReservation;
+
+const FALLBACK_BUFFER_SIZE: usize = 8 * 1024;
+
+/// Returns the capacity to allocate. A refused size falls back to a small
buffer charged with
+/// `grow`, so the task carries on and the reservation still matches what is
held.
+fn reserve_buffer(reservation: &MemoryReservation, size: usize) -> usize {
+ if reservation.try_grow(size).is_ok() {
Review Comment:
[P2] `try_grow(size).is_ok()` treats a Spark exception as permission to use
the 8 KiB fallback. If a peer's `spill()` throws `IOException`, the map can
publish offsets while Spark still holds a partial grant. `restore_buffer`
discards the same error, and `sync_buffer_reservation` can hide it through
infallible `resize` even after the file-buffer request succeeds. Use fallible
growth for reusable capacity, recover only from `ResourcesExhausted`, and
propagate other errors through attachment, writing, and the post-spill hook.
Evidence: With Spark 4.1.3, a 2 MiB off-heap pool, and a peer holding 1.5
MiB that threw on its first spill, both `fair_unified` and `greedy_unified`
published `[0, 262, 536]` and left 524,288 bytes charged until task cleanup.
The base made no peer-spill call and left zero. In a separate run the 1 MiB
file grant and 256-byte input grant both succeeded. Spark threw on the 516-byte
scratch charge, yet the head finished and stranded 256 bytes. Injecting
`External` only during spill-buffer restoration also allowed publication.
--
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]
