refine code

Author: xudong963

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -3305,8 +3305,7 @@ fn preserve_ordering_for_streaming_sorted_aggregate() -> Result<()> {
 
     let test_config = TestConfig::default().with_query_execution_partitions(2);
 
-    let plan_distrib =
-        test_config.to_plan(physical_plan.clone(), &DISTRIB_DISTRIB_SORT);
+    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
       RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2, preserve_order=true, sort_exprs=a@0 ASC
@@ -3340,8 +3339,7 @@ fn preserve_ordering_for_streaming_partially_sorted_aggregate() -> Result<()> {
 
     let test_config = TestConfig::default().with_query_execution_partitions(2);
 
-    let plan_distrib =
-        test_config.to_plan(physical_plan.clone(), &DISTRIB_DISTRIB_SORT);
+    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])
       RepartitionExec: partitioning=Hash([a@0, b@1], 2), input_partitions=2, preserve_order=true, sort_exprs=a@0 ASC

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -1381,6 +1381,9 @@ pub fn ensure_distribution(
                 }
             };
 
+            let streaming_benefit = child.data
+                && preserving_order_enables_streaming(&plan, &child.plan);
+
             // There is an ordering requirement of the operator:
             if let Some(required_input_ordering) = required_input_ordering {
                 // Either:
@@ -1392,8 +1395,6 @@ pub fn ensure_distribution(
                     .equivalence_properties()
                     .ordering_satisfy_requirement(sort_req.clone())?;
 
-                let streaming_benefit = child.data
-                    && preserving_order_enables_streaming(&plan, &child.plan);
                 if (!ordering_satisfied || !order_preserving_variants_desirable)
                     && !streaming_benefit
                     && child.data
@@ -1420,11 +1421,6 @@ pub fn ensure_distribution(
                 match requirement {
                     // Operator requires specific distribution.
                     Distribution::SinglePartition | Distribution::HashPartitioned(_) => {
-                        let streaming_benefit = child.data
-                            && preserving_order_enables_streaming(
-                                &plan,
-                                &child.plan,
-                            );
                         if !streaming_benefit {
                             child =
                                 replace_order_preserving_variants(child)?;