sunchao commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4097144385
##########
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:
You are right: an unchanged pool reservation did not account for those extra
live batches. Simplified in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2):
the PR now only reduces an intermediate merge's selected inputs when fewer
suffice to target the admitted final fan-in. I removed widening and its
source/output budget. Ordinary per-run admission, read-ahead, and reservation
ownership now match the base implementation; the change makes no new
total-allocation or RSS guarantee. The reduced scope and final-revision
evidence are in the [updated PR
description](https://github.com/apache/datafusion/pull/25428).
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -855,16 +990,34 @@ fn effective_spill_merge_fan_in(configured_fan_in: usize)
-> usize {
}
}
+/// Cumulative buffer costs, shared by admission and fixed-budget widening.
+fn spill_merge_memory_requirements<'a>(
+ spills: impl Iterator<Item = &'a SortedSpillFile>,
+ buffer_len: usize,
+ max_spill_files: usize,
+) -> impl Iterator<Item = usize> {
+ spills.take(max_spill_files).scan(0, move |total, spill| {
+ *total += get_reserved_bytes_for_record_batch_size(
+ spill.max_record_batch_memory,
+ spill.max_record_batch_memory,
+ ) * buffer_len;
Review Comment:
Removed the buffer-one widening path in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2),
so this PR no longer needs a separate read-ahead pricing rule. The original
admission calculation and buffer fallback remain unchanged. Sizing keeps the
admitted reservation while returning an unnecessary suffix of selected files
for a later merge. See the [updated PR
description](https://github.com/apache/datafusion/pull/25428) for the final
scope and evidence.
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -284,6 +295,13 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
);
drop(timer);
+ if self.flush_on_input_batch_boundary {
Review Comment:
Removed the PR's `MergeBatchMemoryBudget`, input-boundary flushing, and
output-estimator changes in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2).
`BatchBuilder` and the output-batching code in `merge.rs` now match the base.
The remaining optimization only changes eligible intermediate input selection;
no type-specific estimator or new full-batch guarantee remains. The [PR
description](https://github.com/apache/datafusion/pull/25428) reports
final-revision diagnostics separately.
##########
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:
This is now the entire optimization in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2).
With nine uniform runs and admission for seven, it merges three and leaves
seven for final replay, retaining the original grant. The guard requires
matching recorded maximum batch memory and batch-size limits across all
admitted and pending runs; heterogeneous layouts keep the original selection.
Only `AggregateSpill` opts in, leaving legacy aggregation unchanged. Final
admission still runs normally, so different output sizes or pool availability
can require another pass. Tests and exact-base/final-head measurements are in
the [updated description](https://github.com/apache/datafusion/pull/25428).
##########
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:
Removed the wider-write retry path entirely in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2).
The existing `StreamAttachedReservation` owns the ordinary
`MemoryReservation`; this PR adds no shared reservation owner or retry
lifetime. Sizing retains that admitted reservation and only returns unneeded
selected files before constructing the stream. The [updated
description](https://github.com/apache/datafusion/pull/25428) reflects this
smaller scope.
##########
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:
Removed the boundary-flushing policy and its documentation in
[`c17e315d0f`](https://github.com/apache/datafusion/commit/c17e315d0fe71cfc1db2c5a376dd1f5e70e14ed2),
including the stale "as before" wording. The remaining change is intermediate
input sizing, described with the nine-run/seven-input example in the [updated
PR description](https://github.com/apache/datafusion/pull/25428).
--
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]
