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

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

@@ -51,8 +51,8 @@ use crate::{InputOrderMode, SendableRecordBatchStream};
 ///   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 OrderedFinalAggregateStream {

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

@@ -37,6 +37,16 @@ use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
 use crate::spill::spill_manager::SpillManager;
 use crate::{InputOrderMode, SendableRecordBatchStream};
 
+/// 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.
 ///
 /// Every aggregation stream that spills does so the same way. Each spill event
@@ -222,12 +232,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(
@@ -280,6 +296,7 @@ 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))

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

@@ -98,8 +98,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
@@ -112,7 +112,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

@@ -86,6 +86,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>>>,
 
@@ -146,11 +156,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,
@@ -184,15 +205,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
     }
 
@@ -250,14 +280,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

@@ -154,6 +154,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>>,
@@ -198,6 +200,7 @@ impl MultiLevelMergeBuilder {
             expr,
             metrics,
             batch_size,
+            max_batch_bytes: None,
             reservation,
             merge_pool: None,
             reserve_replay_headroom: false,
@@ -212,6 +215,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 {
@@ -542,7 +552,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