[SPARK-59905][SQL][4.2] Fix storage-partitioned join sorting partitions the wrong way for an ORDER BY over a join-key expression

[peter-toth]

Differences from #59189

This is the branch-4.2 backport of #59189. The description below is the original one. On this branch:

• The change to EnsureRequirements does not cherry-pick cleanly, since the OrderedDistribution arm sits one level deeper here. The code is the same as on master, re-indented, and its comment is re-wrapped.
• The two comment changes in reorderJoinKeysRecursively are left out. Branch-4.2 has no such comment there.
• The test is the same as on master. Without the change, it fails on branch-4.2 the same way as on master.
• Ran KeyGroupedPartitioningSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite, ProjectedOrderingAndPartitioningSuite and PlannerSuite on branch-4.2, 306 tests in all, plus dev/lint-scala.

What changes were proposed in this pull request?

When EnsureRequirements serves an ORDER BY by putting a keyed child's partitions in key order, it now compares the key rows position by position. Sort order i reads field i of the key row. It used to bind the ordering to the references of the partition expressions instead. The binding takes the types the key rows hold (keyDataTypes), and an assertion pins the precondition it relies on next to it.

Why are the changes needed?

A key row holds the values of the partition expressions. Binding the ordering to their references is right only when each expression is a bare column, the only thing a scan reports that an ORDER BY can name. A join can report the other side's join key instead. A one-side shuffle does, and so does an inner broadcast hash join, whose expandOutputPartitioning reports the streamed side's layout over the build side's join key.

With spark.sql.sources.v2.bucketing.shuffle.enabled and spark.sql.sources.v2.bucketing.sorting.enabled on, this query returns b as 0 to 6 instead of 6 to 0. ident(id) is identity-partitioned and holds 94 to 100, and plain(b) holds 0 to 6:

SELECT p.b FROM ident i JOIN plain p ON i.id = 100 - p.b ORDER BY 100 - p.b

The same happens with only spark.sql.sources.v2.bucketing.sorting.enabled when plain is broadcast.

1. The join shuffles plain onto ident's layout and reports 100 - b as plain's partition expression. Its key rows hold 94 to 100.
2. KeyedPartitioning.keysSatisfy admits that partitioning for ORDER BY 100 - p.b, since its expression is the ordering's. So no range shuffle is planned, and the global sort is dropped as redundant.
3. EnsureRequirements binds 100 - b to b and evaluates it on the key rows, so it computes 100 - (100 - b). The partitions end up in the reverse order.

Two other join keys fail planning instead:

• For a join on i.id = p.b + p.c, ORDER BY p.b + p.c binds two references to a key row with one field. Planning fails with an ArrayIndexOutOfBoundsException.
• For a join on i.id = p.s.a, ORDER BY p.s.a binds the struct s to the long key and reads a field of it. Planning fails with a ClassCastException.

keysSatisfy admits a partitioning here only when its expressions are the ordering's, position by position (OrderedDistribution.areAllClusterKeysMatched). So reading field i for sort order i is exact.

Does this PR introduce any user-facing change?

Yes. Such a query returns its rows in the right order, and the two other forms no longer fail planning. They all need spark.sql.sources.v2.bucketing.sorting.enabled, which is off by default. A broadcast join needs nothing else, and a one-side shuffle also needs spark.sql.sources.v2.bucketing.shuffle.enabled. No range shuffle is added.

How was this patch tested?

New test in KeyGroupedPartitioningSuite, "SPARK-59905: ORDER BY a join-key expression a join reports as a partition expression". It runs five queries:

• ORDER BY 100 - p.b, ascending and descending;
• ORDER BY p.s.a;
• ORDER BY 100 - p.b, p.c DESC over a two-key identity side, which needs a mixed-direction GroupPartitionsExec;
• ORDER BY p.b + p.c.

Each runs through a one-side shuffle and through a broadcast join, with AQE on and off. The p.b + p.c one runs with AQE off only, since AQE's plan validation fails on it for SPARK-59901, which #59165 fixes. Each run checks the order of the rows, the number of shuffles, and whether a GroupPartitionsExec sorts the partitions.

All 18 runs fail on master:

• the 100 - p.b and two-key queries return their rows in the wrong order;
• the p.s.a one fails with the ClassCastException;
• the p.b + p.c one fails with the ArrayIndexOutOfBoundsException.

Also ran the KeyGroupedPartitioning* suites, EnsureRequirementsSuite, GroupPartitionsExecSuite, ProjectedOrderingAndPartitioningSuite and PlannerSuite, 427 tests in all, plus dev/lint-scala.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

[dongjoon-hyun]

+1, LGTM. Thank you, @peter-toth .

[dongjoon-hyun]

Merge Summary:

• merged into branch-4.2 26e6dc0

Posted by merge_spark_pr.py