refactor: separate aggregate group completion from input ordering

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

@@ -28,14 +28,13 @@ use datafusion_common::assert_or_internal_err;
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_expr::EmitTo;
 
-use crate::InputOrderMode;
 use crate::PhysicalExpr;
 use crate::aggregates::group_values::{
     AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics,
     GroupByMetrics, GroupValues, new_group_values,
 };
 use crate::aggregates::grouped_hash_stream::create_group_accumulator;
-use crate::aggregates::order::GroupOrdering;
+use crate::aggregates::order::{GroupCompletionMode, GroupOrdering};
 use crate::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions,
     evaluate_group_by,
@@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics {
 /// `OrderedAggrMode` selects the aggregate semantics. For example,
 /// `OrderedAggregateTable::<PartialMarker>::new(...)` consumes raw rows
 /// and emits partial states, while
-/// `OrderedAggregateTable::<FinalMarker>::new_with_input_order(...)`
+/// `OrderedAggregateTable::<FinalMarker>::new_with_group_completion(...)`
 /// consumes partial states and emits final values.
 ///
 /// Shared methods live on `impl<T>`; single/partial/final behavior lives on
@@ -184,7 +183,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         output_schema: SchemaRef,
         state_schema: SchemaRef,
         batch_size: usize,
-        input_order_mode: &InputOrderMode,
+        group_completion_mode: &GroupCompletionMode,
         aggregate_mode: &AggregateMode,
         filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
         metrics: OrderedAggregateTableMetrics,
@@ -194,7 +193,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
             "OrderedAggregateTable requires config batch_size >= 1"
         );
 
-        let group_ordering = GroupOrdering::try_new(input_order_mode)?;
+        let group_ordering = GroupOrdering::try_new_for_mode(group_completion_mode)?;
         let group_schema = agg.group_by.group_schema(input_schema)?;
         let group_values = new_group_values(group_schema, &group_ordering)?;
         let aggregate_arguments = aggregate_expressions(

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

@@ -25,8 +25,8 @@ 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::order::GroupCompletionMode;
 use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase};
 
 use super::common::HashAggregateAccumulator;
@@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}
 ///
 /// See comments at [`OrderedAggregateTable`] for details.
 impl OrderedAggregateTable<FinalMarker> {
-    pub(in crate::aggregates) fn new_with_input_order(
+    pub(in crate::aggregates) fn new_with_group_completion(
         agg: &AggregateExec,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
         batch_size: usize,
-        input_order_mode: &InputOrderMode,
+        group_completion_mode: &GroupCompletionMode,
         metrics: OrderedAggregateTableMetrics,
     ) -> Result<Self> {
         Self::new_for_mode(
@@ -56,7 +56,7 @@ impl OrderedAggregateTable<FinalMarker> {
             output_schema,
             Arc::clone(input_schema),
             batch_size,
-            input_order_mode,
+            group_completion_mode,
             &AggregateMode::Final,
             vec![None; agg.aggr_expr.len()],
             metrics,

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

@@ -68,7 +68,7 @@ impl OrderedAggregateTable<PartialMarker> {
             output_schema,
             state_schema,
             batch_size,
-            &agg.input_order_mode,
+            &agg.group_completion_mode,
             &AggregateMode::Partial,
             agg.filter_expr.iter().cloned().collect(),
             metrics,

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

@@ -59,7 +59,7 @@ impl OrderedAggregateTable<SingleMarker> {
             output_schema,
             state_schema,
             batch_size,
-            &agg.input_order_mode,
+            &agg.group_completion_mode,
             &agg.mode,
             agg.filter_expr.iter().cloned().collect(),
             metrics,

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

@@ -518,7 +518,7 @@ impl GroupedHashAggregateStream {
             .collect::<Vec<_>>()
             .join(", ");
         let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})");
-        let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?;
+        let group_ordering = GroupOrdering::try_new_for_mode(&agg.group_completion_mode)?;
         let oom_mode = match (agg.mode, &group_ordering) {
             // In partial aggregation mode, always prefer to emit incomplete results early.
             (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly,

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

@@ -45,6 +45,7 @@ use super::aggregate_hash_table::{
     AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker,
     PartialSkipMarker,
 };
+use super::order::GroupCompletionMode;
 use super::ordered_final_stream::OrderedFinalAggregateStream;
 use super::skip_partial::SkipAggregationProbe;
 use crate::metrics::{
@@ -326,6 +327,7 @@ impl FinalSpillContext {
 
         let mut final_agg = agg.clone();
         final_agg.input_order_mode = InputOrderMode::Sorted;
+        final_agg.group_completion_mode = GroupCompletionMode::Full;
 
         Ok(Self {
             final_agg,
@@ -414,7 +416,7 @@ impl FinalSpillContext {
             &context,
             partition,
             merged,
-            &InputOrderMode::Sorted,
+            &GroupCompletionMode::Full,
             baseline_metrics.clone(),
             metrics,
             None,

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

@@ -203,6 +203,7 @@ use datafusion_physical_expr_common::sort_expr::{
 use datafusion_expr::utils::AggregateOrderSensitivity;
 use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
 use itertools::Itertools;
+use order::GroupCompletionMode;
 use topk::hash_table::is_supported_hash_key_type;
 use topk::heap::is_supported_heap_type;
 
@@ -877,6 +878,8 @@ pub struct AggregateExec {
     required_input_ordering: Option<OrderingRequirements>,
     /// Describes how the input is ordered relative to the group by columns
     input_order_mode: InputOrderMode,
+    /// Describes how the executor can determine that groups are complete.
+    group_completion_mode: GroupCompletionMode,
     cache: Arc<PlanProperties>,
     /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]),
     /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node.
@@ -901,6 +904,7 @@ impl AggregateExec {
             required_input_ordering: self.required_input_ordering.clone(),
             metrics: ExecutionPlanMetricsSet::new(),
             input_order_mode: self.input_order_mode.clone(),
+            group_completion_mode: self.group_completion_mode.clone(),
             cache: Arc::clone(&self.cache),
             mode: self.mode,
             group_by: Arc::clone(&self.group_by),
@@ -921,6 +925,7 @@ impl AggregateExec {
             required_input_ordering: self.required_input_ordering.clone(),
             metrics: ExecutionPlanMetricsSet::new(),
             input_order_mode: self.input_order_mode.clone(),
+            group_completion_mode: self.group_completion_mode.clone(),
             cache: Arc::clone(&self.cache),
             mode: self.mode,
             group_by: Arc::clone(&self.group_by),
@@ -1043,6 +1048,8 @@ impl AggregateExec {
             input_order_mode = InputOrderMode::Linear;
         }
 
+        let group_completion_mode = GroupCompletionMode::from(&input_order_mode);
+
         // construct a map from the input expression to the output expression of the Aggregation group by
         let group_expr_mapping =
             ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?;
@@ -1073,6 +1080,7 @@ impl AggregateExec {
             required_input_ordering,
             limit_options: None,
             input_order_mode,
+            group_completion_mode,
             cache: Arc::new(cache),
             dynamic_filter: None,
         };
@@ -1281,7 +1289,7 @@ impl AggregateExec {
 
     fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool {
         self.mode == AggregateMode::Partial
-            && self.input_order_mode == InputOrderMode::Linear
+            && self.group_completion_mode == GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
             && self.limit_options_supported_by_hash_stream()
@@ -1292,7 +1300,7 @@ impl AggregateExec {
         _context: &TaskContext,
     ) -> bool {
         self.mode == AggregateMode::Partial
-            && self.input_order_mode != InputOrderMode::Linear
+            && self.group_completion_mode != GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
             && self.limit_options_supported_by_hash_stream()
@@ -1303,7 +1311,7 @@ impl AggregateExec {
             self.mode,
             AggregateMode::Final | AggregateMode::FinalPartitioned
         ) && self.limit_options_supported_by_hash_stream()
-            && self.input_order_mode == InputOrderMode::Linear
+            && self.group_completion_mode == GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
     }
@@ -1316,7 +1324,7 @@ impl AggregateExec {
 
         self.mode == AggregateMode::PartialReduce
             && self.limit_options.is_none()
-            && self.input_order_mode == InputOrderMode::Linear
+            && self.group_completion_mode == GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
     }
@@ -1326,7 +1334,7 @@ impl AggregateExec {
             self.mode,
             AggregateMode::Single | AggregateMode::SinglePartitioned
         ) && self.limit_options.is_none()
-            && self.input_order_mode == InputOrderMode::Linear
+            && self.group_completion_mode == GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
     }
@@ -1336,7 +1344,7 @@ impl AggregateExec {
             self.mode,
             AggregateMode::Single | AggregateMode::SinglePartitioned
         ) && self.limit_options.is_none()
-            && self.input_order_mode != InputOrderMode::Linear
+            && self.group_completion_mode != GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
     }
@@ -1346,7 +1354,7 @@ impl AggregateExec {
             self.mode,
             AggregateMode::Final | AggregateMode::FinalPartitioned
         ) && self.limit_options_supported_by_hash_stream()
-            && self.input_order_mode != InputOrderMode::Linear
+            && self.group_completion_mode != GroupCompletionMode::None
             && !self.group_by.is_true_no_grouping()
             && self.group_by.is_single()
     }
@@ -2354,6 +2362,8 @@ impl ExecutionPlan for AggregateExec {
             required_input_ordering: _,
             // Derived at construction from the input ordering and `group_by`.
             input_order_mode: _,
+            // Derived at construction from `input_order_mode`.
+            group_completion_mode: _,
             // Derived at construction by `Self::compute_properties`.
             cache: _,
             dynamic_filter,
@@ -4539,6 +4549,61 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn unsorted_contiguous_groups_use_final_emission() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("key", DataType::Int32, false),
+            Field::new("time_bin", DataType::Int64, false),
+            Field::new("value", DataType::Int64, false),
+        ]));
+        // Two sorted logical runs are concatenated into one DataFusion partition.
+        // Every distinct grouping tuple occupies one contiguous range, but tuple
+        // order resets between runs, so (key, time_bin) is not globally sorted.
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 1, 2, 2, 1, 1, 2, 2])),
+                Arc::new(Int64Array::from(vec![20, 20, 20, 20, 0, 0, 0, 0])),
+                Arc::new(Int64Array::from(vec![10, 20, 30, 40, 50, 60, 70, 80])),
+            ],
+        )?;
+        let group_by = PhysicalGroupBy::new_single(vec![
+            (col("key", &schema)?, "key".to_string()),
+            (col("time_bin", &schema)?, "time_bin".to_string()),
+        ]);
+        let aggr_expr = Arc::new(
+            AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
+                .schema(Arc::clone(&schema))
+                .alias("SUM(value)")
+                .build()?,
+        );
+        let input: Arc<dyn ExecutionPlan> =
+            TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?;
+        assert_eq!(input.output_partitioning().partition_count(), 1);
+
+        let aggregate = AggregateExec::try_new(
+            AggregateMode::Single,
+            group_by,
+            vec![aggr_expr],
+            vec![None],
+            input,
+            schema,
+        )?;
+
+        assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
+        assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::None);
+        // This captures the behavior before #24438. When the source can declare
+        // `(key, time_bin)` group-contiguous, the corresponding case can use
+        // `EmissionType::Incremental`.
+        assert_eq!(aggregate.cache().emission_type, EmissionType::Final);
+
+        let task_ctx = new_migrated_hash_ctx(1024);
+        let stream = aggregate.execute_typed(0, &task_ctx)?;
+        assert!(matches!(stream, StreamType::SingleHash(_)));
+
+        Ok(())
+    }
+
     /// Ensures for ordered input, `OrderedPartialAggregateStream` is used.
     #[tokio::test]
     async fn ordered_partial_aggregate_planning() -> Result<()> {

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

@@ -28,6 +28,33 @@ use crate::InputOrderMode;
 pub use full::GroupOrderingFull;
 pub use partial::GroupOrderingPartial;
 
+/// Describes how an aggregate can determine that groups are complete.
+///
+/// This is distinct from [`InputOrderMode`], which describes the ordering of
+/// the input relative to the grouping expressions. Input ordering is one way
+/// to establish a group-completion mode, but the execution machinery only
+/// needs to know when it can safely emit completed groups.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) enum GroupCompletionMode {
+    /// Groups cannot be completed before the input ends.
+    None,
+    /// Groups sharing the values at these grouping-expression indices form a
+    /// contiguous range.
+    Partial(Vec<usize>),
+    /// Every complete grouping tuple forms a contiguous range.
+    Full,
+}
+
+impl From<&InputOrderMode> for GroupCompletionMode {
+    fn from(value: &InputOrderMode) -> Self {
+        match value {
+            InputOrderMode::Linear => Self::None,
+            InputOrderMode::PartiallySorted(indices) => Self::Partial(indices.clone()),
+            InputOrderMode::Sorted => Self::Full,
+        }
+    }
+}
+
 /// Ordering information for each group in the hash table
 #[derive(Debug)]
 pub enum GroupOrdering {
@@ -40,15 +67,22 @@ pub enum GroupOrdering {
 }
 
 impl GroupOrdering {
-    /// Create a `GroupOrdering` for the specified ordering
+    /// Create a `GroupOrdering` for the specified input order mode.
     pub fn try_new(mode: &InputOrderMode) -> Result<Self> {
+        Self::try_new_for_mode(&GroupCompletionMode::from(mode))
+    }
+
+    /// Create a `GroupOrdering` for the specified group-completion mode.
+    pub(crate) fn try_new_for_mode(mode: &GroupCompletionMode) -> Result<Self> {
         match mode {
-            InputOrderMode::Linear => Ok(GroupOrdering::None),
-            InputOrderMode::PartiallySorted(order_indices) => {
+            GroupCompletionMode::None => Ok(GroupOrdering::None),
+            GroupCompletionMode::Partial(order_indices) => {
                 GroupOrderingPartial::try_new(order_indices.clone())
                     .map(GroupOrdering::Partial)
             }
-            InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())),
+            GroupCompletionMode::Full => {
+                Ok(GroupOrdering::Full(GroupOrderingFull::new()))
+            }
         }
     }