Implement Streaming Aggregation: Do not break pipeline in aggregation if group by columns are ordered

datafusion-examples/examples/custom_datasource.rs

@@ -217,7 +217,7 @@ impl ExecutionPlan for CustomExec {
         datafusion::physical_plan::Partitioning::UnknownPartitioning(1)
     }
 
-    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+    fn output_ordering(&self) -> Option<Vec<PhysicalSortExpr>> {
         None
     }
 

datafusion/core/src/physical_optimizer/repartition.rs

@@ -1147,7 +1147,7 @@ mod tests {
             self.input.output_partitioning()
         }
 
-        fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        fn output_ordering(&self) -> Option<Vec<PhysicalSortExpr>> {
             self.input.output_ordering()
         }
 
@@ -1159,6 +1159,7 @@ mod tests {
         fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
             vec![self
                 .output_ordering()
+                .as_deref()
                 .map(PhysicalSortRequirement::from_sort_exprs)]
         }
 

datafusion/core/src/physical_optimizer/sort_enforcement.rs

@@ -468,7 +468,7 @@ fn ensure_sorting(
         match (required_ordering, physical_ordering) {
             (Some(required_ordering), Some(physical_ordering)) => {
                 if !ordering_satisfy_requirement_concrete(
-                    physical_ordering,
+                    &physical_ordering,
                     &required_ordering,
                     || child.equivalence_properties(),
                 ) {
@@ -534,8 +534,8 @@ fn analyze_immediate_sort_removal(
         let sort_input = sort_exec.input().clone();
         // If this sort is unnecessary, we should remove it:
         if ordering_satisfy(
-            sort_input.output_ordering(),
-            sort_exec.output_ordering(),
+            sort_input.output_ordering().as_deref(),
+            sort_exec.output_ordering().as_deref(),
             || sort_input.equivalence_properties(),
         ) {
             // Since we know that a `SortExec` has exactly one child,
@@ -785,7 +785,7 @@ fn can_skip_sort(
         return Ok(None);
     };
     let orderby_exprs = convert_to_expr(orderby_keys);
-    let physical_ordering_exprs = convert_to_expr(physical_ordering);
+    let physical_ordering_exprs = convert_to_expr(&physical_ordering);
     let equal_properties = || input.equivalence_properties();
     // indices of the order by expressions among input ordering expressions
     let ob_indices = get_indices_of_matching_exprs(
@@ -826,7 +826,7 @@ fn can_skip_sort(
     if first_n <= furthest_ob_index {
         return Ok(None);
     }
-    let input_orderby_columns = get_at_indices(physical_ordering, &unique_ob_indices)?;
+    let input_orderby_columns = get_at_indices(&physical_ordering, &unique_ob_indices)?;
     let expected_orderby_columns =
         get_at_indices(orderby_keys, find_indices(&ob_indices, unique_ob_indices)?)?;
     let should_reverse = if let Some(should_reverse) = check_alignments(

datafusion/core/src/physical_optimizer/sort_pushdown.rs

@@ -125,9 +125,11 @@ pub(crate) fn pushdown_sorts(
     let err = || DataFusionError::Plan(ERR_MSG.to_string());
     if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
         let mut new_plan = plan.clone();
-        if !ordering_satisfy_requirement(plan.output_ordering(), parent_required, || {
-            plan.equivalence_properties()
-        }) {
+        if !ordering_satisfy_requirement(
+            plan.output_ordering().as_deref(),
+            parent_required,
+            || plan.equivalence_properties(),
+        ) {
             // If the current plan is a SortExec, modify it to satisfy parent requirements:
             let parent_required_expr = PhysicalSortRequirement::to_sort_exprs(
                 parent_required.ok_or_else(err)?.iter().cloned(),
@@ -137,6 +139,7 @@ pub(crate) fn pushdown_sorts(
         };
         let required_ordering = new_plan
             .output_ordering()
+            .as_deref()
             .map(PhysicalSortRequirement::from_sort_exprs);
         // Since new_plan is a SortExec, we can safely get the 0th index.
         let child = &new_plan.children()[0];
@@ -155,9 +158,11 @@ pub(crate) fn pushdown_sorts(
         }
     } else {
         // Executors other than SortExec
-        if ordering_satisfy_requirement(plan.output_ordering(), parent_required, || {
-            plan.equivalence_properties()
-        }) {
+        if ordering_satisfy_requirement(
+            plan.output_ordering().as_deref(),
+            parent_required,
+            || plan.equivalence_properties(),
+        ) {
             // Satisfies parent requirements, immediately return.
             return Ok(Transformed::Yes(SortPushDown {
                 required_ordering: None,

datafusion/core/src/physical_optimizer/utils.rs

@@ -41,7 +41,7 @@ pub fn add_sort_above(
     sort_expr: Vec<PhysicalSortExpr>,
 ) -> Result<()> {
     // If the ordering requirement is already satisfied, do not add a sort.
-    if !ordering_satisfy(node.output_ordering(), Some(&sort_expr), || {
+    if !ordering_satisfy(node.output_ordering().as_deref(), Some(&sort_expr), || {
         node.equivalence_properties()
     }) {
         let new_sort = SortExec::new(sort_expr, node.clone());