comphead commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4052429956
##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -198,6 +198,11 @@ FROM (
statement ok
SET datafusion.execution.target_partitions = 4
+# Bound merge buffers so each partition can allocate replay state while the
+# other partitions retain aggregate state in the shared greedy memory pool.
+statement ok
+SET datafusion.runtime.max_spill_merge_fan_in = 2
Review Comment:
**P1.** This is the part I would push back on. Reserving replay headroom
also implicitly bounded how much of a *shared* pool a single partition's merge
could hold, and at `target_partitions = 4` all four partitions merge
concurrently against one greedy pool. Widening removes that bound for
intermediate passes, and the change here caps fan-in rather than capping the
share.
Two costs:
- Case G no longer covers default fan-in at 4 partitions, which is the
configuration most likely to hit the new behaviour.
- The condition is timing dependent, so the cap hides it rather than closing
it. I removed both statements and ran the file 12 times on a 16-core machine:
12/12 green. It reproduces somewhere but not everywhere, which is the shape of
a contention window, not of a deterministic budget error.
Could widening instead be bounded by a share of the pool, or skipped while
other consumers hold state, so the SLT can keep default fan-in? If the cap is
genuinely the right fix, the commit message should say which failure it
addresses, because `test: bound merge fan-in for parallel aggregate spilling`
with an empty body does not record that.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -274,15 +278,46 @@ impl MultiLevelMergeBuilder {
return Ok(stream);
}
- // Need to sort to a spill file
- let Some((spill_file, max_record_batch_memory)) = self
+ // A wider merge can reduce total writes while increasing peak disk
+ // usage. Keep its inputs and admitted buffers until the writer has
+ // finished, so a failed write can retry the original smaller
merge.
+ let mut result = self
.spill_manager
.spill_record_batch_stream_and_return_max_batch_memory(
&mut stream,
"MultiLevelMergeBuilder intermediate spill",
)
- .await?
- else {
+ .await;
+ drop(stream);
+ // A successful write drops the backups before the next selection.
+ if let Some(retry) = retry
Review Comment:
The retry fires on any `Err`, but only a disk-capacity failure motivates it.
A memory, I/O or cancellation error pays a full re-read plus re-merge and then
surfaces the *second* error, with the first silently dropped. Worth gating on
the quota error specifically, or at least chaining the original error as
context so the real cause survives.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -373,95 +413,161 @@ 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, 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_memory = memory_reservation.size();
+ if self.reserve_replay_headroom
+ && self.widen_intermediate_merges
+ && self.sorted_streams.is_empty()
+ && !spills.is_empty()
+ {
+ // Extend the admitted selection without releasing its
grant
+ // or changing read-ahead. Keep one run for the final
merge,
+ // which must still leave space for aggregate replay.
+ let max_fan_in = effective_spill_merge_fan_in(
+ self.spill_manager
+ .env()
+ .disk_manager
+ .max_spill_merge_fan_in(),
+ );
+ let mut extra = 0;
+ for (spill, _) in self
+ .sorted_spill_files
+ .iter()
+ .take(self.sorted_spill_files.len().saturating_sub(1))
+ {
+ if spills.len() + extra >= max_fan_in {
+ break;
+ }
+ let additional =
get_reserved_bytes_for_record_batch_size(
+ spill.max_record_batch_memory,
+ spill.max_record_batch_memory,
+ ) * buffer_size;
+ if memory_reservation.try_grow(additional).is_err() {
+ break;
+ }
+ extra += 1;
+ }
+ spills.extend(self.sorted_spill_files.drain(..extra));
}
+ let reservation = Arc::new(memory_reservation);
+ let retry =
+ (spills.len() > original_count).then(||
IntermediateMergeRetry {
+ spills: spills
Review Comment:
`read_spill_as_stream` takes `Arc<dyn SpillFile>`, so `merge_selected_runs`
can take `&[(SortedSpillFile, usize)]` and `Arc::clone` internally. The caller
then keeps `spills` and moves it straight into the guard, and this hand-written
clone plus the duplicate `Vec` both go away. Deriving `Clone` on
`SortedSpillFile` (an `Arc` and a `usize`) would also do it.
Minor, same area: `retry` is `None` at four of the five `MergeStep::Stream`
construction sites. Holding it on the builder as `self.pending_retry` would
leave those four unrelated arms untouched.
##########
datafusion/physical-plan/src/sorts/multi_level_merge/replay_headroom_tests.rs:
##########
@@ -0,0 +1,552 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::*;
+
+use crate::expressions::PhysicalSortExpr;
+use arrow::array::{AsArray, Int64Array, StringArray};
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Field, Int64Type, Schema};
+use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer,
MemoryPool};
+use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::metrics::{ExecutionPlanMetricsSet,
SpillMetrics};
+use std::sync::atomic::{AtomicBool, Ordering};
+
+struct ReplayMergeFixture {
+ builder: MultiLevelMergeBuilder,
+ env: Arc<RuntimeEnv>,
+ pool: Arc<dyn MemoryPool>,
+ pool_size: usize,
+ metrics: SpillMetrics,
+ input_bytes: usize,
+}
+
+fn replay_merge_builder(
Review Comment:
`mod tests` in the parent file already provides `test_schema`,
`build_spill_manager`, `make_sorted_spill_file` and `build_merge_builder`, and
already holds replay-headroom tests
(`replay_headroom_splits_an_oversized_first_run`,
`replay_headroom_allows_only_an_indivisible_minimum`). This module re-creates
all four helpers. Either put the new cases in that module, or make the helpers
`pub(super)` and import them. Two replay-headroom test modules in one file will
be confusing to the next person touching this.
##########
datafusion/physical-plan/src/sorts/multi_level_merge/replay_headroom_tests.rs:
##########
@@ -0,0 +1,552 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::*;
+
+use crate::expressions::PhysicalSortExpr;
+use arrow::array::{AsArray, Int64Array, StringArray};
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Field, Int64Type, Schema};
+use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer,
MemoryPool};
+use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::metrics::{ExecutionPlanMetricsSet,
SpillMetrics};
+use std::sync::atomic::{AtomicBool, Ordering};
+
+struct ReplayMergeFixture {
+ builder: MultiLevelMergeBuilder,
+ env: Arc<RuntimeEnv>,
+ pool: Arc<dyn MemoryPool>,
+ pool_size: usize,
+ metrics: SpillMetrics,
+ input_bytes: usize,
+}
+
+fn replay_merge_builder(
+ spill_manager: SpillManager,
+ schema: SchemaRef,
+ spills: Vec<SortedSpillFile>,
+ pool: &Arc<dyn MemoryPool>,
+ batch_size: usize,
+) -> MultiLevelMergeBuilder {
+ MultiLevelMergeBuilder::new(
+ spill_manager,
+ schema,
+ spills,
+ vec![],
+ [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(),
+ BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ batch_size,
+ MemoryConsumer::new("replay headroom test").register(pool),
+ None,
+ false,
+ )
+ .with_replay_headroom(true)
+}
+
+/// Create equally sized, interleaved input runs.
+fn replay_merge_fixture(
+ run_count: usize,
+ rows_per_run: usize,
+ memory_batches: usize,
+ max_fan_in: usize,
+) -> Result<ReplayMergeFixture> {
+ let env = RuntimeEnvBuilder::new()
+ .with_max_spill_merge_fan_in(max_fan_in)
+ .build_arc()?;
+ let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64,
false)]));
+ let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+ let spill_manager =
+ SpillManager::new(Arc::clone(&env), metrics.clone(),
Arc::clone(&schema));
+ let mut spills = Vec::with_capacity(run_count);
+ for run in 0..run_count {
+ let values = Int64Array::from_iter_values(
+ (0..rows_per_run).map(|row| (row * run_count + run) as i64),
+ );
+ let batch = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(values)])?;
+ let (file, max_record_batch_memory) = spill_manager
+ .spill_record_batch_iter_and_return_max_batch_memory(
+ std::iter::once(Ok(batch)),
+ "replay headroom test input",
+ )?
+ .expect("a nonempty input must spill");
+ spills.push(SortedSpillFile {
+ file,
+ max_record_batch_memory,
+ });
+ }
+ let input_bytes = metrics.spilled_bytes.value();
+ let pool_size = memory_batches * spills[0].max_record_batch_memory;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(pool_size));
+ let builder =
+ replay_merge_builder(spill_manager, schema, spills, &pool,
rows_per_run);
+ Ok(ReplayMergeFixture {
+ builder,
+ env,
+ pool,
+ pool_size,
+ metrics,
+ input_bytes,
+ })
+}
+
+/// Return additional spill metrics, checking final replay space and cleanup.
+async fn merge_replay_runs(
+ run_count: usize,
+ rows_per_run: usize,
+ memory_batches: usize,
+) -> Result<(usize, usize, usize)> {
+ let ReplayMergeFixture {
+ builder,
+ env,
+ pool,
+ pool_size,
+ metrics,
+ input_bytes,
+ } = replay_merge_fixture(run_count, rows_per_run, memory_batches, 0)?;
+ let schema = Arc::clone(&builder.schema);
+ let replay = MemoryConsumer::new("replay consumer").register(&pool);
+ let mut stream = builder.create_spillable_merge_stream();
+ let mut batches = Vec::new();
+ while let Some(batch) = stream.try_next().await? {
+ assert!(pool.reserved() <= pool_size / 2);
+ assert!(crate::spill::get_record_batch_memory_size(&batch) <=
pool.reserved());
+ // Check that another consumer can actually claim the replay allowance.
+ replay.try_grow(pool_size / 2)?;
+ replay.free();
+ batches.push(batch);
+ }
+ let merged = concat_batches(&schema, &batches)?;
+ let expected = Int64Array::from_iter_values(0..(run_count * rows_per_run)
as i64);
+ assert_eq!(merged.column(0).as_primitive::<Int64Type>(), &expected);
+ assert_eq!(pool.reserved(), 0);
+ drop(stream);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+
+ Ok((
+ metrics.spill_file_count.value() - run_count,
+ metrics.spilled_rows.value() - run_count * rows_per_run,
+ metrics.spilled_bytes.value() - input_bytes,
+ ))
+}
+
+#[rstest::rstest]
+#[case::intermediate_uses_full_pool(6, 16, 0, 2, 16)]
+#[case::intermediate_holds_back_a_run(3, 16, 0, 1, 8)]
+#[case::final_disables_read_ahead(3, 12, 0, 0, 6)]
+#[case::fan_in_limited_intermediate(4, 8, 2, 2, 4)]
+#[case::fan_in_limited_final(2, 8, 2, 0, 4)]
+#[tokio::test]
+async fn replay_headroom_depends_on_merge_phase(
+ #[case] run_count: usize,
+ #[case] memory_batches: usize,
+ #[case] max_fan_in: usize,
+ #[case] remaining_runs: usize,
+ #[case] reserved_batches: usize,
+) -> Result<()> {
+ let ReplayMergeFixture {
+ mut builder,
+ env,
+ pool,
+ ..
+ } = replay_merge_fixture(run_count, 128, memory_batches, max_fan_in)?;
+ let batch_memory = builder.sorted_spill_files[0].0.max_record_batch_memory;
+ let MergeStep::Stream { stream, .. } =
+ builder.merge_sorted_runs_within_mem_limit(false)?
+ else {
+ panic!("the merge should fit without splitting a run");
+ };
+ assert_eq!(builder.sorted_spill_files.len(), remaining_runs);
+ assert_eq!(pool.reserved(), reserved_batches * batch_memory);
+ let batches: Vec<RecordBatch> = stream.try_collect().await?;
+ assert_eq!(
+ batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
+ (run_count - remaining_runs) * 128
+ );
+ drop(builder);
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_keeps_split_retries_before_intermediate_merges() ->
Result<()> {
+ let ReplayMergeFixture {
+ mut builder, pool, ..
+ } = replay_merge_fixture(3, 128, 6, 0)?;
+ assert!(matches!(
+ builder.merge_sorted_runs_within_mem_limit(false)?,
+ MergeStep::SplitThenRetry(_)
+ ));
+ assert_eq!(builder.sorted_spill_files.len(), 3);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_preserves_indivisible_run_batch_limits() ->
Result<()> {
+ let env = Arc::new(RuntimeEnv::default());
+ let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Utf8,
false)]));
+ let spill_manager = SpillManager::new(
+ Arc::clone(&env),
+ SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ Arc::clone(&schema),
+ );
+ let values = (0..96)
+ .map(|value| format!("{value:03}{}", "x".repeat(1024)))
+ .collect::<Vec<_>>();
+ let mut spills = Vec::new();
+ for run in 0..3 {
+ // Every batch is already indivisible, but the configured merge batch
+ // size is much larger. The split/retry path must discover the one-row
+ // limit before an intermediate pass can concatenate these batches.
+ let batches = values.iter().skip(run).step_by(3).map(|value| {
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(StringArray::from(vec![value.as_str()]))],
+ )
+ .map_err(Into::into)
+ });
+ let (file, max_record_batch_memory) = spill_manager
+ .spill_record_batch_iter_and_return_max_batch_memory(
+ batches,
+ "indivisible replay input",
+ )?
+ .expect("a nonempty input must spill");
+ spills.push(SortedSpillFile {
+ file,
+ max_record_batch_memory,
+ });
+ }
+ let pool_size = 6 * spills[0].max_record_batch_memory;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(pool_size));
+ let builder =
+ replay_merge_builder(spill_manager, Arc::clone(&schema), spills,
&pool, 8192);
+ let mut stream = builder.create_spillable_merge_stream();
+ let mut batches = Vec::new();
+ while let Some(batch) = stream.try_next().await? {
+ assert_eq!(batch.num_rows(), 1);
+ assert!(crate::spill::get_record_batch_memory_size(&batch) <=
pool.reserved());
+ assert!(pool.reserved() <= pool_size);
+ batches.push(batch);
+ }
+ let merged = concat_batches(&schema, &batches)?;
+ assert_eq!(
+ merged.column(0).as_string::<i32>(),
+ &StringArray::from(values)
+ );
+ drop(stream);
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_does_not_rewrite_intermediate_runs_twice() ->
Result<()> {
+ // The pool holds eight read-ahead inputs, or four plus equal replay
headroom.
+ // Four intermediate merges of eight runs leave four runs for the final
merge.
+ // Reserving headroom for every pass instead writes ten intermediate files,
+ // rewriting every input row twice before returning the final merge.
+ let (spill_count, spilled_rows, spilled_bytes) =
+ merge_replay_runs(32, 256, 32).await?;
+ println!(
+ "Intermediate spill: {spill_count} files, {spilled_rows} rows,
{spilled_bytes} bytes"
+ );
+ assert_eq!(spill_count, 4, "additional spill bytes: {spilled_bytes}");
+ assert_eq!(spilled_rows, 32 * 256);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_is_restored_after_intermediate_split_retries() ->
Result<()> {
+ // Two inputs need four batches of workspace. Three available batches force
+ // an intermediate split; the final merge then needs further splitting to
+ // leave replay headroom. Both phases must preserve every row and release
+ // all reservations and temporary files when the stream finishes.
+ let (spill_count, _, _) = merge_replay_runs(3, 128, 3).await?;
+ assert!(spill_count > 1, "the merge must spill and split runs");
+ Ok(())
+}
+
+#[tokio::test]
+async fn intermediate_merge_preserves_short_batch_replay() -> Result<()> {
Review Comment:
This and `replay_headroom_preserves_indivisible_run_batch_limits` above
share a fixture (1-row Utf8 runs of 1024-char values) and differ only in run
count and pool size, and the latter in turn overlaps the pre-existing
`replay_headroom_allows_only_an_indivisible_minimum`. One `rstest` matrix over
(run_count, memory_batches, expect_split) would cover all three at a third of
the length.
Wider point:
`datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs`
already has `RunTestWithLimitedMemoryArgs` with
`MemoryBehavior::TakeAllMemoryAtTheBeginning`, varying and oversized record
batches, and `assert_all_output_batches_roughly_match_batch_size_conf`. Those
are the same shapes this module hand-builds, including the 90-line
`HandoffPool` in `intermediate_merge_keeps_admitted_buffers`, and extending it
would add randomized coverage this PR currently has none of.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -373,95 +413,161 @@ 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, 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_memory = memory_reservation.size();
+ if self.reserve_replay_headroom
Review Comment:
This re-implements `get_sorted_spill_files_to_merge`.
`effective_spill_merge_fan_in`, the
`get_reserved_bytes_for_record_batch_size(x, x) * buffer_len` formula, the
fan-in cap and the drain all already exist there. It is the same loop with
`check_headroom = false` plus a "leave one run" stop condition, and that
function already carries an `allow_minimum_without_headroom` knob for exactly
this kind of variation. Folding it in keeps admission accounting in one place
and removes the second counter, the second drain and the duplicated fan-in
lookup.
Style drift worth noting while the two exist side by side: the original loop
grows to an absolute `total_needed`, this one grows by a delta. They agree only
because the shrink just above left the reservation at `accepted_memory`.
Separately: this also runs when `allow_minimum_without_headroom` is true,
that is, immediately after the pool failed to seat two streams with headroom.
Trying to grow past the just-admitted minimum in that state is contradictory.
Suggest adding `&& !allow_minimum_without_headroom`.
##########
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:
**P2.** This flushes on *every* input batch boundary, so a cut happens
whenever any one of N inputs exhausts a batch. In steady state the output batch
size converges to the *input* batch size, independent of `batch_size`.
Measured by driving `StreamingMergeBuilder` directly, 8 runs x 10 batches,
`batch_size = 8192`:
| input batch | flush off | flush on |
| --- | --- | --- |
| 1000 rows | 10 batches, avg 8000 rows | 80 batches, avg 1000 rows |
| 8192 rows | 80 batches, avg 8192 rows | 150 batches, avg 4369 rows |
So even when runs are written at the full `batch_size`, the intermediate run
ends up with about 1.9x the batches at about 53% the rows, and the final pass
re-reads all of them. That is the IPC framing cost the description mentions,
quantified, and it is paid on every widened pass.
The invariant you need is "pending output must fit the reserved workspace",
not "at most one batch per input". `BatchBuilder` already tracks
`batches_mem_used` and the per-merge grant is known at construction, so the
flush could be conditional on approaching that budget and keep full batches the
rest of the time. That would also make the short-wide-batch case the only one
that pays.
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -129,6 +129,42 @@ 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) {
+ assert!(self.indices.is_empty());
Review Comment:
Bare `assert!` with no message on a per-input-batch path. `merge.rs` uses
`assert_or_internal_err!` a few lines from the new call site, and
`debug_assert!` would also fit here since the invariant is private and
established by the flush immediately before.
Also: `release_unused_memory()` below is a no-op in practice.
`flush_on_input_batch_boundary` is only ever set on merges built with
`with_bypass_mempool()`, so the reservation is an `UnboundedMemoryPool` mock.
Fine to keep for internal consistency, but it reads as if it is returning bytes
to the real pool.
##########
datafusion/physical-plan/src/sorts/streaming_merge.rs:
##########
@@ -310,6 +323,181 @@ mod tests {
ExecutionPlanMetricsSet, SpillMetrics,
};
+ #[rstest::rstest]
+ #[case::primitive(DataType::Int32)]
+ #[case::strings(DataType::Utf8)]
+ #[case::views(DataType::Utf8View)]
+ #[tokio::test]
+ async fn intermediate_merge_flushes_before_replacing_input_batch(
+ #[case] data_type: DataType,
+ #[values(false, true)] row_cursor: bool,
+ #[values(false, true)] round_robin: bool,
Review Comment:
`round_robin` only affects tie-break ordering among equal keys and does not
interact with batch-boundary flushing, but it doubles a matrix that is already
3 data types x 2 cursor kinds x 2 fetch values, each running an inner `for
flush in [false, true]`, so 48 merges. Dropping it halves that with no coverage
loss. The `flush == false` branch is also asserting pre-existing batching
behaviour rather than anything this PR changes.
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -129,6 +129,42 @@ 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) {
+ assert!(self.indices.is_empty());
+ self.retain_current_batches(false);
+ self.release_unused_memory();
+ }
+
+ fn retain_current_batches(&mut self, keep_consumed: bool) {
Review Comment:
`retain_current_batches(true)` at the pre-existing call site reads as "keep
the current batches", which is the opposite of what the flag selects. Naming
the parameter for the behaviour it adds (`drop_consumed`) keeps the old site
readable, or split into two small named methods over a shared inner.
##########
datafusion/physical-plan/src/sorts/multi_level_merge/replay_headroom_tests.rs:
##########
@@ -0,0 +1,552 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::*;
+
+use crate::expressions::PhysicalSortExpr;
+use arrow::array::{AsArray, Int64Array, StringArray};
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Field, Int64Type, Schema};
+use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer,
MemoryPool};
+use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::metrics::{ExecutionPlanMetricsSet,
SpillMetrics};
+use std::sync::atomic::{AtomicBool, Ordering};
+
+struct ReplayMergeFixture {
+ builder: MultiLevelMergeBuilder,
+ env: Arc<RuntimeEnv>,
+ pool: Arc<dyn MemoryPool>,
+ pool_size: usize,
+ metrics: SpillMetrics,
+ input_bytes: usize,
+}
+
+fn replay_merge_builder(
+ spill_manager: SpillManager,
+ schema: SchemaRef,
+ spills: Vec<SortedSpillFile>,
+ pool: &Arc<dyn MemoryPool>,
+ batch_size: usize,
+) -> MultiLevelMergeBuilder {
+ MultiLevelMergeBuilder::new(
+ spill_manager,
+ schema,
+ spills,
+ vec![],
+ [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(),
+ BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ batch_size,
+ MemoryConsumer::new("replay headroom test").register(pool),
+ None,
+ false,
+ )
+ .with_replay_headroom(true)
+}
+
+/// Create equally sized, interleaved input runs.
+fn replay_merge_fixture(
+ run_count: usize,
+ rows_per_run: usize,
+ memory_batches: usize,
+ max_fan_in: usize,
+) -> Result<ReplayMergeFixture> {
+ let env = RuntimeEnvBuilder::new()
+ .with_max_spill_merge_fan_in(max_fan_in)
+ .build_arc()?;
+ let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64,
false)]));
+ let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+ let spill_manager =
+ SpillManager::new(Arc::clone(&env), metrics.clone(),
Arc::clone(&schema));
+ let mut spills = Vec::with_capacity(run_count);
+ for run in 0..run_count {
+ let values = Int64Array::from_iter_values(
+ (0..rows_per_run).map(|row| (row * run_count + run) as i64),
+ );
+ let batch = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(values)])?;
+ let (file, max_record_batch_memory) = spill_manager
+ .spill_record_batch_iter_and_return_max_batch_memory(
+ std::iter::once(Ok(batch)),
+ "replay headroom test input",
+ )?
+ .expect("a nonempty input must spill");
+ spills.push(SortedSpillFile {
+ file,
+ max_record_batch_memory,
+ });
+ }
+ let input_bytes = metrics.spilled_bytes.value();
+ let pool_size = memory_batches * spills[0].max_record_batch_memory;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(pool_size));
+ let builder =
+ replay_merge_builder(spill_manager, schema, spills, &pool,
rows_per_run);
+ Ok(ReplayMergeFixture {
+ builder,
+ env,
+ pool,
+ pool_size,
+ metrics,
+ input_bytes,
+ })
+}
+
+/// Return additional spill metrics, checking final replay space and cleanup.
+async fn merge_replay_runs(
+ run_count: usize,
+ rows_per_run: usize,
+ memory_batches: usize,
+) -> Result<(usize, usize, usize)> {
+ let ReplayMergeFixture {
+ builder,
+ env,
+ pool,
+ pool_size,
+ metrics,
+ input_bytes,
+ } = replay_merge_fixture(run_count, rows_per_run, memory_batches, 0)?;
+ let schema = Arc::clone(&builder.schema);
+ let replay = MemoryConsumer::new("replay consumer").register(&pool);
+ let mut stream = builder.create_spillable_merge_stream();
+ let mut batches = Vec::new();
+ while let Some(batch) = stream.try_next().await? {
+ assert!(pool.reserved() <= pool_size / 2);
+ assert!(crate::spill::get_record_batch_memory_size(&batch) <=
pool.reserved());
+ // Check that another consumer can actually claim the replay allowance.
+ replay.try_grow(pool_size / 2)?;
+ replay.free();
+ batches.push(batch);
+ }
+ let merged = concat_batches(&schema, &batches)?;
+ let expected = Int64Array::from_iter_values(0..(run_count * rows_per_run)
as i64);
+ assert_eq!(merged.column(0).as_primitive::<Int64Type>(), &expected);
+ assert_eq!(pool.reserved(), 0);
+ drop(stream);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+
+ Ok((
+ metrics.spill_file_count.value() - run_count,
+ metrics.spilled_rows.value() - run_count * rows_per_run,
+ metrics.spilled_bytes.value() - input_bytes,
+ ))
+}
+
+#[rstest::rstest]
+#[case::intermediate_uses_full_pool(6, 16, 0, 2, 16)]
+#[case::intermediate_holds_back_a_run(3, 16, 0, 1, 8)]
+#[case::final_disables_read_ahead(3, 12, 0, 0, 6)]
+#[case::fan_in_limited_intermediate(4, 8, 2, 2, 4)]
+#[case::fan_in_limited_final(2, 8, 2, 0, 4)]
+#[tokio::test]
+async fn replay_headroom_depends_on_merge_phase(
+ #[case] run_count: usize,
+ #[case] memory_batches: usize,
+ #[case] max_fan_in: usize,
+ #[case] remaining_runs: usize,
+ #[case] reserved_batches: usize,
+) -> Result<()> {
+ let ReplayMergeFixture {
+ mut builder,
+ env,
+ pool,
+ ..
+ } = replay_merge_fixture(run_count, 128, memory_batches, max_fan_in)?;
+ let batch_memory = builder.sorted_spill_files[0].0.max_record_batch_memory;
+ let MergeStep::Stream { stream, .. } =
+ builder.merge_sorted_runs_within_mem_limit(false)?
+ else {
+ panic!("the merge should fit without splitting a run");
+ };
+ assert_eq!(builder.sorted_spill_files.len(), remaining_runs);
+ assert_eq!(pool.reserved(), reserved_batches * batch_memory);
+ let batches: Vec<RecordBatch> = stream.try_collect().await?;
+ assert_eq!(
+ batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
+ (run_count - remaining_runs) * 128
+ );
+ drop(builder);
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_keeps_split_retries_before_intermediate_merges() ->
Result<()> {
+ let ReplayMergeFixture {
+ mut builder, pool, ..
+ } = replay_merge_fixture(3, 128, 6, 0)?;
+ assert!(matches!(
+ builder.merge_sorted_runs_within_mem_limit(false)?,
+ MergeStep::SplitThenRetry(_)
+ ));
+ assert_eq!(builder.sorted_spill_files.len(), 3);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_preserves_indivisible_run_batch_limits() ->
Result<()> {
+ let env = Arc::new(RuntimeEnv::default());
+ let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Utf8,
false)]));
+ let spill_manager = SpillManager::new(
+ Arc::clone(&env),
+ SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ Arc::clone(&schema),
+ );
+ let values = (0..96)
+ .map(|value| format!("{value:03}{}", "x".repeat(1024)))
+ .collect::<Vec<_>>();
+ let mut spills = Vec::new();
+ for run in 0..3 {
+ // Every batch is already indivisible, but the configured merge batch
+ // size is much larger. The split/retry path must discover the one-row
+ // limit before an intermediate pass can concatenate these batches.
+ let batches = values.iter().skip(run).step_by(3).map(|value| {
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(StringArray::from(vec![value.as_str()]))],
+ )
+ .map_err(Into::into)
+ });
+ let (file, max_record_batch_memory) = spill_manager
+ .spill_record_batch_iter_and_return_max_batch_memory(
+ batches,
+ "indivisible replay input",
+ )?
+ .expect("a nonempty input must spill");
+ spills.push(SortedSpillFile {
+ file,
+ max_record_batch_memory,
+ });
+ }
+ let pool_size = 6 * spills[0].max_record_batch_memory;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(pool_size));
+ let builder =
+ replay_merge_builder(spill_manager, Arc::clone(&schema), spills,
&pool, 8192);
+ let mut stream = builder.create_spillable_merge_stream();
+ let mut batches = Vec::new();
+ while let Some(batch) = stream.try_next().await? {
+ assert_eq!(batch.num_rows(), 1);
+ assert!(crate::spill::get_record_batch_memory_size(&batch) <=
pool.reserved());
+ assert!(pool.reserved() <= pool_size);
+ batches.push(batch);
+ }
+ let merged = concat_batches(&schema, &batches)?;
+ assert_eq!(
+ merged.column(0).as_string::<i32>(),
+ &StringArray::from(values)
+ );
+ drop(stream);
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0);
+ assert_eq!(env.disk_manager.used_disk_space(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn replay_headroom_does_not_rewrite_intermediate_runs_twice() ->
Result<()> {
+ // The pool holds eight read-ahead inputs, or four plus equal replay
headroom.
+ // Four intermediate merges of eight runs leave four runs for the final
merge.
+ // Reserving headroom for every pass instead writes ten intermediate files,
+ // rewriting every input row twice before returning the final merge.
+ let (spill_count, spilled_rows, spilled_bytes) =
+ merge_replay_runs(32, 256, 32).await?;
+ println!(
Review Comment:
Leftover debug output. All three values are already in scope for the asserts
below, so folding them into the existing assert message keeps the diagnostic
without the noise on every run.
--
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]
