dongjoon-hyun commented on code in PR #59189:
URL: https://github.com/apache/spark/pull/59189#discussion_r4158266417
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5461,6 +5461,41 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59905: ORDER BY the join-key expression of a one-side shuffle") {
Review Comment:
**The same bug is also reachable through a broadcast hash join, with only
`sorting.enabled`.** `BroadcastHashJoinExec.expandOutputPartitioning`
(BroadcastHashJoinExec.scala:119-121) substitutes the build-side join key into
the partitioning of the streamed side, so an inner broadcast join turns
`KeyedPartitioning([i.id])` into a collection that also holds
`KeyedPartitioning([100 - p.b])` over the same layout. With these tables,
`spark.sql.sources.v2.bucketing.sorting.enabled=true` and `shuffle.enabled`
left at its default (false), this query reaches the same `OrderedDistribution`
arm:
```sql
SELECT /*+ BROADCAST(p) */ p.b FROM testcat.ns.ident i JOIN testcat.ns.plain
p ON i.id = 100 - p.b ORDER BY 100 - p.b
```
On master it should return `b` as 0 to 6 instead of 6 to 0, and the `p.b +
p.c` and `p.s.a` forms should fail the same way as in this test. I traced this
by reading the code and did not run it. This PR fixes that route too, but no
test covers it, and the PR description says all the cases need
`spark.sql.sources.v2.bucketing.shuffle.enabled`. Could you add a broadcast
variant of these queries and update the description?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5461,6 +5461,41 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59905: ORDER BY the join-key expression of a one-side shuffle") {
+ 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(", "))
+
+ // The join shuffles `plain` onto `ident`'s keys and reports the join key,
e.g. `100 - b`, as
+ // its partition expression. The ORDER BY over it is served by the key
order, so the one-side
+ // shuffle is the only one. The keys come out ascending, so only a
descending order needs a
+ // `GroupPartitionsExec` to sort the partitions. AQE stays off, since its
plan validation fails
+ // on `b + c` until SPARK-59901 is fixed.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
Review Comment:
**AQE is turned off for all four queries, but the stated reason only applies
to `p.b + p.c`.** The SPARK-59901 failure is the `assert(refs.size == 1)` in
`KeyedShuffleSpec.keyPositions` (partitioning.scala:1954), which fires only for
a key with two references. `100 - p.b` and `p.s.a` have one reference each and
map to an empty position set, so they should plan under AQE too. As written,
the default AQE-on path is never tested for those three queries, e.g. the
`ValidateRequirements` check in `OptimizeSkewedJoin` and the re-planning over
the keyed `ShuffleQueryStageExec`. Could you run the three single-reference
queries with AQE both off and on, and keep only `p.b + p.c` AQE-off until
#59165 lands?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5461,6 +5461,41 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59905: ORDER BY the join-key expression of a one-side shuffle") {
+ 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(", "))
+
+ // The join shuffles `plain` onto `ident`'s keys and reports the join key,
e.g. `100 - b`, as
+ // its partition expression. The ORDER BY over it is served by the key
order, so the one-side
+ // shuffle is the only one. The keys come out ascending, so only a
descending order needs a
+ // `GroupPartitionsExec` to sort the partitions. AQE stays off, since its
plan validation fails
+ // on `b + c` until SPARK-59901 is fixed.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+ val ids = (94 to 100).map(i => Row(i.toLong))
+ Seq(
Review Comment:
Minor: Every case here has a single partition key. Two-key positional
reading is already covered on a scan layout by SPARK-48655, but no test covers
a multi-key member built by the one-side shuffle, where key-row field `i` has
to line up with the `i`-th expression that
`KeyedShuffleSpec.createPartitioning` emits and the direction differs per
position. For example, `ident2(x, y)` partitioned by `identity(x), identity(y)`
with keys (94, 1), (94, 2), (95, 1), joined `ON i.x = 100 - p.b AND i.y = p.c`
and ordered by `100 - p.b, p.c DESC`, which needs a mixed-direction
`GroupPartitionsExec`. Could you add such a case?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -139,10 +139,16 @@ 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, e.g. a shuffled side's join key
`100 - b`.
+ val keyRowOrdering = RowOrdering.create(
+ o.ordering.zipWithIndex.map { case (order, i) =>
+ SortOrder(BoundReference(i, order.child.dataType, nullable =
true),
Review Comment:
Minor: The type here is `order.child.dataType`, which is the partition
expression type, not the type the key rows are compared at. The
`KeyedPartitioning` doc says `expressionDataTypes` "Need not be what the
`partitionKeys` rows hold. See `keyDataTypes`", and that `ShuffleExchangeExec`
"is the one reader that stays on `expressionDataTypes`" (partitioning.scala:704
and :723). Today the two differ only in struct field names and nullability,
which the generated ordering ignores, so the result is right. But that depends
on join type coercion and the reduce gates keeping them in step. Could you zip
`o.ordering` with `satisfyingKeyedPartitioning.keyDataTypes` and use that type?
The lengths match because `areAllClusterKeysMatched` uses `corresponds`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -139,10 +139,16 @@ 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, e.g. a shuffled side's join key
`100 - b`.
+ val keyRowOrdering = RowOrdering.create(
Review Comment:
Minor: This ordering is only right because `KeyedPartitioning.keysSatisfy`
admits the partitioning through `o.areAllClusterKeysMatched(expressions)`
(partitioning.scala:963), in another module, and nothing here checks it. If
that gate is ever relaxed, e.g. to admit an ordering prefix the way
`RangePartitioning.satisfies0` does, then `KeyedPartitioning([a])` under `ORDER
BY a, b` would read field 1 of a one-field key row, and a permuted match would
sort the partitions in the wrong order without any error. Could you add
`assert(o.areAllClusterKeysMatched(satisfyingKeyedPartitioning.expressions))`
here, or build this ordering in `KeyedPartitioning` next to `keysSatisfy`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -139,10 +139,16 @@ 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, e.g. a shuffled side's join key
`100 - b`.
Review Comment:
Minor: This new comment is right, but two copies of the comment it replaces
are still in this file, at L536 and L543 in `reorderJoinKeysRecursively`: "The
single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
attribute per partition expression." `supportsExpressions` is only checked
where a scan reports its partitioning (DataSourceV2ScanExecBase.scala:105). For
the `p.b + p.c` join in this test, `kp.expressions.flatMap(_.references)` there
gives two attributes for one expression. The code is harmless (`reorder` gives
up on a size mismatch), but this assumption is what caused SPARK-59905, and
#59165 only rewords the third copy, at L1130. Could you update these two as
well?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -139,10 +139,16 @@ 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, e.g. a shuffled side's join key
`100 - b`.
+ val keyRowOrdering = RowOrdering.create(
+ o.ordering.zipWithIndex.map { case (order, i) =>
+ SortOrder(BoundReference(i, order.child.dataType, nullable =
true),
+ order.direction, order.nullOrdering, Seq.empty)
Review Comment:
Nit: `order.copy(child = BoundReference(i, ..., nullable = true),
sameOrderExpressions = Seq.empty)` would say "the same sort order, pointed at
key position `i`" without restating `direction` and `nullOrdering`.
ShuffleExchangeExec.scala:393-395 uses that form for the `RangePartitioner`
sample ordering. Keep `nullable = true` rather than copying `order.nullable` as
that code does, since a key row can hold a null where the `ORDER BY` expression
is not nullable. Could you use `copy` here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
