address comments

Author: shujingyang-db

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

@@ -638,7 +638,6 @@ case class ShufflePartitionIdPassThrough(
     expr: DirectShufflePartitionID,
     numPartitions: Int) extends Expression with Partitioning with Unevaluable {
 
-  // We don't support creating partitioning for ShufflePartitionIdPassThrough.
   override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
     ShufflePartitionIdPassThroughSpec(this, distribution)
   }
@@ -1009,6 +1008,7 @@ case class ShufflePartitionIdPassThroughSpec(
       false
   }
 
+  // We don't support creating partitioning for ShufflePartitionIdPassThrough.
   override def canCreatePartitioning: Boolean = false
 
   override def numPartitions: Int = partitioning.numPartitions

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

@@ -516,7 +516,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     val dist3 = ClusteredDistribution(Seq($"e", $"b"))
     checkCompatible(
       p1.createShuffleSpec(dist3),
-      p2.createShuffleSpec(dist2),
+      p2.createShuffleSpec(dist),
       expected = false
     )
   }

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

@@ -165,27 +165,28 @@ case class EnsureRequirements(
       // Check if the following conditions are satisfied:
       //   1. There are exactly two children (e.g., join). Note that Spark doesn't support
       //      multi-way join at the moment, so this check should be sufficient.
-      //   2. All children are of `KeyGroupedPartitioning`, and they are compatible with each other
+      //   2. All children are of the same partitioning, and they are compatible with each other
       // If both are true, skip shuffle.
-      val isKeyGroupCompatible = parent.isDefined &&
+      val areChildrenCompatible = parent.isDefined &&
           children.length == 2 && childrenIndexes.length == 2 && {
         val left = children.head
         val right = children(1)
+
+        // key group compatibility check
         val newChildren = checkKeyGroupCompatible(
           parent.get, left, right, requiredChildDistributions)
         if (newChildren.isDefined) {
           children = newChildren.get
+          true
+        } else {
+          // If key group check fails, check ShufflePartitionIdPassThrough compatibility
+          checkShufflePartitionIdPassThroughCompatible(
+            left, right, requiredChildDistributions)
         }
-        newChildren.isDefined
       }
 
-      val isShufflePassThroughCompatible = !isKeyGroupCompatible &&
-        parent.isDefined && children.length == 2 && childrenIndexes.length == 2 &&
-        checkShufflePartitionIdPassThroughCompatible(
-          children.head, children(1), requiredChildDistributions)
-
       children = children.zip(requiredChildDistributions).zipWithIndex.map {
-        case ((child, _), idx) if isKeyGroupCompatible || isShufflePassThroughCompatible ||
+        case ((child, _), idx) if areChildrenCompatible ||
             !childrenIndexes.contains(idx) =>
           child
         case ((child, dist), idx) =>

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

@@ -1197,7 +1197,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     TransformExpression(BucketFunction, expr, Some(numBuckets))
   }
 
-  test("ShufflePartitionIdPassThrough - avoid necessary shuffle when they are compatible") {
+  test("ShufflePartitionIdPassThrough - avoid unnecessary shuffle when children are compatible") {
     withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
       val plan1 = DummySparkPlan(
         outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprA), 5))
@@ -1253,7 +1253,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
       val plan2 = DummySparkPlan(
         outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprC), 5))
       // Join on different keys than partitioning keys
-      val smjExec = SortMergeJoinExec(exprB :: Nil, exprD :: Nil, Inner, None, plan1, plan2)
+      val smjExec = SortMergeJoinExec(exprA :: exprB :: Nil, exprD :: exprC :: Nil, Inner, None,
+        plan1, plan2)
 
       EnsureRequirements.apply(smjExec) match {
         case SortMergeJoinExec(_, _, _, _,
@@ -1306,32 +1307,6 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
-  test("ShufflePartitionIdPassThrough - incompatible due to different expressions " +
-       "with same base column") {
-    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
-      // Even though both use exprA as base and have same numPartitions,
-      // different Pmod operations make them incompatible
-      val plan1 = DummySparkPlan(
-        outputPartitioning = ShufflePartitionIdPassThrough(
-          DirectShufflePartitionID(Pmod(exprA, Literal(10))), 5))
-      val plan2 = DummySparkPlan(
-        outputPartitioning = ShufflePartitionIdPassThrough(
-          DirectShufflePartitionID(Pmod(exprA, Literal(5))), 5))
-      val smjExec = SortMergeJoinExec(exprA :: Nil, exprA :: Nil, Inner, None, plan1, plan2)
-
-      EnsureRequirements.apply(smjExec) match {
-        case SortMergeJoinExec(_, _, _, _,
-          SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _),
-          SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) =>
-          // Both sides should be shuffled due to expression mismatch
-          assert(p1.numPartitions == 10)
-          assert(p2.numPartitions == 10)
-          assert(p1.expressions == Seq(exprA))
-          assert(p2.expressions == Seq(exprA))
-        case other => fail(s"Expected shuffles on both sides, but got: $other")
-      }
-    }
-  }
 
   test("ShufflePartitionIdPassThrough - compatible with multiple clustering keys") {
     withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
@@ -1359,6 +1334,33 @@ class EnsureRequirementsSuite extends SharedSparkSession {
         case other => fail(s"We don't expect shuffle on neither sides with multiple " +
           s"clustering keys, but got: $other")
       }
+
+      // Test case 2: partition key matches at position 1
+      // Both sides partitioned by exprB and join on (exprA, exprB)
+      // Should be compatible because partition key exprB matches at position 1 in join keys
+      val plan3 = DummySparkPlan(
+        outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprB), 5))
+      val plan4 = DummySparkPlan(
+        outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprB), 5))
+      val smjExec2 = SortMergeJoinExec(exprA :: exprB :: Nil, exprA :: exprB :: Nil, Inner, None,
+        plan3, plan4)
+
+      EnsureRequirements.apply(smjExec2) match {
+        case SortMergeJoinExec(
+            leftKeys,
+            rightKeys,
+            _,
+            _,
+            SortExec(_, _, DummySparkPlan(_, _, _: ShufflePartitionIdPassThrough, _, _), _),
+            SortExec(_, _, DummySparkPlan(_, _, _: ShufflePartitionIdPassThrough, _, _), _),
+            _
+            ) =>
+          // No shuffles because exprB (partition key) appears at position 1 in join keys
+          assert(leftKeys === Seq(exprA, exprB))
+          assert(rightKeys === Seq(exprA, exprB))
+        case other => fail(s"Expected no shuffles due to position overlap at position 1, " +
+          s"but got: $other")
+      }
     }
   }
 
@@ -1414,35 +1416,6 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
-  test("ShufflePartitionIdPassThrough - compatible when partition key matches at any position") {
-    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
-      // Both sides partitioned by exprB and join on (exprA, exprB)
-      // Should be compatible because partition key exprB matches at position 1 in join keys
-      val plan1 = DummySparkPlan(
-        outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprB), 5))
-      val plan2 = DummySparkPlan(
-        outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprB), 5))
-      val smjExec = SortMergeJoinExec(exprA :: exprB :: Nil, exprA :: exprB :: Nil, Inner, None,
-        plan1, plan2)
-
-      EnsureRequirements.apply(smjExec) match {
-        case SortMergeJoinExec(
-            leftKeys,
-            rightKeys,
-            _,
-            _,
-            SortExec(_, _, DummySparkPlan(_, _, _: ShufflePartitionIdPassThrough, _, _), _),
-            SortExec(_, _, DummySparkPlan(_, _, _: ShufflePartitionIdPassThrough, _, _), _),
-            _
-            ) =>
-          // No shuffles because exprB (partition key) appears at position 1 in join keys
-          assert(leftKeys === Seq(exprA, exprB))
-          assert(rightKeys === Seq(exprA, exprB))
-        case other => fail(s"Expected no shuffles due to position overlap at position 1, " +
-          s"but got: $other")
-      }
-    }
-  }
 
   def years(expr: Expression): TransformExpression = {
     TransformExpression(YearsFunction, Seq(expr))