[SPARK-59256][SQL] Choose a ShuffleSpecCollection member by pairing, not by enumeration order

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -318,7 +318,7 @@ trait HashPartitioningLike extends Expression with Partitioning with Unevaluable
 case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int)
   extends HashPartitioningLike {
 
-  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+  override def createShuffleSpec(distribution: ClusteredDistribution): HashShuffleSpec =
     HashShuffleSpec(this, distribution)
 
   /**
@@ -364,7 +364,7 @@ case class NullAwareHashPartitioning(expressions: Seq[Expression], numPartitions
     }
   }
 
-  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+  override def createShuffleSpec(distribution: ClusteredDistribution): NullAwareHashShuffleSpec =
     NullAwareHashShuffleSpec(this, distribution)
 
   override protected def withNewChildrenInternal(
@@ -791,7 +791,7 @@ case class KeyedPartitioning(
     if (isGrouped) keysSatisfy(required) else mayGroupToSatisfy(required)
   }
 
-  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
+  override def createShuffleSpec(distribution: ClusteredDistribution): KeyedShuffleSpec = {
     val result = KeyedShuffleSpec(this, distribution)
     if (SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys) {
       // If allowing operation keys to be a subset of partition keys, create a new
@@ -1198,18 +1198,6 @@ case class BroadcastPartitioning(mode: BroadcastMode) extends Partitioning {
   }
 }
 
-/**
- * This is used in the scenario where an operator has multiple children (e.g., join) and one or more
- * of which have their own requirement regarding whether its data can be considered as
- * co-partitioned from others. This offers APIs for:
- *
- *   - Comparing with specs from other children of the operator and check if they are compatible.
- *      When two specs are compatible, we can say their data are co-partitioned, and Spark will
- *      potentially be able to eliminate shuffle if necessary.
- *   - Creating a partitioning that can be used to re-partition another child, so that to make it
- *      having a compatible partitioning as this node.
- */
-
 /**
  * Represents a partitioning where partition IDs are passed through directly from the
  * DirectShufflePartitionID expression. This partitioning scheme is used when users
@@ -1253,12 +1241,16 @@ case class ShufflePartitionIdPassThrough(
     copy(expr = newChildren.head.asInstanceOf[DirectShufflePartitionID])
 }
 
-trait ShuffleSpec {
-  /**
-   * Returns the number of partitions of this shuffle spec
-   */
-  def numPartitions: Int
-
+/**
+ * Describes how a child's data is laid out, for the purpose of deciding whether two children are
+ * co-partitioned and, if not, what to shuffle the other one onto.
+ *
+ * A [[LeafShuffleSpec]] is one concrete layout. A [[ShuffleSpecCollection]] stands for a choice
+ * between several. A collection can answer [[isCompatibleWith]], which succeeds when any member
+ * matches. It cannot answer anything that needs one member: which one is right depends on what the
+ * other side matched, and only the caller comparing the two sides can see that.
+ */
+sealed trait ShuffleSpec {
   /**
    * Returns true iff this spec is compatible with the provided shuffle spec.
    *
@@ -1271,10 +1263,28 @@ trait ShuffleSpec {
   def isCompatibleWith(other: ShuffleSpec): Boolean
 
   /**
-   * Whether this shuffle spec can be used to create partitionings for the other children.
+   * Whether this shuffle spec can be used to create partitionings for the other children. A
+   * [[ShuffleSpecCollection]] answers for the whole choice, since the planner asks it of a child's
+   * spec as a whole. Building the partitioning is [[LeafShuffleSpec.createPartitioning]], and that
+   * is always one member's job.
    */
   def canCreatePartitioning: Boolean
 
+  /**
+   * This spec's leaf specs: a [[ShuffleSpecCollection]] yields its members recursively, and a
+   * [[LeafShuffleSpec]] yields itself. A caller that needs one member picks from these. Never
+   * empty, since a collection has at least one member.
+   */
+  def flatten: Seq[LeafShuffleSpec]
+}
+
+/** A [[ShuffleSpec]] describing one layout, as opposed to a choice between several. */
+trait LeafShuffleSpec extends ShuffleSpec {
+  /**
+   * Returns the number of partitions of this shuffle spec
+   */
+  def numPartitions: Int
+
   /**
    * Creates a partitioning that can be used to re-partition the other side with the given
    * clustering expressions.
@@ -1284,11 +1294,19 @@ trait ShuffleSpec {
    */
   def createPartitioning(clustering: Seq[Expression]): Partitioning =
     throw SparkUnsupportedOperationException()
+
+  override final def flatten: Seq[LeafShuffleSpec] = this +: Nil
 }
 
-case object SinglePartitionShuffleSpec extends ShuffleSpec {
-  override def isCompatibleWith(other: ShuffleSpec): Boolean = {
-    other.numPartitions == 1
+case object SinglePartitionShuffleSpec extends LeafShuffleSpec {
+  override def isCompatibleWith(other: ShuffleSpec): Boolean = other match {
+    case leaf: LeafShuffleSpec => leaf.numPartitions == 1
+    // `forall`, not the `exists` the other specs use for a collection. They ask whether *some*
+    // member matches them, and the caller then plans on that member. This asks a property of the
+    // child itself, and the child is whichever member the other side picks, so a single partition
+    // has to be the answer for every one of them. The members can only disagree when the subset
+    // config projects them onto different key sets.
+    case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith)
   }
 
   override def canCreatePartitioning: Boolean = false
@@ -1301,7 +1319,7 @@ case object SinglePartitionShuffleSpec extends ShuffleSpec {
 
 case class RangeShuffleSpec(
     numPartitions: Int,
-    distribution: ClusteredDistribution) extends ShuffleSpec {
+    distribution: ClusteredDistribution) extends LeafShuffleSpec {
 
   // `RangePartitioning` is not compatible with any other partitioning since it can't guarantee
   // data are co-partitioned for all the children, as range boundaries are randomly sampled. We
@@ -1338,7 +1356,7 @@ private object HashShuffleSpecCompatibility {
 
 case class HashShuffleSpec(
     partitioning: HashPartitioning,
-    distribution: ClusteredDistribution) extends ShuffleSpec {
+    distribution: ClusteredDistribution) extends LeafShuffleSpec {
 
   /**
    * A sequence where each element is a set of positions of the hash partition key to the cluster
@@ -1424,7 +1442,7 @@ case class HashShuffleSpec(
  */
 case class NullAwareHashShuffleSpec(
     partitioning: NullAwareHashPartitioning,
-    distribution: ClusteredDistribution) extends ShuffleSpec {
+    distribution: ClusteredDistribution) extends LeafShuffleSpec {
 
   lazy val hashKeyPositions: Seq[mutable.BitSet] = {
     val distKeyToPos = mutable.Map.empty[Expression, mutable.BitSet]
@@ -1481,8 +1499,8 @@ case class NullAwareHashShuffleSpec(
 }
 
 case class CoalescedHashShuffleSpec(
-    from: ShuffleSpec,
-    partitions: Seq[CoalescedBoundary]) extends ShuffleSpec {
+    from: LeafShuffleSpec,
+    partitions: Seq[CoalescedBoundary]) extends LeafShuffleSpec {
 
   override def isCompatibleWith(other: ShuffleSpec): Boolean = other match {
     case SinglePartitionShuffleSpec =>
@@ -1558,7 +1576,7 @@ case class IdentityReducer(transform: TransformExpression) extends Reducer[Any,
 case class KeyedShuffleSpec(
     partitioning: KeyedPartitioning,
     distribution: ClusteredDistribution,
-    joinKeyPositions: Option[Seq[Int]] = None) extends ShuffleSpec {
+    joinKeyPositions: Option[Seq[Int]] = None) extends LeafShuffleSpec {
 
   /**
    * A sequence where each element is a set of positions of the partition expression to the cluster
@@ -1595,7 +1613,7 @@ case class KeyedShuffleSpec(
     //  4. the partition values from both sides are following the same order.
     case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _) =>
       distribution.clustering.length == otherDistribution.clustering.length &&
-        numPartitions == other.numPartitions && areKeysCompatible(otherSpec) &&
+        numPartitions == otherSpec.numPartitions && areKeysCompatible(otherSpec) &&
           partitioning.partitionKeys == otherPartitioning.partitionKeys
     case ShuffleSpecCollection(specs) =>
       specs.exists(isCompatibleWith)
@@ -1775,7 +1793,7 @@ case class KeyedShuffleSpec(
 
 case class ShufflePartitionIdPassThroughSpec(
     partitioning: ShufflePartitionIdPassThrough,
-    distribution: ClusteredDistribution) extends ShuffleSpec {
+    distribution: ClusteredDistribution) extends LeafShuffleSpec {
 
   /**
    * A sequence where each element is a set of positions of the partition key to the cluster
@@ -1818,24 +1836,20 @@ case class ShufflePartitionIdPassThroughSpec(
   override def numPartitions: Int = partitioning.numPartitions
 }
 
+/**
+ * A choice between several layouts, produced by [[PartitioningCollection.createShuffleSpec]].
+ *
+ * `specs` can hold a nested collection, since a [[PartitioningCollection]] can hold a nested one.
+ */
 case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec {
+  require(specs.nonEmpty, "expected specs to be non-empty")
+
   override def isCompatibleWith(other: ShuffleSpec): Boolean = {
     specs.exists(_.isCompatibleWith(other))
   }
 
   override def canCreatePartitioning: Boolean =
     specs.forall(_.canCreatePartitioning)
 
-  override def createPartitioning(clustering: Seq[Expression]): Partitioning = {
-    // as we only consider # of partitions as the cost now, it doesn't matter which one we choose
-    // since they should all have the same # of partitions.
-    require(specs.map(_.numPartitions).toSet.size == 1, "expected all specs in the collection " +
-      "to have the same number of partitions")
-    specs.head.createPartitioning(clustering)
-  }
-
-  override def numPartitions: Int = {
-    require(specs.nonEmpty, "expected specs to be non-empty")
-    specs.head.numPartitions
-  }
+  override def flatten: Seq[LeafShuffleSpec] = specs.flatMap(_.flatten)
 }

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst
 
 import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
 import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Attribute, DirectShufflePartitionID, Expression, TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression}
 import org.apache.spark.sql.catalyst.plans.SQLHelper
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
@@ -59,7 +59,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
   }
 
   protected def checkCreatePartitioning(
-      spec: ShuffleSpec,
+      spec: LeafShuffleSpec,
       dist: ClusteredDistribution,
       expected: Partitioning): Unit = {
     val actual = spec.createPartitioning(dist.clustering)
@@ -506,7 +506,6 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
 
     withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
       val spec = reduced.createShuffleSpec(ClusteredDistribution(Seq(a)))
-        .asInstanceOf[KeyedShuffleSpec]
       assert(spec.joinKeyPositions === Some(Seq(0)))
       assert(spec.partitioning.partitionKeys.map(_.row.getInt(0)) === Seq(2020, 2021))
     }
@@ -613,13 +612,6 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
       SinglePartition
     )
 
-    checkCreatePartitioning(ShuffleSpecCollection(Seq(
-      HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution),
-        RangeShuffleSpec(10, distribution))),
-      ClusteredDistribution(Seq($"c", $"d")),
-      HashPartitioning(Seq($"c"), 10)
-    )
-
     // unsupported cases
 
     checkError(
@@ -629,7 +621,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
       condition = "UNSUPPORTED_CALL.WITHOUT_SUGGESTION",
       parameters = Map(
         "methodName" -> "createPartitioning$",
-        "className" -> "org.apache.spark.sql.catalyst.plans.physical.ShuffleSpec"))
+        "className" -> "org.apache.spark.sql.catalyst.plans.physical.LeafShuffleSpec"))
   }
 
   test("compatibility: ShufflePartitionIdPassThroughSpec on both sides") {
@@ -690,4 +682,73 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
       expected = false
     )
   }
+
+  test("SPARK-59080: a collection whose members cover different key subsets disagrees") {
+    val id = AttributeReference("id", IntegerType)()
+    val t1 = AttributeReference("t1", IntegerType)()
+    val t2 = AttributeReference("t2", IntegerType)()
+    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.fromPartitionings(Seq(
+      KeyedPartitioning(Seq(id, t1), keys),
+      KeyedPartitioning(Seq(id, t2), keys)))
+
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1)))
+
+      // 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, so the
+      // collection has no partition count of its own to read.
+      val memberPartitions = spec.flatten.map(_.numPartitions)
+      assert(memberPartitions.toSet === Set(3, 2))
+      assert(spec.isCompatibleWith(spec), "every member stays available for matching")
+    }
+  }
+
+  test("SPARK-59256: a single-partition side needs every member of a collection to be one") {
+    val id = AttributeReference("id", IntegerType)()
+    val t1 = AttributeReference("t1", IntegerType)()
+    val t2 = AttributeReference("t2", IntegerType)()
+    // Clustering on (id, t1) again. The first member matches `id` only and collapses to one
+    // partition, the second projects onto both positions and keeps three. The child itself has
+    // three, so it is not co-partitioned with a single partition, whatever the head says.
+    val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3))
+    val collection = PartitioningCollection.fromPartitionings(Seq(
+      KeyedPartitioning(Seq(id, t2), keys),
+      KeyedPartitioning(Seq(id, t1), keys)))
+
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1)))
+      val memberPartitions = spec.flatten.map(_.numPartitions)
+      // Ordered, not just as a set: the whole point is that the head is the wrong one to read.
+      assert(memberPartitions === Seq(1, 3))
+
+      // Reading the head's count answers one partition and claims a co-partitioning that does not
+      // hold. Only this direction is asserted: `KeyedShuffleSpec` has no
+      // `SinglePartitionShuffleSpec` case at all, so the reverse is false whatever the member
+      // counts are. That asymmetry is pre-existing and is not what this change is about.
+      assert(!SinglePartitionShuffleSpec.isCompatibleWith(spec))
+    }
+  }
+
+  test("SPARK-59256: an empty collection is rejected at construction, not at the first read") {
+    // `numPartitions` carried this `require` and threw the same way, but only once something asked.
+    // `flatten` and the ranking's `max` both rely on a collection being non-empty.
+    val e = intercept[IllegalArgumentException](ShuffleSpecCollection(Nil))
+    assert(e.getMessage.contains("expected specs to be non-empty"))
+  }
+
+  test("SPARK-59256: flattening reaches the members of a nested collection") {
+    val distribution = ClusteredDistribution(Seq($"a", $"b"))
+    val buried = HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution)
+    val direct = HashShuffleSpec(HashPartitioning(Seq($"b"), 10), distribution)
+    // A `PartitioningCollection` can hold another one, so a spec collection can nest too.
+    val collection = ShuffleSpecCollection(Seq(ShuffleSpecCollection(Seq(buried)), direct))
+
+    assert(collection.flatten === Seq(buried, direct))
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

@@ -87,9 +87,10 @@ case class GroupPartitionsExec(
                 assert(projectedExpressions.length == exprs.length)
                 projectedExpressions.zip(exprs).map {
                   case (expr, Some(KeyReducer(_, reduced))) =>
-                    // `reduced` was stored from the single spec that `createKeyedShuffleSpec`
-                    // picked (`collectFirst`); re-target it at this `KeyedPartitioning`'s own key
-                    // attribute so that every `KeyedPartitioning` in a collection keeps its own.
+                    // `reduced` came from the one member `checkKeyGroupCompatible` paired this
+                    // side on, which need not be the member being rewritten. The keys are reduced
+                    // once, from the shared key rows, so `reduced` describes them whichever member
+                    // this is, and only the key attribute has to be re-targeted.
                     reduced.withReference(expr.references.head)
                   case (expr, None) => expr
                 }