feat: stream aggregates over grouped input

datafusion/core/tests/dataframe/mod.rs

@@ -3422,9 +3422,9 @@ async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_reparti
     assert_snapshot!(
         union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(false).await?,
         @r"
-    AggregateExec: mode=Final, gby=[id@0 as id], aggr=[], ordering_mode=Sorted
+    AggregateExec: mode=Final, gby=[id@0 as id], aggr=[], group_clustering_mode=Full
       SortPreservingMergeExec: [id@0 ASC NULLS LAST]
-        AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[], ordering_mode=Sorted
+        AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[], group_clustering_mode=Full
           UnionExec
             DataSourceExec: file_groups={1 group: [[{testdata}/alltypes_tiny_pages.parquet]]}, projection=[id], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet
             SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false]
@@ -3440,9 +3440,9 @@ async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_reparti
     assert_snapshot!(
         union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(true).await?,
         @r"
-    AggregateExec: mode=Final, gby=[id@0 as id], aggr=[], ordering_mode=Sorted
+    AggregateExec: mode=Final, gby=[id@0 as id], aggr=[], group_clustering_mode=Full
       SortPreservingMergeExec: [id@0 ASC NULLS LAST]
-        AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[], ordering_mode=Sorted
+        AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[], group_clustering_mode=Full
           UnionExec
             DataSourceExec: file_groups={1 group: [[{testdata}/alltypes_tiny_pages.parquet]]}, projection=[id], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet
             SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false]

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -41,15 +41,14 @@ use datafusion_common_runtime::JoinSet;
 use datafusion_functions_aggregate::sum::sum_udaf;
 use datafusion_physical_expr::PhysicalSortExpr;
 use datafusion_physical_expr::expressions::{Column, col, lit};
-use datafusion_physical_plan::InputOrderMode;
 use test_utils::{StringBatchGenerator, add_empty_batches};
 
 use datafusion_execution::TaskContext;
 use datafusion_execution::memory_pool::FairSpillPool;
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_physical_expr::aggregate::AggregateExprBuilder;
 use datafusion_physical_plan::aggregates::{
-    AggregateExec, AggregateMode, PhysicalGroupBy,
+    AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy,
 };
 use datafusion_physical_plan::metrics::MetricValue;
 use datafusion_physical_plan::{ExecutionPlan, collect, displayable};
@@ -303,7 +302,7 @@ async fn streaming_aggregate_test() {
 /// two `AggregateExec` variants produce the same result: the pipeline breaking
 /// one over unordered input (`PartialHashAggregateStream`) and the
 /// non-pipeline breaking one over ordered input
-/// (`OrderedPartialAggregateStream`).
+/// (`ClusteredPartialAggregateStream`).
 async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str>) {
     let schema = input1[0].schema();
     let session_config = SessionConfig::new().with_batch_size(50);
@@ -354,8 +353,8 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
         .unwrap(),
     );
     assert_ne!(
-        aggregate_exec_running.input_order_mode(),
-        &InputOrderMode::Linear,
+        aggregate_exec_running.group_clustering_mode(),
+        &GroupClusteringMode::None,
         "running aggregate should observe ordered input for group_by: {group_by:?}"
     );
 
@@ -555,13 +554,17 @@ async fn verify_ordered_aggregate(frame: &DataFrame, expected_sort: bool) {
 
         fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
             if let Some(exec) = node.downcast_ref::<AggregateExec>() {
+                assert_eq!(
+                    exec.properties().output_ordering().is_some(),
+                    self.expected_sort
+                );
                 if self.expected_sort {
                     assert!(matches!(
-                        exec.input_order_mode(),
-                        InputOrderMode::PartiallySorted(_) | InputOrderMode::Sorted
+                        exec.group_clustering_mode(),
+                        GroupClusteringMode::Partial(_) | GroupClusteringMode::Full
                     ));
                 } else {
-                    assert_eq!(*exec.input_order_mode(), InputOrderMode::Linear);
+                    assert_eq!(*exec.group_clustering_mode(), GroupClusteringMode::None);
                 }
             }
             Ok(TreeNodeRecursion::Continue)

datafusion/core/tests/fuzz_cases/aggregation_fuzzer/query_builder.rs

@@ -93,9 +93,9 @@ pub struct QueryBuilder {
     ///   ...
     /// ```
     ///
-    /// More details can see [`GroupOrdering`].
+    /// More details can see [`GroupClustering`].
     ///
-    /// [`GroupOrdering`]:  datafusion_physical_plan::aggregates::order::GroupOrdering
+    /// [`GroupClustering`]:  datafusion_physical_plan::aggregates::order::GroupClustering
     dataset_sort_keys: Vec<Vec<String>>,
 
     /// If we will also test the no grouping case like:

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -4867,9 +4867,9 @@ fn preserve_ordering_for_streaming_sorted_aggregate() -> Result<()> {
 
     let plan_distrib = test_config.to_plan(physical_plan.clone(), &DISTRIB_DISTRIB_SORT);
     assert_plan!(plan_distrib, @r"
-    AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[COUNT(b)], ordering_mode=Sorted
+    AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[COUNT(b)], group_clustering_mode=Full
       RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2, preserve_order=true, sort_exprs=a@0 ASC
-        AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[COUNT(b)], ordering_mode=Sorted
+        AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[COUNT(b)], group_clustering_mode=Full
           DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
     ");
 
@@ -4901,9 +4901,9 @@ fn preserve_ordering_for_streaming_partially_sorted_aggregate() -> Result<()> {
 
     let plan_distrib = test_config.to_plan(physical_plan.clone(), &DISTRIB_DISTRIB_SORT);
     assert_plan!(plan_distrib, @r"
-    AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[COUNT(c)], ordering_mode=PartiallySorted([0])
+    AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[COUNT(c)], group_clustering_mode=Partial([0])
       RepartitionExec: partitioning=Hash([a@0, b@1], 2), input_partitions=2, preserve_order=true, sort_exprs=a@0 ASC
-        AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[COUNT(c)], ordering_mode=PartiallySorted([0])
+        AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[COUNT(c)], group_clustering_mode=Partial([0])
           DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
     ");
 

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -520,7 +520,7 @@ fn test_has_order_by() -> Result<()> {
         actual,
         @r"
     LocalLimitExec: fetch=10
-      AggregateExec: mode=Single, gby=[a@0 as a], aggr=[], ordering_mode=Sorted
+      AggregateExec: mode=Single, gby=[a@0 as a], aggr=[], group_clustering_mode=Full
         DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
     "
     );

datafusion/core/tests/physical_optimizer/pushdown_sort.rs

@@ -731,13 +731,13 @@ fn test_pushdown_through_blocking_node() {
     OptimizationTest:
       input:
         - SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
-        -   AggregateExec: mode=Final, gby=[a@0 as a], aggr=[COUNT(b)], ordering_mode=Sorted
+        -   AggregateExec: mode=Final, gby=[a@0 as a], aggr=[COUNT(b)], group_clustering_mode=Full
         -     SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
         -       DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
       output:
         Ok:
           - SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
-          -   AggregateExec: mode=Final, gby=[a@0 as a], aggr=[COUNT(b)], ordering_mode=Sorted
+          -   AggregateExec: mode=Final, gby=[a@0 as a], aggr=[COUNT(b)], group_clustering_mode=Full
           -     SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
           -       DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet, sort_order_for_reorder=[a@0 DESC NULLS LAST], reverse_row_groups=true
     "

datafusion/expr-common/src/sort_properties.rs

@@ -37,16 +37,24 @@ use arrow::datatypes::DataType;
 pub enum SortProperties {
     /// Use the ordinary [`SortOptions`] struct to represent ordered data:
     Ordered(SortOptions),
-    // This alternative represents unordered data:
+    /// Within each partition, all rows with the same value for this expression
+    /// form one contiguous run. The runs may occur in any order.
+    Grouped,
+    /// This alternative represents unordered data:
     #[default]
     Unordered,
-    // Singleton is used for single-valued literal numbers:
+    /// Singleton is used for single-valued literal numbers:
     Singleton,
 }
 
 impl SortProperties {
     pub fn add(&self, rhs: &Self) -> Self {
         match (self, rhs) {
+            // Addition can collapse distinct values (for example through
+            // floating-point rounding), which may join non-adjacent groups.
+            (Self::Grouped, Self::Singleton) | (Self::Singleton, Self::Grouped) => {
+                Self::Unordered
+            }
             (Self::Singleton, _) => *rhs,
             (_, Self::Singleton) => *self,
             (Self::Ordered(lhs), Self::Ordered(rhs))
@@ -65,6 +73,11 @@ impl SortProperties {
     pub fn sub(&self, rhs: &Self) -> Self {
         match (self, rhs) {
             (Self::Singleton, Self::Singleton) => Self::Singleton,
+            // Subtraction can collapse distinct values (for example through
+            // floating-point rounding), which may join non-adjacent groups.
+            (Self::Grouped, Self::Singleton) | (Self::Singleton, Self::Grouped) => {
+                Self::Unordered
+            }
             (Self::Singleton, Self::Ordered(rhs)) => Self::Ordered(SortOptions {
                 descending: !rhs.descending,
                 nulls_first: rhs.nulls_first,
@@ -89,6 +102,9 @@ impl SortProperties {
                 descending: !rhs.descending,
                 nulls_first: rhs.nulls_first,
             }),
+            // Comparisons can map several non-adjacent grouped values to the
+            // same boolean value, so they do not preserve grouping.
+            (Self::Grouped, Self::Singleton) => Self::Unordered,
             (_, Self::Singleton) => *self,
             (Self::Ordered(lhs), Self::Ordered(rhs))
                 if lhs.descending != rhs.descending
@@ -147,6 +163,7 @@ mod sort_properties_test {
     const ASC_NL: SortProperties = ordered(false, false);
     const DESC_NF: SortProperties = ordered(true, true);
     const DESC_NL: SortProperties = ordered(true, false);
+    const GROUPED: SortProperties = SortProperties::Grouped;
     const UNORDERED: SortProperties = SortProperties::Unordered;
     const SINGLETON: SortProperties = SortProperties::Singleton;
 
@@ -197,6 +214,27 @@ mod sort_properties_test {
                 SINGLETON,
                 SINGLETON,
             ),
+            (
+                "add: may collapse grouped values",
+                SortProperties::add,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
+            (
+                "sub: may collapse grouped values",
+                SortProperties::sub,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
+            (
+                "comparison does not preserve grouping",
+                SortProperties::gt_or_gteq,
+                GROUPED,
+                SINGLETON,
+                UNORDERED,
+            ),
             // `and` keeps ASC NULLS LAST / DESC NULLS FIRST, `or` keeps ASC
             // NULLS FIRST / DESC NULLS LAST. Both are commutative.
             (
@@ -482,7 +520,7 @@ impl Neg for SortProperties {
 #[derive(Debug, Clone)]
 pub struct ExprProperties {
     /// Properties that describe the sorting behavior of the expression,
-    /// such as whether it is ordered, unordered, or a singleton value.
+    /// such as whether it is ordered, grouped, unordered, or a singleton value.
     pub sort_properties: SortProperties,
     /// A closed interval representing the range of possible values for
     /// the expression. Used to compute reliable bounds.

datafusion/expr/src/udf.rs

@@ -360,8 +360,20 @@ impl ScalarUDF {
 
     /// Calculates the [`SortProperties`] of this function based on its
     /// children's properties.
+    ///
+    /// [`SortProperties::Grouped`] is retained only when the implementation
+    /// also reports that the transformation is strictly order-preserving. A
+    /// many-to-one function can otherwise make an output value occur in
+    /// multiple non-adjacent runs.
     pub fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
-        self.inner.output_ordering(inputs)
+        let sort_properties = self.inner.output_ordering(inputs)?;
+        if sort_properties == SortProperties::Grouped
+            && !self.inner.strictly_order_preserving(inputs)?
+        {
+            Ok(SortProperties::Unordered)
+        } else {
+            Ok(sort_properties)
+        }
     }
 
     pub fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result<bool> {
@@ -957,7 +969,12 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         Ok(Some(vec![]))
     }
 
-    /// Calculates the [`SortProperties`] of this function based on its children's properties.
+    /// Calculates the [`SortProperties`] of this function based on its children's
+    /// properties.
+    ///
+    /// The [`ScalarUDF`] wrapper retains a [`SortProperties::Grouped`] result
+    /// only when [`Self::strictly_order_preserving`] also returns `true` for
+    /// the same inputs.
     fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
         if !self.preserves_lex_ordering(inputs)? {
             return Ok(SortProperties::Unordered);

datafusion/ffi/src/expr/expr_properties.rs

@@ -67,13 +67,15 @@ pub enum FFI_SortProperties {
     Ordered(FFI_SortOptions),
     Unordered,
     Singleton,
+    Grouped,
 }
 
 impl From<&SortProperties> for FFI_SortProperties {
     fn from(value: &SortProperties) -> Self {
         match value {
             SortProperties::Unordered => FFI_SortProperties::Unordered,
             SortProperties::Singleton => FFI_SortProperties::Singleton,
+            SortProperties::Grouped => FFI_SortProperties::Grouped,
             SortProperties::Ordered(o) => FFI_SortProperties::Ordered(o.into()),
         }
     }
@@ -84,6 +86,7 @@ impl From<&FFI_SortProperties> for SortProperties {
         match value {
             FFI_SortProperties::Unordered => SortProperties::Unordered,
             FFI_SortProperties::Singleton => SortProperties::Singleton,
+            FFI_SortProperties::Grouped => SortProperties::Grouped,
             FFI_SortProperties::Ordered(o) => SortProperties::Ordered(o.into()),
         }
     }

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -21,7 +21,7 @@ use arrow_schema::DataType;
 use datafusion_catalog::TableFunctionImpl;
 use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
-use datafusion_expr::sort_properties::ExprProperties;
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
     AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF,
     ScalarUDFImpl, Signature, Volatility, WindowUDF,
@@ -172,6 +172,11 @@ impl ScalarUDFImpl for PlacementUDF {
         &self,
         inputs: &[ExprProperties],
     ) -> datafusion_common::Result<bool> {
+        // This test-only sentinel verifies that the new `Grouped` variant
+        // survives the cross-library `ExprProperties` conversion.
+        if matches!(inputs, [input] if input.sort_properties == SortProperties::Grouped) {
+            return Ok(true);
+        }
         Ok(inputs.iter().all(|input| input.preserves_lex_ordering))
     }
 }

datafusion/ffi/src/udf/mod.rs

@@ -553,6 +553,7 @@ impl ScalarUDFImpl for ForeignScalarUDF {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use datafusion_expr::sort_properties::SortProperties;
 
     #[derive(Debug, PartialEq, Eq, Hash)]
     struct PlacementUDF {
@@ -594,6 +595,13 @@ mod tests {
                 return internal_err!("preserves_lex_ordering requires an input");
             }
 
+            // This test-only sentinel verifies that the new `Grouped` variant
+            // travels through the foreign path.
+            if matches!(inputs, [input] if input.sort_properties == SortProperties::Grouped)
+            {
+                return Ok(true);
+            }
+
             Ok(inputs.iter().all(|input| input.preserves_lex_ordering))
         }
 
@@ -692,6 +700,8 @@ mod tests {
                 .preserves_lex_ordering(&[preserves, does_not_preserve])
                 .unwrap()
         );
+        let grouped = ExprProperties::new_unknown().with_order(SortProperties::Grouped);
+        assert!(foreign_udf.preserves_lex_ordering(&[grouped]).unwrap());
         assert!(foreign_udf.preserves_lex_ordering(&[]).is_err());
 
         let updated = foreign_udf

datafusion/ffi/tests/ffi_udf.rs

@@ -27,7 +27,7 @@ mod tests {
     use datafusion::prelude::{SessionContext, col};
     use datafusion_execution::config::SessionConfig;
     use datafusion_expr::lit;
-    use datafusion_expr::sort_properties::ExprProperties;
+    use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
     use datafusion_ffi::tests::create_record_batch;
     use datafusion_ffi::tests::utils::get_module;
     use std::sync::Arc;
@@ -119,6 +119,9 @@ mod tests {
         assert!(foreign_func.preserves_lex_ordering(std::slice::from_ref(&preserves))?);
         assert!(!foreign_func.preserves_lex_ordering(&[preserves, does_not_preserve])?);
 
+        let grouped = ExprProperties::new_unknown().with_order(SortProperties::Grouped);
+        assert!(foreign_func.preserves_lex_ordering(&[grouped])?);
+
         Ok(())
     }