refactor(hash-aggr): Support spilling for ordered aggregation

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs

@@ -81,6 +81,10 @@ pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> {
     /// Output schema: group columns followed by aggregate state or final values.
     pub(super) output_schema: SchemaRef,
 
+    /// Intermediate-state schema used when memory pressure requires the table
+    /// to pass through or spill its current state.
+    pub(super) state_schema: SchemaRef,
+
     /// Maximum rows per emitted output batch, from config `batch_size`.
     pub(super) batch_size: usize,
 
@@ -129,13 +133,14 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     )]
     pub(super) fn new_for_mode(
         agg: &AggregateExec,
-        partition: usize,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
+        state_schema: SchemaRef,
         batch_size: usize,
         input_order_mode: &InputOrderMode,
         aggregate_mode: &AggregateMode,
         filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
+        group_by_metrics: GroupByMetrics,
     ) -> Result<Self> {
         assert_or_internal_err!(
             batch_size > 0,
@@ -168,8 +173,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
 
         Ok(Self {
             output_schema,
+            state_schema,
             batch_size,
-            group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
+            group_by_metrics,
             buffer: OrderedAggregateTableBuffer {
                 group_by: Arc::clone(&agg.group_by),
                 group_ordering,
@@ -217,9 +223,19 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         self.buffer.group_ordering.input_done();
     }
 
+    /// Returns the ordering state used to decide how memory pressure is handled.
+    pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering {
+        &self.buffer.group_ordering
+    }
+
+    /// Number of groups currently buffered.
+    pub(in crate::aggregates) fn num_groups(&self) -> usize {
+        self.buffer.group_values.len()
+    }
+
     /// Check if there is zero groups accumulated so far.
     pub(in crate::aggregates) fn is_empty(&self) -> bool {
-        self.buffer.group_values.is_empty()
+        self.num_groups() == 0
     }
 
     /// All internal buffer's memory size.
@@ -234,6 +250,43 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
             + self.buffer.group_indices.allocated_size()
     }
 
+    pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {
+        self.group_by_metrics.clone()
+    }
+
+    /// Takes every intermediate aggregate state and resets the table so it can
+    /// continue with a new ordered input segment.
+    ///
+    /// Unlike normal ordered emission, this operation is allowed to take the
+    /// active (incomplete) groups. Partial aggregation can pass those states to
+    /// its final stage, while final aggregation sorts and spills them before
+    /// replay.
+    pub(in crate::aggregates) fn take_state_batch(
+        &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        if self.buffer.group_values.is_empty() {
+            return Ok(None);
+        }
+
+        let mut output = self.buffer.group_values.emit(EmitTo::All)?;
+        for acc in &mut self.buffer.accumulators {
+            output.extend(acc.state(EmitTo::All)?);
+        }
+
+        let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?;
+        debug_assert!(batch.num_rows() > 0);
+
+        // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
+        // key/index buffers too so the memory reservation can be released
+        // before the batch is passed downstream or sorted for spilling.
+        self.buffer.group_values.clear_shrink(0);
+        self.buffer.group_indices.clear();
+        self.buffer.group_indices.shrink_to_fit();
+        self.buffer.group_ordering.reset();
+
+        Ok(Some(batch))
+    }
+
     /// Returns the [`EmitTo`], clamped to the specified batch size
     ///
     /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number

datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs

@@ -19,12 +19,15 @@
 //!
 //! See comments in [`super::ordered_partial_table`] for details.
 
+use std::sync::Arc;
+
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 
 use crate::InputOrderMode;
 use crate::aggregates::aggregate_hash_table::FinalMarker;
+use crate::aggregates::group_values::GroupByMetrics;
 use crate::aggregates::{AggregateExec, AggregateMode};
 
 use super::common_ordered::OrderedAggregateTable;
@@ -41,21 +44,22 @@ use super::common_ordered::OrderedAggregateTable;
 impl OrderedAggregateTable<FinalMarker> {
     pub(in crate::aggregates) fn new_with_input_order(
         agg: &AggregateExec,
-        partition: usize,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
         batch_size: usize,
         input_order_mode: &InputOrderMode,
+        group_by_metrics: GroupByMetrics,
     ) -> Result<Self> {
         Self::new_for_mode(
             agg,
-            partition,
             input_schema,
             output_schema,
+            Arc::clone(input_schema),
             batch_size,
             input_order_mode,
             &AggregateMode::Final,
             vec![None; agg.aggr_expr.len()],
+            group_by_metrics,
         )
     }
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs

@@ -29,12 +29,15 @@
 //! The implementation is separated from other aggregate tables because this
 //! execution path is likely to be optimized further in the future.
 
+use std::sync::Arc;
+
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 
 use crate::aggregates::{
     AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker,
+    group_values::GroupByMetrics,
 };
 
 use super::common_ordered::OrderedAggregateTable;
@@ -56,15 +59,18 @@ impl OrderedAggregateTable<PartialMarker> {
         batch_size: usize,
     ) -> Result<Self> {
         let input_schema = agg.input().schema();
+        let state_schema = Arc::clone(&output_schema);
+        let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
         Self::new_for_mode(
             agg,
-            partition,
             &input_schema,
             output_schema,
+            state_schema,
             batch_size,
             &agg.input_order_mode,
             &AggregateMode::Partial,
             agg.filter_expr.iter().cloned().collect(),
+            group_by_metrics,
         )
     }
 

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

@@ -1235,12 +1235,10 @@ impl AggregateExec {
             && self.limit_options_supported_by_hash_stream()
     }
 
-    fn should_use_ordered_partial_aggregate_stream(&self, context: &TaskContext) -> bool {
-        // TODO: implement memory-limited path and remove this limitation
-        if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
-            return false;
-        }
-
+    fn should_use_ordered_partial_aggregate_stream(
+        &self,
+        _context: &TaskContext,
+    ) -> bool {
         self.mode == AggregateMode::Partial
             && self.input_order_mode != InputOrderMode::Linear
             && !self.group_by.is_true_no_grouping()
@@ -1291,12 +1289,7 @@ impl AggregateExec {
             && self.group_by.is_single()
     }
 
-    fn should_use_ordered_final_aggregate_stream(&self, context: &TaskContext) -> bool {
-        // TODO: implement memory-limited path and remove this limitation
-        if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
-            return false;
-        }
-
+    fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool {
         matches!(
             self.mode,
             AggregateMode::Final | AggregateMode::FinalPartitioned
@@ -3930,10 +3923,10 @@ mod tests {
 +----------+-----------+-------------------------+
 ");
 
-        // Ordered streams don't implement memory limits yet.
+        // Ordered partial aggregation supports finite memory.
         let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?;
         let stream = aggregate.execute_typed(0, &finite_memory_task_ctx)?;
-        assert!(matches!(stream, StreamType::GroupedHash(_)));
+        assert!(matches!(stream, StreamType::OrderedPartialAggregate(_)));
 
         Ok(())
     }
@@ -4007,10 +4000,10 @@ mod tests {
 +-----+--------------+
 ");
 
-        // Ordered streams don't implement memory limits yet.
+        // Ordered final aggregation supports finite memory.
         let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?;
         let stream = final_aggregate.execute_typed(0, &finite_memory_task_ctx)?;
-        assert!(matches!(stream, StreamType::GroupedHash(_)));
+        assert!(matches!(stream, StreamType::OrderedFinalAggregate(_)));
 
         Ok(())
     }

datafusion/physical-plan/src/aggregates/order/full.rs

@@ -115,6 +115,11 @@ impl GroupOrderingFull {
         self.state = State::Complete;
     }
 
+    /// Starts tracking a new fully ordered input segment.
+    pub fn reset(&mut self) {
+        self.state = State::Start;
+    }
+
     /// Called when new groups are added in a batch. See documentation
     /// on [`super::GroupOrdering::new_groups`]
     pub fn new_groups(&mut self, total_num_groups: usize) {

datafusion/physical-plan/src/aggregates/order/mod.rs

@@ -93,6 +93,20 @@ impl GroupOrdering {
         }
     }
 
+    /// Resets the ordering state while preserving the configured ordering mode.
+    ///
+    /// Ordered partial aggregation uses this after passing intermediate states
+    /// downstream, and ordered final aggregation uses it after spilling a run.
+    /// In both cases the hash table is empty and can start tracking the next
+    /// input batch from a fresh ordering state.
+    pub fn reset(&mut self) {
+        match self {
+            GroupOrdering::None => {}
+            GroupOrdering::Partial(partial) => partial.reset(),
+            GroupOrdering::Full(full) => full.reset(),
+        }
+    }
+
     /// Removes the first `n` groups from the internal state, shifting all
     /// existing indexes down by `n`.
     pub fn remove_groups(&mut self, n: usize) {

datafusion/physical-plan/src/aggregates/order/partial.rs

@@ -186,6 +186,12 @@ impl GroupOrderingPartial {
         };
     }
 
+    /// Starts tracking a new ordered input segment with the same sort-key
+    /// columns.
+    pub fn reset(&mut self) {
+        self.state = State::Start;
+    }
+
     fn updated_sort_key(
         current_sort: usize,
         sort_key: Option<Vec<ScalarValue>>,