diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index 4cc8cf22b124..f12e68b3dd16 100644 --- 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: