dongjoon-hyun commented on code in PR #59703:
URL: https://github.com/apache/spark/pull/59703#discussion_r4349390218
##########
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)
Review Comment:
Minor: This holds only because `CombineAdjacentAggregation` merges the
partial and final aggregates over join1 into one `Complete` aggregate. Without
that merge, both aggregates sit above join1, and the nested `collect` at
L1939-1942 returns join1 once for each. So with the internal
`spark.sql.execution.combineAdjacentAggregation` off, or after any change that
stops that merge, this fails although the rule did the right thing. I checked
this by reading the code and did not run it. Could this take only the topmost
aggregate, e.g. `collectFirst(adaptivePlan) { case agg: BaseAggregateExec =>
findTopLevelSortMergeJoin(agg) }.getOrElse(Nil)`, or dedupe by identity?
##########
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
+ * and contains nothing else, but a broadcast-hash join only constrains its
streamed side, so
+ * another candidate can sit in its build side. Replacing the outer join
first keeps that inner
+ * candidate an instance `plan` holds, and the walk then splits it inside
the replacement.
+ */
+ private[adaptive] def applySkewJoins(
+ plan: SparkPlan,
+ candidates: Seq[SkewJoinCandidate]): SparkPlan = {
+ if (candidates.isEmpty) {
return plan
}
+ plan.transformDown { operator =>
+ candidates.find(_.join eq
operator).map(_.replacement).getOrElse(operator)
+ }
+ }
- // We try to optimize every skewed sort-merge/shuffle-hash joins in the
query plan. If this
- // introduces extra shuffles, we give up the optimization and return the
original query plan, or
- // accept the extra shuffles if the force-apply config is true.
- // TODO: It's possible that only one skewed join in the query plan leads
to extra shuffles and
- // we only need to skip optimizing that join. We should make the
strategy smarter here.
- val optimized = optimizeSkewJoin(plan)
- val requirementSatisfied = if
(ensureRequirements.requiredDistribution.isDefined) {
- ValidateRequirements.validate(optimized,
ensureRequirements.requiredDistribution.get)
- } else {
- ValidateRequirements.validate(optimized)
+ private def requirementsSatisfied(plan: SparkPlan): Boolean = {
+ ensureRequirements.requiredDistribution match {
+ case Some(distribution) => ValidateRequirements.validate(plan,
distribution)
+ case None => ValidateRequirements.validate(plan)
}
- if (requirementSatisfied) {
- optimized.transform {
- case SkewJoinChildWrapper(child) => child
+ }
+
+ private def stripSkewJoinChildWrappers(plan: SparkPlan): SparkPlan = {
+ plan.transform { case SkewJoinChildWrapper(child) => child }
+ }
+
+ /**
+ * Applies the candidates one at a time, in the order a bottom-up pass
visits them, keeping a
+ * candidate when the plan holding the kept ones and this one still
satisfies every requirement.
+ * A split the plan cannot take therefore costs only its own join, where the
whole-plan fast path
+ * would have cost every join.
+ *
+ * Each trial is rebuilt from `plan` rather than from the previous trial, so
acceptance never
+ * depends on how a rebuild kept node instances; the handful of skewed joins
keeps that cheap.
+ */
+ private def applySkewJoinsSatisfyingRequirements(
+ plan: SparkPlan,
+ candidates: Seq[SkewJoinCandidate]): SparkPlan = {
+ var accepted = Seq.empty[SkewJoinCandidate]
+ candidates.foreach { candidate =>
+ val trial = applySkewJoins(plan, accepted :+ candidate)
+ if (requirementsSatisfied(trial)) {
+ accepted = accepted :+ candidate
+ } else {
+ logDebug(s"Skip splitting the skewed ${candidate.joinType} join " +
+ s"${candidate.join.nodeName}: with it the plan fails
ValidateRequirements")
}
+ }
+ logDebug(s"Optimized ${accepted.length} of ${candidates.length} skewed
joins")
+ stripSkewJoinChildWrappers(applySkewJoins(plan, accepted))
+ }
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.getConf(SQLConf.SKEW_JOIN_ENABLED)) {
+ return plan
+ }
+
+ // Try every split at once first: one pass to find the candidates and one
validation cover the
+ // common case. When that breaks a requirement of the plan, either accept
the extra shuffles if
+ // the force-apply config is true, or keep only the joins whose splits the
plan can take. The
+ // candidates are collected once, so the fallback reuses the splits the
fast path built.
+ val candidates = collectSplitSkewJoins(plan)
+ val optimized = applySkewJoins(plan, candidates)
+ if (requirementsSatisfied(optimized)) {
+ stripSkewJoinChildWrappers(optimized)
} else if (conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN)) {
- ensureRequirements.apply(optimized).transform {
- case SkewJoinChildWrapper(child) => child
- }
+ stripSkewJoinChildWrappers(ensureRequirements.apply(optimized))
} else {
Review Comment:
Minor: This branch used to return `plan`, and nothing turns the new fallback
off. Every AQE query whose all-at-once split fails validation with
`spark.sql.adaptive.forceOptimizeSkewedJoin` off now gets a different plan at
runtime, with trade-offs such as the coalescing one at L345. If the fallback
misbehaves, the only way back is `spark.sql.adaptive.skewJoin.enabled=false`,
which also drops the splits the old rule made, and
`spark.sql.adaptive.optimizer.excludedRules` cannot exclude this rule because
it filters only `AQEOptimizer`'s logical rules (AQEOptimizer.scala:50-61).
Neighboring rules have such a switch, e.g.
`spark.sql.execution.combineAdjacentAggregation`, and SPARK-52921 added
`spark.sql.unionOutputPartitioning` for a plan-shape change like this one.
Could an internal conf, e.g.
`spark.sql.adaptive.skewJoin.perJoinFallback.enabled` with default true and
version 4.4.0, guard this branch and return `plan` when it is off?
##########
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
+ * and contains nothing else, but a broadcast-hash join only constrains its
streamed side, so
+ * another candidate can sit in its build side. Replacing the outer join
first keeps that inner
+ * candidate an instance `plan` holds, and the walk then splits it inside
the replacement.
+ */
+ private[adaptive] def applySkewJoins(
+ plan: SparkPlan,
+ candidates: Seq[SkewJoinCandidate]): SparkPlan = {
+ if (candidates.isEmpty) {
return plan
}
+ plan.transformDown { operator =>
+ candidates.find(_.join eq
operator).map(_.replacement).getOrElse(operator)
+ }
+ }
- // We try to optimize every skewed sort-merge/shuffle-hash joins in the
query plan. If this
- // introduces extra shuffles, we give up the optimization and return the
original query plan, or
- // accept the extra shuffles if the force-apply config is true.
- // TODO: It's possible that only one skewed join in the query plan leads
to extra shuffles and
- // we only need to skip optimizing that join. We should make the
strategy smarter here.
- val optimized = optimizeSkewJoin(plan)
- val requirementSatisfied = if
(ensureRequirements.requiredDistribution.isDefined) {
- ValidateRequirements.validate(optimized,
ensureRequirements.requiredDistribution.get)
- } else {
- ValidateRequirements.validate(optimized)
+ private def requirementsSatisfied(plan: SparkPlan): Boolean = {
+ ensureRequirements.requiredDistribution match {
+ case Some(distribution) => ValidateRequirements.validate(plan,
distribution)
+ case None => ValidateRequirements.validate(plan)
}
- if (requirementSatisfied) {
- optimized.transform {
- case SkewJoinChildWrapper(child) => child
+ }
+
+ private def stripSkewJoinChildWrappers(plan: SparkPlan): SparkPlan = {
+ plan.transform { case SkewJoinChildWrapper(child) => child }
+ }
+
+ /**
+ * Applies the candidates one at a time, in the order a bottom-up pass
visits them, keeping a
+ * candidate when the plan holding the kept ones and this one still
satisfies every requirement.
+ * A split the plan cannot take therefore costs only its own join, where the
whole-plan fast path
+ * would have cost every join.
+ *
+ * Each trial is rebuilt from `plan` rather than from the previous trial, so
acceptance never
+ * depends on how a rebuild kept node instances; the handful of skewed joins
keeps that cheap.
+ */
+ private def applySkewJoinsSatisfyingRequirements(
+ plan: SparkPlan,
+ candidates: Seq[SkewJoinCandidate]): SparkPlan = {
+ var accepted = Seq.empty[SkewJoinCandidate]
+ candidates.foreach { candidate =>
+ val trial = applySkewJoins(plan, accepted :+ candidate)
+ if (requirementsSatisfied(trial)) {
+ accepted = accepted :+ candidate
+ } else {
+ logDebug(s"Skip splitting the skewed ${candidate.joinType} join " +
+ s"${candidate.join.nodeName}: with it the plan fails
ValidateRequirements")
}
+ }
+ logDebug(s"Optimized ${accepted.length} of ${candidates.length} skewed
joins")
+ stripSkewJoinChildWrappers(applySkewJoins(plan, accepted))
+ }
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.getConf(SQLConf.SKEW_JOIN_ENABLED)) {
+ return plan
+ }
+
+ // Try every split at once first: one pass to find the candidates and one
validation cover the
+ // common case. When that breaks a requirement of the plan, either accept
the extra shuffles if
+ // the force-apply config is true, or keep only the joins whose splits the
plan can take. The
+ // candidates are collected once, so the fallback reuses the splits the
fast path built.
+ val candidates = collectSplitSkewJoins(plan)
+ val optimized = applySkewJoins(plan, candidates)
+ if (requirementsSatisfied(optimized)) {
+ stripSkewJoinChildWrappers(optimized)
} else if (conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN)) {
- ensureRequirements.apply(optimized).transform {
- case SkewJoinChildWrapper(child) => child
- }
+ stripSkewJoinChildWrappers(ensureRequirements.apply(optimized))
} else {
- plan
+ applySkewJoinsSatisfyingRequirements(plan, candidates)
Review Comment:
**When this fallback keeps one union branch's split and rejects another,
`CoalesceShufflePartitions` skips coalescing for every stage under that union,
while before this PR nothing there was split and all of them were coalesced.**
Take the first new test's query. `StampUnionDecisions` runs before this rule
(AdaptiveSparkPlanExec.scala:155 and 175), when both branches still report
`hashpartitioning(key1)`, so the union is stamped non-plain. Here join2's split
is kept and join1's is rejected, so the union then reports
`UnknownPartitioning` and `unionRDDs` concatenates. But
`childrenNeedCompatiblePartitioning` reads the stamp (`case u: UnionExec =>
!u.isPlainUnion` at CoalesceShufflePartitions.scala:189, which reads it since
SPARK-59122). So join1's two stages, which carry no partition specs, and
join2's two skew reads form one coalesce group, and
`coalescePartitionsWithSkew` logs "Could not apply partition coalescing because
of missing MapOutputStatistics or shuffle partition specs.
" and gives up (ShufflePartitionsUtil.scala:135-140). Results are unchanged,
but join1's branch and join2's non-skewed partitions now run at the full
shuffle partition count, while join1, still unsplit, keeps its straggler. The
third case of SPARK-37652 takes the same path. Before this PR, a union could
already reach this with a split join next to a branch that was never skewed,
but now the shape this PR targets reaches it too. Only master and branch-4.x
are affected, since branch-4.2 and earlier have `case _: UnionExec => false`
there. I traced this by reading the code and did not run it. Could
`CoalesceShufflePartitions` group a union's children only while the union still
reports a real partitioning, e.g. `case u: UnionExec => !u.isPlainUnion &&
!u.outputPartitioning.isInstanceOf[UnknownPartitioning]` as its comment at
CoalesceShufflePartitions.scala:185-188 describes, here or in a separate JIRA,
with the first test asserting that join1's children get an `AQEShuffleReadExec`?
##########
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) {
Review Comment:
Nit: These lines are
`findTopLevelAggregate(adaptivePlan).flatMap(findTopLevelSortMergeJoin)`
(L181-185 and L157-161), and `collect(result) { case j: SortMergeJoinExec => j
}` at L2038 is `findTopLevelSortMergeJoin(result)`. L1945 and L1976 already use
the helper. Could both use the helpers, here in the form my item 3 suggests?
##########
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])))
Review Comment:
Nit: This is `AQEShuffleReadExec.hasSkewedPartition`
(AQEShuffleReadExec.scala:144-145), which this suite already uses, e.g. at
L2605 and L2660. Could this be `}.exists(_.hasSkewedPartition))`?
##########
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:
Minor: The PR description predates this test and 351074c. It still says "Two
new tests in `AdaptiveQueryExecSuite`" and "Full suite 148/148 on the rebased
head", while this commit adds this third test and the suite now has 149 tests,
and its implementation notes do not say that `applySkewJoins` now walks
top-down so that a join nested in a broadcast-hash join's build side keeps its
split. It also says only the first arm of the union test passes on the base
code, but the plain arm passes there too, as the control. Could the description
be updated, since it becomes the commit message?
##########
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:
Nit: `joinType` repeats `join.joinType`, since all three joins extend
`BaseJoinExec` (BaseJoinExec.scala:27-28). Its only reader is the debug log at
L321-322, and the broadcast arms changed `_` to a `joinType` binder (L243 and
L249) only to fill it. That log also cannot tell joins apart: in the union
test's windowed arm it prints "Skip splitting the skewed Inner join
SortMergeJoin: ..." twice. Could this be `SkewJoinCandidate(join: BaseJoinExec,
replacement: SparkPlan)`, with the log printing
`candidate.join.simpleString(conf.maxToStringFields)`, which shows the keys and
the join type?
##########
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:
Nit: A sort-merge candidate's children are the `SortExec`s at L223-224, with
the stages under them; only a shuffled-hash candidate has stages as children.
Two more new doc comments are off: L260 still says "the bottom-up order a
`transformUp` visits them", though this file no longer uses `transformUp` (the
order is `foreachUp`'s), and L215, "The candidate `plan` stands for, when it is
...", does not say that it returns `None` otherwise. Could these read "a
shuffled-hash candidate has shuffle stages as children and a sort-merge one has
sorts over them, so neither contains another candidate", "in the bottom-up
order `foreachUp` visits them", and "Returns the candidate for `plan` when it
is ..., and `None` otherwise"?
##########
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:
Nit: `runAdaptiveAndVerifyResult` already collects the AQE answer and
compares it with the non-AQE one (L119-123), and this rule changes only the AQE
plan, so a wrong split already fails there. This line, L1999 and L2008 run each
query once more and collect 334,001 or 668,000 rows to the driver in a
`@SlowSQLTest` suite, against constants derived by hand. Could these three
blocks go, keeping the plan assertions?
--
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]
