[SPARK-59289][SQL] Decide a co-partitioned child's GroupPartitionsExec once, where the join's merged keys are known

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

@@ -20,14 +20,23 @@ package org.apache.spark.sql.catalyst
 import org.apache.spark.SparkFunSuite
 /* Implicit conversions */
 import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod, SortOrder, 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
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.IntegerType
+import org.apache.spark.sql.types.{DataType, IntegerType}
 
 class DistributionSuite extends SparkFunSuite with SQLHelper {
 
+  /** A bound function with a stable canonical name, enough to build a `TransformExpression`. */
+  private object TestYearsFunction extends ScalarFunction[Int] {
+    override def inputTypes(): Array[DataType] = Array(IntegerType)
+    override def resultType(): DataType = IntegerType
+    override def name(): String = "years"
+    override def canonicalName(): String = "test.years"
+  }
+
   protected def checkSatisfied(
       inputPartitioning: Partitioning,
       requiredDistribution: Distribution,
@@ -413,6 +422,56 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true)
   }
 
+  test("SPARK-59289: satisfies is strict about a projection that merges partitions") {
+    val a = AttributeReference("a", IntegerType)()
+    val b = AttributeReference("b", IntegerType)()
+    val clustered = ClusteredDistribution(Seq(a))
+
+    // Partitioned by [a, b], clustered on [a] alone. Both are grouped, so the only question is
+    // whether rows sharing `a` share a partition.
+    val merging = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(1, 2)))
+    val notMerging = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2, 2)))
+    assert(merging.isGrouped && notMerging.isGrouped)
+
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      // a = 1 sits on two partitions, so this satisfies nothing until a GroupPartitionsExec
+      // projects the keys onto [a]. `keysMaySatisfy` is the question that says so.
+      checkSatisfied(merging, clustered, false)
+      assert(merging.keysMaySatisfy(clustered))
+
+      // Projecting here would merge no partition, so every `a` already sits on one and the
+      // partitioning satisfies as it stands. Keeping it beats projecting, since it still names `b`
+      // for a downstream operator to co-partition on.
+      checkSatisfied(notMerging, clustered, true)
+    }
+  }
+
+  test("SPARK-59289: satisfies accepts a partition expression that is itself a cluster key") {
+    val ts = AttributeReference("ts", IntegerType)()
+    val other = AttributeReference("other", IntegerType)()
+    val years = TransformExpression(TestYearsFunction, Seq(ts))
+    val partitioning = KeyedPartitioning(Seq(years), Seq(InternalRow(1), InternalRow(2)))
+    assert(partitioning.isGrouped)
+
+    // The clustering names the transform rather than the column it reads, so rows sharing
+    // `years(ts)` already sit on one partition and nothing has to be re-partitioned. The test that
+    // decides this reads the expression as well as its references, and it runs on the default
+    // configuration.
+    checkSatisfied(partitioning, ClusteredDistribution(Seq(years)), true)
+
+    // The `requireAllClusterKeys` arm has always accepted this shape, since it compares
+    // expressions rather than references, so the two arms agree now.
+    checkSatisfied(
+      partitioning, ClusteredDistribution(Seq(years), requireAllClusterKeys = true), true)
+
+    // Clustered on the column the transform reads, rows sharing `ts` share a partition too, since
+    // the partition expression is a function of it.
+    checkSatisfied(partitioning, ClusteredDistribution(Seq(ts)), true)
+
+    // Neither the expression nor its reference is a cluster key here, so this is refused.
+    checkSatisfied(partitioning, ClusteredDistribution(Seq(other)), false)
+  }
+
   test("SPARK-56877: fromPartitionings reuses already-consistent nested collections") {
     val x = AttributeReference("x", IntegerType)()
     val y = AttributeReference("y", IntegerType)()

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

@@ -27,6 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StructType}
 
 class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
+
   private val passThrough_a_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"a"), 10)
   private val passThrough_b_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"b"), 10)
   private val passThrough_c_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"c"), 10)
@@ -39,6 +40,19 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     override def canonicalName(): String = "test.bucket"
   }
 
+  test("SPARK-59289: createShuffleSpec keeps members that only satisfy after grouping") {
+    val a = AttributeReference("a", IntegerType)()
+    // Filtering the members on the strict `satisfies` would leave nothing here, and an empty
+    // `ShuffleSpecCollection` is rejected by its own `require`.
+    val ungroupedKeyed = KeyedPartitioning(
+      Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
+    assert(!ungroupedKeyed.isGrouped, "test setup: no member serves the distribution as it stands")
+    val collection = PartitioningCollection(Seq(ungroupedKeyed, ungroupedKeyed))
+    val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(a)))
+    assert(spec.asInstanceOf[ShuffleSpecCollection].specs.size == 2)
+    assert(spec.flatten.map(_.numPartitions) == Seq(3, 3))
+  }
+
   protected def checkCompatible(
       left: ShuffleSpec,
       right: ShuffleSpec,

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala

@@ -45,7 +45,7 @@ object AQEUtils {
     case p: ProjectExec =>
       getRequiredDistribution(p.child).flatMap {
         case h: ClusteredDistribution =>
-          if (h.clustering.forall(e => p.projectList.exists(_.semanticEquals(e)))) {
+          if (h.allClusterKeysAmong(p.projectList)) {
             Some(h)
           } else {
             // It's possible that the user-specified repartition is effective but the output