viirya commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4054345925
##########
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:
The doc explains what the estimate covers but not how it behaves when the
budget is tight, which is the question a reader debugging a fragmented
intermediate run will arrive with.
The floor is safe, and saying so is reassuring: the check runs only at input
batch boundaries, never per row, and returns `false` on `is_empty()`, so the
worst case degenerates to one output batch per input batch — exactly the
unconditional flushing this replaces — and can never produce a zero-row batch
or fail to make progress.
Something like: "When the allowance is small this degrades to flushing at
every input boundary, which is the previous behaviour; it never blocks
progress."
##########
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:
Worth pinning this invariant here with a `debug_assert`.
`merge.rs:302` indexes `budget.input_batch_sizes[winner_stream]`, but the
stream vector built below starts with the in-memory streams from
`mem::take(&mut self.sorted_streams)` and appends the spill streams after them,
while `input_batch_sizes` is built from `sorted_spill_files` alone. The indices
line up only because widening requires `self.sorted_streams.is_empty()` at line
449, three layers up from here.
Not a live bug — I traced every `sorted_streams` mutation between line 449
and line 507 and there is none, and the `input_batch_sizes.len() ==
streams.len()` check in `streaming_merge.rs` would catch a mismatch as an
internal error rather than a silent misattribution. But nothing at this
signature says `bound_batch_memory` implies "no in-memory streams", so a later
change that admits in-memory streams alongside a widened pass would be
attributing the wrong per-input size to each stream.
A `debug_assert!(!bound_batch_memory || sorted_streams.is_empty())` after
the `mem::take`, or a line on `bound_batch_memory` in the doc comment, would
make the coupling local.
##########
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:
`memory_limit` here is `sum(max_record_batch_memory) * 2`, and the comment
says it is "the same two-batch allowance used for spill admission". That checks
out — `get_reserved_bytes_for_record_batch_size(x, x)` is `x + x`, so admission
reserves `2 * max_batch_memory` per run per buffer, and this equals the
admitted per-run allowance at `buffer_len == 1`, which is what a widened pass
runs at.
That equivalence is load-bearing and only holds while widening forces
`buffer_size = 1`. Worth naming the `buffer_len == 1` assumption in the comment
so a future change to the widened read-ahead does not silently leave this
budget over-generous.
--
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]
