feat: stream exact group-contiguous aggregates

datafusion/datasource/src/source.rs

@@ -165,6 +165,18 @@ pub trait DataSource: Any + Send + Sync + Debug {
 
     fn output_partitioning(&self) -> Partitioning;
     fn eq_properties(&self) -> EquivalenceProperties;
+
+    /// Expressions whose complete tuple is contiguous within each output
+    /// partition.
+    ///
+    /// See
+    /// [`ExecutionPlanProperties::group_contiguous_exprs`](datafusion_physical_plan::ExecutionPlanProperties::group_contiguous_exprs)
+    /// for the full correctness contract. Expressions must refer to the schema
+    /// returned by [`Self::eq_properties`].
+    fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &[]
+    }
+
     fn scheduling_type(&self) -> SchedulingType {
         SchedulingType::NonCooperative
     }
@@ -674,6 +686,7 @@ impl DataSourceExec {
             EmissionType::Incremental,
             Boundedness::Bounded,
         )
+        .with_group_contiguous_exprs(data_source.group_contiguous_exprs().to_vec())
         .with_scheduling_type(data_source.scheduling_type())
     }
 
@@ -705,3 +718,90 @@ where
         Self::new(Arc::new(source))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion_physical_expr::expressions::Column;
+    use datafusion_physical_plan::EmptyRecordBatchStream;
+
+    #[derive(Debug)]
+    struct GroupContiguousSource {
+        schema: SchemaRef,
+        group_contiguous_exprs: Vec<Arc<dyn PhysicalExpr>>,
+    }
+
+    impl DataSource for GroupContiguousSource {
+        fn open(
+            &self,
+            _partition: usize,
+            _context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(
+                &self.schema,
+            ))))
+        }
+
+        fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
+            write!(f, "GroupContiguousSource")
+        }
+
+        fn output_partitioning(&self) -> Partitioning {
+            Partitioning::UnknownPartitioning(1)
+        }
+
+        fn eq_properties(&self) -> EquivalenceProperties {
+            EquivalenceProperties::new(Arc::clone(&self.schema))
+        }
+
+        fn group_contiguous_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+            &self.group_contiguous_exprs
+        }
+
+        fn partition_statistics(
+            &self,
+            _partition: Option<usize>,
+        ) -> Result<Arc<Statistics>> {
+            Ok(Arc::new(Statistics::new_unknown(&self.schema)))
+        }
+
+        fn with_fetch(&self, _limit: Option<usize>) -> Option<Arc<dyn DataSource>> {
+            None
+        }
+
+        fn fetch(&self) -> Option<usize> {
+            None
+        }
+
+        fn try_swapping_with_projection(
+            &self,
+            _projection: &ProjectionExprs,
+        ) -> Result<Option<Arc<dyn DataSource>>> {
+            Ok(None)
+        }
+
+        fn apply_expressions(
+            &self,
+            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+        ) -> Result<TreeNodeRecursion> {
+            Ok(TreeNodeRecursion::Continue)
+        }
+    }
+
+    #[test]
+    fn data_source_exec_exposes_group_contiguous_exprs() -> Result<()> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
+        let key = Arc::new(Column::new("key", 0)) as Arc<dyn PhysicalExpr>;
+        let exec = DataSourceExec::new(Arc::new(GroupContiguousSource {
+            schema,
+            group_contiguous_exprs: vec![Arc::clone(&key)],
+        }));
+
+        assert_eq!(exec.properties().group_contiguous_exprs().len(), 1);
+        assert!(exec.properties().group_contiguous_exprs()[0].eq(&key));
+        Ok(())
+    }
+}

datafusion/ffi/src/session/mod.rs

@@ -985,7 +985,7 @@ mod tests {
         let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?;
         assert_eq!(
             format!("{physical_plan:?}"),
-            "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None } }"
+            "EmptyExec { schema: Schema { fields: [], metadata: {} }, partitions: 1, cache: PlanProperties { eq_properties: EquivalenceProperties { eq_group: EquivalenceGroup { map: {}, classes: [] }, oeq_class: OrderingEquivalenceClass { orderings: [] }, oeq_cache: OrderingEquivalenceCache { normal_cls: OrderingEquivalenceClass { orderings: [] }, leading_map: {} }, constraints: Constraints { inner: [] }, schema: Schema { fields: [], metadata: {} } }, partitioning: UnknownPartitioning(1), emission_type: Incremental, boundedness: Bounded, evaluation_type: Lazy, scheduling_type: Cooperative, output_ordering: None, group_contiguous_exprs: [] } }"
         );
 
         assert_eq!(

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,8 @@ 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_group_completion(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

@@ -15,12 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Aggregate table for partial aggregation when input is ordered by group keys.
+//! Aggregate table for partial aggregation when input groups can be completed
+//! before end of input.
 //!
 //! See the [`super::common_ordered`] comments for the high-level ideas.
 //!
-//! This operator handles input that is ordered by group keys:
-//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b`
+//! This operator handles both complete-key contiguity and input ordered by a
+//! subset of the group keys:
+//! - Fully contiguous: every distinct `(a, b)` tuple occurs in one input range
 //! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a`
 //!
 //! When a group key combination is exhausted, this table eagerly flushes the
@@ -68,7 +70,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,
@@ -90,7 +92,7 @@ impl OrderedAggregateTable<PartialMarker> {
     }
 
     /// Emits the next batch of partial state rows for groups proven complete by
-    /// the input ordering.
+    /// the input's group-completion guarantee.
     ///
     /// For example, when the query is `GROUP BY a` and the input is ordered by
     /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3`

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,8 @@ 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_group_completion(&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,
@@ -530,10 +531,10 @@ impl GroupedHashAggregateStream {
             {
                 OutOfMemoryMode::Spill
             }
-            // For `GroupOrdering::Full`, the incoming stream is already sorted. This ensures the
-            // number of incomplete groups can be kept small at all times. If we still hit
-            // an out-of-memory condition, spilling to disk would not be beneficial since the same
-            // situation is likely to reoccur when reading back the spilled data.
+            // For `GroupOrdering::Full`, every complete group tuple is contiguous. This keeps the
+            // number of incomplete groups small at all times. If we still hit an out-of-memory
+            // condition, spilling to disk would not be beneficial since the same situation is
+            // likely to reoccur when reading back the spilled data.
             // Therefore, we fall back to simply reporting the error immediately.
             // This mode will also be used if the `DiskManager` is not configured to allow spilling
             // to disk.

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,