sunchao commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4055100674
##########
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(
Review Comment:
Added the local spill-only `debug_assert!` immediately after `mem::take`,
and documented the `bound_batch_memory` contract in `merge_selected_runs`. This
keeps the input-size/stream-index coupling visible at the helper itself.
Included in rebased head `ede48fcdd`.
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -129,6 +129,201 @@ impl BatchBuilder {
&self.schema
}
+ /// Release fully consumed batches after a merge drains at an input
boundary.
+ /// Keeping their dictionaries can otherwise enlarge the next output even
+ /// though none of its rows refer to those batches.
+ pub(super) fn discard_consumed_batches(&mut self) -> Result<()> {
+ assert_or_internal_err!(
+ self.indices.is_empty(),
+ "pending merge rows must be emitted before discarding source
batches"
+ );
+ self.retain_current_batches(true);
+ // Bypassed spill merges only update their local accounting here; their
+ // real pool reservation remains attached to the outer merge stream.
+ self.release_unused_memory();
+ Ok(())
+ }
+
+ /// Whether replacing an exhausted input would exceed the allowance for
+ /// retained source batches and materializing output together. This
preserves
+ /// the caller's existing source/output estimate; cursor, read-ahead and
IPC
+ /// allocations still depend on the merge's heuristic workspace
reservation.
+ pub(super) fn should_flush_before_input(
Review Comment:
Documented the fallback to flushing at each input boundary with pending
rows, and that an empty builder skips the check so the policy cannot emit empty
batches or wait for a larger allowance. I avoided promising one output batch
per input because normal batch-size limits and overflow recovery can also split
output. Included in `ede48fcdd`.
##########
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:
The comment now names `buffer_size == 1`, and a local `debug_assert!`
enforces that assumption whenever the batch budget is enabled. It also
distinguishes the selected runs' two-batch estimate from the retained original
grant: the former can be smaller. Included in `ede48fcdd`.
--
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]
