ulysses-you commented on code in PR #59703:
URL: https://github.com/apache/spark/pull/59703#discussion_r4373287001
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala:
##########
@@ -211,65 +211,137 @@ case class OptimizeSkewedJoin(ensureRequirements:
EnsureRequirements)
}
}
- def optimizeSkewJoin(plan: SparkPlan): SparkPlan = {
- val localShuffleReaderEnabled =
conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
- plan.transformUp {
+ /**
+ * Returns the candidate for `plan` when it is a sort-merge, shuffled-hash
or broadcast-hash join
+ * with a skewed shuffled side this rule can split, and `None` otherwise.
+ */
+ 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, 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, 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, _,
ShuffleStage(left: ShuffleQueryStageExec), _, _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(left).map {
- case newLeft => bhj.copy(left = newLeft, isSkewJoin = true)
- }.getOrElse(bhj)
+ case newLeft => SkewJoinCandidate(bhj, bhj.copy(left = newLeft,
isSkewJoin = true))
+ }
case bhj @ BroadcastHashJoinExec(_, _, _, BuildLeft, _, _,
ShuffleStage(right: ShuffleQueryStageExec), _, false) if
!localShuffleReaderEnabled =>
tryOptimizeBroadcastHashJoinStreamedPlan(right).map {
- case newRight => bhj.copy(right = newRight, isSkewJoin = true)
- }.getOrElse(bhj)
+ case newRight => SkewJoinCandidate(bhj, bhj.copy(right = newRight,
isSkewJoin = true))
+ }
+
+ case _ => None
}
+
+ /**
+ * Every join in `plan` this rule can split, in the bottom-up order
`foreachUp` 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 shuffled-hash candidate has shuffle stages as
children and a
+ * sort-merge one has sorts over them, so neither contains another
candidate; 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 =>
Review Comment:
Applied in 73738e1d935: with every earlier candidate kept, the last trial is
the plan the fast path rejected, so it is skipped without rebuilding or
validating, and nothing is rebuilt when no candidate is kept. The two shapes
you named, `SPARK-30524`'s GROUP BY query and `SPARK-37652`'s third query,
still pass.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1559,6 +1559,17 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val SKEW_JOIN_PER_JOIN_FALLBACK_ENABLED =
+ buildConf("spark.sql.adaptive.skewJoin.perJoinFallback.enabled")
+ .internal()
+ .doc("When true, a plan whose all-at-once skew join optimization breaks
a requirement " +
+ "keeps the splits the plan can still satisfy, instead of giving up the
optimization for " +
+ "every join. When false, such a plan is left unchanged.")
Review Comment:
Added in 73738e1d935: the doc now ends with 'Only takes effect when
'spark.sql.adaptive.forceOptimizeSkewedJoin' is false.' with the key as a
literal.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala:
##########
@@ -285,6 +357,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.
Review Comment:
Updated in 73738e1d935 as suggested.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -1901,6 +1901,154 @@ 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 rejectedJoins = collect(adaptivePlan) {
+ case agg: BaseAggregateExec => findTopLevelSortMergeJoin(agg)
+ }.flatten.distinct
+ assert(rejectedJoins.size == 1, adaptivePlan.toString)
+ val rejected = rejectedJoins.head
+ assert(!rejected.isSkewJoin)
+ // The rejected join does not lose partition coalescing to the kept
split.
+ val rejectedReads = collect(rejected) { case r: AQEShuffleReadExec =>
r }
+ assert(rejectedReads.size == 2 &&
rejectedReads.forall(_.hasCoalescedPartition),
+ adaptivePlan.toString)
+
+ val keptJoins = findTopLevelSortMergeJoin(adaptivePlan).filterNot(_ eq
rejected)
+ assert(keptJoins.size == 1, adaptivePlan.toString)
+ assert(keptJoins.head.isSkewJoin)
+ assert(collect(keptJoins.head) {
+ case r: AQEShuffleReadExec => r
+ }.exists(_.hasSkewedPartition))
+ assert(ValidateRequirements.validate(adaptivePlan))
+
+ // Turning the fallback off gives the plan back to the all-or-nothing
path.
+ withSQLConf(SQLConf.SKEW_JOIN_PER_JOIN_FALLBACK_ENABLED.key ->
"false") {
+ val (_, withoutFallback) = runAdaptiveAndVerifyResult(query)
+
assert(findTopLevelSortMergeJoin(withoutFallback).forall(!_.isSkewJoin))
Review Comment:
Fixed in 73738e1d935: the arm asserts the two joins are present before
checking that neither is split.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala:
##########
@@ -186,7 +186,11 @@ case class CoalesceShufflePartitions(session:
SparkSession) extends AQEShuffleRe
// partitioning relied on by downstream operators, so its children must
remain co-partitioned
// -- coalesce them as a single group. When it reports UnknownPartitioning
(config off, or
// children not compatibly partitioned), the children may still be
coalesced independently.
- case u: UnionExec => !u.isPlainUnion
+ // The stamped decision alone is not enough: a rule that runs after the
stamp, such as
+ // `OptimizeSkewedJoin` splitting one branch, can take the layout away,
and the union then
+ // concatenates whatever its children produce.
+ case u: UnionExec =>
+ !u.isPlainUnion &&
!u.outputPartitioning.isInstanceOf[UnknownPartitioning]
Review Comment:
Simplified in 73738e1d935 as suggested. The comment now says a union that
reports a concrete layout needs its children co-partitioned, and that the
answer is read now because a rule after the stamp can take the layout away.
--
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]
