sunchao commented on code in PR #25428: URL: https://github.com/apache/datafusion/pull/25428#discussion_r4054156943
########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). I removed the Case G fan-in override. Widening now trades read-ahead from two batches to one within the original admitted reservation; it never grows that reservation. This preserves the pool capacity that the narrow selection leaves for peer replay, without a racy pool-usage snapshot or guessed partition share. The new deterministic test compares both policies with a competing consumer. On `6987df2ec`, widening retains 32 KiB rather than 16 KiB and the peer's next 4 KiB allocation fails with no space remaining. Both selections pass on the revised code. Seeded concurrent Greedy/Fair aggregate tests and the default-fan-in SQL case are also covered. The commit body records the original CI allocation failure and why the test cap was insufficient. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). The caller now supplies an explicit source/output allowance and per-input maximum sizes. Before replacing a source, the merge checks retained inputs plus a conservative bound on output that can accumulate before the next boundary; it no longer flushes unconditionally. This uses the admitted spill budget, not the bypass pool's reservation size. For your 8-run × 10-batch cases, the revised code produces 80 full output batches for 8,192-row inputs (previously 150), and 38 batches for 1,000-row inputs (previously 80). Both cursor paths have regression coverage. The short-input case still pays for conservative memory bounds; dictionary cleanup and overflow draining remain covered. The PR description now reports these counts and the isolated base/head benchmark. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc) by retaining the original error as context if either rebuilding or writing the narrower retry fails. The insufficient-quota tests verify that both quota errors survive. I kept the single bounded retry because quota failures, OS disk-full errors, and custom backend failures do not share a portable typed quota variant here. Dropping the merge future cancels and releases it; that cancellation does not enter this retry branch. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). Admission and widening now share the cumulative spill-cost/fan-in iterator and fan-in lookup. The widening phase performs no pool allocation: it chooses a larger selection at read-buffer capacity one within the already admitted grant, so its selection remains separate from fallible admission. It also explicitly excludes `allow_minimum_without_headroom` and capacity-one selections. The original grant stays live throughout. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc): `merge_selected_runs` borrows the spill slice and clones the file Arcs internally, then the original vector moves into the retry guard. This removes the duplicate metadata vector. I kept the guard in `MergeStep::Stream` so the stream, backup inputs, and retained reservation remain one owned result; storing it separately on the builder would loosen that lifetime coupling. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). Cleanup now returns an internal error with an invariant message rather than panicking. The comment also clarifies that bypassed spill merges update local accounting here; their actual pool grant remains attached to the outer merge stream and is not released by this call. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). The parameter is now `drop_consumed`; the ordinary retention path passes `false`, and the boundary cleanup path passes `true`. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). The four existing parent helpers are now `pub(super)` and imported by the replay tests. The remaining small wrapper only enables replay headroom. Shared fixture construction stays in one place. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc) with a shared singleton/short-run fixture matrix and seeded cases in the existing constrained-memory fuzz harness. Two aggregate pipelines are constructed before polling and run concurrently against one Greedy or Fair pool, with varied/oversized payloads and default fan-in. The cases verify every key and ARRAY_AGG payload, spilling, reservation bounds, and cleanup. I retained the direct minimum-admission probe and the deterministic HandoffPool test: they force specific selection and release/re-admission transitions that randomized scheduling does not guarantee. The new matrix checks the distinct end-to-end propagation of short-run limits. ########## 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: Removed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). Spill rows and bytes are now included in the assertion diagnostic. ########## 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: Addressed in [0ee7a4d06](https://github.com/apache/datafusion/commit/0ee7a4d06aad89bd699059d9e0cf8a613120a7dc). I removed the redundant tie-mode axis from the unique-key matrix and added separate duplicate-key tests with source tags and Pending input boundaries, covering both tie modes and cursor paths. The ordinary/no-budget control remains to verify that the new policy does not change ordinary batching. -- 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]
