fix: size aggregate spill batches by bytes as well as rows

datafusion/physical-plan/src/aggregates/clustered_final_stream.rs

@@ -40,8 +40,8 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter};
 ///   extra index array is used in later sorting before spilling.
 /// - On memory pressure, materialize all group states into one batch.
 /// - Use [`IncrementalSortIterator`](crate::sorts::IncrementalSortIterator) to compute the full-batch index, then
-///   materialize and write one sorted `batch_size` slice at a time. The original
-///   batch and full index remain live until the run is written.
+///   materialize and write one sorted slice at a time. The original batch and
+///   full index remain live until the run is written.
 /// - After input ends, merge the sorted runs and replay them through a fully
 ///   ordered final aggregate stream.
 pub(crate) struct ClusteredFinalAggregateStream {

datafusion/physical-plan/src/aggregates/spill.rs

@@ -14,6 +14,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
 use datafusion_physical_expr::PhysicalSortExpr;
 use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use futures::StreamExt;
 
 use super::aggregate_hash_table::ClusteredAggregateTableMetrics;
 use super::clustered_final_stream::ClusteredFinalAggregateStream;
@@ -23,7 +24,19 @@ use crate::SendableRecordBatchStream;
 use crate::metrics::{BaselineMetrics, SpillMetrics};
 use crate::sorts::IncrementalSortIterator;
 use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
+use crate::spill::gc_nested_arrays;
 use crate::spill::spill_manager::SpillManager;
+use crate::stream::RecordBatchStreamAdapter;
+
+/// Target size in bytes of the batches that spilling writes and replay
+/// merges. A batch can be larger when it holds a larger group.
+///
+/// Replay holds the groups of the batch it merges in memory, and merging
+/// several spill files reserves memory for a few batches of each. A limit of
+/// `batch_size` rows alone does not bound the size of a batch: a few groups
+/// with large state, such as `array_agg` over a low-cardinality key, would fit
+/// in a single batch.
+const SPILL_BATCH_TARGET_BYTES: usize = 1024 * 1024;
 
 /// Spill configuration and accumulated runs of one grouped aggregation stream.
 ///
@@ -210,12 +223,18 @@ impl AggregateSpill {
             return Ok(());
         };
 
-        let max_batch_rows = state_batch.num_rows().min(self.batch_size);
+        let mut max_batch_rows = 0;
         let sorted_iter = IncrementalSortIterator::new(
             state_batch,
             self.spill_expr.clone(),
             self.batch_size,
-        );
+        )
+        .with_max_batch_bytes(SPILL_BATCH_TARGET_BYTES)
+        .inspect(|batch| {
+            if let Ok(batch) = batch {
+                max_batch_rows = max_batch_rows.max(batch.num_rows());
+            }
+        });
         let spill_file = self
             .spill_manager
             .spill_record_batch_iter_and_return_max_batch_memory(
@@ -289,10 +308,20 @@ impl AggregateSpill {
             .with_expressions(&spill_expr)
             .with_metrics(baseline_metrics.intermediate())
             .with_batch_size(batch_size)
+            .with_max_batch_bytes(Some(SPILL_BATCH_TARGET_BYTES))
             .with_reservation(merge_reservation)
             .with_replay_headroom()
             .with_intermediate_merge_sizing(Some(min_spill_batch_rows))
             .build()?;
+        // A merged batch references the whole data buffers of the spilled
+        // batches its rows come from, and can hold all their dictionary
+        // values. Replay keeps the list states of a group, such as those of
+        // `array_agg`, until the group completes, so compact them to hold only
+        // the bytes of their own rows.
+        let merged: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(
+            merged.schema(),
+            merged.map(|batch| gc_nested_arrays(&batch?)),
+        ));
         let replay = ClusteredFinalAggregateStream::new_with_input_and_metrics(
             &replay_agg,
             &context,

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

@@ -85,8 +85,8 @@ impl BatchBuilder {
         }
     }
 
-    /// Append a new batch in `stream_idx`
-    pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> {
+    /// Append a new batch in `stream_idx`, returning its memory size
+    pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<usize> {
         let size = get_record_batch_memory_size(&batch);
         self.batches_mem_used += size;
         // Only request additional memory from the pool when actual batch
@@ -99,7 +99,7 @@ impl BatchBuilder {
             batch_idx,
             row_idx: 0,
         };
-        Ok(())
+        Ok(size)
     }
 
     /// Append the next row from `stream_idx`

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

@@ -73,6 +73,16 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// Target batch size
     batch_size: usize,
 
+    /// See [`Self::with_max_batch_bytes`]
+    max_batch_bytes: Option<usize>,
+
+    /// The average memory size of the rows of each input's current batch
+    avg_row_bytes: Vec<usize>,
+
+    /// The estimated memory size of the in-progress rows, tracked when
+    /// `max_batch_bytes` is set
+    in_progress_bytes: usize,
+
     /// Cursors for each input partition. `None` means the input is exhausted
     cursors: Vec<Option<Cursor<C>>>,
 
@@ -133,11 +143,22 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             poll_reset_epochs: vec![0; stream_count],
             loser_tree: vec![],
             batch_size,
+            max_batch_bytes: None,
+            avg_row_bytes: vec![0; stream_count],
+            in_progress_bytes: 0,
             fetch,
             produced: 0,
         }
     }
 
+    /// Also emits a batch once its rows take about `max_batch_bytes`, as
+    /// estimated from the average row size of the batches they come from,
+    /// before it has `batch_size` rows.
+    pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option<usize>) -> Self {
+        self.max_batch_bytes = max_batch_bytes;
+        self
+    }
+
     pub(crate) fn into_stream(self) -> SendableRecordBatchStream
     where
         C: 'static,
@@ -171,15 +192,24 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             Some(Err(e)) => Poll::Ready(Err(e)),
             Some(Ok((cursor, batch))) => {
                 self.cursors[idx] = Some(Cursor::new(cursor));
-                Poll::Ready(self.in_progress.push_batch(idx, batch))
+                let num_rows = batch.num_rows();
+                let size = match self.in_progress.push_batch(idx, batch) {
+                    Ok(size) => size,
+                    Err(e) => return Poll::Ready(Err(e)),
+                };
+                self.avg_row_bytes[idx] = size.div_ceil(num_rows.max(1));
+                Poll::Ready(Ok(()))
             }
         }
     }
 
     fn emit_in_progress_batch(&mut self) -> Result<Option<RecordBatch>> {
         let rows_before = self.in_progress.len();
         let result = self.in_progress.build_record_batch();
-        self.produced += rows_before - self.in_progress.len();
+        let rows_after = self.in_progress.len();
+        self.produced += rows_before - rows_after;
+        // Rows kept after an offset overflow are assumed to be of average size
+        self.in_progress_bytes = self.in_progress_bytes * rows_after / rows_before.max(1);
         result
     }
 
@@ -237,14 +267,21 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                 // 3.1. add loser_tree[0] (minimum) stream to pending record batch
                 let winner_stream = self.loser_tree[0];
                 self.in_progress.push_row(winner_stream);
+                if self.max_batch_bytes.is_some() {
+                    self.in_progress_bytes += self.avg_row_bytes[winner_stream];
+                }
 
                 // 3.2. If the new row reached the limit
                 if self.fetch_reached() {
                     break;
                 }
 
                 // 3.3. if there is enough to emit for a full record batch
-                if self.in_progress.len() >= self.batch_size {
+                if self.in_progress.len() >= self.batch_size
+                    || self
+                        .max_batch_bytes
+                        .is_some_and(|max_bytes| self.in_progress_bytes >= max_bytes)
+                {
                     // 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/multi_level_merge.rs

@@ -141,6 +141,8 @@ pub(crate) struct MultiLevelMergeBuilder {
     expr: LexOrdering,
     metrics: BaselineMetrics,
     batch_size: usize,
+    /// See [`StreamingMergeBuilder::with_max_batch_bytes`]
+    max_batch_bytes: Option<usize>,
     reservation: MemoryReservation,
     /// Workspace retained across retries and intermediate spill passes.
     merge_pool: Option<Arc<MergeMemoryPool>>,
@@ -185,6 +187,7 @@ impl MultiLevelMergeBuilder {
             expr,
             metrics,
             batch_size,
+            max_batch_bytes: None,
             reservation,
             merge_pool: None,
             reserve_replay_headroom: false,
@@ -199,6 +202,13 @@ impl MultiLevelMergeBuilder {
         self
     }
 
+    /// Applies [`StreamingMergeBuilder::with_max_batch_bytes`] to every pass,
+    /// including the passes whose output is spilled again.
+    pub(super) fn with_max_batch_bytes(mut self, max_batch_bytes: Option<usize>) -> Self {
+        self.max_batch_bytes = max_batch_bytes;
+        self
+    }
+
     /// Leave replay headroom while selecting merge buffers. Temporary splitting
     /// workspace can still use the full pool because replay has not started.
     pub(super) fn with_replay_headroom(mut self, reserve: bool) -> Self {
@@ -529,7 +539,8 @@ impl MultiLevelMergeBuilder {
                 self.metrics.intermediate()
             })
             .with_round_robin_tie_breaker(self.enable_round_robin_tie_breaker)
-            .with_streams(streams);
+            .with_streams(streams)
+            .with_max_batch_bytes(self.max_batch_bytes);
 
         if !all_in_memory {
             // Don't track memory used by this stream as we reserve that memory by worst case sceneries