From 40e5b36bc32f7efd1d0a0a0c7ac679d813b0752b Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 1 Nov 2026 13:56:57 +0800 Subject: [PATCH 1/3] feat: let SpillManager bound the size of the batches it writes --- Cargo.lock | 1 + datafusion/physical-plan/Cargo.toml | 1 + .../src/spill/in_progress_spill_file.rs | 45 ++- datafusion/physical-plan/src/spill/mod.rs | 1 + .../physical-plan/src/spill/spill_manager.rs | 31 ++ .../physical-plan/src/spill/spill_pieces.rs | 364 ++++++++++++++++++ 6 files changed, 428 insertions(+), 15 deletions(-) create mode 100644 datafusion/physical-plan/src/spill/spill_pieces.rs diff --git a/Cargo.lock b/Cargo.lock index b581b3c8d42d3..fa9c59c14357d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2540,6 +2540,7 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", + "arrow-select", "async-trait", "bytes", "criterion", diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 82a2215581f22..7b943dd531737 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,6 +64,7 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } +arrow-select = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } diff --git a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs index 0f83ef1a4b441..b98b33ed7b230 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -12,8 +12,8 @@ use datafusion_common::exec_datafusion_err; use datafusion_execution::spill_file::SpillFile; use super::{ - AsyncIPCStreamWriter, IPCStreamEncoder, IPCStreamWriter, gc_view_arrays, - spill_manager::{GetSlicedSize, SpillManager}, + AsyncIPCStreamWriter, IPCStreamEncoder, IPCStreamWriter, spill_manager::SpillManager, + spill_pieces::SpillPieces, }; enum InProgressWriter { @@ -50,7 +50,12 @@ impl InProgressSpillFile { /// buffers. When a view array is sliced, it still references the original full buffers, /// causing massive spill files without GC (see issue #19414: 820MB → 33MB after GC). /// - /// Returns the post-GC sliced memory size of the batch for memory accounting. + /// When the spill manager bounds batch bytes (see + /// [`SpillManager::with_max_batch_bytes`]), a larger batch is written as + /// several smaller batches in the same row order. + /// + /// Returns the post-GC sliced memory size of the largest batch written, which + /// is the largest batch a reader of the file decodes. /// /// # Errors /// - Returns an error if the file is not active (has been finalized) @@ -63,7 +68,7 @@ impl InProgressSpillFile { )); } - let gc_batch = gc_view_arrays(batch)?; + let pieces = SpillPieces::new(batch, self.spill_manager.max_batch_bytes())?; if self.writer.is_none() { // Use the SpillManager's declared schema rather than the batch's schema. @@ -84,18 +89,23 @@ impl InProgressSpillFile { self.spill_manager.metrics.spill_file_count.add(1); } } + let mut max_size = 0; if let Some(InProgressWriter::Sync(writer)) = &mut self.writer { - // The writer calculates how many serialized bytes were emitted - let (spilled_rows, delta_bytes) = writer.write(&gc_batch)?; + for piece in pieces { + let (piece, size) = piece?; + // The writer calculates how many serialized bytes were emitted + let (spilled_rows, delta_bytes) = writer.write(&piece)?; - self.spill_manager.metrics.spilled_rows.add(spilled_rows); - self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + max_size = max_size.max(size); + } } else if self.writer.is_some() { return Err(exec_datafusion_err!( "Cannot use synchronous append after asynchronous spill writing has started" )); } - gc_batch.get_sliced_size() + Ok(max_size) } /// Appends a `RecordBatch` using the asynchronous spill writer. @@ -106,7 +116,7 @@ impl InProgressSpillFile { )); } - let gc_batch = gc_view_arrays(batch)?; + let pieces = SpillPieces::new(batch, self.spill_manager.max_batch_bytes())?; if self.writer.is_none() { let schema = self.spill_manager.schema(); @@ -128,12 +138,17 @@ impl InProgressSpillFile { } } + let mut max_size = 0; match &mut self.writer { Some(InProgressWriter::Async(writer)) => { - let (spilled_rows, delta_bytes) = writer.write(&gc_batch).await?; - - self.spill_manager.metrics.spilled_rows.add(spilled_rows); - self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + for piece in pieces { + let (piece, size) = piece?; + let (spilled_rows, delta_bytes) = writer.write(&piece).await?; + + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + max_size = max_size.max(size); + } } Some(InProgressWriter::Sync(_)) => { return Err(exec_datafusion_err!( @@ -146,7 +161,7 @@ impl InProgressSpillFile { )); } } - gc_batch.get_sliced_size() + Ok(max_size) } pub fn flush(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index ab57d87dceecb..bd60543936a92 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -7,6 +7,7 @@ pub(crate) mod in_progress_spill_file; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; +pub(crate) mod spill_pieces; pub mod spill_pool; use datafusion_execution::spill_file::{AsyncSpillWriter, SpillWriter}; // Moved for refactor, re-export to keep the public API stable diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 57e819b1408c4..a828fb20f597c 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -32,6 +32,9 @@ pub struct SpillManager { batch_read_buffer_capacity: usize, /// general-purpose compression options pub(crate) compression: SpillCompression, + /// Upper bound on the in-memory size of each written batch, see + /// [`Self::with_max_batch_bytes`]. + max_batch_bytes: Option, } impl SpillManager { @@ -42,6 +45,7 @@ impl SpillManager { schema, batch_read_buffer_capacity: 2, compression: SpillCompression::default(), + max_batch_bytes: None, } } @@ -58,6 +62,33 @@ impl SpillManager { self } + /// Bounds the in-memory size of every batch written through this manager. + /// + /// Meant for managers that write sorted runs for a spill merge + /// ([`StreamingMergeBuilder::with_sorted_spill_files`]): the merge reserves + /// memory for the largest batch of each run, so a bound keeps the smallest + /// merge within a known amount whatever the row width. DataFusion's sort uses + /// `sort_spill_reservation_bytes / 4`; its aggregate also uses at most an + /// eighth of the state it held when it first spilled. + /// + /// A larger appended batch is written as several batches in the same row + /// order, so a reader gets more batches back than were appended. Readers that + /// expect one batch per append must not use a bounded manager. One row can + /// still exceed the bound, and so can batches with a column that cannot be + /// split by rows (dictionaries or run-end encoded arrays below the top level, + /// unions and list views). + /// + /// [`StreamingMergeBuilder::with_sorted_spill_files`]: crate::sorts::streaming_merge::StreamingMergeBuilder::with_sorted_spill_files + pub fn with_max_batch_bytes(mut self, max_batch_bytes: usize) -> Self { + self.max_batch_bytes = Some(max_batch_bytes); + self + } + + /// See [`Self::with_max_batch_bytes`]. + pub fn max_batch_bytes(&self) -> Option { + self.max_batch_bytes + } + /// Returns the schema for batches managed by this SpillManager pub fn schema(&self) -> &SchemaRef { &self.schema diff --git a/datafusion/physical-plan/src/spill/spill_pieces.rs b/datafusion/physical-plan/src/spill/spill_pieces.rs new file mode 100644 index 0000000000000..7ddc9a637a470 --- /dev/null +++ b/datafusion/physical-plan/src/spill/spill_pieces.rs @@ -0,0 +1,364 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Splits a batch into byte-bounded pieces before it is written to a spill file. + +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, AsArray, RecordBatch, RecordBatchOptions, UInt32Array, make_array, +}; +use arrow::compute::take; +use arrow::datatypes::{DataType, Schema}; +use arrow_data::transform::MutableArrayData; +use arrow_select::dictionary::garbage_collect_any_dictionary; +use datafusion_common::Result; + +use super::gc_view_arrays; +use super::spill_manager::GetSlicedSize; + +/// The pieces that one appended batch is written as, in row order. +/// +/// Without a bound, or when the batch is already within it, there is one piece: the batch +/// after [`gc_view_arrays`], as it is written without a bound. Its dictionaries are not +/// changed, so the batches of a spill file can keep sharing one dictionary. +/// +/// When the batch is larger than `max_batch_bytes`, it is cut into equal row ranges by its +/// average row size. Each range is copied out with [`compact`], so that it holds only its +/// own rows, and measured with [`GetSlicedSize`]. A range that is still too large (the rows +/// differ in size) is halved until it fits or has one row. Pieces are produced one at a +/// time, so that the caller can write each one before the next is copied. +/// +/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written whole. +pub(crate) struct SpillPieces { + batch: RecordBatch, + max_batch_bytes: usize, + /// Row ranges still to produce, the next one last. + pending: Vec>, + /// The whole batch, when it is written as one piece. + whole: Option<(RecordBatch, usize)>, +} + +impl SpillPieces { + pub(crate) fn new( + batch: &RecordBatch, + max_batch_bytes: Option, + ) -> Result { + let batch = gc_view_arrays(batch)?; + let size = batch.get_sliced_size()?; + let rows = batch.num_rows(); + let whole = |batch, size| Self { + whole: Some((batch, size)), + batch: RecordBatch::new_empty(Arc::new(Schema::empty())), + max_batch_bytes: 0, + pending: vec![], + }; + let Some(max_batch_bytes) = max_batch_bytes + .filter(|max| size > *max && rows > 1 && can_split(batch.schema_ref())) + else { + return Ok(whole(batch, size)); + }; + // A shared dictionary can hold many values that no row of this batch uses. Drop + // them once before choosing the row ranges, so that the ranges are sized by what + // the batch uses and the pieces do not each prune the whole dictionary again. + let batch = prune_dictionaries(&batch)?; + let size = batch.get_sliced_size()?; + if size <= max_batch_bytes { + return Ok(whole(batch, size)); + } + // The sliced size over-counts nested columns, whose children are not narrowed by + // a slice, so this can give more pieces than needed but not fewer. + let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows); + let rows_per_piece = rows.div_ceil(pieces); + let pending = (0..rows) + .step_by(rows_per_piece) + .map(|start| start..(start + rows_per_piece).min(rows)) + .rev() + .collect(); + Ok(Self { + batch, + max_batch_bytes, + pending, + whole: None, + }) + } +} + +impl Iterator for SpillPieces { + /// A piece and its size, as [`GetSlicedSize`] measures it. + type Item = Result<(RecordBatch, usize)>; + + fn next(&mut self) -> Option { + if let Some(whole) = self.whole.take() { + return Some(Ok(whole)); + } + while let Some(range) = self.pending.pop() { + let piece = match compact(&self.batch.slice(range.start, range.len())) { + Ok(piece) => piece, + Err(e) => return Some(Err(e)), + }; + let size = match piece.get_sliced_size() { + Ok(size) => size, + Err(e) => return Some(Err(e)), + }; + if size > self.max_batch_bytes && range.len() > 1 { + let mid = range.start + range.len() / 2; + self.pending.push(mid..range.end); + self.pending.push(range.start..mid); + continue; + } + return Some(Ok((piece, size))); + } + None + } +} + +/// Copies a slice so that it holds only its own rows: list children are narrowed to the +/// range the slice uses, view arrays are garbage collected, top-level dictionaries drop +/// the values that no key uses, and top-level run-end encoded arrays keep only the runs +/// the slice covers. +fn compact(slice: &RecordBatch) -> Result { + let columns = slice + .columns() + .iter() + .map(|column| { + if let DataType::RunEndEncoded(_, _) = column.data_type() { + // `MutableArrayData` does not support run-end encoded arrays; `take` + // re-encodes the runs the indices select, with their values copied. + let indices = UInt32Array::from_iter_values(0..column.len() as u32); + return Ok(take(column.as_ref(), &indices, None)?); + } + let data = column.to_data(); + let mut copy = MutableArrayData::new(vec![&data], false, data.len()); + copy.try_extend(0, 0, data.len())?; + let copied = make_array(copy.freeze()); + Ok(match copied.data_type() { + DataType::Dictionary(_, _) => { + garbage_collect_any_dictionary(copied.as_any_dictionary())? + } + _ => copied, + }) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(slice.num_rows())); + let piece = RecordBatch::try_new_with_options(slice.schema(), columns, &options)?; + gc_view_arrays(&piece) +} + +/// Drops the values of top-level dictionaries that no key uses. +fn prune_dictionaries(batch: &RecordBatch) -> Result { + if !batch + .schema_ref() + .fields() + .iter() + .any(|field| matches!(field.data_type(), DataType::Dictionary(_, _))) + { + return Ok(batch.clone()); + } + let columns = batch + .columns() + .iter() + .map(|column| { + Ok(match column.data_type() { + DataType::Dictionary(_, _) => { + garbage_collect_any_dictionary(column.as_any_dictionary())? + } + _ => Arc::clone(column), + }) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + Ok(RecordBatch::try_new_with_options( + batch.schema(), + columns, + &options, + )?) +} + +/// Whether [`compact`] narrows every column of `schema` to its rows. It does not for list +/// views and unions, whose children a row range does not bound, run-end encoded arrays +/// below the top level, or dictionaries below the top level, whose values it does not +/// prune. A batch with such a column is written whole. +fn can_split(schema: &Schema) -> bool { + fn supported(data_type: &DataType, top_level: bool) -> bool { + match data_type { + DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Union(_, _) => false, + DataType::RunEndEncoded(_, values) => { + top_level && supported(values.data_type(), false) + } + DataType::Dictionary(_, values) => top_level && supported(values, false), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => supported(field.data_type(), false), + DataType::Struct(fields) => fields + .iter() + .all(|field| supported(field.data_type(), false)), + _ => true, + } + } + schema + .fields() + .iter() + .all(|field| supported(field.data_type(), true)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + Array, DictionaryArray, Int32Array, ListBuilder, RunArray, StringArray, + StringBuilder, StringViewArray, StructArray, + }; + use arrow::datatypes::{Field, Int32Type}; + use arrow::util::display::{ArrayFormatter, FormatOptions}; + + fn batch_of(column: ArrayRef) -> RecordBatch { + RecordBatch::try_from_iter(vec![("c", column)]).unwrap() + } + + /// Each row of every column, formatted, so pieces can be compared with the + /// batch they came from whatever their encoding. + fn rows(batches: &[RecordBatch]) -> Vec { + let options = FormatOptions::default(); + batches + .iter() + .flat_map(|batch| { + let formatters = batch + .columns() + .iter() + .map(|c| ArrayFormatter::try_new(c.as_ref(), &options).unwrap()) + .collect::>(); + (0..batch.num_rows()) + .map(|row| { + formatters + .iter() + .map(|f| f.value(row).to_string()) + .collect::>() + .join("|") + }) + .collect::>() + }) + .collect() + } + + fn pieces(batch: &RecordBatch, bound: Option) -> Vec<(RecordBatch, usize)> { + SpillPieces::new(batch, bound) + .unwrap() + .collect::>>() + .unwrap() + } + + /// Splits `batch` and checks that the pieces keep every row in order and + /// that each is within `bound` or holds one row. Returns the piece count. + fn check_split(batch: &RecordBatch, bound: usize) -> usize { + let pieces = pieces(batch, Some(bound)); + for (piece, size) in &pieces { + assert_eq!(*size, piece.get_sliced_size().unwrap()); + assert!( + *size <= bound || piece.num_rows() == 1, + "{} rows, {size} bytes, bound {bound}", + piece.num_rows() + ); + } + let batches = pieces.into_iter().map(|(p, _)| p).collect::>(); + assert_eq!(rows(&batches), rows(std::slice::from_ref(batch))); + batches.len() + } + + fn wide_strings(rows: usize, width: usize) -> Vec { + (0..rows) + .map(|i| format!("{i}{}", "x".repeat(width))) + .collect() + } + + #[test] + fn without_a_bound_the_batch_is_one_piece() { + let batch = batch_of(Arc::new(StringArray::from(wide_strings(100, 1000)))); + let pieces = pieces(&batch, None); + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].0.num_rows(), 100); + } + + #[test] + fn wide_rows_are_split_in_row_order() { + let batch = batch_of(Arc::new(StringArray::from(wide_strings(100, 1000)))); + assert!(check_split(&batch, 10 * 1024) >= 10); + } + + #[test] + fn view_rows_are_split_in_row_order() { + let values = wide_strings(100, 1000); + let batch = batch_of(Arc::new(StringViewArray::from_iter_values(&values))); + assert!(check_split(&batch, 10 * 1024) >= 10); + } + + #[test] + fn skewed_rows_are_halved_until_they_fit() { + let values = (0..256) + .map(|i| "x".repeat(if i < 4 { 8 * 1024 } else { 8 })) + .collect::>(); + let batch = batch_of(Arc::new(StringArray::from(values))); + assert!(check_split(&batch, 10 * 1024) > 2); + } + + #[test] + fn one_row_larger_than_the_bound_is_one_piece() { + let batch = batch_of(Arc::new(StringArray::from(wide_strings(1, 100 * 1024)))); + assert_eq!(check_split(&batch, 1024), 1); + } + + #[test] + fn top_level_dictionaries_keep_only_their_values() { + let values = StringArray::from(wide_strings(256, 1000)); + let keys = Int32Array::from_iter_values(0..256); + let dictionary = DictionaryArray::::try_new(keys, Arc::new(values)); + let batch = batch_of(Arc::new(dictionary.unwrap())); + assert!(check_split(&batch, 16 * 1024) >= 16); + } + + #[test] + fn run_end_encoded_rows_are_split() { + let run_ends = Int32Array::from_iter_values((1..=128).map(|i| i * 2)); + let values = StringArray::from(wide_strings(128, 1000)); + let runs = RunArray::::try_new(&run_ends, &values).unwrap(); + let batch = batch_of(Arc::new(runs)); + assert!(check_split(&batch, 16 * 1024) >= 8); + } + + #[test] + fn lists_are_split() { + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in 0..100 { + for value in wide_strings(4, 250) { + builder.values().append_value(format!("{row}{value}")); + } + builder.append(true); + } + let batch = batch_of(Arc::new(builder.finish())); + assert!(check_split(&batch, 10 * 1024) >= 10); + } + + #[test] + fn nested_dictionaries_are_written_whole() { + let values = StringArray::from(wide_strings(64, 1000)); + let keys = Int32Array::from_iter_values(0..64); + let dictionary = + DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let field = Field::new("d", dictionary.data_type().clone(), false); + let nested = StructArray::try_new( + vec![field].into(), + vec![Arc::new(dictionary) as ArrayRef], + None, + ) + .unwrap(); + let batch = batch_of(Arc::new(nested)); + let pieces = pieces(&batch, Some(1024)); + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].0.num_rows(), 64); + } +} From a1566c318ff288e4906f2c4ec6272c1d198dfa91 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 1 Nov 2026 13:58:10 +0800 Subject: [PATCH 2/3] fix: write sorted spill runs in byte-bounded batches --- ...spilling_fuzz_in_memory_constrained_env.rs | 10 +- .../src/aggregates/grouped_hash_stream.rs | 18 ++- .../physical-plan/src/aggregates/spill.rs | 35 ++++- datafusion/physical-plan/src/sorts/builder.rs | 30 ++++ datafusion/physical-plan/src/sorts/merge.rs | 24 +++- datafusion/physical-plan/src/sorts/mod.rs | 2 +- .../src/sorts/multi_level_merge.rs | 20 ++- datafusion/physical-plan/src/sorts/sort.rs | 4 +- .../src/sorts/streaming_merge.rs | 26 +++- .../test_files/spill_batch_bound.slt | 129 ++++++++++++++++++ 10 files changed, 280 insertions(+), 18 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/spill_batch_bound.slt diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index ba566d597bb27..ab8473ffc06f4 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -978,7 +978,7 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( /// Assert baseline metrics are as expected or around that /// /// `output_batch_size` should be `None` when you expect to not get batched at the same size -/// `Some(session conf batch size)` for the rest +/// `Some(session conf batch size)` for the rest, where no batch has more rows than that fn assert_baseline_metrics_for_non_empty_output( metrics: &MetricsSet, expected_output_rows: usize, @@ -1009,9 +1009,11 @@ fn assert_baseline_metrics_for_non_empty_output( let output_batches = get_output_batches_from_metrics(metrics); if let Some(output_batch_size) = output_batch_size { - assert_eq!( - output_batches, - expected_output_rows.div_ceil(output_batch_size) + // A spilled merge also ends a batch once it reaches the spill batch byte + // bound, so wide rows can come in more, smaller batches. + assert!( + output_batches >= expected_output_rows.div_ceil(output_batch_size), + "{output_batches} output batches for {expected_output_rows} rows" ); } else { assert_ne!(output_batches, 0,); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 1ee6d02d13e8e..d713a7b9121c3 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -24,6 +24,7 @@ use std::vec; use super::aggregate_hash_table::{accumulator_phases, create_group_accumulator}; use super::order::GroupClustering; use super::skip_partial::SkipAggregationProbe; +use super::spill::bound_spill_batches_by_first_spill; use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, @@ -36,6 +37,7 @@ use crate::aggregates::{ group_id_array, max_duplicate_ordinal, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; +use crate::sorts::multi_level_merge::max_spill_batch_bytes; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; use crate::stream::EmptyRecordBatchStream; @@ -560,7 +562,14 @@ impl GroupedHashAggregateStream { metrics::SpillMetrics::new(&agg.metrics, partition), Arc::clone(&spill_schema), ) - .with_compression_type(context.session_config().spill_compression()); + .with_compression_type(context.session_config().spill_compression()) + .with_max_batch_bytes(max_spill_batch_bytes( + context + .session_config() + .options() + .execution + .sort_spill_reservation_bytes, + )); let spill_state = SpillState { spills: vec![], @@ -1250,6 +1259,13 @@ impl GroupedHashAggregateStream { ) })?; + if self.spill_state.spills.is_empty() { + self.spill_state.spill_manager = bound_spill_batches_by_first_spill( + &self.spill_state.spill_manager, + &emit, + )?; + } + let sorted_iter = IncrementalSortIterator::new( emit, self.spill_state.spill_expr.clone(), diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 9da8da2eae62a..3fd744179e8ad 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -22,8 +22,9 @@ use super::{AggregateExec, AggregateMode}; use crate::SendableRecordBatchStream; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; +use crate::sorts::multi_level_merge::max_spill_batch_bytes; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::spill::spill_manager::SpillManager; +use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; /// Spill configuration and accumulated runs of one grouped aggregation stream. /// @@ -175,12 +176,19 @@ impl AggregateSpill { return internal_err!("{label}: spill expression is empty"); }; + let session_config = context.session_config(); let spill_manager = SpillManager::new( context.runtime_env(), spill_metrics, Arc::clone(spill_schema), ) - .with_compression_type(context.session_config().spill_compression()); + .with_compression_type(session_config.spill_compression()) + .with_max_batch_bytes(max_spill_batch_bytes( + session_config + .options() + .execution + .sort_spill_reservation_bytes, + )); Ok(Self { replay_agg, @@ -210,6 +218,11 @@ impl AggregateSpill { return Ok(()); }; + if self.spills.is_empty() { + self.spill_manager = + bound_spill_batches_by_first_spill(&self.spill_manager, &state_batch)?; + } + let max_batch_rows = state_batch.num_rows().min(self.batch_size); let sorted_iter = IncrementalSortIterator::new( state_batch, @@ -307,3 +320,21 @@ impl AggregateSpill { Ok(replay.into_stream()) } } + +/// Bounds the batches of an aggregate's sorted spill runs by the state of its +/// first spill. +/// +/// Unlike the sort, the merge of an aggregate's runs holds no workspace of its +/// own. The batches are bounded so that the smallest merge, plus as much again +/// for the replay, fits in what the aggregate held when memory first ran out. +/// Later spills, such as the remainder spilled at the end of the input, can +/// hold much less, so they do not lower the bound. +pub(super) fn bound_spill_batches_by_first_spill( + spill_manager: &SpillManager, + first_spill: &RecordBatch, +) -> Result { + let held = first_spill.get_sliced_size()?; + let bound = max_spill_batch_bytes(held / 2) + .min(spill_manager.max_batch_bytes().unwrap_or(usize::MAX)); + Ok(spill_manager.clone().with_max_batch_bytes(bound)) +} diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 93b7622f40e49..0c6c0115169fb 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -63,6 +63,14 @@ pub struct BatchBuilder { /// The accumulated stream indexes from which to pull rows /// Consists of a tuple of `(batch_idx, row_idx)` indices: Vec<(usize, usize)>, + + /// Average row size of the current batch of each stream, set by + /// [`Self::push_batch`]. Empty unless [`Self::track_pending_bytes`] was called. + row_bytes: Vec, + + /// Estimated size of the rows in `indices`: the sum of their batches' + /// average row sizes. + pending_bytes: usize, } impl BatchBuilder { @@ -79,6 +87,8 @@ impl BatchBuilder { batches: Vec::with_capacity(stream_count * 2), cursors: vec![BatchCursor::RELEASED; stream_count], indices: Vec::with_capacity(batch_size), + row_bytes: vec![], + pending_bytes: 0, reservation, batches_mem_used: 0, initial_reservation, @@ -88,6 +98,9 @@ impl BatchBuilder { /// Append a new batch in `stream_idx` pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> { let size = get_record_batch_memory_size(&batch); + if let Some(row_bytes) = self.row_bytes.get_mut(stream_idx) { + *row_bytes = size.div_ceil(batch.num_rows().max(1)); + } self.batches_mem_used += size; // Only request additional memory from the pool when actual batch // usage exceeds the current reservation (which may include @@ -116,6 +129,9 @@ impl BatchBuilder { let row_idx = cursor.row_idx; cursor.row_idx += 1; self.indices.push((cursor.batch_idx, row_idx)); + if let Some(row_bytes) = self.row_bytes.get(stream_idx) { + self.pending_bytes += row_bytes; + } } /// Returns the number of in-progress rows in this [`BatchBuilder`] @@ -123,6 +139,18 @@ impl BatchBuilder { self.indices.len() } + /// Starts estimating [`Self::pending_bytes`]. Merges that do not cap their + /// output by bytes skip the per-row bookkeeping. + pub fn track_pending_bytes(&mut self) { + self.row_bytes = vec![0; self.cursors.len()]; + } + + /// Estimated in-memory size of the batch [`Self::build_record_batch`] would + /// build now, if [`Self::track_pending_bytes`] was called, otherwise 0. + pub fn pending_bytes(&self) -> usize { + self.pending_bytes + } + /// Returns `true` if this [`BatchBuilder`] contains no in-progress rows pub fn is_empty(&self) -> bool { self.indices.is_empty() @@ -161,7 +189,9 @@ impl BatchBuilder { columns: Vec, ) -> Result { // Remove consumed indices, keeping any remaining for the next call. + let rows_before = self.indices.len(); self.indices.drain(..rows_to_emit); + self.pending_bytes = self.pending_bytes * self.indices.len() / rows_before; if self.indices.is_empty() { self.retain_cursor_batches(); diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8e900372121bd..5ec382801e301 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -73,6 +73,11 @@ pub(crate) struct SortPreservingMergeStream { /// Target batch size batch_size: usize, + /// Emit a batch once its estimated size reaches this many bytes, even if it + /// has fewer than `batch_size` rows. Set for merges of spilled runs, see + /// [`MultiLevelMergeBuilder`](crate::sorts::multi_level_merge::MultiLevelMergeBuilder). + max_batch_bytes: Option, + /// Cursors for each input partition. `None` means the input is exhausted cursors: Vec>>, @@ -133,11 +138,28 @@ impl SortPreservingMergeStream { poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], batch_size, + max_batch_bytes: None, fetch, produced: 0, } } + /// See the `max_batch_bytes` field. + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + if max_batch_bytes.is_some() { + self.in_progress.track_pending_bytes(); + } + self.max_batch_bytes = max_batch_bytes; + self + } + + fn batch_is_full(&self) -> bool { + self.in_progress.len() >= self.batch_size + || self + .max_batch_bytes + .is_some_and(|max| self.in_progress.pending_bytes() >= max) + } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream where C: 'static, @@ -244,7 +266,7 @@ impl SortPreservingMergeStream { } // 3.3. if there is enough to emit for a full record batch - if self.in_progress.len() >= self.batch_size { + if self.batch_is_full() { // 3.3.1 build pending record batch and reset builder let Some(batch) = self.emit_in_progress_batch()? else { return internal_err!("must have batch in progress to emit"); diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index dc65c6e49d2f4..c2e8755bbb8a3 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -9,7 +9,7 @@ mod cursor; mod merge; #[cfg(test)] mod merge_trace; -mod multi_level_merge; +pub(crate) mod multi_level_merge; pub mod partial_sort; pub mod partitioned_topk; pub mod sort; diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 424d1faffcf33..966f68b01d255 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -152,6 +152,23 @@ pub(crate) struct MultiLevelMergeBuilder { enable_round_robin_tie_breaker: bool, } +/// Largest in-memory size of a batch in a sorted spill run, such that the +/// smallest merge (two runs, one batch read ahead from each) fits in +/// `merge_memory`. Each batch is charged twice its size (see +/// [`get_reserved_bytes_for_record_batch_size`]). +/// +/// The sort passes `sort_spill_reservation_bytes`, the workspace it holds for +/// its merge. The aggregate, which holds no workspace, also passes half of the +/// memory it held when it spilled. +pub(crate) fn max_spill_batch_bytes(merge_memory: usize) -> usize { + const MIN_SPILL_BATCH_BYTES: usize = 64 * 1024; + let per_batch = get_reserved_bytes_for_record_batch_size(1, 1); + (merge_memory / (MIN_SPILL_MERGE_FAN_IN * per_batch)).max(MIN_SPILL_BATCH_BYTES) +} + +/// A merge pass needs at least two inputs to make progress. +const MIN_SPILL_MERGE_FAN_IN: usize = 2; + impl Debug for MultiLevelMergeBuilder { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "MultiLevelMergeBuilder") @@ -521,6 +538,7 @@ impl MultiLevelMergeBuilder { .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) .with_batch_size(output_batch_size) + .with_max_batch_bytes(self.spill_manager.max_batch_bytes()) .with_fetch(self.fetch) .with_metrics(if is_output { // Only add the metrics to the last run @@ -898,7 +916,7 @@ fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize { if configured_fan_in == 0 { usize::MAX } else { - configured_fan_in.max(2) + configured_fan_in.max(MIN_SPILL_MERGE_FAN_IN) } } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 8663515c07a28..6e8c25d354b86 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -30,6 +30,7 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; +use crate::sorts::multi_level_merge::max_spill_batch_bytes; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::get_record_batch_memory_size; use crate::spill::in_progress_spill_file::InProgressSpillFile; @@ -348,7 +349,8 @@ impl ExternalSorter { metrics.spill_metrics.clone(), Arc::clone(&schema), ) - .with_compression_type(spill_compression); + .with_compression_type(spill_compression) + .with_max_batch_bytes(max_spill_batch_bytes(sort_spill_reservation_bytes)); let in_mem_batches_counter = (!schema .fields() diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index da341cf0cf047..2edbc1080619d 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -30,7 +30,7 @@ macro_rules! primitive_merge_helper { } macro_rules! merge_helper { - ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ + ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $max_batch_bytes:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); return Ok(SortPreservingMergeStream::new( @@ -42,6 +42,7 @@ macro_rules! merge_helper { $reservation, $enable_round_robin_tie_breaker, ) + .with_max_batch_bytes($max_batch_bytes) .into_stream()); }}; } @@ -86,6 +87,7 @@ pub struct StreamingMergeBuilder<'a> { /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, min_spill_batch_rows: Option, + max_batch_bytes: Option, enable_round_robin_tie_breaker: bool, } @@ -171,6 +173,14 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Emit a batch once it holds about this many bytes, even if it has fewer + /// than `batch_size` rows. Set by [`MultiLevelMergeBuilder`] for merges of + /// spilled runs. + pub(super) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + /// See [SortPreservingMergeExec::with_round_robin_repartition] for more /// information. /// @@ -206,6 +216,7 @@ impl<'a> StreamingMergeBuilder<'a> { merge_pool, reserve_replay_headroom, min_spill_batch_rows, + max_batch_bytes, fetch, expressions, enable_round_robin_tie_breaker, @@ -268,12 +279,12 @@ impl<'a> StreamingMergeBuilder<'a> { let sort = expressions[0].clone(); let data_type = sort.expr.data_type(schema.as_ref())?; downcast_primitive! { - data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker), - DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker), + DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) _ => {} } } @@ -293,6 +304,7 @@ impl<'a> StreamingMergeBuilder<'a> { reservation, enable_round_robin_tie_breaker, ) + .with_max_batch_bytes(max_batch_bytes) .into_stream()) } } diff --git a/datafusion/sqllogictest/test_files/spill_batch_bound.slt b/datafusion/sqllogictest/test_files/spill_batch_bound.slt new file mode 100644 index 0000000000000..42c42019cd04b --- /dev/null +++ b/datafusion/sqllogictest/test_files/spill_batch_bound.slt @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +# Spilled sorts and aggregates over wide rows. +# +# The batches of a sorted spill run are bounded by bytes, so that merging the +# runs needs a known amount of memory whatever the row width. Each query below +# spills rows much wider than `batch_size` rows of narrow data, under a memory +# limit that fits each row but not a whole batch of them. + +statement ok +SET datafusion.execution.target_partitions = 1 + +# --- Sort: 512 rows of 64 KiB, built in small batches so each input batch fits --- + +statement ok +SET datafusion.execution.sort_spill_reservation_bytes = 1048576 + +statement ok +SET datafusion.execution.batch_size = 4 + +statement ok +CREATE TABLE wide AS +SELECT value AS id, concat(CAST(value AS VARCHAR), repeat('x', 65536)) AS payload +FROM generate_series(0, 511) + +statement ok +CREATE TABLE wide_ree AS +SELECT id, arrow_cast(payload, 'RunEndEncoded("run_ends": Int32, "values": Utf8)') AS payload +FROM wide + +statement ok +SET datafusion.execution.batch_size = 512 + +statement ok +SET datafusion.runtime.memory_limit = '16M' + +# The last column counts rows out of order, so 0 means sorted. +query III +SELECT count(*), sum(length(payload)), sum(CASE WHEN prev <= id THEN 1 ELSE 0 END) +FROM (SELECT id, payload, lag(id) OVER (ORDER BY id DESC) AS prev FROM wide) +---- +512 33555858 0 + +# Run-end encoded payloads are split by rows too. +statement ok +SET datafusion.runtime.memory_limit = '32M' + +query III +SELECT count(*), sum(length(CAST(payload AS VARCHAR))), sum(CASE WHEN prev <= id THEN 1 ELSE 0 END) +FROM (SELECT id, payload, lag(id) OVER (ORDER BY id DESC) AS prev FROM wide_ree) +---- +512 33555858 0 + +statement ok +DROP TABLE wide + +statement ok +DROP TABLE wide_ree + +statement ok +RESET datafusion.execution.sort_spill_reservation_bytes + +statement ok +RESET datafusion.execution.batch_size + +# --- Aggregate: a few groups with large states (#25851) --- + +statement ok +SET datafusion.execution.target_partitions = 2 + +statement ok +SET datafusion.runtime.memory_limit = '12M' + +query II +SELECT count(*) AS groups, sum(n) AS values_collected FROM ( + SELECT k, cardinality(array_agg(v)) AS n + FROM ( + SELECT value % 4 AS k, concat(CAST(value AS VARCHAR), repeat('x', 100)) AS v + FROM generate_series(1, 200000) + ) + GROUP BY k +) +---- +4 200000 + +statement ok +SET datafusion.execution.target_partitions = 1 + +# --- Aggregate: many groups whose states are wide relative to the pool --- + +statement ok +SET datafusion.runtime.memory_limit = '4M' + +query II +SELECT count(*), sum(length(s)) FROM ( + SELECT k, string_agg(v, ',') AS s + FROM (SELECT value % 40000 AS k, repeat('x', 600) AS v FROM generate_series(1, 160000)) + GROUP BY k +) +---- +40000 96120000 + +# The legacy aggregate stream bounds its spill batches the same way. +statement ok +SET datafusion.execution.enable_migration_aggregate = false + +query II +SELECT count(*), sum(length(s)) FROM ( + SELECT k, string_agg(v, ',') AS s + FROM (SELECT value % 40000 AS k, repeat('x', 200) AS v FROM generate_series(1, 160000)) + GROUP BY k +) +---- +40000 32120000 + +statement ok +RESET datafusion.execution.enable_migration_aggregate + +# Restore settings to slt runner defaults +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema From 42e0eae91476a10cb3b9b0ee1162b4872df79d8d Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 1 Nov 2026 14:00:49 +0800 Subject: [PATCH 3/3] refactor: remove the spill merge's halving re-spill --- .../src/sorts/multi_level_merge.rs | 1185 +++-------------- .../intermediate_sizing_tests.rs | 6 +- .../library-user-guide/upgrading/56.0.0.md | 31 + 3 files changed, 225 insertions(+), 997 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 966f68b01d255..81b20fbd7540e 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -14,14 +14,12 @@ use std::sync::Arc; use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; -use datafusion_common::{Result, internal_err, resources_err}; +use datafusion_common::{Result, resources_err}; use datafusion_execution::memory_pool::{MemoryReservation, MergeMemoryPool}; use crate::sorts::builder::try_grow_reservation_to_at_least; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::spill::gc_view_arrays; -use crate::spill::spill_manager::GetSlicedSize; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -107,42 +105,39 @@ use futures::{Stream, StreamExt}; /// /// ## Memory Management Strategy /// -/// This multi-level merge make sure that we can handle any amount of data to sort as long as -/// we have enough memory to merge at least 2 streams at a time, even when individual record -/// batches are skewed (very wide). +/// The merge relies on one property of its inputs: every batch of a sorted spill +/// run is at most the bound its [`SpillManager`] sets (see +/// [`max_spill_batch_bytes`]), unless it holds a single row or a column that +/// cannot be split by rows (see [`SpillManager::with_max_batch_bytes`]). The writers of sorted runs (the sort, the aggregate, and this merge for +/// its intermediate runs) get it from their [`SpillManager`], which splits larger +/// batches as it writes them. The merge keeps its own output to the same bound, +/// so an intermediate run, and the stream it returns, have it too. /// -/// 1. **Worst-Case Memory Reservation**: Reserves memory based on the largest -/// batch size encountered in each spill file to merge, ensuring sufficient memory is always -/// available during merge operations. -/// 2. **Adaptive Buffer Sizing**: Reduces buffer sizes when memory is constrained -/// 3. **Spill-to-Disk**: Spill to disk when we cannot merge all files in memory -/// 4. **Re-spilling Skewed Runs**: If even at the smallest read-buffer size we still cannot -/// reserve memory for the minimum of 2 streams - because a single run's largest batch is so -/// wide that two streams' worth of reservation exceeds the budget - the larger of the two -/// runs is re-spilled with each batch sliced in half. This shrinks its largest batch, -/// lowering the per-stream reservation, and the merge pass is retried. The re-spilled run -/// is tracked alongside a per-run batch-size limit equal to half the batch size it was -/// written with, so any later merge that includes it caps its output batch size to match - -/// otherwise the merged run could rebuild a full-size batch and reintroduce the skew. -/// Crucially the global merge batch size is *not* lowered, so re-spilling more than one run -/// does not compound the reduction. If a batch cannot be split any further (a single row -/// wider than the budget), the merge surfaces `ResourcesExhausted` instead of looping -/// forever. +/// 1. **Worst-Case Memory Reservation**: Reserves memory for the largest batch of +/// each run it reads, times the number of batches read ahead. +/// 2. **Adaptive Buffer Sizing**: Reads fewer batches ahead when memory is +/// constrained. +/// 3. **Replay Headroom**: When the caller asks for it (the aggregate replaying +/// its spills), a run is admitted only if the pool could give as much again. +/// Only the two runs a pass needs, with one batch read ahead, may use the +/// whole pool. +/// 4. **Spill-to-Disk**: Spill to disk when we cannot merge all runs in memory. +/// When the caller asks for it, an intermediate pass merges only as many runs +/// as needed for the remaining runs to fit one final pass. +/// +/// If two runs do not fit even with one batch read ahead (single rows or columns +/// that cannot be split, or a pool smaller than the bound was chosen for), the +/// merge returns `ResourcesExhausted`. pub(crate) struct MultiLevelMergeBuilder { spill_manager: SpillManager, schema: SchemaRef, - /// Sorted runs still to be merged. Each run is paired with the batch-size limit a - /// merge consuming it must cap its output at. Runs written at the full batch size - /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved - /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit - /// out of the public [`SortedSpillFile`], so no external caller has to set it. - sorted_spill_files: Vec<(SortedSpillFile, usize)>, + sorted_spill_files: Vec, sorted_streams: Vec, expr: LexOrdering, metrics: BaselineMetrics, batch_size: usize, reservation: MemoryReservation, - /// Workspace retained across retries and intermediate spill passes. + /// Workspace retained across intermediate spill passes. merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, @@ -192,12 +187,7 @@ impl MultiLevelMergeBuilder { Self { spill_manager, schema, - // Initial runs are written at the full batch size, so they impose no cap - // on later merges - record `batch_size` as their (unconstrained) limit. - sorted_spill_files: sorted_spill_files - .into_iter() - .map(|file| (file, batch_size)) - .collect(), + sorted_spill_files, sorted_streams, expr, metrics, @@ -216,8 +206,7 @@ impl MultiLevelMergeBuilder { self } - /// Leave replay headroom while selecting merge buffers. Temporary splitting - /// workspace can still use the full pool because replay has not started. + /// Leave replay headroom while selecting merge buffers. pub(super) fn with_replay_headroom(mut self, reserve: bool) -> Self { self.reserve_replay_headroom = reserve; self @@ -242,38 +231,8 @@ impl MultiLevelMergeBuilder { } async fn create_stream(mut self) -> Result { - let mut allow_minimum_without_headroom = false; loop { - let (mut stream, batch_size_limit) = match self - .merge_sorted_runs_within_mem_limit(allow_minimum_without_headroom)? - { - MergeStep::Stream { - stream, - batch_size_limit, - } => (stream, batch_size_limit), - MergeStep::SplitThenRetry(index) => { - // Couldn't reserve memory for the minimum of 2 streams. Re-spill - // the larger of the two we're trying to merge with half its batch - // size so its largest batch shrinks, lowering the per-stream - // reservation, then retry. Makes the merge resilient to skewed - // (very wide) rows. - let retry_unsplittable = - self.reserve_replay_headroom && !allow_minimum_without_headroom; - if !self - .split_spill_file_in_half(index, retry_unsplittable) - .await? - { - // A single row may prevent leaving replay headroom while - // the minimum merge still fits the actual shared pool. - allow_minimum_without_headroom = true; - } - continue; - } - }; - allow_minimum_without_headroom = false; - - // TODO - add a threshold for number of files to disk even if empty and reading from disk so - // we can avoid the memory reservation + let mut stream = self.merge_sorted_runs_within_mem_limit()?; // If no spill files are left, we can return the stream as this is the last sorted run // TODO - We can write to disk before reading it back to avoid having multiple streams in memory @@ -293,7 +252,7 @@ impl MultiLevelMergeBuilder { } // Need to sort to a spill file - let Some((spill_file, max_record_batch_memory, max_batch_rows)) = self + let Some((file, max_record_batch_memory, max_batch_rows)) = self .spill_manager .spill_record_batch_stream_and_return_max_batch_stats( &mut stream, @@ -304,21 +263,14 @@ impl MultiLevelMergeBuilder { continue; }; - // Offset overflow can flush partial batches even without a fetch limit. - // Keep sizing disabled once a run cannot fill its output batch limit. - self.size_intermediate_merges &= max_batch_rows >= batch_size_limit; - - // Add the spill file paired with the batch-size limit of the merge that - // produced it: if that merge consumed a shrunk (skew-resolved) run, its - // output was capped and this intermediate run is likewise capped, so a - // later pass that re-merges it won't rebuild an oversized batch. - self.sorted_spill_files.push(( - SortedSpillFile { - file: spill_file, - max_record_batch_memory, - }, - batch_size_limit, - )); + // Offset overflow or the byte bound can end batches early. Keep + // sizing disabled once a run cannot fill its output batch. + self.size_intermediate_merges &= max_batch_rows >= self.batch_size; + + self.sorted_spill_files.push(SortedSpillFile { + file, + max_record_batch_memory, + }); } } @@ -326,59 +278,42 @@ impl MultiLevelMergeBuilder { /// as possible within the memory limit. fn merge_sorted_runs_within_mem_limit( &mut self, - allow_minimum_without_headroom: bool, - ) -> Result { + ) -> Result { match (self.sorted_spill_files.len(), self.sorted_streams.len()) { // No data so empty batch (0, 0) => { let empty_stream = Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); - Ok(MergeStep::Stream { - stream: self.observe_output(empty_stream), - batch_size_limit: self.batch_size, - }) + Ok(self.observe_output(empty_stream)) } // Only in-memory stream, return that (0, 1) => { let output_stream = self.sorted_streams.remove(0); - Ok(MergeStep::Stream { - stream: self.observe_output(output_stream), - batch_size_limit: self.batch_size, - }) + Ok(self.observe_output(output_stream)) } // Only single sorted spill file so return it (1, 0) => { - let (spill_file, batch_size) = self.sorted_spill_files.remove(0); + let spill_file = self.sorted_spill_files.remove(0); // Not reserving any memory for this disk as we are not holding it in memory let output_stream = self .spill_manager .read_spill_as_stream(spill_file.file, None)?; - - Ok(MergeStep::Stream { - stream: self.observe_output(output_stream), - batch_size_limit: batch_size, - }) + Ok(self.observe_output(output_stream)) } - // Only in memory streams, so merge them all in a single pass. In-memory - // runs are never shrunk for skew, so this merge runs at the full batch - // size and its output carries no limit. + // Only in memory streams, so merge them all in a single pass (0, _) => { let sorted_stream = mem::take(&mut self.sorted_streams); // No need to wrap with observed stream since merge sort will update the observed metrics - Ok(MergeStep::Stream { - stream: self.create_new_merge_sort( - sorted_stream, - // If we have no sorted spill files left, this is the last run - true, - true, - self.batch_size, - )?, - batch_size_limit: self.batch_size, - }) + self.create_new_merge_sort( + sorted_stream, + // If we have no sorted spill files left, this is the last run + true, + true, + ) } // Need to merge multiple streams @@ -390,30 +325,13 @@ impl MultiLevelMergeBuilder { // allocation. let mut memory_reservation = self.reservation.take(); - // Compute the minimum before taking the in-memory streams so that, if we - // need to re-spill and retry, `self.sorted_streams` is left untouched. - let minimum_number_of_required_streams = - 2_usize.saturating_sub(self.sorted_streams.len()); - - let (mut sorted_spill_files, buffer_size) = match self + let (mut sorted_spill_files, buffer_size) = self .get_sorted_spill_files_to_merge( 2, // we must have at least 2 streams to merge - minimum_number_of_required_streams, + MIN_SPILL_MERGE_FAN_IN.saturating_sub(self.sorted_streams.len()), &mut memory_reservation, - allow_minimum_without_headroom, - )? { - SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => { - (sorted_spill_files, 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)); - } - }; + )?; if self.size_intermediate_merges && self.sorted_streams.is_empty() @@ -429,21 +347,20 @@ impl MultiLevelMergeBuilder { } else { final_fan_in }; - let needed = self.sorted_spill_files.len() + sorted_spill_files.len() - - final_fan_in - + 1; + let needed = + (self.sorted_spill_files.len() + sorted_spill_files.len() + 1) + .saturating_sub(final_fan_in); if needed < sorted_spill_files.len() { - let (first, first_limit) = &sorted_spill_files[0]; - // A wider pending run or a smaller output batch limit can - // change the next pass's fan-in and add intermediate work. - // Apply this sizing heuristic only to uniform run budgets. + let first = &sorted_spill_files[0]; + // A wider pending run can change the next pass's + // fan-in and add intermediate work. Apply this sizing + // heuristic only to uniform run budgets. let uniform_runs = sorted_spill_files .iter() .chain(&self.sorted_spill_files) - .all(|(file, limit)| { + .all(|file| { file.max_record_batch_memory == first.max_record_batch_memory - && limit == first_limit }); if uniform_runs { let unused = sorted_spill_files.split_off(needed); @@ -471,15 +388,7 @@ impl MultiLevelMergeBuilder { mem::swap(&mut self.reservation, &mut memory_reservation); } - // 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 { + for spill in sorted_spill_files { let stream = self .spill_manager .clone() @@ -488,7 +397,6 @@ impl MultiLevelMergeBuilder { 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( @@ -496,7 +404,6 @@ impl MultiLevelMergeBuilder { // 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 @@ -508,20 +415,14 @@ impl MultiLevelMergeBuilder { "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory" ); - Ok(MergeStep::Stream { - stream: merge_sort_stream, - batch_size_limit: output_batch_size, - }) + Ok(merge_sort_stream) } 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, - }) + Ok(Box::pin(StreamAttachedReservation::new( + merge_sort_stream, + memory_reservation, + ))) } } } @@ -532,12 +433,11 @@ impl MultiLevelMergeBuilder { streams: Vec, is_output: bool, all_in_memory: bool, - output_batch_size: usize, ) -> Result { let mut builder = StreamingMergeBuilder::new() .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) - .with_batch_size(output_batch_size) + .with_batch_size(self.batch_size) .with_max_batch_bytes(self.spill_manager.max_batch_bytes()) .with_fetch(self.fetch) .with_metrics(if is_output { @@ -566,56 +466,41 @@ impl MultiLevelMergeBuilder { builder.build() } - /// Return the sorted spill files to use for the next phase, and the buffer size - /// This will try to get as many spill files as possible to merge, and if we don't have enough streams - /// it will try to reduce the buffer size until we have enough streams to merge - /// otherwise it will return an error + /// Returns the sorted spill files to read in the next pass, and how many + /// batches to read ahead from each. + /// + /// Takes as many runs as fit, at least `minimum_number_of_required_streams`, + /// leaving replay headroom if the caller asked for it (see the type docs). If + /// the minimum does not fit, retries with fewer batches read ahead; with one + /// batch read ahead the minimum only has to fit the pool, or the pool's error + /// is returned. fn get_sorted_spill_files_to_merge( &mut self, buffer_len: usize, minimum_number_of_required_streams: usize, reservation: &mut MemoryReservation, - allow_minimum_without_headroom: bool, - ) -> Result { + ) -> Result<(Vec, usize)> { assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); - let mut number_of_spills_to_read_for_current_phase = 0; - let configured_fan_in = self - .spill_manager - .env() - .disk_manager - .max_spill_merge_fan_in(); - let max_spill_files = effective_spill_merge_fan_in(configured_fan_in); - // Track total memory needed for spill file buffers. When the - // reservation has pre-reserved bytes (from sort_spill_reservation_bytes), - // those bytes cover the first N spill files without additional pool - // allocation, preventing starvation under memory pressure. - let mut total_needed: usize = 0; + let max_spill_files = self.max_fan_in(); + let mut number_of_spills = 0; let mut accepted_memory: usize = 0; - for (spill, _) in &self.sorted_spill_files { - if number_of_spills_to_read_for_current_phase >= max_spill_files - || (allow_minimum_without_headroom - && number_of_spills_to_read_for_current_phase - >= minimum_number_of_required_streams) - { + for spill in &self.sorted_spill_files { + if number_of_spills >= max_spill_files { break; } - let per_spill = get_reserved_bytes_for_record_batch_size( - spill.max_record_batch_memory, - // Size will be the same as the sliced size, bc it is a spilled batch. - spill.max_record_batch_memory, - ) * buffer_len; - total_needed += per_spill; - - // If a run cannot shrink, allow only the minimum merge without - // replay headroom. Disable read-ahead and still ask the pool for - // every byte used by the merge. - let skip_headroom = allow_minimum_without_headroom - && buffer_len == 1 - && number_of_spills_to_read_for_current_phase - < minimum_number_of_required_streams; - let check_headroom = self.reserve_replay_headroom && !skip_headroom; + let total_needed = accepted_memory + + get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + // Size will be the same as the sliced size, bc it is a spilled batch. + spill.max_record_batch_memory, + ) * buffer_len; + // Only the minimum merge with one batch read ahead may use the + // whole pool: without it, the merge cannot make progress at all. + let check_headroom = self.reserve_replay_headroom + && !(buffer_len == 1 + && number_of_spills < minimum_number_of_required_streams); let admission = if check_headroom { // Ask the pool for merge buffers plus equal replay space, then // return the spare bytes before exposing the merge stream. @@ -629,62 +514,36 @@ impl MultiLevelMergeBuilder { try_grow_reservation_to_at_least(reservation, total_needed) }; match admission { - Ok(_) => { - number_of_spills_to_read_for_current_phase += 1; + Ok(()) => { + number_of_spills += 1; accepted_memory = total_needed; } - // If we can't grow the reservation, we need to stop - Err(err) => { - // We must have at least 2 streams to merge, so if we don't have enough memory - // fail - if minimum_number_of_required_streams - > number_of_spills_to_read_for_current_phase - { - // Free the memory we reserved for this merge as we either try again or fail - reservation.free(); - if buffer_len > 1 { - // Try again with smaller buffer size, it will be slower but at least we can merge - return self.get_sorted_spill_files_to_merge( - buffer_len - 1, - minimum_number_of_required_streams, - reservation, - allow_minimum_without_headroom, - ); - } - - // buffer_len == 1 and we still can't seat the minimum of 2 streams. - if number_of_spills_to_read_for_current_phase == 0 { - if check_headroom { - // Replay has not started, so splitting may use the - // full pool to leave room for replay afterward. - return Ok(SpillFilesToMerge::SplitThenRetry(0)); - } - // We couldn't even reserve a single stream - one record batch - // is larger than the whole merge budget. That's the lone-batch - // case, not the 2-stream merge skew we rescue here - surface it. - return Err(err); - } - - // We seated one stream (index 0) but not the second (index 1, the - // batch that just failed to reserve). Those are by definition the - // only two streams we are trying to merge, so re-spill the larger - // of them with a smaller batch size and retry, the smaller max - // batch lowers the per-stream reservation enough to seat both. - let split_index = usize::from( - self.sorted_spill_files[1].0.max_record_batch_memory - > self.sorted_spill_files[0].0.max_record_batch_memory, + Err(err) if number_of_spills < minimum_number_of_required_streams => { + // Free the memory we reserved for this merge as we either try again or fail + reservation.free(); + if buffer_len > 1 { + // Try again with smaller buffer size, it will be slower but at least we can merge + return self.get_sorted_spill_files_to_merge( + buffer_len - 1, + minimum_number_of_required_streams, + reservation, ); - return Ok(SpillFilesToMerge::SplitThenRetry(split_index)); } - - // We reached the maximum amount of memory we can use - // for this merge - break; + return Err(err.context(format!( + "Cannot merge two sorted spill runs: each needs {} bytes \ + to read one batch; consider increasing the memory limit", + get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + spill.max_record_batch_memory + ) + ))); } + // We reached the maximum amount of memory we can use for this merge + Err(_) => break, } } - if self.reserve_replay_headroom { + if self.reserve_replay_headroom && reservation.size() > accepted_memory { // `total_needed` may include a rejected candidate. Keep only the // buffers that were admitted, releasing temporary replay headroom. reservation.shrink(reservation.size() - accepted_memory); @@ -692,180 +551,20 @@ impl MultiLevelMergeBuilder { let spills = self .sorted_spill_files - .drain(..number_of_spills_to_read_for_current_phase) + .drain(..number_of_spills) .collect::>(); - Ok(SpillFilesToMerge::Ready(spills, buffer_len)) + Ok((spills, buffer_len)) } - /// Re-spill the spill file at `index` with half its batch size, putting it back - /// at the same position. We read the file back and re-spill it through the normal - /// spill API (which owns batch layout), slicing every batch in two, which halves - /// the largest written batch and so lowers the per-stream merge reservation enough - /// for the next attempt to seat both streams. One stream's worth of memory is - /// reserved for the duration and freed afterwards. Makes the merge resilient to skew. - /// - /// Instead of halving the *global* merge batch size (which would compound when more - /// than one run is re-spilled), the shrunk run records its own smaller batch-size - /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that - /// actually consume it pay the reduced batch size. - /// - /// Returns whether the largest batch shrank. If `retry_unsplittable` is true, - /// restore an unchanged run for one minimum-merge admission attempt. - async fn split_spill_file_in_half( - &mut self, - index: usize, - retry_unsplittable: bool, - ) -> Result { - log::debug!( - "2 spilled streams could not be loaded into memory for merge \ - (requires 2x of the largest batch from both), re-spilling the larger of the two with half \ - the batch size to reduce memory needs for the next merge attempt. the shrunk run carries \ - a halved batch-size limit so only merges consuming it use the smaller batch size" - ); - - // Extract the target in O(1) instead of `remove(index)`, which would shift - // every following spill file. Swap it to the back and pop it; the matching - // swap after re-spilling restores the original order, so the vec ends up - // exactly as it started, just with the target file shrunk. - // `old_batch_size` is the batch size this run was written with (the full merge - // batch size unless it was already shrunk once). Halving it caps the next merge - // that reads this run so the merged output can't rebuild a full-size batch. - let last = self.sorted_spill_files.len() - 1; - self.sorted_spill_files.swap(index, last); - let (mut target, old_batch_size) = self - .sorted_spill_files - .pop() - .expect("index is in bounds, so the vec is non-empty"); - let old_max = target.max_record_batch_memory; - let mut max_batch_rows = old_batch_size; - - // Reserve enough to hold a single stream of this file while we re-spill it. - let reservation = self.reservation.new_empty(); - reservation - .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?; - - if self.reserve_replay_headroom { - // A maximum-sized singleton cannot shrink. Find it without writing - // another file: the original runs may already fill the disk quota. - // Use an unbuffered reader so no background read outlives this guard. - let mut source = self.spill_manager.read_spill_as_stream_unbuffered( - Arc::clone(&target.file), - Some(old_max), - )?; - let mut all_singletons = true; - let mut max_is_singleton = false; - let mut decoded_max = 0; - max_batch_rows = 0; - while let Some(batch) = source.next().await { - let batch = batch?; - max_batch_rows = max_batch_rows.max(batch.num_rows()); - all_singletons &= batch.num_rows() == 1; - decoded_max = decoded_max.max(batch.get_sliced_size()?); - if batch.num_rows() == 1 - && gc_view_arrays(&batch)?.get_sliced_size()? >= old_max - { - max_is_singleton = true; - break; - } - } - // IPC can discard spare view-buffer capacity included in `old_max`. - // A complete scan can correct that estimate without rewriting the - // file. Use decoded buffers without GC, as the merge retains them. - let shrank = !max_is_singleton && decoded_max < old_max; - if shrank || all_singletons || max_is_singleton { - if !shrank && !retry_unsplittable { - return resources_err!( - "Cannot merge sorted runs: a single record batch of {old_max} bytes \ - exceeds the available merge memory and cannot be split further" - ); - } - if shrank { - target.max_record_batch_memory = decoded_max; - } - let batch_size_limit = if all_singletons || max_is_singleton { - 1 - } else { - old_batch_size.min(max_batch_rows).max(1) - }; - self.sorted_spill_files.push((target, batch_size_limit)); - self.sorted_spill_files.swap(index, last); - return Ok(shrank); - } - } - - let source = self - .spill_manager - .read_spill_as_stream(target.file, Some(old_max))?; - // Re-spill with half the batch size: slice every batch in two. The spill - // writer owns the batch layout, we only change how many rows per batch. - let mut halved: SendableRecordBatchStream = - Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&self.schema), - source.flat_map(|batch| { - futures::stream::iter(match batch { - Ok(batch) => split_batch_in_half(batch) - .into_iter() - .map(Ok) - .collect::>(), - Err(e) => vec![Err(e)], - }) - }), - )); - - let result = self - .spill_manager - .spill_record_batch_stream_and_return_max_batch_stats( - &mut halved, - "MultiLevelMergeBuilder split skewed spill", - ) - .await?; - - reservation.free(); - - let Some((file, new_max, new_max_batch_rows)) = result else { - return internal_err!("re-spilling a skewed spill file produced no data"); - }; - - // If halving cannot reduce the largest batch, only a requested retry - // against the actual pool can make progress. The caller permits that - // retry once before surfacing this error. - let shrank = new_max < old_max; - if !shrank && !retry_unsplittable { - return resources_err!( - "Cannot merge sorted runs: a single record batch of {old_max} bytes \ - exceeds the available merge memory and cannot be split further" - ); - } - - // Record the halved batch size as a *per-run* limit rather than lowering the - // global batch size. Merges that don't touch this run keep the full batch - // size. a merge that reads it caps its output at this limit so the merged run - // can't rebuild a full-size batch and reintroduce the skew. Actual batches - // can be shorter than the configured limit; use their size so the merge - // cannot accumulate extra split batches outside its reservation. - let new_batch_size_limit = if shrank { - (old_batch_size / 2).min(max_batch_rows.div_ceil(2)).max(1) - } else { - // Skipping headroom only covers indivisible input rows, not a - // larger output batch formed by concatenating those rows. - 1 - }; - self.size_intermediate_merges &= new_max_batch_rows >= new_batch_size_limit; - - // Push the re-spilled (smaller) file and swap it back into `index`, undoing - // the swap-to-back above so the order is preserved. - self.sorted_spill_files.push(( - SortedSpillFile { - file, - max_record_batch_memory: new_max, - }, - new_batch_size_limit, - )); - let last = self.sorted_spill_files.len() - 1; - self.sorted_spill_files.swap(index, last); - - Ok(shrank) + /// The most runs one pass may merge (`max_spill_merge_fan_in`). + fn max_fan_in(&self) -> usize { + effective_spill_merge_fan_in( + self.spill_manager + .env() + .disk_manager + .max_spill_merge_fan_in(), + ) } fn observe_output( @@ -876,42 +575,6 @@ impl MultiLevelMergeBuilder { } } -/// Outcome of trying to reserve memory for one multi-level merge pass. -enum SpillFilesToMerge { - /// Enough memory: the spill files to read this pass (each paired with its - /// batch-size limit) and the read-ahead buffer size. - Ready(Vec<(SortedSpillFile, usize)>, usize), - /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index - /// with a smaller (halved) batch size, then retry the pass. - SplitThenRetry(usize), -} - -/// What one iteration of the multi-level merge loop should do next. -enum MergeStep { - /// A merged stream is ready to be consumed (and possibly spilled back). - Stream { - stream: SendableRecordBatchStream, - /// The batch-size limit to stamp on the run if this stream is re-spilled as an - /// intermediate result: the batch size its merge ran at. It equals the full - /// merge batch size unless the merge consumed a skew-resolved run, in which - /// case it is that run's smaller limit so the re-spilled result stays capped - /// and can't rebuild an oversized batch. - batch_size_limit: usize, - }, - /// Re-spill the spill file at this index smaller, then retry the merge step. - SplitThenRetry(usize), -} - -/// Slice `batch` into two row-halves so a re-spill writes batches half the size. -fn split_batch_in_half(batch: RecordBatch) -> Vec { - let num_rows = batch.num_rows(); - if num_rows <= 1 { - return vec![batch]; - } - let mid = num_rows / 2; - vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)] -} - fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize { if configured_fan_in == 0 { usize::MAX @@ -1201,122 +864,28 @@ mod tests { Ok(()) } - /// Two sorted runs whose largest batches are too big to both - /// be seated in the merge budget at once are re-spilled (halved) until they - /// fit, and the merge then completes with fully sorted, complete output. - #[tokio::test] - async fn skewed_runs_reuse_retained_workspace_across_retries() -> Result<()> { - // These budgets require one and two re-spills, respectively. - for (budget_halves, expected_splits) in [(7, 1), (5, 2)] { - let capacity = 1024 * 1024; - let parent: Arc = Arc::new(GreedyMemoryPool::new(capacity)); - let env = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&parent)) - .build_arc()?; - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let spill_count = spill_manager.metrics.spill_file_count.clone(); - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - let workspace = m * budget_halves / 2; - assert!(workspace < capacity); - let merge_pool = Arc::new(MergeMemoryPool::new( - Arc::clone(&parent), - MemoryConsumer::new("merge workspace"), - )); - let pool: Arc = - Arc::clone(&merge_pool) as Arc; - let batch_size = 8192; - let mut builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - batch_size, - ) - .with_merge_pool(Some(Arc::clone(&merge_pool))); - builder.reservation.try_grow(workspace)?; - merge_pool.retain(workspace); - let contender = MemoryConsumer::new("contender").register(&parent); - contender.try_grow(capacity - workspace)?; - - // Drive each retry so its freed child reservation cannot silently - // release the workspace and reacquire it from the parent pool. - for _ in 0..expected_splits { - let MergeStep::SplitThenRetry(index) = - builder.merge_sorted_runs_within_mem_limit(false)? - else { - panic!("the merge must re-spill a skewed run"); - }; - assert_eq!(parent.reserved(), capacity); - assert!(contender.try_grow(1).is_err()); - builder.split_spill_file_in_half(index, false).await?; - assert_eq!(parent.reserved(), capacity); - assert!(contender.try_grow(1).is_err()); - } - - let final_bytes: usize = builder - .sorted_spill_files - .iter() - .map(|(file, _)| { - get_reserved_bytes_for_record_batch_size( - file.max_record_batch_memory, - file.max_record_batch_memory, - ) - }) - .sum(); - assert!(final_bytes < workspace && workspace < 2 * final_bytes); - - let mut stream = builder.create_spillable_merge_stream(); - let first = stream.try_next().await?.expect("nonempty merge"); - assert_eq!(spill_count.value(), 2 + expected_splits); - assert_eq!(merge_pool.reserved(), final_bytes); - assert_eq!(parent.reserved(), contender.size() + final_bytes); - - let mut batches = vec![first]; - while let Some(batch) = stream.try_next().await? { - batches.push(batch); - } - assert_eq!( - batches.iter().map(|batch| batch.num_rows()).max(), - Some(batch_size / 2), - "splitting both runs must halve, not quarter, the output batch size" - ); - let merged = concat_batches(&schema, &batches)?; - let expected = - Int64Array::from_iter_values((0..n).flat_map(|value| [value, value])); - assert_eq!(merged.column(0).as_primitive::(), &expected); - - // The pool stays alive: EOF must release live bytes and the idle floor. - assert_eq!(merge_pool.reserved(), 0); - assert_eq!(parent.reserved(), contender.size()); - drop(stream); - assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); - assert_eq!(env.disk_manager.used_disk_space(), 0); - drop(contender); - assert_eq!(parent.reserved(), 0); - } - Ok(()) - } - + /// Runs written through a manager with a byte bound merge in a pool that + /// could not hold two of their batches had they been written whole. #[tokio::test] - async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> { + async fn bounded_runs_merge_in_a_pool_smaller_than_two_whole_batches() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let n: i64 = 16384; + let whole = build_spill_manager(&env, &schema); + let whole_batch = make_sorted_spill_file(&whole, &schema, (0..n).collect()) + .max_record_batch_memory; + let bound = whole_batch / 4; + let spill_manager = + build_spill_manager(&env, &schema).with_max_batch_bytes(bound); let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // Seating two streams needs ~4*m (2*m each), which does NOT fit, but the - // budget is large enough once a run is halved. The rescue keeps halving - // the blocking run until two streams fit (here, after one halving). - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + assert!(f0.max_record_batch_memory <= bound); + assert!(f1.max_record_batch_memory <= bound); + // Merging the whole batches needs 4 * whole_batch with one batch read + // ahead, more than this pool. + let pool: Arc = + Arc::new(GreedyMemoryPool::new(whole_batch * 7 / 2)); let builder = build_merge_builder( spill_manager, Arc::clone(&schema), @@ -1324,83 +893,27 @@ mod tests { &pool, 8192, ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!( - total_rows, - (2 * n) as usize, - "the merge must emit every input row" - ); - - let merged = concat_batches(&schema, &batches)?; - let col = merged.column(0).as_primitive::(); - for i in 1..col.len() { - assert!( - col.value(i - 1) <= col.value(i), - "merge output must be sorted: {} > {} at {i}", - col.value(i - 1), - col.value(i), - ); - } - - Ok(()) - } - - /// A run can fit the pool during splitting but leave no room for replay. - /// Splitting must shrink the merge buffers to half the available memory. - #[rstest::rstest] - #[tokio::test] - async fn replay_headroom_splits_an_oversized_first_run( - #[values(128, 129, 4096)] rows: i64, - ) -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let first = make_sorted_spill_file(&spill_manager, &schema, (0..rows).collect()); - let second = - make_sorted_spill_file(&spill_manager, &schema, (rows..2 * rows).collect()); - let merge_limit = first.max_record_batch_memory; - // One original run fills the pool. Splitting may use that space before - // replay starts, but the final merge must leave half the pool for replay. - let pool_size = 2 * merge_limit; - let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_size)); - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![first, second], - &pool, - 8192, - ) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - let batch = batch?; - // Short input runs must not rebuild a larger output batch than - // the attached merge reservation can hold. - assert!( - crate::spill::get_record_batch_memory_size(&batch) <= pool.reserved() - ); - batches.push(batch); - assert!(pool.reserved() <= merge_limit); + let batches: Vec = builder + .create_spillable_merge_stream() + .try_collect() + .await?; + for batch in &batches { + // The merge output keeps to the same bound, give or take a row. + assert!(crate::spill::get_record_batch_memory_size(batch) <= bound + 64); } let merged = concat_batches(&schema, &batches)?; - let values = merged.column(0).as_primitive::(); - assert_eq!(values.len(), (2 * rows) as usize); - for (expected, value) in values.values().iter().enumerate() { - assert_eq!(*value, expected as i64); - } - drop(stream); + let col = merged.column(0).as_primitive::(); + assert_eq!(col.len(), (2 * n) as usize); + assert!(col.values().windows(2).all(|w| w[0] <= w[1])); + drop(batches); assert_eq!(pool.reserved(), 0); - assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); - assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); Ok(()) } + /// With one batch read ahead, the two runs a pass needs only have to fit + /// the pool, not the replay headroom. #[tokio::test] - async fn replay_headroom_allows_only_an_indivisible_minimum() -> Result<()> { + async fn replay_headroom_allows_the_minimum_without_headroom() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); let schema = test_schema(); let spill_manager = build_spill_manager(&env, &schema); @@ -1414,16 +927,9 @@ mod tests { build_merge_builder(spill_manager, Arc::clone(&schema), spills, &pool, 8192) .with_replay_headroom(true); - assert!(!builder.split_spill_file_in_half(0, true).await?); - assert_eq!(builder.sorted_spill_files.len(), 3); - // Actual batches contain one row even though the nominal size is 8192. - assert_eq!(builder.sorted_spill_files[0].1, 1); let mut reservation = builder.reservation.new_empty(); - let SpillFilesToMerge::Ready(spills, buffer_len) = - builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation, true)? - else { - panic!("minimum merge should fit the pool"); - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation)?; assert_eq!(buffer_len, 1); assert_eq!(spills.len(), 2); assert_eq!(builder.sorted_spill_files.len(), 1); @@ -1432,11 +938,10 @@ mod tests { reservation.free(); let mut stream = builder.create_spillable_merge_stream(); - let mut values = Vec::new(); + let mut values: Vec = Vec::new(); while let Some(batch) = stream.next().await { let batch = batch?; - assert_eq!(batch.num_rows(), 1); - values.push(batch.column(0).as_primitive::().value(0)); + values.extend(batch.column(0).as_primitive::().values()); } assert_eq!(values, vec![0, 1, 2]); drop(stream); @@ -1446,6 +951,32 @@ mod tests { Ok(()) } + /// With replay headroom, reading two batches ahead must also leave as much + /// again for the replay; otherwise the merge reads one batch ahead. + #[test] + fn read_ahead_needs_headroom() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let spills = (0..2) + .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value])) + .collect::>(); + let batch_memory = spills[0].max_record_batch_memory; + // Two batches read ahead from two runs take 8 batches of the pool's 9. + let pool: Arc = Arc::new(GreedyMemoryPool::new(9 * batch_memory)); + let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_replay_headroom(true); + let mut reservation = builder.reservation.new_empty(); + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation)?; + assert_eq!(buffer_len, 1); + assert_eq!(spills.len(), 2); + assert_eq!(reservation.size(), 4 * batch_memory); + drop((reservation, spills, builder)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[test] fn replay_headroom_is_released_after_rejected_candidate() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); @@ -1460,11 +991,8 @@ mod tests { let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) .with_replay_headroom(true); let mut reservation = builder.reservation.new_empty(); - let SpillFilesToMerge::Ready(spills, buffer_len) = - builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? - else { - panic!("two streams and replay headroom should fit the pool"); - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation)?; assert_eq!(buffer_len, 1); assert_eq!(spills.len(), 2); assert_eq!(builder.sorted_spill_files.len(), 1); @@ -1495,7 +1023,13 @@ mod tests { .with_replay_headroom(true); let mut stream = builder.create_spillable_merge_stream(); let error = stream.next().await.unwrap().unwrap_err(); - assert!(error.to_string().contains("cannot be split further")); + assert!( + matches!( + error.find_root(), + datafusion_common::DataFusionError::ResourcesExhausted(_) + ), + "{error}" + ); drop(stream); assert_eq!(pool.reserved(), 0); assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); @@ -1503,330 +1037,6 @@ mod tests { Ok(()) } - #[rstest::rstest] - #[case(1, false, DataType::Utf8)] - #[case(8192, true, DataType::Utf8)] - #[case(8192, true, DataType::Utf8View)] - #[tokio::test] - async fn indivisible_spill_merge_needs_no_extra_disk_space( - #[case] batch_size: usize, - #[case] mixed_batches: bool, - #[case] data_type: DataType, - ) -> Result<()> { - use arrow::array::{ArrayRef, StringArray, StringViewArray}; - - const KEY_BYTES: usize = 300_000; - const POOL_BYTES: usize = 2 * 1024 * 1024; - let schema = Arc::new(Schema::new(vec![Field::new("x", data_type, false)])); - let make_runs = |manager: &SpillManager| -> Result> { - (0..2) - .map(|run| { - let batches = (0..8).map(|batch| { - // A small two-row batch before the large singleton - // prevents treating the first batch as representative. - let keys = if mixed_batches && batch == 0 { - vec![format!("{run:04}"), format!("{:04}", run + 2)] - } else { - let key = run + 2 * (batch + usize::from(mixed_batches)); - vec![format!("{key:04}{}", "x".repeat(KEY_BYTES - 4))] - }; - let values: ArrayRef = match schema.field(0).data_type() { - DataType::Utf8 => Arc::new(StringArray::from(keys)), - DataType::Utf8View => Arc::new(StringViewArray::from(keys)), - _ => unreachable!(), - }; - RecordBatch::try_new(Arc::clone(&schema), vec![values]) - .map_err(Into::into) - }); - let (file, max_record_batch_memory) = manager - .spill_record_batch_iter_and_return_max_batch_memory( - batches, - "indivisible input run", - )? - .unwrap(); - Ok(SortedSpillFile { - file, - max_record_batch_memory, - }) - }) - .collect() - }; - - // Calibrate the quota to exactly the original IPC files, with no room - // for even a replacement header. Keep several batches in each run so - // read-ahead cannot retire the original before the unnecessary write. - let calibration = Arc::new(RuntimeEnv::default()); - let runs = make_runs(&build_spill_manager(&calibration, &schema))?; - let quota = runs.iter().map(|run| run.file.size().unwrap()).sum(); - drop(runs); - let env = RuntimeEnvBuilder::new() - .with_max_temp_directory_size(quota) - .build_arc()?; - let manager = build_spill_manager(&env, &schema); - let runs = make_runs(&manager)?; - assert_eq!(env.disk_manager.spilling_progress().current_bytes, quota); - assert!(4 * runs[0].max_record_batch_memory > POOL_BYTES / 2); - assert!(4 * runs[0].max_record_batch_memory <= POOL_BYTES); - let pool: Arc = Arc::new(GreedyMemoryPool::new(POOL_BYTES)); - let builder = build_merge_builder(manager, schema, runs, &pool, batch_size) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut expected = 0; - while let Some(batch) = stream.next().await { - let batch = batch?; - assert_eq!(batch.num_rows(), 1); - let values = arrow::compute::cast(batch.column(0), &DataType::Utf8)?; - let key = values.as_string::().value(0); - assert_eq!(key[..4].parse::().unwrap(), expected); - expected += 1; - } - assert_eq!(expected, if mixed_batches { 18 } else { 16 }); - drop(stream); - assert_eq!(pool.reserved(), 0); - let progress = env.disk_manager.spilling_progress(); - assert_eq!(progress.current_bytes, 0); - assert_eq!(progress.active_files_count, 0); - Ok(()) - } - - #[rstest::rstest] - #[tokio::test] - async fn small_view_spill_merge_needs_no_extra_disk_space( - #[values(3, 6)] pool_batches: usize, - #[values(false, true)] mixed_batches: bool, - ) -> Result<()> { - use arrow::array::StringViewBuilder; - - let schema = Arc::new(Schema::new(vec![Field::new( - "x", - DataType::Utf8View, - false, - )])); - let make_runs = |manager: &SpillManager| -> Result> { - (0..2) - .map(|run| { - let batches = (0..8).map(|batch| { - // A non-inline value retains the builder's 8 KiB block, - // but IPC stores only its used bytes. Small view buffers - // are below the spill writer's compaction threshold. - let key = format!("{:04}xxxxxxxxx", run + 2 * batch); - let mut values = StringViewBuilder::with_capacity(1) - .with_fixed_block_size(8192); - values.append_value(key); - // A larger final batch requires scanning the entire run - // before replacing its maximum decoded size. - if mixed_batches && batch == 7 { - values.append_value(format!( - "{:04}xxxxxxxxx", - run + 2 * (batch + 1) - )); - } - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(values.finish())], - ) - .map_err(Into::into) - }); - let (file, max_record_batch_memory) = manager - .spill_record_batch_iter_and_return_max_batch_memory( - batches, - "small view singleton run", - )? - .unwrap(); - Ok(SortedSpillFile { - file, - max_record_batch_memory, - }) - }) - .collect() - }; - - // Fill the quota with the original runs. Multiple batches prevent - // read-ahead from retiring a run before an unnecessary replacement write. - let calibration = Arc::new(RuntimeEnv::default()); - let runs = make_runs(&build_spill_manager(&calibration, &schema))?; - let quota = runs.iter().map(|run| run.file.size().unwrap()).sum(); - drop(runs); - let env = RuntimeEnvBuilder::new() - .with_max_temp_directory_size(quota) - .build_arc()?; - let manager = build_spill_manager(&env, &schema); - let runs = make_runs(&manager)?; - assert_eq!(env.disk_manager.spilling_progress().current_bytes, quota); - let stored_max = runs[0].max_record_batch_memory; - let mut source = manager.read_spill_as_stream_unbuffered( - Arc::clone(&runs[0].file), - Some(stored_max), - )?; - let batch = source.next().await.unwrap()?; - assert_eq!(batch.num_rows(), 1); - assert!(gc_view_arrays(&batch)?.get_sliced_size()? < stored_max); - drop((source, batch)); - - // The stale estimate needs four batches for the minimum merge. Test - // limits below and above that estimate: both can fit the decoded rows. - let pool_bytes = pool_batches * stored_max; - let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_bytes)); - let builder = build_merge_builder(manager, schema, runs, &pool, 8192) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut expected = 0; - while let Some(batch) = stream.next().await { - let batch = batch?; - if mixed_batches { - assert!(batch.num_rows() <= 2); - } else { - assert_eq!(batch.num_rows(), 1); - } - for key in batch.column(0).as_string_view().iter() { - assert_eq!(key.unwrap(), format!("{expected:04}xxxxxxxxx")); - expected += 1; - } - } - assert_eq!(expected, if mixed_batches { 18 } else { 16 }); - drop(stream); - assert_eq!(pool.reserved(), 0); - let progress = env.disk_manager.spilling_progress(); - assert_eq!(progress.current_bytes, 0); - assert_eq!(progress.active_files_count, 0); - Ok(()) - } - - /// Tests the `new_max >= old_max` guard: a single-row run cannot be split - /// any smaller, so re-spilling it does not shrink the largest batch and the - /// rescue surfaces `ResourcesExhausted` rather than looping forever. - #[tokio::test] - async fn respilling_an_unsplittable_run_surfaces_resources_exhausted() -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - // A one-row run: `split_batch_in_half` returns it unchanged, so the - // re-spilled file's largest batch cannot drop below the original. - let f0 = make_sorted_spill_file(&spill_manager, &schema, vec![42]); - - // Ample budget so the only possible failure is the un-splittable guard, - // not the single-stream reservation itself. - let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let mut builder = - build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024); - - let err = builder - .split_spill_file_in_half(0, false) - .await - .expect_err("re-spilling a one-row run cannot shrink it"); - assert!( - err.to_string().contains("cannot be split further"), - "expected the un-splittable guard error, got: {err}" - ); - - Ok(()) - } - - /// Proves the re-spill also halves the merge output batch size: after one - /// re-spill the merged run is emitted in 4096-row batches (not the original - /// 8192), so it cannot rebuild a full-size batch and reintroduce the skew. - #[tokio::test] - async fn respill_halves_the_merge_output_batch_size() -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // 3.5*m forces exactly one re-spill (split one run, then both fit), which - // halves the merge output batch size. - let initial_batch_size = 8192; - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); - - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - initial_batch_size, - ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - // All rows are still present. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, (2 * n) as usize); - - // The largest emitted batch is the halved size, not the original 8192: the - // shrunk run carries a halved batch-size limit, and the final pass consumes - // it, so the merge output is capped there. Without the per-run limit the merge - // would rebuild 8192-row batches. - let expected_batch_size = initial_batch_size / 2; - let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); - assert_eq!( - max_batch_rows, expected_batch_size, - "after one re-spill the merge must emit {expected_batch_size}-row \ - batches, got a largest batch of {max_batch_rows} rows" - ); - - Ok(()) - } - - /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight - /// enough that *both* runs must be re-spilled before the merge fits - the scenario - /// where the batch-size reduction could compound. Because the reduction is tracked - /// per-run (each run capped at half) rather than by halving the global batch size on - /// every split, the merged output is emitted in 4096-row batches - half, not a - /// quarter. A global-halving implementation would have halved once per re-spill and - /// emitted 2048-row batches. - #[tokio::test] - async fn respilling_two_skewed_runs_halves_the_output_without_compounding() - -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // 2.5*m is tight enough that even after halving one run the two still don't - // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m, - // as in the single-split test, would let the pair fit after one split.) This - // is exactly the scenario where a compounding, global-halving implementation - // would drive the output batch size down to a quarter. - let initial_batch_size = 8192; - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 5 / 2)); - - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - initial_batch_size, - ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - // All rows are still present. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, (2 * n) as usize); - - // Each run was re-spilled once, so each is capped at half the original batch - // size and the merge caps its output at that half - NOT a quarter. A global - // halving-per-split implementation would have emitted 2048-row batches here. - let expected_batch_size = initial_batch_size / 2; - let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); - assert_eq!( - max_batch_rows, expected_batch_size, - "two re-spills must halve (not quarter) the output: expected \ - {expected_batch_size}-row batches, got a largest batch of \ - {max_batch_rows} rows" - ); - - Ok(()) - } - #[test] fn spill_merge_fan_in_is_unlimited_by_default() { assert_eq!(effective_spill_merge_fan_in(0), usize::MAX); @@ -1881,17 +1091,8 @@ mod tests { let mut merge_reservation = MemoryConsumer::new("spill_merge_fan_in_phase") .register(&runtime.memory_pool); - let (spills, buffer_len) = match builder.get_sorted_spill_files_to_merge( - 1, - 2, - &mut merge_reservation, - false, - )? { - SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len), - SpillFilesToMerge::SplitThenRetry(index) => { - panic!("expected ready spill files, got retry for index {index}") - } - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut merge_reservation)?; assert_eq!(spills.len(), 2); assert_eq!(buffer_len, 1); diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs b/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs index 35b2f6486a6a5..03c726b8a2974 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs @@ -213,11 +213,7 @@ async fn sized_merge_retains_reservation_until_builder_drop() -> Result<()> { build_merge_builder(spill_manager, Arc::clone(&schema), files, &pool, ROWS) .with_replay_headroom(true) .with_intermediate_merge_sizing(Some(ROWS)); - let MergeStep::Stream { stream, .. } = - builder.merge_sorted_runs_within_mem_limit(false)? - else { - panic!("six inputs should be selected for the intermediate merge"); - }; + let stream = builder.merge_sorted_runs_within_mem_limit()?; assert_eq!(builder.sorted_spill_files.len(), 2); let retained = 4 * 7 * batch_memory; assert_eq!(builder.reservation.size(), retained); diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 6b33b5ededd17..7cef10dcf99af 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1480,3 +1480,34 @@ fn is_nullable(&self) -> bool { ``` See [issue #18882](https://github.com/apache/datafusion/issues/18882) for details. + +### Spill merges no longer re-spill runs with halved batches + +`StreamingMergeBuilder::with_sorted_spill_files` merges sorted spill runs by +reserving memory for the largest batch of each run. When two runs did not fit, +the merge used to re-spill the larger run with every batch halved and retry. +It no longer does: it returns `ResourcesExhausted` instead. + +DataFusion's sort and aggregate now write their runs in batches of a bounded +size, so that two runs fit the memory their merges expect; only a single row, +or a column that cannot be split by rows, can be larger. The new +`SpillManager::with_max_batch_bytes` sets that bound on a spill manager: +every batch written through it is split into batches of at most that many +bytes, unless a batch holds one row or a column that cannot be split by rows. +`SpillManager::max_batch_bytes` returns the bound. + +**Migration guide:** + +If you write your own runs for `with_sorted_spill_files`, write them through a +`SpillManager` with a bound, chosen so that two runs fit the memory you expect +the merge to have. Each batch is charged twice its size, and the merge reads +at least one batch ahead from each run: + +```rust,ignore +// Two runs, one batch ahead each, charged twice: 4 batches fit `merge_memory`. +let spill_manager = SpillManager::new(env, metrics, schema) + .with_max_batch_bytes(merge_memory / 4); +``` + +Do not set a bound on a manager whose reader expects one batch back for each +appended batch.