peter-toth commented on code in PR #60058:
URL: https://github.com/apache/spark/pull/60058#discussion_r4469233265
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4154,6 +4154,65 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-60460: nested SPJ retargets a one-side shuffle onto merged
partition values") {
+ createTable("pa", Array(Column.create("id", IntegerType)),
Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pa VALUES (1), (2), (3)")
+ createTable("pd", Array(Column.create("id", IntegerType)), Array.empty)
+ sql("INSERT INTO testcat.ns.pd VALUES (1), (2), (3)")
+ createTable("pc", Array(Column.create("id", IntegerType)),
Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pc VALUES (1), (2), (3), (4), (5)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(pa, pd, pc) */ pa.id, pd.id, pc.id
+ |FROM testcat.ns.pa JOIN testcat.ns.pd ON pa.id = pd.id
+ |JOIN testcat.ns.pc ON pa.id = pc.id
+ |""".stripMargin)
+ val shuffles = collectAllShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.size == 1, "only pd's one-side shuffle should remain")
+ assert(shuffles.head.outputPartitioning.numPartitions == 5,
+ "that shuffle should be retargeted onto the merged keys")
+ checkAnswer(df, Seq(Row(1, 1, 1), Row(2, 2, 2), Row(3, 3, 3)))
+ }
+ }
+ }
+
+ test("SPARK-60460: nested SPJ shuffles when a one-side shuffle is wider than
the outer key") {
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("dept", IntegerType),
+ Column.create("data", StringType))
+ createTable("sa", cols, Array(identity("store"), identity("dept")))
+ createTable("sb", cols, Array.empty)
+ createTable("sc", cols, Array(identity("dept")))
+ sql("INSERT INTO testcat.ns.sa VALUES (1, 10, 'a10'), (2, 20, 'a20')")
+ sql("INSERT INTO testcat.ns.sb VALUES (1, 10, 'b10'), (2, 20, 'b20')")
+ sql("INSERT INTO testcat.ns.sc VALUES (0, 10, 'c10'), (0, 20, 'c20'), (0,
30, 'c30')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(sa, sb, sc) */ sa.dept, sb.data, sc.data
+ |FROM testcat.ns.sa JOIN testcat.ns.sb
+ |ON sa.store = sb.store AND sa.dept = sb.dept
+ |JOIN testcat.ns.sc ON sa.dept = sc.dept
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(10, "b10", "c10"), Row(20, "b20", "c20")))
Review Comment:
**Finding 3.** This test passes on branch-4.1 (`336826d`) as well, AQE on
and off. The query selects only `sa.dept`, so the `Project` above the inner
join drops `store`. The inner join's partitioning on `(store, dept)` goes with
it. branch-4.1 then plans an exchange above the inner join (`Exchange
keygroupedpartitioning(dept#28, 3, [10], [20], [30], ...)`), so the nested push
never happens. The data would hide the bug anyway. Two depts and two `(store,
dept)` pairs give both sides two partitions in the same order.
With `sa.store` selected and three pairs, branch-4.1 takes the nested SPJ.
It pushes `dept` into `sa`'s scan and leaves `sb`'s exchange on the three
`(store, dept)` partitions. It silently drops `[2, 10, b2, c10]`. This head
returns all three rows, AQE on and off:
```suggestion
sql("INSERT INTO testcat.ns.sa VALUES (1, 10, 'a1'), (2, 10, 'a2'), (1,
20, 'a3')")
sql("INSERT INTO testcat.ns.sb VALUES (1, 10, 'b1'), (2, 10, 'b2'), (1,
20, 'b3')")
sql("INSERT INTO testcat.ns.sc VALUES (0, 10, 'c10'), (0, 20, 'c20'),
(0, 30, 'c30')")
Seq(true, false).foreach { aqe =>
withSQLConf(
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
val df = sql(
"""
|SELECT /*+ MERGE(sa, sb, sc) */ sa.store, sa.dept, sb.data,
sc.data
|FROM testcat.ns.sa JOIN testcat.ns.sb
|ON sa.store = sb.store AND sa.dept = sb.dept
|JOIN testcat.ns.sc ON sa.dept = sc.dept
|""".stripMargin)
checkAnswer(df,
Seq(Row(1, 10, "b1", "c10"), Row(2, 10, "b2", "c10"), Row(1, 20,
"b3", "c20")))
```
The description says this test fails on branch-4.1 "the same way". For this
shape branch-4.1 returns wrong rows rather than failing, which belongs under
the user-facing changes too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -902,6 +917,92 @@ case class EnsureRequirements(
case other => other.children.exists(hasReducedScanOutsideJoin)
}
+ /** Partition rows a push expands to, one per output partition. */
+ private def pushedPartitionRows(values: Seq[(InternalRow, Int)]):
Seq[InternalRow] =
+ values.flatMap { case (row, n) => Seq.fill(n)(row) }
+
+ private def samePartitionValues(
+ partitioning: KeyGroupedPartitioning,
+ rows: Seq[InternalRow]): Boolean =
+ partitioning.partitionValues.length == rows.length &&
+ partitioning.partitionValues.zip(rows).forall { case (left, right) =>
+ left.numFields == partitioning.expressions.length &&
+ left.numFields == right.numFields &&
+ InternalRowComparableWrapper(left, partitioning.expressions) ==
+ InternalRowComparableWrapper(right, partitioning.expressions)
+ }
+
+ /**
+ * A one-side shuffle can take `values` when it already has them, or when it
evaluates the same
+ * keys and can be rewritten onto them. A reducer changes the value domain,
and partial
+ * clustering replicates partitions, neither of which the shuffle's
partitioner can follow.
+ * Anything other than [[ShuffleExchangeExec]] (a query stage, a reused
exchange) cannot be
+ * rewritten, so it follows only when it already has `values`.
+ */
+ private def retargetKeyGroupedShuffle(
+ exchange: ShuffleExchangeExec,
+ values: Seq[(InternalRow, Int)],
+ reducers: Option[Seq[Option[Reducer[_, _]]]],
+ applyPartialClustering: Boolean): Option[ShuffleExchangeExec] = {
+ exchange.outputPartitioning match {
+ case k: KeyGroupedPartitioning if samePartitionValues(k,
pushedPartitionRows(values)) =>
+ Some(exchange)
+ case k: KeyGroupedPartitioning if !applyPartialClustering &&
Review Comment:
**Finding 1.** For branch-4.1 I'd keep only the refusal and drop this
rewrite. 4.2 and later handle these shapes with `GroupPartitionsExec` above the
inner join, so the retarget is a mechanism only 4.1 would have. It also needs
its own care. With no merged values, every guard of this case holds, since
`forall` and `headOption.forall` are true on `Nil`. The exchange then gets zero
partitions, and `KeyGroupedPartitioner` does `nonNegativeMod(hash, 0)`. With
`partition.filter.enabled`, `pa` and `pd` holding 1 to 3 and `pc` holding 7:
```sql
SELECT /*+ MERGE(pa, pd, pc) */ pa.id, pd.id, pc.id
FROM pa JOIN pd ON pa.id = pd.id
JOIN pc ON pa.id = pc.id
```
- This head, AQE on: `ArithmeticException: / by zero`. AQE off returns the
empty result.
- branch-4.1 (`336826d`): `PartitioningCollection requires all of its
partitionings have the same numPartitions` with AQE on, `Can't zip RDDs with
unequal numbers of partitions: List(0, 3)` with AQE off.
Refusing costs one shuffle, of the inner join's output, only where the
retarget fires. I measured it on this head with the rewrite removed, AQE on and
off:
- The motivating query, the reordered values and the partition-filter `RIGHT
JOIN` return the right rows with 2 shuffles instead of 1. branch-4.1 fails or
returns wrong rows on all three.
- The query above returns the empty result in both modes.
- Every other shape from this review keeps its rows and its shuffle count,
once the change from finding 2 is in.
The prototype pushes only when every key-grouped shuffle the push reaches
already has exactly the merged values, and no reducer is pushed. It removes
`retargetKeyGroupedShuffle`, `outputFollowsPush` and the new
`ShuffleExchangeExec` case in `populateSpjParams`, which goes back to `case s:
ShuffleExchangeLike => s`:
```scala
/**
* Whether every key-grouped shuffle [[populateSpjParams]] would reach
below `plan` already has
* `values`. The push stops at such a shuffle, which was planned at its
own values, so the push is
* sound only when it changes nothing there: the same values in the same
order, and no reducer.
* The walk stops where the push stops.
*/
private def nestedShufflesFollowPush(
plan: SparkPlan,
values: Seq[(InternalRow, Int)],
reducers: Option[Seq[Option[Reducer[_, _]]]]): Boolean = plan match {
case _: BroadcastExchangeLike => true
case s @ (_: ShuffleExchangeLike | _: ShuffleQueryStageExec | _:
ReusedExchangeExec) =>
s.outputPartitioning match {
case k: KeyGroupedPartitioning =>
reducers.forall(_.forall(_.isEmpty)) &&
samePartitionValues(k, pushedPartitionRows(values))
case _ => false
}
case other => other.children.forall(nestedShufflesFollowPush(_, values,
reducers))
}
```
With it and the change from finding 2, `KeyGroupedPartitioningSuite` and
`EnsureRequirementsSuite` pass, 155 tests. The first new test then expects two
shuffles, of 3 and 5 partitions.
The reducer condition comes from reading the code. With a pushed reducer,
the scan regroups on the reduced values while the exchange routes on the raw
ones. Equal value lists keep the two aligned only when the reducer maps each of
those values to itself. The test catalog's bucket reducers always do, so I
couldn't build a failing case, but the `ReducibleFunction` contract doesn't
promise it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -230,7 +230,11 @@ case class EnsureRequirements(
!childrenIndexes.contains(idx) =>
child
case ((child, dist), idx) =>
- if (bestMemberOpt.isDefined && matchedIndexes.contains(idx)) {
+ // A nested one-side shuffle is already routed at its own partition
count. Pushing join
+ // key positions into the scan beside it would regroup that scan and
leave the shuffle
+ // behind, so this child is shuffled onto the best spec instead.
+ if (bestMemberOpt.isDefined && matchedIndexes.contains(idx) &&
+ !hasKeyGroupedShuffleUnder(child)) {
Review Comment:
**Finding 2.** This also catches a matched child that the push leaves alone:
- The matching member has no join key positions, which is the default
(`allowJoinKeysSubsetOfPartitionKeys.enabled=false`). The `case _ => child`
below keeps such a child as it is.
- The subset conf is on, and the positions select every partition
expression. Pushing them doesn't regroup the scan.
In both cases the nested shuffle still lines up with the scan beside it. But
the child now takes the `else` branch and is shuffled onto its own layout.
Measured with `pa` partitioned by `id`, `pd` and `pe` not partitioned, each
holding 1 to 3, and `shuffle.enabled=true`, AQE on and off:
```sql
SELECT /*+ MERGE(pa, pd, pe) */ pa.id, pd.id, pe.id
FROM pa JOIN pd ON pa.id = pd.id
JOIN pe ON pa.id = pe.id
```
- branch-4.1 (`336826d`): 2 shuffles, for `pd` and `pe`.
- This head: 3 shuffles. The inner join's output is shuffled again, onto
`keygroupedpartitioning(id, 3, [1], [2], [3])`.
The rows are right on both, and the subset conf gives the same counts. A
fact table joined to two unpartitioned tables on its partition key is a common
shape. The reducer shape from the ticket's [second
comment](https://issues.apache.org/jira/browse/SPARK-60460?focusedCommentId=18136631#comment-18136631),
which fails on branch-4.1, has 3 shuffles here for the same reason.
The check only has to fire when the positions narrow the layout the nested
shuffle routes on. Comparing them with the matching member's own expressions is
not enough, because those are already projected onto the positions. An outer
join on `store`, a prefix of `(store, dept)`, then pushes `[0]` and fails at
planning (the `store` test in finding 4). Comparing them with the nested
shuffle's expressions works. I ran this version:
- the query above has 2 shuffles again, with and without the subset conf;
- the reducer shape has 2 shuffles instead of 3;
- the narrowed-positions shapes from finding 4 still return the right rows;
- `KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite` pass, 155
tests.
```scala
// A nested one-side shuffle keeps its own layout. Join key
positions that narrow it would
// regroup the scan beside it and leave the shuffle behind, so
this child is shuffled
// onto the best spec instead.
if (bestMemberOpt.isDefined && matchedIndexes.contains(idx) &&
!positionsRegroupBesideShuffle(child, specs(idx),
bestMemberOpt.get)) {
```
```scala
/**
* Whether pushing the join key positions of the member of `spec` that
matches `bestMember`
* would regroup a scan beside a one-side shuffle below `plan`. The
shuffle keeps its own layout,
* so only positions that select all of its partition expressions keep the
two aligned. The walk
* stops where the push stops.
*/
private def positionsRegroupBesideShuffle(
plan: SparkPlan,
spec: ShuffleSpec,
bestMember: ShuffleSpec): Boolean = {
def regroups(p: SparkPlan, positions: Seq[Int]): Boolean = p match {
case _: BroadcastExchangeLike => false
case s @ (_: ShuffleExchangeLike | _: ShuffleQueryStageExec | _:
ReusedExchangeExec) =>
s.outputPartitioning match {
case k: KeyGroupedPartitioning => positions !=
k.expressions.indices
case _ => false
}
case other => other.children.exists(regroups(_, positions))
}
flattenSpec(spec).find(bestMember.isCompatibleWith) match {
case Some(KeyGroupedShuffleSpec(_, _, Some(positions))) =>
regroups(plan, positions)
case _ => false
}
}
```
`hasKeyGroupedShuffleUnder` and `containsKeyGroupedShuffle` have no callers
after that.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4154,6 +4154,65 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-60460: nested SPJ retargets a one-side shuffle onto merged
partition values") {
Review Comment:
**Finding 4.** Both new tests stay green when the positions check is
removed, and when the values are compared without their order. The ticket's
comments and a partition-filter `RIGHT JOIN` give shapes that catch both. I ran
each one on this head, with the refusal from finding 1, and on branch-4.1
(`336826d`), AQE on and off:
- **Narrowed positions:** the corrected test from finding 3, with `sc` on
depts 10 and 20 only. Also the same tables with `sc` partitioned by `store`
(stores 1 and 2) and the outer join on `store`. branch-4.1 fails both at
planning with `PartitioningCollection requires all of its partitionings have
the same numPartitions`. So does this head without the
`hasKeyGroupedShuffleUnder` check. The ticket's [`da`/`db`/`de`
query](https://issues.apache.org/jira/browse/SPARK-60460?focusedCommentId=18132917#comment-18132917)
narrows the same way.
- **Reordered values**
([comment](https://issues.apache.org/jira/browse/SPARK-60460?focusedCommentId=18132917#comment-18132917)):
branch-4.1 returns no rows instead of 2, with no error. `sb`'s exchange keeps
`[20], [10]` while `sa`'s scan is sorted to `[10], [20]`. A check that compares
the two value lists as sets returns no rows too.
- **Partition filter:** branch-4.1 returns `[null,null,2]`, `[null,null,3]`,
`[null,null,4]` instead of `[2,2,2]`, `[3,3,3]`, `[null,null,4]`. `pa`'s scan
is filtered to `pc`'s keys, while `pd`'s exchange keeps `[1], [2], [3]`.
- **Reducer**
([comment](https://issues.apache.org/jira/browse/SPARK-60460?focusedCommentId=18136631#comment-18136631)):
branch-4.1 fails at planning or execution.
```sql
-- pa PARTITIONED BY (id): 1, 2, 3 pd not partitioned: 1, 2, 3
-- pc PARTITIONED BY (id): 2, 3, 4
-- shuffle.enabled=true, partition.filter.enabled=true
SELECT /*+ MERGE(pa, pd, pc) */ pa.id, pd.id, pc.id
FROM pa JOIN pd ON pa.id = pd.id
RIGHT JOIN pc ON pa.id = pc.id
```
All of them return the right rows here, with and without the retarget. The
reordered and partition-filter shapes returned wrong rows on branch-4.1 with no
error, so they belong in the user-facing section too.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4154,6 +4154,65 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-60460: nested SPJ retargets a one-side shuffle onto merged
partition values") {
+ createTable("pa", Array(Column.create("id", IntegerType)),
Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pa VALUES (1), (2), (3)")
+ createTable("pd", Array(Column.create("id", IntegerType)), Array.empty)
+ sql("INSERT INTO testcat.ns.pd VALUES (1), (2), (3)")
+ createTable("pc", Array(Column.create("id", IntegerType)),
Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pc VALUES (1), (2), (3), (4), (5)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(pa, pd, pc) */ pa.id, pd.id, pc.id
+ |FROM testcat.ns.pa JOIN testcat.ns.pd ON pa.id = pd.id
+ |JOIN testcat.ns.pc ON pa.id = pc.id
+ |""".stripMargin)
+ val shuffles = collectAllShuffles(df.queryExecution.executedPlan)
Review Comment:
**Finding 5.** Under AQE, `executedPlan` before execution is the initial
adaptive plan. So with `aqe = true` these assertions don't see what
re-optimization does once the stages materialize. Reading the shuffles after
`checkAnswer`, which runs the query, would check the final plan.
--
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]
