fix: write sorted spill runs in byte-bounded batches and drop the halving re-spill

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,);

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 }

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(),

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<SpillManager> {
+    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))
+}

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<usize>,
+
+    /// 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,13 +129,28 @@ 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`]
     pub fn len(&self) -> usize {
         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<ArrayRef>,
     ) -> Result<RecordBatch> {
         // 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();

datafusion/physical-plan/src/sorts/merge.rs

@@ -73,6 +73,11 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// 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<usize>,
+
     /// Cursors for each input partition. `None` means the input is exhausted
     cursors: Vec<Option<Cursor<C>>>,
 
@@ -133,11 +138,28 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             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<usize>) -> 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<C: CursorValues> SortPreservingMergeStream<C> {
                 }
 
                 // 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");

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;