sunchao commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4125261149
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -373,95 +424,163 @@ impl MultiLevelMergeBuilder {
let minimum_number_of_required_streams =
2_usize.saturating_sub(self.sorted_streams.len());
- let (sorted_spill_files, buffer_size) = match self
- .get_sorted_spill_files_to_merge(
- 2,
- // we must have at least 2 streams to merge
- minimum_number_of_required_streams,
- &mut memory_reservation,
- allow_minimum_without_headroom,
- )? {
- SpillFilesToMerge::Ready(sorted_spill_files, buffer_size)
=> {
- (sorted_spill_files, buffer_size)
+ let selection = self.get_sorted_spill_files_to_merge(
+ 2,
+ minimum_number_of_required_streams,
+ &mut memory_reservation,
+ allow_minimum_without_headroom,
+ )?;
+ let (mut spills, mut buffer_size) = match selection {
+ SpillFilesToMerge::Ready(spills, buffer_size) => {
+ (spills, buffer_size)
}
- // Not enough memory to seat 2 streams. Re-spill the
blocking file
- // smaller and retry. `get_sorted_spill_files_to_merge`
already freed
- // the reservation and `self.sorted_streams` is untouched,
so the
- // retry starts clean.
SpillFilesToMerge::SplitThenRetry(index) => {
return Ok(MergeStep::SplitThenRetry(index));
}
};
- // Don't account for existing streams memory
- // as we are not holding the memory for them
- let mut sorted_streams = mem::take(&mut self.sorted_streams);
-
- let is_only_merging_memory_streams =
sorted_spill_files.is_empty();
-
- // If no spill files were selected (e.g. all too large for
- // available memory but enough in-memory streams exist),
- // return the pre-reserved bytes to self.reservation so
- // create_new_merge_sort can transfer them to the merge
- // stream's BatchBuilder.
- if is_only_merging_memory_streams {
- mem::swap(&mut self.reservation, &mut memory_reservation);
+ let original_count = spills.len();
+ let original_buffer_size = buffer_size;
+ let original_memory = memory_reservation.size();
+ if self.reserve_replay_headroom
+ && self.widen_intermediate_merges
+ && !allow_minimum_without_headroom
+ && buffer_size > 1
+ && self.sorted_streams.is_empty()
+ && !spills.is_empty()
+ {
+ // Trade read-ahead for fan-in without taking any more pool
+ // memory. Other partitions retain exactly the space left
by
+ // the original admission, even while they replay
aggregates.
+ // Keep one run for the final merge and its replay
headroom.
+ let candidates = spills
+ .iter()
+ .chain(&self.sorted_spill_files)
+ .take(spills.len() + self.sorted_spill_files.len() - 1)
+ .map(|(spill, _)| spill);
+ let widened_count = spill_merge_memory_requirements(
Review Comment:
The final sizing rule in `78c4b7b` also accounts for the next admission's
replay headroom within the retained grant. The earlier nine-run/merge-three
example above is superseded: with headroom enabled, eight uniform full-batch
runs and an initial fan-in of seven now merge six and leave three for final
replay; nine runs keep the ordinary selection. The current description includes
fresh exact-base/head SF10 diagnostics, including actual trimmed passes and
unchanged/adverse cases.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -373,95 +424,163 @@ impl MultiLevelMergeBuilder {
let minimum_number_of_required_streams =
2_usize.saturating_sub(self.sorted_streams.len());
- let (sorted_spill_files, buffer_size) = match self
- .get_sorted_spill_files_to_merge(
- 2,
- // we must have at least 2 streams to merge
- minimum_number_of_required_streams,
- &mut memory_reservation,
- allow_minimum_without_headroom,
- )? {
- SpillFilesToMerge::Ready(sorted_spill_files, buffer_size)
=> {
- (sorted_spill_files, buffer_size)
+ let selection = self.get_sorted_spill_files_to_merge(
+ 2,
+ minimum_number_of_required_streams,
+ &mut memory_reservation,
+ allow_minimum_without_headroom,
+ )?;
+ let (mut spills, mut buffer_size) = match selection {
+ SpillFilesToMerge::Ready(spills, buffer_size) => {
+ (spills, buffer_size)
}
- // Not enough memory to seat 2 streams. Re-spill the
blocking file
- // smaller and retry. `get_sorted_spill_files_to_merge`
already freed
- // the reservation and `self.sorted_streams` is untouched,
so the
- // retry starts clean.
SpillFilesToMerge::SplitThenRetry(index) => {
return Ok(MergeStep::SplitThenRetry(index));
}
};
- // Don't account for existing streams memory
- // as we are not holding the memory for them
- let mut sorted_streams = mem::take(&mut self.sorted_streams);
-
- let is_only_merging_memory_streams =
sorted_spill_files.is_empty();
-
- // If no spill files were selected (e.g. all too large for
- // available memory but enough in-memory streams exist),
- // return the pre-reserved bytes to self.reservation so
- // create_new_merge_sort can transfer them to the merge
- // stream's BatchBuilder.
- if is_only_merging_memory_streams {
- mem::swap(&mut self.reservation, &mut memory_reservation);
+ let original_count = spills.len();
+ let original_buffer_size = buffer_size;
+ let original_memory = memory_reservation.size();
+ if self.reserve_replay_headroom
+ && self.widen_intermediate_merges
+ && !allow_minimum_without_headroom
+ && buffer_size > 1
+ && self.sorted_streams.is_empty()
+ && !spills.is_empty()
+ {
+ // Trade read-ahead for fan-in without taking any more pool
+ // memory. Other partitions retain exactly the space left
by
+ // the original admission, even while they replay
aggregates.
+ // Keep one run for the final merge and its replay
headroom.
+ let candidates = spills
+ .iter()
+ .chain(&self.sorted_spill_files)
+ .take(spills.len() + self.sorted_spill_files.len() - 1)
+ .map(|(spill, _)| spill);
+ let widened_count = spill_merge_memory_requirements(
+ candidates,
+ 1,
+ self.max_spill_merge_fan_in(),
+ )
+ .take_while(|needed| *needed <= original_memory)
+ .count();
+ if widened_count > original_count {
+ buffer_size = 1;
+ spills.extend(
+ self.sorted_spill_files
+ .drain(..widened_count - original_count),
+ );
+ }
}
+ let reservation = Arc::new(memory_reservation);
+ let widened = spills.len() > original_count;
+ let retry_reservation = widened.then(||
Arc::clone(&reservation));
+ let (stream, batch_size_limit) =
+ self.merge_selected_runs(&spills, buffer_size,
reservation, widened)?;
+ let retry = retry_reservation.map(|reservation|
IntermediateMergeRetry {
+ spills,
+ original_count,
+ buffer_size: original_buffer_size,
+ reservation,
+ });
+ Ok(MergeStep::Stream {
+ stream,
+ batch_size_limit,
+ retry,
+ })
+ }
+ }
+ }
- // Cap the merge output at the smallest limit among the runs
we're
- // about to merge. Runs that were shrunk for skew carry a
smaller limit,
- // if none do, every run carries `self.batch_size` and the
merge runs at
- // the full batch size. The output stream is tagged with the
same limit
- // (see the `MergeStep::Stream` returns below) so a re-spilled
- // intermediate run stays shrunk and won't rebuild an
oversized batch on
- // a later pass.
- let mut output_batch_size = self.batch_size;
- for (spill, batch_size_limit) in sorted_spill_files {
- let stream = self
- .spill_manager
- .clone()
- .with_batch_read_buffer_capacity(buffer_size)
- .read_spill_as_stream(
- spill.file,
- Some(spill.max_record_batch_memory),
- )?;
- output_batch_size =
output_batch_size.min(batch_size_limit);
- sorted_streams.push(stream);
- }
- let merge_sort_stream = self.create_new_merge_sort(
+ /// Build a stream from an already admitted selection. The reservation can
+ /// also be held by a retry guard until an intermediate writer finishes.
+ /// `bound_batch_memory` requires spill-only inputs and a one-batch read
buffer.
+ fn merge_selected_runs(
+ &mut self,
+ sorted_spill_files: &[(SortedSpillFile, usize)],
+ buffer_size: usize,
+ memory_reservation: Arc<MemoryReservation>,
+ bound_batch_memory: bool,
+ ) -> Result<(SendableRecordBatchStream, usize)> {
+ // Don't account for existing streams memory
+ // as we are not holding the memory for them
+ let mut sorted_streams = mem::take(&mut self.sorted_streams);
+ debug_assert!(!bound_batch_memory || sorted_streams.is_empty());
+ debug_assert!(!bound_batch_memory || buffer_size == 1);
+
+ let is_only_merging_memory_streams = sorted_spill_files.is_empty();
+
+ // If no spill files were selected (e.g. all too large for
+ // available memory but enough in-memory streams exist),
+ // return the pre-reserved bytes to self.reservation so
+ // create_new_merge_sort can transfer them to the merge
+ // stream's BatchBuilder.
+ if is_only_merging_memory_streams {
+ self.reservation = Arc::try_unwrap(memory_reservation)
Review Comment:
Confirming the final ownership model at `78c4b7b`: there is no widened-write
retry or shared reservation wrapper. Only an actually trimmed intermediate pass
keeps the existing plain `MemoryReservation` in the builder through EOF and
spill completion, then transfers it to the next admission. Other passes retain
the ordinary stream-attached ownership. The contention and builder-drop tests
cover the retained grant and cleanup.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -373,95 +424,158 @@ impl MultiLevelMergeBuilder {
let minimum_number_of_required_streams =
2_usize.saturating_sub(self.sorted_streams.len());
- let (sorted_spill_files, buffer_size) = match self
- .get_sorted_spill_files_to_merge(
- 2,
- // we must have at least 2 streams to merge
- minimum_number_of_required_streams,
- &mut memory_reservation,
- allow_minimum_without_headroom,
- )? {
- SpillFilesToMerge::Ready(sorted_spill_files, buffer_size)
=> {
- (sorted_spill_files, buffer_size)
+ let selection = self.get_sorted_spill_files_to_merge(
+ 2,
+ minimum_number_of_required_streams,
+ &mut memory_reservation,
+ allow_minimum_without_headroom,
+ )?;
+ let (mut spills, mut buffer_size) = match selection {
+ SpillFilesToMerge::Ready(spills, buffer_size) => {
+ (spills, buffer_size)
}
- // Not enough memory to seat 2 streams. Re-spill the
blocking file
- // smaller and retry. `get_sorted_spill_files_to_merge`
already freed
- // the reservation and `self.sorted_streams` is untouched,
so the
- // retry starts clean.
SpillFilesToMerge::SplitThenRetry(index) => {
return Ok(MergeStep::SplitThenRetry(index));
}
};
- // Don't account for existing streams memory
- // as we are not holding the memory for them
- let mut sorted_streams = mem::take(&mut self.sorted_streams);
-
- let is_only_merging_memory_streams =
sorted_spill_files.is_empty();
-
- // If no spill files were selected (e.g. all too large for
- // available memory but enough in-memory streams exist),
- // return the pre-reserved bytes to self.reservation so
- // create_new_merge_sort can transfer them to the merge
- // stream's BatchBuilder.
- if is_only_merging_memory_streams {
- mem::swap(&mut self.reservation, &mut memory_reservation);
+ let original_count = spills.len();
+ let original_buffer_size = buffer_size;
+ let original_memory = memory_reservation.size();
+ if self.reserve_replay_headroom
+ && self.widen_intermediate_merges
+ && !allow_minimum_without_headroom
+ && buffer_size > 1
+ && self.sorted_streams.is_empty()
+ && !spills.is_empty()
+ {
+ // Trade read-ahead for fan-in without taking any more pool
+ // memory. Other partitions retain exactly the space left
by
+ // the original admission, even while they replay
aggregates.
+ // Keep one run for the final merge and its replay
headroom.
+ let candidates = spills
+ .iter()
+ .chain(&self.sorted_spill_files)
+ .take(spills.len() + self.sorted_spill_files.len() - 1)
+ .map(|(spill, _)| spill);
+ let widened_count = spill_merge_memory_requirements(
+ candidates,
+ 1,
+ self.max_spill_merge_fan_in(),
+ )
+ .take_while(|needed| *needed <= original_memory)
+ .count();
+ if widened_count > original_count {
+ buffer_size = 1;
+ spills.extend(
+ self.sorted_spill_files
+ .drain(..widened_count - original_count),
+ );
+ }
}
+ let reservation = Arc::new(memory_reservation);
+ let widened = spills.len() > original_count;
+ let retry_reservation = widened.then(||
Arc::clone(&reservation));
+ let (stream, batch_size_limit) =
+ self.merge_selected_runs(&spills, buffer_size,
reservation, widened)?;
+ let retry = retry_reservation.map(|reservation|
IntermediateMergeRetry {
+ spills,
+ original_count,
+ buffer_size: original_buffer_size,
+ reservation,
+ });
+ Ok(MergeStep::Stream {
+ stream,
+ batch_size_limit,
+ retry,
+ })
+ }
+ }
+ }
- // Cap the merge output at the smallest limit among the runs
we're
- // about to merge. Runs that were shrunk for skew carry a
smaller limit,
- // if none do, every run carries `self.batch_size` and the
merge runs at
- // the full batch size. The output stream is tagged with the
same limit
- // (see the `MergeStep::Stream` returns below) so a re-spilled
- // intermediate run stays shrunk and won't rebuild an
oversized batch on
- // a later pass.
- let mut output_batch_size = self.batch_size;
- for (spill, batch_size_limit) in sorted_spill_files {
- let stream = self
- .spill_manager
- .clone()
- .with_batch_read_buffer_capacity(buffer_size)
- .read_spill_as_stream(
- spill.file,
- Some(spill.max_record_batch_memory),
- )?;
- output_batch_size =
output_batch_size.min(batch_size_limit);
- sorted_streams.push(stream);
- }
- let merge_sort_stream = self.create_new_merge_sort(
+ /// Build a stream from an already admitted selection. The reservation can
+ /// also be held by a retry guard until an intermediate writer finishes.
+ fn merge_selected_runs(
+ &mut self,
+ sorted_spill_files: &[(SortedSpillFile, usize)],
+ buffer_size: usize,
+ memory_reservation: Arc<MemoryReservation>,
+ bound_batch_memory: bool,
+ ) -> Result<(SendableRecordBatchStream, usize)> {
+ // Don't account for existing streams memory
+ // as we are not holding the memory for them
+ let mut sorted_streams = mem::take(&mut self.sorted_streams);
+
+ let is_only_merging_memory_streams = sorted_spill_files.is_empty();
+
+ // If no spill files were selected (e.g. all too large for
+ // available memory but enough in-memory streams exist),
+ // return the pre-reserved bytes to self.reservation so
+ // create_new_merge_sort can transfer them to the merge
+ // stream's BatchBuilder.
+ if is_only_merging_memory_streams {
+ self.reservation = Arc::try_unwrap(memory_reservation)
+ .expect("in-memory merges do not retain a spill retry
reservation");
+ return Ok((
+ self.create_new_merge_sort(
sorted_streams,
- // If we have no sorted spill files left, this is the last
run
self.sorted_spill_files.is_empty(),
- is_only_merging_memory_streams,
- output_batch_size,
- )?;
-
- // If we're only merging memory streams, we don't need to
attach the memory reservation
- // as it's empty
- if is_only_merging_memory_streams {
- assert_eq!(
- memory_reservation.size(),
- 0,
- "when only merging memory streams, we should not have
any memory reservation and let the merge sort handle the memory"
- );
+ true,
+ self.batch_size,
+ None,
+ )?,
+ self.batch_size,
+ ));
+ }
- Ok(MergeStep::Stream {
- stream: merge_sort_stream,
- batch_size_limit: output_batch_size,
- })
- } else {
- // Attach the memory reservation to the stream to make
sure we have enough memory
- // throughout the merge process as we bypassed the memory
pool for the merge sort stream
- Ok(MergeStep::Stream {
- stream: Box::pin(StreamAttachedReservation::new(
- merge_sort_stream,
- memory_reservation,
- )),
- batch_size_limit: output_batch_size,
- })
- }
- }
+ // Cap the merge output at the smallest limit among the runs we're
+ // about to merge. Runs that were shrunk for skew carry a smaller
limit,
+ // if none do, every run carries `self.batch_size` and the merge runs
at
+ // the full batch size. The output stream is tagged with the same limit
+ // (see the `MergeStep::Stream` returns below) so a re-spilled
+ // intermediate run stays shrunk and won't rebuild an oversized batch
on
+ // a later pass.
+ let mut output_batch_size = self.batch_size;
+ for (spill, batch_size_limit) in sorted_spill_files {
+ let stream = self
+ .spill_manager
+ .clone()
+ .with_batch_read_buffer_capacity(buffer_size)
+ .read_spill_as_stream(
+ Arc::clone(&spill.file),
+ Some(spill.max_record_batch_memory),
+ )?;
+ output_batch_size = output_batch_size.min(*batch_size_limit);
+ sorted_streams.push(stream);
}
+ let batch_memory_budget = bound_batch_memory.then(|| {
Review Comment:
One update to the earlier reply for the final head, `78c4b7b`: ordinary
admission pricing and read-ahead remain unchanged, and widening, custom batch
budgets, and boundary flushing are gone. Reservation ownership now differs
deliberately for actually trimmed intermediate passes: the builder retains the
original plain grant through the write and transfers it to the next admission,
without taking extra pool memory. This fixes the subsequent contention finding;
it does not assert that baseline admission accounting bounds allocator usage or
RSS.
--
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]
