fix: leave memory for aggregate spill replay

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

@@ -1367,6 +1367,7 @@ impl GroupedHashAggregateStream {
                 .with_metrics(self.baseline_metrics.clone())
                 .with_batch_size(self.batch_size)
                 .with_reservation(self.reservation.new_empty())
+                .with_replay_headroom()
                 .build()?;
             self.input_done = false;
 
@@ -1394,6 +1395,9 @@ impl GroupedHashAggregateStream {
             // to ensure we don't spill the spilled data to disk again.
             self.oom_mode = OutOfMemoryMode::ReportError;
 
+            // Release unused initial capacity from recreated group values so it
+            // does not consume the memory available for spill replay.
+            self.group_values.clear_shrink(0);
             self.update_memory_reservation()?;
 
             ExecutionState::ReadingInput

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

@@ -348,6 +348,7 @@ impl FinalSpillContext {
             .with_metrics(baseline_metrics.intermediate())
             .with_batch_size(batch_size)
             .with_reservation(merge_reservation)
+            .with_replay_headroom()
             .build()?;
         let replay = OrderedFinalAggregateStream::new_with_input_and_metrics(
             &final_agg,

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

@@ -243,6 +243,7 @@ impl OrderedFinalSpillContext {
             .with_metrics(baseline_metrics.intermediate())
             .with_batch_size(batch_size)
             .with_reservation(merge_reservation)
+            .with_replay_headroom()
             .build()?;
         let replay = OrderedFinalAggregateStream::new_with_input_and_metrics(
             &agg,

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

@@ -303,6 +303,7 @@ impl OrderedSingleSpillContext {
             .with_metrics(baseline_metrics.intermediate())
             .with_batch_size(batch_size)
             .with_reservation(merge_reservation)
+            .with_replay_headroom()
             .build()?;
         let replay = OrderedFinalAggregateStream::new_with_input_and_metrics(
             &final_agg,

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

@@ -329,6 +329,7 @@ impl SingleSpillContext {
             .with_metrics(baseline_metrics.intermediate())
             .with_batch_size(batch_size)
             .with_reservation(merge_reservation)
+            .with_replay_headroom()
             .build()?;
         let replay = OrderedFinalAggregateStream::new_with_input_and_metrics(
             &final_agg,

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

@@ -96,6 +96,8 @@ pub struct StreamingMergeBuilder<'a> {
     fetch: Option<usize>,
     reservation: Option<MemoryReservation>,
     merge_pool: Option<Arc<MergeMemoryPool>>,
+    /// Leave memory for the aggregate consuming the merged spill rows.
+    reserve_replay_headroom: bool,
     enable_round_robin_tie_breaker: bool,
 }
 
@@ -161,6 +163,13 @@ impl<'a> StreamingMergeBuilder<'a> {
         self
     }
 
+    /// Leave room for aggregate replay by checking that the pool can admit
+    /// merge buffers plus equal headroom. Release the headroom before replay.
+    pub(crate) fn with_replay_headroom(mut self) -> Self {
+        self.reserve_replay_headroom = true;
+        self
+    }
+
     /// See [SortPreservingMergeExec::with_round_robin_repartition] for more
     /// information.
     ///
@@ -194,6 +203,7 @@ impl<'a> StreamingMergeBuilder<'a> {
             batch_size,
             reservation,
             merge_pool,
+            reserve_replay_headroom,
             fetch,
             expressions,
             enable_round_robin_tie_breaker,
@@ -235,6 +245,7 @@ impl<'a> StreamingMergeBuilder<'a> {
                 enable_round_robin_tie_breaker,
             )
             .with_merge_pool(merge_pool)
+            .with_replay_headroom(reserve_replay_headroom)
             .create_spillable_merge_stream());
         }