[SPARK-59025][SQL] Fix subset join key grouping when the keyed side reports a PartitioningCollection

[dongjoon-hyun]

What changes were proposed in this pull request?

Fix the post-shuffle regrouping in EnsureRequirements for storage-partitioned joins with spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys=true. The unwrap that pushes join key positions down to the keyed side only matches a bare KeyedShuffleSpec:

bestSpecOpt match {
  case Some(KeyedShuffleSpec(_, _, Some(joinKeyPositions))) =>
    withJoinKeyPositions(child, joinKeyPositions)
  case _ => child
}

When the keyed child's outputPartitioning is a PartitioningCollection (e.g. an inner SPJ join result, or a projection duplicating a partition column under multiple aliases like SELECT a AS a1, a AS a2, b), its spec is a ShuffleSpecCollection wrapping the KeyedShuffleSpec, so the match falls through and no GroupPartitionsExec is inserted. This PR unwraps a (possibly nested) ShuffleSpecCollection to its head spec — the same spec ShuffleSpecCollection.createPartitioning delegates to — before the match.

Why are the changes needed?

Without the fix, the keyed side keeps its N ungrouped partitions while the other side is shuffled into the projected layout with M partitions, and planning fails with:

java.lang.IllegalArgumentException: requirement failed:
All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys

Does this PR introduce any user-facing change?

The affected queries used to fail with the error above; they now run correctly, shuffling only the non-keyed side.

How was this patch tested?

Pass the CIs.

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

Generated-by: Claude Fable 5

[dongjoon-hyun]

CI is running here.

• https://github.com/dongjoon-hyun/spark/actions/runs/32989186623

[dongjoon-hyun]

cc @peter-toth , too

[peter-toth]

This is independent from other SPJ fixes, I'm reviewing it now.

[peter-toth]

Thanks for the PR, @dongjoon-hyun!

I traced the failure independently and landed on the same diagnosis. With allowKeysSubsetOfPartitionKeys on, a keyed child whose outputPartitioning is a PartitioningCollection reports a ShuffleSpecCollection, so case Some(KeyedShuffleSpec(...)) never fires - the keyed side keeps its 4 partitions while the other side is re-shuffled into the 3-partition projected layout, and the join's PartitioningCollection then rejects the pair. Unwrapping to the head spec is the right minimal shape, since that is the spec ShuffleSpecCollection.numPartitions and createPartitioning already delegate to. I also checked the new test is effective against the current base rather than the PR's: with SPARK-58988's partitioning.scala applied on top of 14e21716 (the only commit touching these files since), it still fails with the reported require, and it passes at e2117ce. KeyGroupedPartitioningSuite + EnsureRequirementsSuite + ProjectedOrderingAndPartitioningSuite + GroupPartitionsExecSuite are 175/175 green at head. Nothing blocking from me.

Non-blocking

• 1. JIRA affects-version: SPARK-59025 says 5.0.0, but EnsureRequirements.scala is byte-identical between this PR's base and apache/branch-4.x, so 4.4.0 is affected too. The same unguarded match also exists back to branch-4.0 under the older KeyGroupedShuffleSpec name, which may be worth a look before you pick fix-versions.

Minor

• 2. Helper placement: the new helper lands between hasKeyedPartitioning and tryEnableSortedMerge, which is the only thing that uses it.
• 3. Test assertions: the test comment states the keyed side must drop from 4 partitions to 3, but nothing asserts it.
• 4. Description: "How was this patch tested?" is just "Pass the CIs.", so the added regression test goes unmentioned in the commit message that dev/merge_spark_pr.py builds. And in the user-facing-change section, "No behavior change." contradicts the sentence right after it - queries that used to throw now return rows. Suggest naming the test in the first, and dropping "No behavior change." from the second.

[peter-toth]

Finding 2. hasKeyedPartitioning above (:307-310) exists only to serve tryEnableSortedMerge's pruning predicate at :339, and the new helper is inserted between the two, splitting that pair.

It reads better next to the other shuffle-spec helpers - withJoinKeyPositions (:778) and createKeyedShuffleSpec (:793) - or directly after ensureDistributionAndOrdering, its only caller.

[peter-toth]

Finding 3. joinKeyPositions.isDefined would also hold if the positions covered both partition columns, in which case the keyed side would stay at 4 partitions and the fix's point would be lost. The comment at :2580-2581 states the number, so it is worth pinning:

Suggested change

	"the keyed side should be grouped by the join keys")
	"the keyed side should be grouped by the join keys")
	assert(groupPartitions.head.outputPartitioning.numPartitions == 3,
	"the keyed side should be grouped down to 3 partitions")
	assert(shuffles.head.outputPartitioning.numPartitions == 3,
	"the shuffled side should match the 3 grouped partitions")

Both assertions pass at e2117ce - I ran them.

[dongjoon-hyun]

Thank you, @peter-toth . I addressed your comments. For old release branches, I'll try to backport and update the JIRA info.

[peter-toth]

Re-checked through 103dad17 - findings 1-4 resolved (affects-version now 4.4.0, unwrapSpecCollection moved next to withJoinKeyPositions, both partition-count assertions in, description corrected), nothing new. Re-measured at the new head rather than carrying the old verdict: with SPARK-58988's partitioning.scala on top of the base, the test still fails with the same require without the EnsureRequirements change, and KeyGroupedPartitioningSuite + EnsureRequirementsSuite + ProjectedOrderingAndPartitioningSuite + GroupPartitionsExecSuite are 175/175 green.

Thanks for working through all of these, @dongjoon-hyun - nothing left open from my side.

[dongjoon-hyun]

Could you approve this PR, @peter-toth ?

[dongjoon-hyun]

Thank you always, @peter-toth !

[peter-toth]

Could you approve this PR, @peter-toth ?

Sorry, I forgot to do it. 😄 Approved now.

[dongjoon-hyun]

Merge Summary:

• merged into master 25bae64
• merged into branch-4.x 3809fc6
• merged into branch-4.3 7cc9194

Posted by merge_spark_pr.py

[dongjoon-hyun]

I backported this to branch-4.2/branch-4.1/branch-4.0.

• branch-4.2: 91cd9a3
• branch-4.1: 48c19fa
• branch-4.0: 26fb64f