peter-toth commented on code in PR #59703:
URL: https://github.com/apache/spark/pull/59703#discussion_r4358209039
##########
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:
**Finding 3.** This config is only read when
`spark.sql.adaptive.forceOptimizeSkewedJoin` is false
(`OptimizeSkewedJoin.scala:341-344`). With the force config on, the plan gets
the extra shuffles, so "such a plan is left unchanged" doesn't hold there.
`ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN` is defined further down this object, so
the key needs to be a literal here.
```suggestion
"every join. When false, such a plan is left unchanged. Only takes
effect when " +
"'spark.sql.adaptive.forceOptimizeSkewedJoin' is false.")
```
##########
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:
**Finding 6.** The `isPlainUnion` clause is redundant. A plain union reports
`super.outputPartitioning`, which is `UnknownPartitioning(0)`
(`basicPhysicalOperators.scala:1377`, `SparkPlan.scala:163`). So the second
clause already returns `false` for it.
It doesn't save a `rawPartitioning` derivation either:
- A stamped plain union answers `outputPartitioning` from the stamp, without
deriving it.
- An unstamped union derives it in `isPlainUnion` already. A non-plain one
then derives it a second time in `outputPartitioning`.
```suggestion
case u: UnionExec =>
!u.outputPartitioning.isInstanceOf[UnknownPartitioning]
```
##########
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:
**Finding 1.** `accepted` keeps the candidates in order. So when every
earlier candidate has been kept, the last trial is `applySkewJoins(plan,
candidates)`, which is the `optimized` plan L339 just rejected. With a single
candidate that is the only trial, and one skewed join under an aggregate on its
key is the usual way for the fast path to fail. Such a plan now gets rebuilt
and validated twice per re-optimization, and then `plan` comes back as before.
A probe that logs the skip fires on the `GROUP BY` query of `SPARK-30524`
(trial 1 of 1). It also fires on the third query of `SPARK-37652` (trial 2 of
2), in three re-optimizations. With the change below, the three new tests,
`SPARK-30524` and `SPARK-37652` still pass:
```scala
candidates.zipWithIndex.foreach { case (candidate, i) =>
// With every earlier candidate kept, the last trial is the plan the
fast path rejected.
val rejectedAlready = i == candidates.length - 1 && accepted.length ==
i
if (!rejectedAlready &&
requirementsSatisfied(applySkewJoins(plan, accepted :+
candidate))) {
accepted = accepted :+ candidate
} else {
logDebug(...)
}
}
logDebug(...)
if (accepted.isEmpty) plan else
stripSkewJoinChildWrappers(applySkewJoins(plan, accepted))
```
The [spark-test
review](https://github.com/apache/spark/pull/59703#pullrequestreview-5613593139)
already suggests reusing the last accepted trial instead of rebuilding it at
L325.
##########
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:
**Finding 4.** This still lists "its join type", but `98db648` dropped
`joinType` from the case class.
```suggestion
// A skewed join `OptimizeSkewedJoin` can split: the join node in the plan,
and the join it is
// replaced by, with the skewed side split.
```
##########
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:
**Finding 5.** `forall` is true on an empty result, so this arm would still
pass if the query stopped planning sort-merge joins. The other arms check the
count first (L1942, L1951, and `splitJoins` at L1987). The plan with the
fallback off holds two joins, and this passes at the current head:
```suggestion
val joins = findTopLevelSortMergeJoin(withoutFallback)
assert(joins.size == 2 && joins.forall(!_.isSkewJoin),
withoutFallback.toString)
```
--
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]
