[SPARK-58996][SQL] Read the pre-alignment plan and its positions from one descent

The pre-alignment plan of each join side and the grouping inserted over it
are read once, and the statistics, the original partition keys and the
positions projecting them all derive from that one value. The partition
-count fallback compares the pre-alignment split counts for the same
reason the statistics read does.

Assisted-by: Claude Code

Author: ulysses-you

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

@@ -1133,6 +1133,14 @@ object PartitioningCollection {
     case _ => None
   }
 
+  /**
+   * The number of partitions of the [[KeyedPartitioning]] representing the keyed members of
+   * `partitioning`, if any. Collections validate on construction that their keyed members agree,
+   * so the representative's count stands for all of them.
+   */
+  def numKeyedPartitions(partitioning: Partitioning): Option[Int] =
+    representativeOf(partitioning).map(_.numPartitions)
+
   /**
    * Builds a [[PartitioningCollection]], unifying the `partitionKeys` reference across all
    * [[KeyedPartitioning]]s (including those in nested collections). Use this when combining

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

@@ -614,8 +614,13 @@ case class EnsureRequirements(
             logInfo(log"Skipping partially clustered distribution as it cannot be applied for " +
               log"join type '${MDC(LogKeys.JOIN_TYPE, joinType)}'")
           } else {
-            val unwrappedLeft = unwrapGroupPartitions(left)
-            val unwrappedRight = unwrapGroupPartitions(right)
+            // The pre-alignment plan of each side and the grouping this rule inserted over it,
+            // read once: the statistics and the original partition keys below come from the
+            // plan, the positions projecting them from the grouping.
+            val leftGrouping = innermostGroupPartition(left)
+            val rightGrouping = innermostGroupPartition(right)
+            val unwrappedLeft = leftGrouping.map(_._1.child).getOrElse(left)
+            val unwrappedRight = rightGrouping.map(_._1.child).getOrElse(right)
 
             val leftLink = unwrappedLeft.logicalLink
             val rightLink = unwrappedRight.logicalLink
@@ -637,11 +642,16 @@ case class EnsureRequirements(
                    |""".stripMargin)
               leftLink.get.stats.sizeInBytes < rightLink.get.stats.sizeInBytes
             } else {
-              // As a simple heuristic, we pick the side with fewer number of partitions
-              // to apply the grouping & replication of partitions
+              // As a simple heuristic, we pick the side with fewer number of partitions to
+              // apply the grouping & replication of partitions. The counts read the
+              // pre-alignment plans, for the same reason the statistics do: on a re-run both
+              // aligned reports hold the same number of keys, so comparing them decides nothing.
               logInfo("Using number of partitions to determine which side of join " +
                   "to fully cluster partition values")
-              leftPartKeys.size < rightPartKeys.size
+              PartitioningCollection.numKeyedPartitions(unwrappedLeft.outputPartitioning)
+                .getOrElse(leftPartKeys.size) <
+                PartitioningCollection.numKeyedPartitions(unwrappedRight.outputPartitioning)
+                .getOrElse(rightPartKeys.size)
             }
 
             replicateRightSide = !replicateLeftSide
@@ -661,18 +671,19 @@ case class EnsureRequirements(
               replicateRightSide = false
             } else {
               // In partially clustered distribution, we should use un-grouped partition values.
-              // The positions projecting them come from the innermost grouping when there is
-              // one: like in `applyGroupPartitions`, they were computed against the raw
-              // partition keys, while the spec's were computed against the node's already
-              // projected report on a re-run.
+              // The child and the positions projecting its keys come from the same grouping:
+              // the keys from the node's child, the positions from the node itself, falling back
+              // to the spec's when there is no grouping. Like in `applyGroupPartitions`, the
+              // node's positions were computed against the raw partition keys, while the spec's
+              // were computed against the node's already projected report on a re-run.
               val (partiallyClusteredChild, partiallyClusteredPositions) =
                 if (replicateLeftSide) {
                   (unwrappedRight,
-                    innermostGroupPartition(right).flatMap(_._1.joinKeyPositions)
+                    rightGrouping.flatMap(_._1.joinKeyPositions)
                       .orElse(rightSpec.joinKeyPositions))
                 } else {
                   (unwrappedLeft,
-                    innermostGroupPartition(left).flatMap(_._1.joinKeyPositions)
+                    leftGrouping.flatMap(_._1.joinKeyPositions)
                       .orElse(leftSpec.joinKeyPositions))
                 }
               // The pre-alignment plan of the side that keeps its splits: its partitioning
@@ -801,12 +812,9 @@ case class EnsureRequirements(
     }
 
   /**
-   * Unwraps the `GroupPartitionsExec` nodes this rule inserted over a join child, down to the
-   * pre-alignment plan, per the descent of [[innermostGroupPartition]].
-   *
-   * The statistics-based replicate-side choice and the original partition keys below must read
-   * from the pre-alignment plan on every pass: the local sort one level down carries no
-   * `logicalLink` and reports the aligned layout instead.
+   * Unwraps the groupings and local sorts this rule inserted over a child, down to the
+   * pre-alignment plan, per the descent of [[innermostGroupPartition]]. Peeling one level stops
+   * at the local sort this rule added, leaving the earlier pass's alignment in place.
    */
   private def unwrapGroupPartitions(plan: SparkPlan): SparkPlan =
     innermostGroupPartition(plan).map(_._1.child).getOrElse(plan)
@@ -827,6 +835,9 @@ case class EnsureRequirements(
       g.copy(
         joinKeyPositions = g.joinKeyPositions.orElse(joinKeyPositions),
         expectedPartitionKeys = Some(mergedPartitionKeys),
+        // Unlike `joinKeyPositions`, these need no `orElse`. A re-run with reducers never reaches
+        // here. Both sides then report the same reduced keys, so `isCompatible` above is true and
+        // the whole block is skipped.
         reducers = reducers,
         distributePartitions = distributePartitions)
     }.getOrElse {

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

@@ -2067,6 +2067,52 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
+  test("SPARK-58996: the replicate-side fallback counts pre-alignment partitions") {
+    // The replicate side is chosen by plan statistics; when there is none (the dummy plans carry
+    // no `logicalLink`, which forces the fallback), the side with fewer partitions is picked. On
+    // a re-run the join children arrive already aligned, and both aligned reports hold the same
+    // keys, so counting them decides nothing: the counts must come from the pre-alignment
+    // partitioning. Each arm hands the rule the output shape of a first pass -- a local sort
+    // over an aligned `GroupPartitionsExec` -- and reads the replicate-side choice back off the
+    // `distributePartitions` flags.
+    def distributeFlags(
+        leftKeys: Seq[InternalRow], leftSplits: Int,
+        rightKeys: Seq[InternalRow], rightSplits: Int): (Boolean, Boolean) = {
+      def alignedChild(expr: AttributeReference, keys: Seq[InternalRow], splits: Int) = {
+        val leaf = DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(expr), keys))
+        val gpe = GroupPartitionsExec(leaf,
+          expectedPartitionKeys = Some(Seq(
+            (InternalRowComparableWrapper(InternalRow(1), Seq(expr)), splits))))
+        SortExec(Seq(SortOrder(expr, Ascending)), global = false, gpe)
+      }
+      val smj = SortMergeJoinExec(Seq(exprA), Seq(exprB), Inner, None,
+        alignedChild(exprA, leftKeys, leftSplits), alignedChild(exprB, rightKeys, rightSplits))
+      withSQLConf(
+          SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+          SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true") {
+        val Seq(newLeft, newRight) = EnsureRequirements.apply(smj).children
+        def distribute(child: SparkPlan): Boolean = child.collectFirst {
+          case g: GroupPartitionsExec => g
+        }.get.distributePartitions
+        (distribute(newLeft), distribute(newRight))
+      }
+    }
+
+    // The left side holds one partition against three on the right: it is replicated, so it
+    // does not distribute. Counting the aligned reports instead sees one key on both sides and
+    // always picks the right side.
+    assert(distributeFlags(
+      leftKeys = Seq(InternalRow(1)), leftSplits = 1,
+      rightKeys = Seq(InternalRow(1), InternalRow(1), InternalRow(1)), rightSplits = 3) ===
+        ((false, true)))
+
+    // The mirrored control: the right side is replicated.
+    assert(distributeFlags(
+      leftKeys = Seq(InternalRow(1), InternalRow(1), InternalRow(1)), leftSplits = 3,
+      rightKeys = Seq(InternalRow(1)), rightSplits = 1) ===
+        ((true, false)))
+  }
+
   test("SPARK-58996: a single-child operator over a partially clustered layout still gets " +
       "grouped") {
     // A partially clustered `GroupPartitionsExec` reports a non-grouped partitioning by design,