Skip to content

Commit 266ee89

Browse files
peter-tothdongjoon-hyun
authored andcommitted
[SPARK-59905][4.3][SQL] Fix storage-partitioned join sorting partitions the wrong way for an ORDER BY over a join-key expression
### Differences from #59189 This is the branch-4.3 backport of #59189. The description below is the original one. On this branch: - The change to `EnsureRequirements` does not cherry-pick cleanly, since the `OrderedDistribution` arm sits one level deeper here. The code is the same as on master, re-indented, and its comment is re-wrapped. - The test is the same as on master. Without the change, it fails on branch-4.3 the same way as on master. - Ran `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite` and `PlannerSuite` on branch-4.3, 360 tests in all, plus `dev/lint-scala`. ### What changes were proposed in this pull request? When `EnsureRequirements` serves an `ORDER BY` by putting a keyed child's partitions in key order, it now compares the key rows position by position. Sort order `i` reads field `i` of the key row. It used to bind the ordering to the references of the partition expressions instead. The binding takes the types the key rows hold (`keyDataTypes`), and an assertion pins the precondition it relies on next to it. ### Why are the changes needed? A key row holds the values of the partition expressions. Binding the ordering to their references is right only when each expression is a bare column, the only thing a scan reports that an `ORDER BY` can name. A join can report the other side's join key instead. A one-side shuffle does, and so does an inner broadcast hash join, whose `expandOutputPartitioning` reports the streamed side's layout over the build side's join key. With `spark.sql.sources.v2.bucketing.shuffle.enabled` and `spark.sql.sources.v2.bucketing.sorting.enabled` on, this query returns `b` as 0 to 6 instead of 6 to 0. `ident(id)` is identity-partitioned and holds 94 to 100, and `plain(b)` holds 0 to 6: ```sql SELECT p.b FROM ident i JOIN plain p ON i.id = 100 - p.b ORDER BY 100 - p.b ``` The same happens with only `spark.sql.sources.v2.bucketing.sorting.enabled` when `plain` is broadcast. 1. The join shuffles `plain` onto `ident`'s layout and reports `100 - b` as `plain`'s partition expression. Its key rows hold 94 to 100. 2. `KeyedPartitioning.keysSatisfy` admits that partitioning for `ORDER BY 100 - p.b`, since its expression is the ordering's. So no range shuffle is planned, and the global sort is dropped as redundant. 3. `EnsureRequirements` binds `100 - b` to `b` and evaluates it on the key rows, so it computes `100 - (100 - b)`. The partitions end up in the reverse order. Two other join keys fail planning instead: - For a join on `i.id = p.b + p.c`, `ORDER BY p.b + p.c` binds two references to a key row with one field. Planning fails with an `ArrayIndexOutOfBoundsException`. - For a join on `i.id = p.s.a`, `ORDER BY p.s.a` binds the struct `s` to the long key and reads a field of it. Planning fails with a `ClassCastException`. `keysSatisfy` admits a partitioning here only when its expressions are the ordering's, position by position (`OrderedDistribution.areAllClusterKeysMatched`). So reading field `i` for sort order `i` is exact. ### Does this PR introduce _any_ user-facing change? Yes. Such a query returns its rows in the right order, and the two other forms no longer fail planning. They all need `spark.sql.sources.v2.bucketing.sorting.enabled`, which is off by default. A broadcast join needs nothing else, and a one-side shuffle also needs `spark.sql.sources.v2.bucketing.shuffle.enabled`. No range shuffle is added. ### How was this patch tested? New test in `KeyGroupedPartitioningSuite`, "SPARK-59905: ORDER BY a join-key expression a join reports as a partition expression". It runs five queries: - `ORDER BY 100 - p.b`, ascending and descending; - `ORDER BY p.s.a`; - `ORDER BY 100 - p.b, p.c DESC` over a two-key identity side, which needs a mixed-direction `GroupPartitionsExec`; - `ORDER BY p.b + p.c`. Each runs through a one-side shuffle and through a broadcast join, with AQE on and off. The `p.b + p.c` one runs with AQE off only, since AQE's plan validation fails on it for SPARK-59901, which #59165 fixes. Each run checks the order of the rows, the number of shuffles, and whether a `GroupPartitionsExec` sorts the partitions. All 18 runs fail on master: - the `100 - p.b` and two-key queries return their rows in the wrong order; - the `p.s.a` one fails with the `ClassCastException`; - the `p.b + p.c` one fails with the `ArrayIndexOutOfBoundsException`. Also ran the `KeyGroupedPartitioning*` suites, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite` and `PlannerSuite`, 427 tests in all, plus `dev/lint-scala`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) Closes #59203 from peter-toth/SPARK-59905-order-by-expression-key-4.3. Authored-by: Peter Toth <peter.toth@gmail.com> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
1 parent 13e88a0 commit 266ee89

3 files changed

Lines changed: 83 additions & 13 deletions

File tree

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -805,11 +805,11 @@ case class KeyedPartitioning(
805805
}
806806

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

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

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -97,10 +97,20 @@ case class EnsureRequirements(
9797
// OrderedDistribution requires grouped KeyedPartitioning with sorted keys
9898
// according to the distribution's ordering.
9999
val satisfyingKeyedPartitioning = resolution.fold(identity, _._1)
100-
// The single-column invariant in KeyedPartitioning.supportsExpressions guarantees
101-
// one attribute per partition expression.
102-
val attrs = satisfyingKeyedPartitioning.expressions.flatMap(_.references)
103-
val keyRowOrdering = RowOrdering.create(o.ordering, attrs)
100+
// A key row holds the values of the partition expressions. `keysSatisfy` admits
101+
// a partitioning here only when those expressions are the ordering's, position
102+
// by position, so each sort order reads its position of the key row. The
103+
// expression need not be a bare column. A join can report the other side's join
104+
// key, e.g. `100 - b`.
105+
assert(o.areAllClusterKeysMatched(satisfyingKeyedPartitioning.expressions),
106+
"the partition expressions must be the ordering's, position by position")
107+
val keyRowOrdering = RowOrdering.create(
108+
o.ordering.zip(satisfyingKeyedPartitioning.keyDataTypes).zipWithIndex.map {
109+
case ((order, dataType), i) =>
110+
order.copy(child = BoundReference(i, dataType, nullable = true),
111+
sameOrderExpressions = Seq.empty)
112+
},
113+
Nil)
104114
val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row)
105115
val keys = satisfyingKeyedPartitioning.partitionKeys
106116
// An empty zip is vacuously sorted, which is the answer for a single key.
@@ -440,15 +450,17 @@ case class EnsureRequirements(
440450
.orElse(reorderJoinKeysRecursively(
441451
leftKeys, rightKeys, leftPartitioning, None))
442452
case (Some(KeyedPartitioning(clustering, _, _, _, _)), _) =>
443-
// The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
444-
// attribute per partition expression.
453+
// A scan reports one column per partition expression. A side a join laid out on another
454+
// side's keys can report an expression over several, e.g. `b + c`. `reorder` then sees
455+
// more columns than keys and gives up.
445456
val leafExprs = clustering.flatMap(_.references)
446457
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, leftKeys)
447458
.orElse(reorderJoinKeysRecursively(
448459
leftKeys, rightKeys, None, rightPartitioning))
449460
case (_, Some(KeyedPartitioning(clustering, _, _, _, _))) =>
450-
// The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
451-
// attribute per partition expression.
461+
// A scan reports one column per partition expression. A side a join laid out on another
462+
// side's keys can report an expression over several, e.g. `b + c`. `reorder` then sees
463+
// more columns than keys and gives up.
452464
val leafExprs = clustering.flatMap(_.references)
453465
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, rightKeys)
454466
.orElse(reorderJoinKeysRecursively(

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4559,6 +4559,64 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
45594559
}
45604560
}
45614561

4562+
test("SPARK-59905: ORDER BY a join-key expression a join reports as a partition expression") {
4563+
createTable("ident", Array(Column.create("id", LongType)), Array(identity("id")))
4564+
sql("INSERT INTO testcat.ns.ident VALUES " + (94 to 100).map(i => s"($i)").mkString(", "))
4565+
createTable("plain", Array(Column.create("b", LongType), Column.create("c", LongType),
4566+
Column.create("s", new StructType().add("a", LongType))), Array.empty)
4567+
sql("INSERT INTO testcat.ns.plain VALUES " +
4568+
(0 to 6).map(b => s"($b, 94, named_struct('a', ${94 + b}L))").mkString(", "))
4569+
createTable("ident2", Array(Column.create("x", LongType), Column.create("y", LongType)),
4570+
Array(identity("x"), identity("y")))
4571+
sql("INSERT INTO testcat.ns.ident2 VALUES (94, 1), (94, 2), (95, 1)")
4572+
createTable("plain2", Array(Column.create("b", LongType), Column.create("c", LongType)),
4573+
Array.empty)
4574+
sql("INSERT INTO testcat.ns.plain2 VALUES (6, 1), (6, 2), (5, 1)")
4575+
4576+
// The join lays `p` out on `i`'s keys and reports the join key, e.g. `100 - b`, as `p`'s
4577+
// partition expression. A one-side shuffle does that under the shuffle conf. An inner
4578+
// broadcast join does it too, by reporting `i`'s layout over `p`'s join key. The ORDER BY over
4579+
// that key is served by the key order, with no range shuffle. The keys come out ascending, so
4580+
// only another order needs a `GroupPartitionsExec` to sort the partitions.
4581+
case class Query(left: String, right: String, on: String, order: String, select: String,
4582+
expected: Seq[Row], sortingNodes: Int)
4583+
val ids = (94 to 100).map(i => Row(i.toLong))
4584+
val singleColumnKeys = Seq(
4585+
Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b", "p.b",
4586+
(6 to 0 by -1).map(b => Row(b.toLong)), 0),
4587+
Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b DESC", "p.b",
4588+
(0 to 6).map(b => Row(b.toLong)), 1),
4589+
Query("ident", "plain", "i.id = p.s.a", "p.s.a", "i.id", ids, 0),
4590+
Query("ident2", "plain2", "i.x = 100 - p.b AND i.y = p.c", "100 - p.b, p.c DESC", "p.b, p.c",
4591+
Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1))
4592+
// AQE's plan validation fails on `b + c` until SPARK-59901 is fixed, so it runs with AQE off.
4593+
val twoColumnKey =
4594+
Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0)
4595+
4596+
for {
4597+
(query, aqe) <-
4598+
singleColumnKeys.flatMap(q => Seq(q -> true, q -> false)) :+ (twoColumnKey -> false)
4599+
broadcast <- Seq(false, true)
4600+
} {
4601+
withSQLConf(
4602+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
4603+
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> (!broadcast).toString,
4604+
SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
4605+
val hint = if (broadcast) "/*+ BROADCAST(p) */ " else ""
4606+
val df = sql(
4607+
s"""SELECT $hint${query.select} FROM testcat.ns.${query.left} i
4608+
|JOIN testcat.ns.${query.right} p ON ${query.on} ORDER BY ${query.order}
4609+
|""".stripMargin)
4610+
val clue = s"${query.order}, aqe = $aqe, broadcast = $broadcast"
4611+
checkAnswer(df, query.expected)
4612+
val plan = df.queryExecution.executedPlan
4613+
// A one-side shuffle is the only shuffle, and a broadcast join needs none.
4614+
assert(collectAllShuffles(plan).size == (if (broadcast) 0 else 1), clue)
4615+
assert(collectAllGroupPartitions(plan).size == query.sortingNodes, clue)
4616+
}
4617+
}
4618+
}
4619+
45624620
test("SPARK-55992: GroupPartitions string in simple and extended explain") {
45634621
val items_partitions = Array(bucket(4, "id"), years("arrive_time"))
45644622
createTable(items, itemsColumns, items_partitions)

0 commit comments

Comments
 (0)