ulysses-you commented on code in PR #59703:
URL: https://github.com/apache/spark/pull/59703#discussion_r4355920996
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -1901,6 +1901,144 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-60458: skip only the skewed join whose split the plan cannot
take") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100") {
+ withTempView("skewData1", "skewData2") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+
+ val query =
+ """
+ |SELECT key1, count(*) FROM skewData1 JOIN skewData2 ON key1 =
key2 GROUP BY key1
+ |UNION ALL
+ |SELECT key1, 1 FROM skewData1 JOIN skewData2 ON key1 = key2
+ |""".stripMargin
+
+ // The aggregate above the first join needs its input on `key1`, which
the split breaks.
+ // The second join feeds the union directly, so its split breaks
nothing: it is kept.
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(query)
+
+ // key2 is always 0, so only the 334 rows of key1 = 0 join, each
against all 1000 rows of
+ // the other side: the aggregated branch collapses them into one row,
the other keeps one
+ // row per joined row.
+ val rows = sql(query).collect()
Review Comment:
Kept on purpose. They pin the rows to constants instead of to the
non-adaptive run, so a regression that affects both engines still fails; they
add about 1 s per test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala:
##########
@@ -211,67 +211,138 @@ case class OptimizeSkewedJoin(ensureRequirements:
EnsureRequirements)
}
}
- def optimizeSkewJoin(plan: SparkPlan): SparkPlan = {
- val localShuffleReaderEnabled =
conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
- plan.transformUp {
+ /**
+ * The candidate `plan` stands for, when it is a sort-merge, shuffled-hash
or broadcast-hash join
+ * with a skewed shuffled side this rule can split.
+ */
+ private def splitSkewJoin(
+ plan: SparkPlan,
+ localShuffleReaderEnabled: Boolean): Option[SkewJoinCandidate] =
+ plan match {
case smj @ SortMergeJoinExec(_, _, joinType, _,
s1 @ SortExec(_, _, ShuffleStage(left: ShuffleQueryStageExec), _),
s2 @ SortExec(_, _, ShuffleStage(right: ShuffleQueryStageExec), _),
false) =>
tryOptimizeJoinChildren(left, right, joinType).map {
case (newLeft, newRight) =>
- smj.copy(
- left = s1.copy(child = newLeft), right = s2.copy(child =
newRight), isSkewJoin = true)
- }.getOrElse(smj)
+ SkewJoinCandidate(smj, joinType, smj.copy(
+ left = s1.copy(child = newLeft), right = s2.copy(child =
newRight),
+ isSkewJoin = true))
+ }
case shj @ ShuffledHashJoinExec(_, _, joinType, _, _,
ShuffleStage(left: ShuffleQueryStageExec),
ShuffleStage(right: ShuffleQueryStageExec), false) =>
tryOptimizeJoinChildren(left, right, joinType).map {
case (newLeft, newRight) =>
- shj.copy(left = newLeft, right = newRight, isSkewJoin = true)
- }.getOrElse(shj)
+ SkewJoinCandidate(
+ shj, joinType, shj.copy(left = newLeft, right = newRight,
isSkewJoin = true))
+ }
// When localShuffleReader is disabled, the streamedPlan of
BroadcastHashJoin may be skewed,
// so we try to optimize it.
- case bhj @ BroadcastHashJoinExec(_, _, _, BuildRight, _,
+ case bhj @ BroadcastHashJoinExec(_, _, joinType, BuildRight, _,
ShuffleStage(left: ShuffleQueryStageExec), _, _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(left).map {
- case newLeft => bhj.copy(left = newLeft, isSkewJoin = true)
- }.getOrElse(bhj)
- case bhj @ BroadcastHashJoinExec(_, _, _, BuildLeft, _, _,
+ case newLeft =>
+ SkewJoinCandidate(bhj, joinType, bhj.copy(left = newLeft,
isSkewJoin = true))
+ }
+ case bhj @ BroadcastHashJoinExec(_, _, joinType, BuildLeft, _, _,
ShuffleStage(right: ShuffleQueryStageExec), _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(right).map {
- case newRight => bhj.copy(right = newRight, isSkewJoin = true)
- }.getOrElse(bhj)
+ case newRight =>
+ SkewJoinCandidate(bhj, joinType, bhj.copy(right = newRight,
isSkewJoin = true))
+ }
+
+ case _ => None
}
+
+ /**
+ * Every join in `plan` this rule can split, in the bottom-up order a
`transformUp` visits them.
+ */
+ private def collectSplitSkewJoins(plan: SparkPlan): Seq[SkewJoinCandidate] =
{
+ val localShuffleReaderEnabled =
conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
+ val candidates = mutable.ArrayBuffer.empty[SkewJoinCandidate]
+ plan.foreachUp { operator =>
+ splitSkewJoin(operator, localShuffleReaderEnabled).foreach(candidates +=
_)
+ }
+ candidates.toSeq
}
- override def apply(plan: SparkPlan): SparkPlan = {
- if (!conf.getConf(SQLConf.SKEW_JOIN_ENABLED)) {
+ /**
+ * `plan` with `candidates` replaced by their splits. Candidates are matched
by node identity,
+ * so structurally equal joins in different places are distinct candidates.
A subset may be
+ * passed, keeping the others out of the result.
+ *
+ * The walk is top-down: a sort-merge or shuffled-hash candidate has shuffle
stages as children
Review Comment:
Reworded in 98db64897e8: the candidate doc states the `None` case, the
collection says `foreachUp`, and the walk says a shuffled-hash candidate has
shuffle stages as children while a sort-merge one has sorts over them.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -1901,6 +1901,144 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-60458: skip only the skewed join whose split the plan cannot
take") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100") {
+ withTempView("skewData1", "skewData2") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+
+ val query =
+ """
+ |SELECT key1, count(*) FROM skewData1 JOIN skewData2 ON key1 =
key2 GROUP BY key1
+ |UNION ALL
+ |SELECT key1, 1 FROM skewData1 JOIN skewData2 ON key1 = key2
+ |""".stripMargin
+
+ // The aggregate above the first join needs its input on `key1`, which
the split breaks.
+ // The second join feeds the union directly, so its split breaks
nothing: it is kept.
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(query)
+
+ // key2 is always 0, so only the 334 rows of key1 = 0 join, each
against all 1000 rows of
+ // the other side: the aggregated branch collapses them into one row,
the other keeps one
+ // row per joined row.
+ val rows = sql(query).collect()
+ assert(rows.length == 334001)
+ assert(rows.count(_ == Row(0L, 334000L)) == 1)
+ assert(rows.count(_ == Row(0L, 1L)) == 334000)
+
+ val aggregateJoins = collect(adaptivePlan) {
+ case agg: BaseAggregateExec =>
+ collect(agg) { case j: SortMergeJoinExec => j }
+ }.flatten
+ assert(aggregateJoins.size == 1)
+ assert(!aggregateJoins.head.isSkewJoin)
+ val unionJoins = findTopLevelSortMergeJoin(adaptivePlan)
+ .filterNot(j => aggregateJoins.exists(_ eq j))
+ assert(unionJoins.size == 1)
+ assert(unionJoins.head.isSkewJoin)
+ assert(collect(unionJoins.head) {
+ case r: AQEShuffleReadExec => r
+
}.exists(_.partitionSpecs.exists(_.isInstanceOf[PartialReducerPartitionSpec])))
+ assert(ValidateRequirements.validate(adaptivePlan))
+ }
+ }
+ }
+
+ test("SPARK-60458: keep the layout a union passes through") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100") {
+ withTempView("skewData1", "skewData2") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+
+ def splitJoins(query: String): Seq[SortMergeJoinExec] = {
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(query)
+ assert(ValidateRequirements.validate(adaptivePlan))
+ val joins = findTopLevelSortMergeJoin(adaptivePlan)
+ assert(joins.size == 2, adaptivePlan.toString)
+ joins
+ }
+
+ val union =
+ """
+ |SELECT key1 AS k, value1 AS v FROM skewData1 JOIN skewData2 ON
key1 = key2
+ |UNION ALL
+ |SELECT key1 AS k, value1 AS v FROM skewData1 JOIN skewData2 ON
key1 = key2
+ |""".stripMargin
+ val windowedQuery =
+ s"SELECT k, rank() OVER (PARTITION BY k ORDER BY v) FROM ($union) u"
+
+ // Both branches report the join's hash layout, so the union claims it
and the window above
+ // depends on it: a split takes the claim away from both branches. The
all-at-once attempt
+ // rejects this plan too; this arm pins that the per-join fallback
keeps no split the claim
+ // cannot take, and the arm below that the same splits are kept when
nothing reads the
+ // claim.
+ val windowed = splitJoins(windowedQuery)
+ assert(windowed.forall(j => !j.isSkewJoin))
+ // Each branch holds the 334000 joined rows of key1 = 0, so across the
union every v of
+ // 0, 3, ..., 999 appears 2000 times and shares its rank with its
duplicates.
+ val windowedRows = sql(windowedQuery).collect()
+ assert(windowedRows.length == 668000)
+ assert(windowedRows.forall(_.getLong(0) == 0L))
+
assert(windowedRows.groupBy(_.getInt(1)).view.mapValues(_.length).toMap ==
+ (1 to 666001 by 2000).map(_ -> 2000).toMap)
+
+ // Nothing above reads the union's layout, so the same splits break
nothing and are kept.
+ val plain = splitJoins(union)
+ assert(plain.forall(_.isSkewJoin))
+ val plainRows = sql(union).collect()
+ assert(plainRows.length == 668000)
+ assert(plainRows.groupBy(_.getLong(1)).view.mapValues(_.length).toMap
==
+ (0L until 1000L by 3).map(_ -> 2000).toMap)
+ }
+ }
+ }
+
+ test("SPARK-60458: split a join nested in a broadcast join's build side") {
Review Comment:
Updated in this push: three tests, 149/149, the top-down walk, the new
config and the coalescing change, and it now says both arms of the union test
pass on the previous rule, with the first arm a guard.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala:
##########
@@ -285,6 +356,10 @@ case class OptimizeSkewedJoin(ensureRequirements:
EnsureRequirements)
}
}
+// A skewed join `OptimizeSkewedJoin` can split: the join node in the plan,
its join type, and the
+// join it is replaced by, with the skewed side split.
+private case class SkewJoinCandidate(join: SparkPlan, joinType: JoinType,
replacement: SparkPlan)
Review Comment:
Changed in 98db64897e8: `SkewJoinCandidate(join: BaseJoinExec, replacement:
SparkPlan)`, the broadcast patterns read `_` again, and the skip log prints
`join.simpleString(conf.maxToStringFields)`, so the two skipped joins read
differently.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala:
##########
@@ -211,67 +211,132 @@ case class OptimizeSkewedJoin(ensureRequirements:
EnsureRequirements)
}
}
- def optimizeSkewJoin(plan: SparkPlan): SparkPlan = {
- val localShuffleReaderEnabled =
conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
- plan.transformUp {
+ /**
+ * The candidate `plan` stands for, when it is a sort-merge, shuffled-hash
or broadcast-hash join
+ * with a skewed shuffled side this rule can split.
+ */
+ private def splitSkewJoin(
+ plan: SparkPlan,
+ localShuffleReaderEnabled: Boolean): Option[SkewJoinCandidate] =
+ plan match {
case smj @ SortMergeJoinExec(_, _, joinType, _,
s1 @ SortExec(_, _, ShuffleStage(left: ShuffleQueryStageExec), _),
s2 @ SortExec(_, _, ShuffleStage(right: ShuffleQueryStageExec), _),
false) =>
tryOptimizeJoinChildren(left, right, joinType).map {
case (newLeft, newRight) =>
- smj.copy(
- left = s1.copy(child = newLeft), right = s2.copy(child =
newRight), isSkewJoin = true)
- }.getOrElse(smj)
+ SkewJoinCandidate(smj, joinType, smj.copy(
+ left = s1.copy(child = newLeft), right = s2.copy(child =
newRight),
+ isSkewJoin = true))
+ }
case shj @ ShuffledHashJoinExec(_, _, joinType, _, _,
ShuffleStage(left: ShuffleQueryStageExec),
ShuffleStage(right: ShuffleQueryStageExec), false) =>
tryOptimizeJoinChildren(left, right, joinType).map {
case (newLeft, newRight) =>
- shj.copy(left = newLeft, right = newRight, isSkewJoin = true)
- }.getOrElse(shj)
+ SkewJoinCandidate(
+ shj, joinType, shj.copy(left = newLeft, right = newRight,
isSkewJoin = true))
+ }
// When localShuffleReader is disabled, the streamedPlan of
BroadcastHashJoin may be skewed,
// so we try to optimize it.
- case bhj @ BroadcastHashJoinExec(_, _, _, BuildRight, _,
+ case bhj @ BroadcastHashJoinExec(_, _, joinType, BuildRight, _,
ShuffleStage(left: ShuffleQueryStageExec), _, _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(left).map {
- case newLeft => bhj.copy(left = newLeft, isSkewJoin = true)
- }.getOrElse(bhj)
- case bhj @ BroadcastHashJoinExec(_, _, _, BuildLeft, _, _,
+ case newLeft =>
+ SkewJoinCandidate(bhj, joinType, bhj.copy(left = newLeft,
isSkewJoin = true))
+ }
+ case bhj @ BroadcastHashJoinExec(_, _, joinType, BuildLeft, _, _,
ShuffleStage(right: ShuffleQueryStageExec), _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(right).map {
- case newRight => bhj.copy(right = newRight, isSkewJoin = true)
- }.getOrElse(bhj)
+ case newRight =>
+ SkewJoinCandidate(bhj, joinType, bhj.copy(right = newRight,
isSkewJoin = true))
+ }
+
+ case _ => None
}
+
+ /**
+ * Every join in `plan` this rule can split, in the bottom-up order a
`transformUp` visits them.
+ */
+ private def collectSplitSkewJoins(plan: SparkPlan): Seq[SkewJoinCandidate] =
{
+ val localShuffleReaderEnabled =
conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
+ val candidates = mutable.ArrayBuffer.empty[SkewJoinCandidate]
+ plan.foreachUp { operator =>
+ splitSkewJoin(operator, localShuffleReaderEnabled).foreach(candidates +=
_)
+ }
+ candidates.toSeq
}
- override def apply(plan: SparkPlan): SparkPlan = {
- if (!conf.getConf(SQLConf.SKEW_JOIN_ENABLED)) {
+ /**
+ * `plan` with `candidates` replaced by their splits. A candidate's children
are shuffle stages,
+ * so no candidate contains another and every candidate is a node `plan`
itself holds, which is
+ * how they are matched here: structurally equal joins in different places
are distinct
+ * candidates. A subset may be passed, keeping the others out of the result.
+ */
+ private def applySkewJoins(plan: SparkPlan, candidates:
Seq[SkewJoinCandidate]): SparkPlan = {
+ if (candidates.isEmpty) {
return plan
}
+ plan.transformUp { operator =>
+ candidates.find(_.join eq
operator).map(_.replacement).getOrElse(operator)
Review Comment:
Fixed in 351074cae68: `applySkewJoins` walks top-down, so the outer
broadcast join is replaced while it is still the node the plan holds and the
inner split is applied inside the replacement. A plan-level test nests a
sort-merge join in a broadcast join's build side; it fails on the previous walk
(`result.isSkewJoin was false`).
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -1901,6 +1901,117 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-60458: skip only the skewed join whose split the plan cannot
take") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100") {
+ withTempView("skewData1", "skewData2") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+
+ val query =
+ """
+ |SELECT key1, count(*) FROM skewData1 JOIN skewData2 ON key1 =
key2 GROUP BY key1
+ |UNION ALL
+ |SELECT key1, 1 FROM skewData1 JOIN skewData2 ON key1 = key2
+ |""".stripMargin
+
+ // The aggregate above the first join needs its input on `key1`, which
the split breaks.
+ // The second join feeds the union directly, so its split breaks
nothing: it is kept.
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(query)
+
+ // key2 is always 0, so only the 334 rows of key1 = 0 join, each
against all 1000 rows of
+ // the other side: the aggregated branch collapses them into one row,
the other keeps one
+ // row per joined row.
+ val rows = sql(query).collect()
+ assert(rows.length == 334001)
+ assert(rows.count(_ == Row(0L, 334000L)) == 1)
+ assert(rows.count(_ == Row(0L, 1L)) == 334000)
+
+ val aggregateJoins = collect(adaptivePlan) {
+ case agg: BaseAggregateExec =>
+ collect(agg) { case j: SortMergeJoinExec => j }
+ }.flatten
+ assert(aggregateJoins.size == 1)
+ assert(!aggregateJoins.head.isSkewJoin)
+ val unionJoins = findTopLevelSortMergeJoin(adaptivePlan)
+ .filterNot(j => aggregateJoins.exists(_ eq j))
+ assert(unionJoins.size == 1)
+ assert(unionJoins.head.isSkewJoin)
+ assert(collect(unionJoins.head) {
+ case r: AQEShuffleReadExec => r
+
}.exists(_.partitionSpecs.exists(_.isInstanceOf[PartialReducerPartitionSpec])))
+ assert(ValidateRequirements.validate(adaptivePlan))
+ }
+ }
+ }
+
+ test("SPARK-60458: keep the layout a union passes through") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100") {
+ withTempView("skewData1", "skewData2") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+
+ def splitJoins(query: String): Seq[SortMergeJoinExec] = {
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(query)
+ assert(ValidateRequirements.validate(adaptivePlan))
+ val joins = findTopLevelSortMergeJoin(adaptivePlan)
+ assert(joins.size == 2, adaptivePlan.toString)
+ joins
+ }
+
+ val union =
+ """
+ |SELECT key1 AS k, value1 AS v FROM skewData1 JOIN skewData2 ON
key1 = key2
+ |UNION ALL
+ |SELECT key1 AS k, value1 AS v FROM skewData1 JOIN skewData2 ON
key1 = key2
+ |""".stripMargin
+ val windowedQuery =
+ s"SELECT k, rank() OVER (PARTITION BY k ORDER BY v) FROM ($union) u"
+
+ // Both branches report the join's hash layout, so the union claims it
and the window above
+ // depends on it: a split takes the claim away from both branches.
Base rejects the whole
Review Comment:
Reworded in 351074cae68 to name the all-at-once attempt instead of the
review-time base.
--
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]
