viirya commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4090257247
##########
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:
One way to price the read-ahead slot here without changing admission at read
buffer 2:
```rust
*total += get_reserved_bytes_for_record_batch_size(
spill.max_record_batch_memory,
spill.max_record_batch_memory,
) + spill.max_record_batch_memory * buffer_len;
```
`get_reserved_bytes_for_record_batch_size`'s own doc says its `2x` covers
the batch plus its sorted copies, so the channel slots are a separate term.
This keeps read buffer 2 at `4x` per run, exactly the current `2x * 2`, so
main's default admission does not change by a byte. Read buffer 1 then costs
`3x`, which is the "widened at 3x per run" row in @comphead's table (1.23 of
the grant against 1.29 for main). The widened selection would take fewer runs
than today, but it would stay within the narrow pass's live footprint. Could
you confirm the peak with the same counting-allocator probe once this is in?
The `memory_limit` comment at line 557 would then need revisiting as well.
It currently describes the second `x` as "one materialized output", which is
neither the documented meaning of that term nor something that accounts for the
channel slot.
##########
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:
Retracting my "That checks out" here. Matching the admission arithmetic
isn't enough, because the admission arithmetic doesn't price the read-ahead
slot — see @comphead's reply above. I've left a suggested pricing change on
`spill_merge_memory_requirements` (line 1003).
##########
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 "as before" that @comphead flagged at line 170 comes from my suggested
wording ("which is the previous behaviour"). My mistake: it describes this PR's
history rather than the code, and `main` never flushed at input boundaries. If
the estimator stays, "falls back to flushing at every input boundary that has
pending rows" says the same thing without the history. If pricing the read slot
removes the estimator, this goes away anyway.
--
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]
