[SPARK-59905][SQL] Fix storage-partitioned join sorting partitions the wrong way for an ORDER BY over a join-key expression

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

@@ -953,11 +953,11 @@ case class KeyedPartitioning(
         }
 
       case o @ OrderedDistribution(_) if SQLConf.get.v2BucketingAllowSorting =>
-        // `EnsureRequirements` orders the key rows by the attributes the expressions are over, and
-        // nothing makes a reducer order-preserving, so that ordering says nothing about reduced
-        // keys. This is the local gate. A reduced position always carries a transform and an
-        // `ORDER BY` cannot name one, so the match below already refuses every case a query can
-        // reach. Do not read the first clause as redundant.
+        // `EnsureRequirements` compares the key rows position by position under the required
+        // ordering. Nothing makes a reducer order-preserving, so that comparison says nothing
+        // about reduced keys. This is the local gate. A reduced position always carries a
+        // transform and an `ORDER BY` cannot name one, so the match below already refuses every
+        // case a query can reach. Do not read the first clause as redundant.
         // An out-of-set key can break the ascending sequence of the declared keys, so a marked
         // layout keeps a global ordering claim only for a single partition (see the `@param`).
         expressionsDescribeKeys && o.areAllClusterKeysMatched(expressions) &&

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

@@ -139,10 +139,20 @@ case class EnsureRequirements(
               // OrderedDistribution requires grouped KeyedPartitioning with sorted keys
               // according to the distribution's ordering.
               val satisfyingKeyedPartitioning = resolution.fold(identity, _._1)
-              // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees
-              // one attribute per partition expression.
-              val attrs = satisfyingKeyedPartitioning.expressions.flatMap(_.references)
-              val keyRowOrdering = RowOrdering.create(o.ordering, attrs)
+              // A key row holds the values of the partition expressions. `keysSatisfy` admits a
+              // partitioning here only when those expressions are the ordering's, position by
+              // position, so each sort order reads its position of the key row. The expression
+              // need not be a bare column. A join can report the other side's join key, e.g.
+              // `100 - b`.
+              assert(o.areAllClusterKeysMatched(satisfyingKeyedPartitioning.expressions),
+                "the partition expressions must be the ordering's, position by position")
+              val keyRowOrdering = RowOrdering.create(
+                o.ordering.zip(satisfyingKeyedPartitioning.keyDataTypes).zipWithIndex.map {
+                  case ((order, dataType), i) =>
+                    order.copy(child = BoundReference(i, dataType, nullable = true),
+                      sameOrderExpressions = Seq.empty)
+                },
+                Nil)
               val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row)
               val keys = satisfyingKeyedPartitioning.partitionKeys
               // An empty zip is vacuously sorted, which is the answer for a single key.
@@ -527,15 +537,17 @@ case class EnsureRequirements(
           .orElse(reorderJoinKeysRecursively(
             leftKeys, rightKeys, leftPartitioning, None))
       case (Some(kp: KeyedPartitioning), _) =>
-        // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
-        // attribute per partition expression.
+        // A scan reports one column per partition expression. A side a join laid out on another
+        // side's keys can report an expression over several, e.g. `b + c`. `reorder` then sees
+        // more columns than keys and gives up.
         val leafExprs = kp.expressions.flatMap(_.references)
         reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, leftKeys)
             .orElse(reorderJoinKeysRecursively(
               leftKeys, rightKeys, None, rightPartitioning))
       case (_, Some(kp: KeyedPartitioning)) =>
-        // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
-        // attribute per partition expression.
+        // A scan reports one column per partition expression. A side a join laid out on another
+        // side's keys can report an expression over several, e.g. `b + c`. `reorder` then sees
+        // more columns than keys and gives up.
         val leafExprs = kp.expressions.flatMap(_.references)
         reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, rightKeys)
             .orElse(reorderJoinKeysRecursively(

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

@@ -5461,6 +5461,64 @@ class KeyGroupedPartitioningSuite
     }
   }
 
+  test("SPARK-59905: ORDER BY a join-key expression a join reports as a partition expression") {
+    createTable("ident", Array(Column.create("id", LongType)), Array(identity("id")))
+    sql("INSERT INTO testcat.ns.ident VALUES " + (94 to 100).map(i => s"($i)").mkString(", "))
+    createTable("plain", Array(Column.create("b", LongType), Column.create("c", LongType),
+      Column.create("s", new StructType().add("a", LongType))), Array.empty)
+    sql("INSERT INTO testcat.ns.plain VALUES " +
+      (0 to 6).map(b => s"($b, 94, named_struct('a', ${94 + b}L))").mkString(", "))
+    createTable("ident2", Array(Column.create("x", LongType), Column.create("y", LongType)),
+      Array(identity("x"), identity("y")))
+    sql("INSERT INTO testcat.ns.ident2 VALUES (94, 1), (94, 2), (95, 1)")
+    createTable("plain2", Array(Column.create("b", LongType), Column.create("c", LongType)),
+      Array.empty)
+    sql("INSERT INTO testcat.ns.plain2 VALUES (6, 1), (6, 2), (5, 1)")
+
+    // The join lays `p` out on `i`'s keys and reports the join key, e.g. `100 - b`, as `p`'s
+    // partition expression. A one-side shuffle does that under the shuffle conf. An inner
+    // broadcast join does it too, by reporting `i`'s layout over `p`'s join key. The ORDER BY over
+    // that key is served by the key order, with no range shuffle. The keys come out ascending, so
+    // only another order needs a `GroupPartitionsExec` to sort the partitions.
+    case class Query(left: String, right: String, on: String, order: String, select: String,
+        expected: Seq[Row], sortingNodes: Int)
+    val ids = (94 to 100).map(i => Row(i.toLong))
+    val singleColumnKeys = Seq(
+      Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b", "p.b",
+        (6 to 0 by -1).map(b => Row(b.toLong)), 0),
+      Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b DESC", "p.b",
+        (0 to 6).map(b => Row(b.toLong)), 1),
+      Query("ident", "plain", "i.id = p.s.a", "p.s.a", "i.id", ids, 0),
+      Query("ident2", "plain2", "i.x = 100 - p.b AND i.y = p.c", "100 - p.b, p.c DESC", "p.b, p.c",
+        Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1))
+    // AQE's plan validation fails on `b + c` until SPARK-59901 is fixed, so it runs with AQE off.
+    val twoColumnKey =
+      Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0)
+
+    for {
+      (query, aqe) <-
+        singleColumnKeys.flatMap(q => Seq(q -> true, q -> false)) :+ (twoColumnKey -> false)
+      broadcast <- Seq(false, true)
+    } {
+      withSQLConf(
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+          SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> (!broadcast).toString,
+          SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+        val hint = if (broadcast) "/*+ BROADCAST(p) */ " else ""
+        val df = sql(
+          s"""SELECT $hint${query.select} FROM testcat.ns.${query.left} i
+             |JOIN testcat.ns.${query.right} p ON ${query.on} ORDER BY ${query.order}
+             |""".stripMargin)
+        val clue = s"${query.order}, aqe = $aqe, broadcast = $broadcast"
+        checkAnswer(df, query.expected)
+        val plan = df.queryExecution.executedPlan
+        // A one-side shuffle is the only shuffle, and a broadcast join needs none.
+        assert(collectAllShuffles(plan).size == (if (broadcast) 0 else 1), clue)
+        assert(collectAllGroupPartitions(plan).size == query.sortingNodes, clue)
+      }
+    }
+  }
+
   test("SPARK-55992: GroupPartitions string in simple and extended explain") {
     val items_partitions = Array(bucket(4, "id"), years("arrive_time"))
     createTable(items, itemsColumns, items_partitions)