Align multi-key equi-join conditions to match Range ordering
Is your feature request related to a problem or challenge?
When joining two tables that are both range-partitioned on composite keys (for example, on (a ASC, b ASC) with identical split points), the join is physically co-partitioned and can execute as a task-local partitioned join with zero network shuffle.
However, if the SQL query writes the join condition with the equality pairs in a different sequence than the physical range partition ordering:
SELECT * FROM t1 JOIN t2 ON t1.b = t2.b AND t1.a = t2.aHashJoinExec preserves the syntactic order of the conjuncts, collecting on = [(t1.b, t2.b), (t1.a, t2.a)].
When checking distribution satisfaction: https://github.com/apache/datafusion/blob/c4f72bad34249798182ac7de605eb63ab134f26e/datafusion/physical-expr/src/partitioning.rs#L474-L486
key_satisfaction invokes equivalent_exprs:
https://github.com/apache/datafusion/blob/c4f72bad34249798182ac7de605eb63ab134f26e/datafusion/physical-expr/src/partitioning.rs#L351-L369
equivalent_exprs checks element-by-element equality in sequence. Because required keys [b, a] do not equal partition keys [a, b], key_satisfaction returns NotSatisfied.
As a result, EnsureRequirements fails to recognize that the inputs are co-partitioned and inserts an unnecessary 2-sided repartition shuffle for inputs that are already co-partitioned. Because equi-join conjuncts are commutative (a = a AND b = b is logically identical to b = b AND a = a), the arbitrary syntactic ordering of ON clauses should not dictate whether a 0-shuffle join is planned.
Describe the solution you'd like
In JoinSelection or HashJoinExec (or during distribution satisfaction in EnsureRequirements), detect when the set of equi-join equality keys matches the columns of an input's lexicographical RangePartitioning.
When a permutation of the join keys aligns with the input's range ordering:
- Reorder the
onkey pairs to match the lexicographical sequence of the existingRangePartitioning. - Allow
key_satisfactionto recognize that the inputs satisfy the partitioned distribution requirement. - Preserve task-local execution without inserting a shuffle.
Additional context
Part of epic #25421.
Source: apache/datafusion