ShayanGho opened a new pull request, #25589:
URL: https://github.com/apache/datafusion/pull/25589
## Which issue does this PR close?
- Closes #25435.
Part of epic #25421.
## Rationale for this change
Two inputs that are both range-partitioned on `(a, b)` with identical split
points are co-partitioned, so a partitioned join can run within each partition
with no shuffle. Whether the planner produced that plan depended on the order
the equality conjuncts were written in:
```sql
SELECT ... FROM t1 JOIN t2 ON t1.b = t2.b AND t1.a = t2.a
```
gives `HashJoinExec` `on = [(b, b), (a, a)]`, so the join required
`KeyPartitioned([b, a])`. `Partitioning::satisfaction` compares key lists
positionally, reported `NotSatisfied` against `Range([a, b])`, and
`EnsureRequirements` hash-repartitioned both already co-partitioned inputs.
Before, from `range_join_key_alignment_reversed_keys` (both inputs ranged on
`(a, b)`, the right one aliased to `(a1, b1)`, joined `b = b1 AND a = a1`):
```
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (a@0,
a1@0)]
RepartitionExec: partitioning=Hash([b@1, a@0], 4), input_partitions=4
DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]},
projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC, b@1 ASC], [(10,
0), (20, 0), (30, 0)], 4), file_type=parquet
RepartitionExec: partitioning=Hash([b1@1, a1@0], 4), input_partitions=4
ProjectionExec: expr=[a@0 as a1, b@1 as b1]
DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]},
projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC, b@1 ASC], [(10,
0), (20, 0), (30, 0)], 4), file_type=parquet
```
After:
```
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1,
b1@1)]
DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]},
projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC, b@1 ASC], [(10,
0), (20, 0), (30, 0)], 4), file_type=parquet
ProjectionExec: expr=[a@0 as a1, b@1 as b1]
DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]},
projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC, b@1 ASC], [(10,
0), (20, 0), (30, 0)], 4), file_type=parquet
```
## What changes are included in this PR?
- In `ensure_distribution_with_stats`, before a node's distribution
requirements are read, a new `reorder_join_keys_to_range_inputs` reorders a
`Partitioned` `HashJoinExec`'s or a `SortMergeJoinExec`'s `on` pairs (and, for
sort-merge joins, the matching `sort_options`) to follow the left, else the
right, input's Range key order, when the join keys are a complete permutation
of that input's range keys. Whole pairs move together, so the join condition is
unchanged. It runs in the bottom-up distribution pass so each join sees its
inputs' partitioning after their own requirements have been enforced, and it
applies regardless of `datafusion.optimizer.top_down_join_key_reordering`.
- Matching is per input (`range_join_key_positions`): an input's range key
expressions are compared only with that same input's join keys, using that
input's equivalence classes. The existing `try_reorder` also compares against
the other side's keys, which for an unaliased crossed join (left `Range([a,
b])`, right `Range([b, a])`, `ON l.b = r.a AND l.a = r.b`) reports "already
aligned" and misses the permutation.
- Distribution satisfaction is deliberately left positional; reordering the
pairs is what makes the requirement satisfiable. Making satisfaction
order-insensitive instead would be unsound: a row with `a = 10, b = 20` is in
partition 1 under `Range([a, b])` but in partition 2 under `Range([b, a])`, so
inputs with the same split values in a different key order are still
repartitioned (covered by
`range_join_key_alignment_swapped_range_keys_still_repartition`).
- The existing hash-key reordering in the join-key reordering phase is
unchanged. Doc comments updated to say where each kind of alignment happens.
- Behavior note: when one side's layout is the swapped one, the reference
side changes. Previously the left input was repartitioned to match the right;
now the join is aligned to the left and the right input is repartitioned. Both
plans are correct.
## What is the testing strategy for this PR?
- `datafusion/core/tests/physical_optimizer/enforce_distribution.rs`:
`range_join_key_alignment_reversed_keys`,
`range_join_key_alignment_crossed_keys`,
`range_join_key_alignment_sort_merge_join` (asserts `sort_options` are permuted
with `on` and the projection survives the rebuild), and
`range_join_key_alignment_swapped_range_keys_still_repartition`. Each runs
under both `top_down_join_key_reordering` settings, passes `SanityCheckPlan`,
and is unchanged by a second `EnsureRequirements` pass.
- `datafusion/sqllogictest/test_files/range_partitioning.slt` TEST 49, over
a new composite-key fixture `range_partitioned_composite` registered in
`test_context/range_partitioning.rs`: `EXPLAIN` of the query from the issue
plus result verification, with rows on the split boundaries and duplicate
matching keys.
## Are there any user-facing changes?
No API changes. Plans for multi-key joins over compatible range-partitioned
inputs no longer depend on the order of the `ON` conjuncts.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
