[SPARK-59080][SQL][4.1] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle

### Backport of apache#58527 to `branch-4.2`

Cherry-pick of `4605662b6e4` with two test-only adjustments; the production change applies unchanged.

- The subset opt-in is named `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled` on this branch, so the four `V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS` references in the new tests become `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS`.
- `PartitioningCollection.fromPartitionings` does not exist here, and neither does the invariant it maintains: this branch's `PartitioningCollection` only requires its members to agree on `numPartitions`, with no shared `partitionKeys` reference to intern. The two fixtures use the plain constructor, which the members' equal key lists make sound.

Green on this branch: `ShuffleSpecSuite` 13 tests, `EnsureRequirementsSuite` + `KeyGroupedPartitioningSuite` 162 tests, all passing. `dev/lint-scala` is clean.

---

### What changes were proposed in this pull request?

`EnsureRequirements` stops asking a `ShuffleSpecCollection` for a single answer. It resolves the one member the matched children agreed on, preferring the finest when several qualify, and uses that member to build a re-shuffled child's partitioning.

- `flattenSpec` replaces the head read. It recurses, because `ShuffledJoin.outputPartitioning` builds `PartitioningCollection.fromPartitionings(Seq(left, right))` for an inner join, so a chain of same-key joins nests collections.
- The chosen member has to be compatible with every matched child. When no member is, there is no shared layout, so every child takes the ordinary shuffle. Reaching that needs three or more clustered children, and no operator has three today.
- The `joinKeyPositions` pushed into a compatible child now come from that child's own matching member, because they index into that child's partition expressions.
- `ShuffleSpecCollection.createPartitioning` is untouched. Its `require` stays as a guard, and a new unit test pins it.

### Why are the changes needed?

Under `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`, `KeyedPartitioning.createShuffleSpec` projects each member of a `PartitioningCollection` onto *its own* join-key subset and drops the duplicate keys that projection creates. The members of the resulting `ShuffleSpecCollection` can therefore end up with different `numPartitions`. `EnsureRequirements` then asks the collection for a shuffle template, and `ShuffleSpecCollection.createPartitioning` requires all members to agree:

```
java.lang.IllegalArgumentException: requirement failed: expected all specs in the collection to have the same number of partitions
```

so planning fails outright. With `items` partitioned by `[identity(id), identity(arrive_time)]`, one row per split, `purchases` unpartitioned, and `v2BucketingShuffleEnabled=true`, `partiallyClusteredDistribution=false`, `allowKeysSubsetOfPartitionKeys=true`:

```sql
SELECT /*+ MERGE(i, p) */ id, t1, t2, i.price AS purchase_price, p.price AS sale_price
FROM (SELECT id, arrive_time AS t1, arrive_time AS t2, price FROM testcat.ns.items) i
JOIN testcat.ns.purchases p ON i.id = p.item_id AND i.t1 = p.time
```

Selecting `arrive_time` twice under two aliases makes the alias cross-product produce members that cover different numbers of join keys, which is where the counts diverge.

The collection cannot answer that question locally. `isCompatibleWith` succeeds when *any* member matches, so the collection alone never said which member the two sides agreed on, and `createPartitioning` fell back to `specs.head`, whichever the alias cross-product enumerated first. Narrowing the collection to its finest members would satisfy the `require`, but it would still be a guess: the right member is the one the *other* side matched, and that is only visible in `EnsureRequirements`.

### Does this PR introduce _any_ user-facing change?

Yes. The query above failed to plan and now runs, producing one shuffle and the right rows.

The `joinKeyPositions` half is user-facing too. I originally wrote here that it was latent, on the grounds that a cogroup's grouping key is synthesized so neither side stays keyed. sunchao pointed out that this is only true of the Scala `CoGroupExec`, whose key comes from an `AppendColumns` that no `KeyedPartitioning` satisfies. A Pandas or Arrow cogroup groups on real columns, so two keyed children do reach the per-child branch, and this suite already had a `FlatMapCoGroupsInPandasExec` test with two of them.

So: with two keyed children whose partition expressions are laid out differently, the second side is handed the first side's positions and ends up grouped on its other partition column. The plan test below reproduces it and fails without the fix, reporting `List(Some(List(1)), Some(List(1)))` where `List(Some(List(1)), Some(List(0)))` is right. I have not built an end-to-end query for it.

The only part that stays latent is the three-or-more clustered children case, which no operator has.

### How was this patch tested?

Four new tests. Each was measured against the same commit with only the `EnsureRequirements` change reverted.

| test | on base |
|---|---|
| `KeyGroupedPartitioningSuite`: both sides of the join land on the same collection member | fails with the `require` above |
| `EnsureRequirementsSuite`: the re-shuffled side lands on the member the keyed side was matched on | fails with the `require` above |
| `EnsureRequirementsSuite`: pushed-down positions index into the child's own partition expressions (a Pandas cogroup over two keyed children) | fails |
| `ShuffleSpecSuite`: a collection whose members cover different key subsets disagrees | passes, by design |

The last one pins the guard rather than the fix. It asserts that the members disagree on purpose, that every one of them stays available for `isCompatibleWith`, and that asking the collection for a single partitioning throws. That turns the `require` from untested prose into a pinned contract, which matters now that the method has no production caller.

Green: `ShuffleSpecSuite`, `EnsureRequirementsSuite`, `KeyGroupedPartitioningSuite`, 202 tests in all. `dev/lint-scala` is clean.

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

Generated-by: Claude Code

Closes apache#58565 from peter-toth/SPARK-59080-shufflespec-numpartitions-4.2.

Authored-by: Peter Toth <peter.toth@gmail.com>
Signed-off-by: Peter Toth <peter.toth@gmail.com>

(cherry picked from commit c4b1024)

Author: peter-toth

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

@@ -664,4 +664,34 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
         "identical marked functions remain compatible")
     }
   }
+
+  test("SPARK-59080: a collection whose members cover different key subsets disagrees") {
+    val id = $"id".int
+    val t1 = $"t1".int
+    val t2 = $"t2".int
+    val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1))
+
+    // The shape an alias cross-product produces: same arity, same keys, different expressions. The
+    // operation clusters on (id, t1), so the first member projects onto both positions and keeps
+    // three partitions, while the second matches only `id` and keeps two.
+    val collection = PartitioningCollection(Seq(
+      KeyGroupedPartitioning(Seq(id, t1), keys.size, keys),
+      KeyGroupedPartitioning(Seq(id, t2), keys.size, keys)))
+
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1)))
+        .asInstanceOf[ShuffleSpecCollection]
+
+      // The disagreement is kept rather than resolved here. Every member has to stay for
+      // `isCompatibleWith`, which answers for any of them, and the collection cannot know which one
+      // the other side matched. `EnsureRequirements` resolves that and asks the member, not the
+      // collection.
+      assert(spec.specs.map(_.numPartitions).toSet === Set(3, 2))
+      assert(spec.isCompatibleWith(spec), "every member stays available for matching")
+
+      // So asking the collection for a single answer is the caller's mistake, and it says so.
+      val e = intercept[IllegalArgumentException](spec.createPartitioning(Seq(id, t1)))
+      assert(e.getMessage.contains("expected all specs in the collection to have the same number"))
+    }
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql.execution.exchange
 
-import scala.annotation.tailrec
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
@@ -194,30 +193,50 @@ case class EnsureRequirements(
         }
       }
 
+      // A `ShuffleSpecCollection` answers `isCompatibleWith` if *any* of its members does, so the
+      // collection alone does not say which member the sides agreed on. The projection pushed into
+      // a compatible child and the partitioning built for a re-shuffled child both have to come
+      // from one member, otherwise the sides end up grouped on different keys, or on a key set the
+      // child does not even have. Pick that member once, preferring the finest when several
+      // qualify. Only the branch that shuffles a child reads these, hence `lazy`.
+      lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best =>
+        childrenIndexes.filter(i => best.isCompatibleWith(specs(i)))
+      }
+      lazy val bestMemberOpt = bestSpecOpt.flatMap { best =>
+        val matchedMembers = matchedIndexes.map(i => flattenSpec(specs(i)))
+        // No member serving every matched child means there is no layout to align them on, so they
+        // all take the ordinary shuffle. That needs three or more clustered children, since with
+        // two the member that reported the match serves both, and no operator has three today.
+        flattenSpec(best)
+          .filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith)))
+          .maxByOption(_.numPartitions)
+      }
+
       children = children.zip(requiredChildDistributions).zipWithIndex.map {
         case ((child, _), idx) if areChildrenCompatible ||
             !childrenIndexes.contains(idx) =>
           child
         case ((child, dist), idx) =>
-          if (bestSpecOpt.isDefined && bestSpecOpt.get.isCompatibleWith(specs(idx))) {
-            // If the child's partitioning is a `PartitioningCollection`, its spec is a
-            // `ShuffleSpecCollection` whose `createPartitioning` delegates to the head spec,
-            // so unwrap to the head spec to stay aligned with the re-shuffled side below.
-            unwrapSpecCollection(bestSpecOpt.get) match {
-              // If keyGroupCompatible = false, we can still perform SPJ
+          if (bestMemberOpt.isDefined && matchedIndexes.contains(idx)) {
+            // The positions come from this child's own matching member, since they index into its
+            // own partition expressions -- the chosen best member only says which member of it the
+            // two sides agreed on.
+            val bestMember = bestMemberOpt.get
+            flattenSpec(specs(idx)).find(bestMember.isCompatibleWith) match {
+              // If `areChildrenCompatible` is false, we can still perform SPJ
               // by shuffling the other side based on join keys (see the else case below).
               // Hence we need to ensure that after this call, the outputPartitioning of the
               // partitioned side's BatchScanExec is grouped by join keys to match,
               // and we do that by pushing down the join keys
-              case KeyGroupedShuffleSpec(_, _, Some(joinKeyPositions)) =>
+              case Some(KeyGroupedShuffleSpec(_, _, Some(joinKeyPositions))) =>
                 populateJoinKeyPositions(child, Some(joinKeyPositions))
               case _ => child
             }
           } else {
-            val newPartitioning = bestSpecOpt.map { bestSpec =>
+            val newPartitioning = bestMemberOpt.map { bestMember =>
               // Use the best spec to create a new partitioning to re-shuffle this child
               val clustering = dist.asInstanceOf[ClusteredDistribution].clustering
-              bestSpec.createPartitioning(clustering)
+              bestMember.createPartitioning(clustering)
             }.getOrElse {
               // No best spec available, so we create default partitioning from the required
               // distribution
@@ -721,12 +740,10 @@ case class EnsureRequirements(
     }
   }
 
-  // Unwraps a `ShuffleSpecCollection` (possibly nested) to the spec that its
-  // `createPartitioning` delegates to, i.e. the head spec.
-  @tailrec
-  private def unwrapSpecCollection(spec: ShuffleSpec): ShuffleSpec = spec match {
-    case ShuffleSpecCollection(specs) => unwrapSpecCollection(specs.head)
-    case other => other
+  // Flattens a (possibly nested) `ShuffleSpecCollection` into its member specs.
+  private def flattenSpec(spec: ShuffleSpec): Seq[ShuffleSpec] = spec match {
+    case ShuffleSpecCollection(specs) => specs.flatMap(flattenSpec)
+    case other => Seq(other)
   }
 
   /**

sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

@@ -2295,6 +2295,58 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
     }
   }
 
+  test("SPARK-59080: both sides of the join land on the same collection member") {
+    // `arrive_time` is selected twice under two aliases, so the projected partitioning is a
+    // `PartitioningCollection` whose members cover different numbers of the join keys: one covers
+    // (id, t1), another only id. Each member's spec is projected onto its own subset, so the specs
+    // disagree on `numPartitions`, and asking the collection for one partitioning fails with
+    // "expected all specs in the collection to have the same number of partitions".
+    //
+    // `EnsureRequirements` now resolves the member the two sides agreed on before it asks, so the
+    // keyed side is grouped on both join keys and the shuffled side is laid out on those same keys.
+    val items_partitions = Array(identity("id"), identity("arrive_time"))
+    createTable(items, itemsColumns, items_partitions)
+
+    sql(s"INSERT INTO testcat.ns.$items VALUES " +
+      "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+      "(1, 'ab', 30.0, cast('2020-01-02' as timestamp)), " +
+      "(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+      "(4, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+    createTable(purchases, purchasesColumns, Array.empty)
+    sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+      "(1, 42.0, cast('2020-01-01' as timestamp)), " +
+      "(1, 89.0, cast('2020-01-02' as timestamp)), " +
+      "(3, 19.5, cast('2020-01-01' as timestamp)), " +
+      "(5, 26.0, cast('2023-01-01' as timestamp))")
+
+    withSQLConf(
+      SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+      SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+      SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val df = sql(
+        s"""
+           |${selectWithMergeJoinHint("i", "p")}
+           |id, t1, t2, i.price AS purchase_price, p.price AS sale_price
+           |FROM (SELECT id, arrive_time AS t1, arrive_time AS t2, price FROM testcat.ns.$items) i
+           |JOIN testcat.ns.$purchases p ON i.id = p.item_id AND i.t1 = p.time
+           |""".stripMargin)
+      val plan = df.queryExecution.executedPlan
+      // On this branch the resolved member is recorded on the scan rather than on a grouping node.
+      val positions = collectScans(plan).flatMap(_.spjParams.joinKeyPositions)
+      assert(positions === Seq(Seq(0, 1)),
+        "the keyed side must be grouped on both join keys, the finest granularity available")
+      assert(collectAllShuffles(plan).size == 1, "only the unpartitioned side shuffles")
+      checkAnswer(df, Seq(
+        Row(1, java.sql.Timestamp.valueOf("2020-01-01 00:00:00"),
+          java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), 40.0, 42.0),
+        Row(1, java.sql.Timestamp.valueOf("2020-01-02 00:00:00"),
+          java.sql.Timestamp.valueOf("2020-01-02 00:00:00"), 30.0, 89.0),
+        Row(3, java.sql.Timestamp.valueOf("2020-01-01 00:00:00"),
+          java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), 10.0, 19.5)))
+    }
+  }
+
   test("SPARK-59025: shuffle one side and join keys are less than partition keys " +
       "when the keyed side reports a PartitioningCollection") {
     val items_partitions = Array(identity("id"), identity("name"))

sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala

@@ -1190,6 +1190,82 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
+  test("SPARK-59080: the re-shuffled side lands on the member the keyed side was matched on") {
+    val id = AttributeReference("id", IntegerType)()
+    val t1 = AttributeReference("t1", IntegerType)()
+    val t2 = AttributeReference("t2", IntegerType)()
+    // Same arity and the same keys, different expressions: what an alias cross-product produces
+    // when one column is selected twice. Clustering on (id, t1), the first member covers `id` only
+    // and keeps two partitions, the second covers both positions and keeps three. The coarse
+    // member comes first so that reading the collection's head would pick the wrong one.
+    val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1))
+    val keyed = new DummySparkPlanWithBatchScanChild(
+      outputPartitioning = PartitioningCollection(Seq(
+        KeyGroupedPartitioning(Seq(id, t2), keys.size, keys),
+        KeyGroupedPartitioning(Seq(id, t1), keys.size, keys))))
+    val unpartitioned = DummySparkPlan(outputPartitioning = UnknownPartitioning(0))
+
+    withSQLConf(
+        SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val smj = SortMergeJoinExec(Seq(id, t1), Seq(id, t1), Inner, None, keyed, unpartitioned)
+      val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec]
+
+      // The keyed side is grouped on both join keys, which is the finest member. On this branch
+      // the resolved member is recorded on the scan rather than on a grouping node.
+      assert(planned.left.collect { case b: BatchScanExec => b.spjParams.joinKeyPositions } ===
+        Seq(Some(Seq(0, 1))),
+        "the keyed side must be grouped on the member covering both join keys")
+      // And the shuffled side lands on that same member. Reading the collection instead would take
+      // whichever member came first and could put the two sides on different key sets.
+      val shuffles = planned.right.collect { case s: ShuffleExchangeExec => s }
+      assert(shuffles.map(_.outputPartitioning.numPartitions) === Seq(3),
+        "the re-shuffled side must land on the same member, not on whichever came first")
+    }
+  }
+
+  test("SPARK-59080: pushed-down positions index into the child's own partition expressions") {
+    val nL = AttributeReference("nL", IntegerType)()
+    val iL = AttributeReference("iL", IntegerType)()
+    val iR = AttributeReference("iR", IntegerType)()
+    val nR = AttributeReference("nR", IntegerType)()
+    // Both sides declare the same two key columns, in the opposite order, so the cogroup key sits
+    // at position 1 on the left and at position 0 on the right. Both project onto {1, 2}, so the
+    // sides are co-partitioned and neither is re-shuffled.
+    val leftKeys = Seq(InternalRow(1, 1), InternalRow(2, 1), InternalRow(3, 2))
+    val rightKeys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 3))
+    val left = new DummySparkPlanWithBatchScanChild(
+      outputPartitioning = KeyGroupedPartitioning(Seq(nL, iL), leftKeys.size, leftKeys))
+    val right = new DummySparkPlanWithBatchScanChild(
+      outputPartitioning = KeyGroupedPartitioning(Seq(iR, nR), rightKeys.size, rightKeys))
+
+    val pythonUdf = PythonUDF("pyUDF", null,
+      StructType(Seq(StructField("value", IntegerType))),
+      Seq.empty,
+      PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+      true)
+    // A cogroup requires `ClusteredDistribution` on both children but is not a `ShuffledJoin`, so
+    // `checkKeyGroupCompatible` declines and both children go through the per-child branch that
+    // pushes the positions down. Unlike the Scala `CoGroupExec`, whose key comes from an
+    // `AppendColumns` no `KeyGroupedPartitioning` satisfies, the Pandas one groups on real columns,
+    // so both sides stay keyed.
+    val cogroup = FlatMapCoGroupsInPandasExec(
+      Seq(iL), Seq(iR), pythonUdf,
+      AttributeReference("value", IntegerType)() :: Nil, left, right)
+
+    withSQLConf(
+        SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val result = EnsureRequirements.apply(cogroup)
+
+      assert(result.collect { case s: ShuffleExchangeExec => s }.isEmpty,
+        "the sides are co-partitioned on the cogroup key")
+      assert(result.collect { case b: BatchScanExec => b.spjParams.joinKeyPositions } ===
+        Seq(Some(Seq(1)), Some(Seq(0))),
+        "the positions pushed into a child must index into that child's own expressions")
+    }
+  }
+
   def bucket(numBuckets: Int, expr: Expression): TransformExpression = {
     TransformExpression(BucketFunction, Seq(expr), Some(numBuckets))
   }