Restore add_roundrobin_on_top

viirya · Jan 3, 2024 · 5ea68e6 · 5ea68e6
1 parent ebec14e
commit 5ea68e6
Showing 1 changed file with 14 additions and 14 deletions.
diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
@@ -1212,6 +1212,19 @@ fn ensure_distribution(
                 }
             }
 
+            if enable_round_robin
+                // Operator benefits from partitioning (e.g. filter):
+                && (would_benefit && repartition_beneficial_stats)
+                // Unless partitioning doesn't increase the partition count, it is not beneficial:
+                && child.plan.output_partitioning().partition_count() < target_partitions
+            {
+                // Increase parallelism by adding round-robin repartitioning
+                // on top of the operator. Note that we only do this if the
+                // partition count is not already equal to the desired partition
+                // count.
+                child = add_roundrobin_on_top(child, target_partitions)?;
+            }
+
             // Satisfy the distribution requirement if it is unmet.
             match requirement {
                 Distribution::SinglePartition => {
@@ -1220,20 +1233,7 @@ fn ensure_distribution(
                 Distribution::HashPartitioned(exprs) => {
                     child = add_hash_on_top(child, exprs.to_vec(), target_partitions)?;
                 }
-                Distribution::UnspecifiedDistribution => {
-                    if enable_round_robin
-                        // Operator benefits from partitioning (e.g. filter):
-                        && (would_benefit && repartition_beneficial_stats)
-                        // Unless partitioning doesn't increase the partition count, it is not beneficial:
-                        && child.plan.output_partitioning().partition_count() < target_partitions
-                    {
-                        // Increase parallelism by adding round-robin repartitioning
-                        // on top of the operator. Note that we only do this if the
-                        // partition count is not already equal to the desired partition
-                        // count.
-                        child = add_roundrobin_on_top(child, target_partitions)?;
-                    }
-                }
+                Distribution::UnspecifiedDistribution => {}
             };
 
             // There is an ordering requirement of the operator: