spark-test commented on code in PR #59703:
URL: https://github.com/apache/spark/pull/59703#discussion_r4347882662
##########
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:
**Nit (P3):** This identity lookup assumes the candidate instance survives
the bottom-up rebuild, and the scaladoc above states that no candidate contains
another. That holds for the SMJ/SHJ arms, whose children are query stages. It
does not hold for the BroadcastHashJoinExec arms, which only constrain the
streamed side; the build side is matched by `_`. With
`spark.sql.adaptive.localShuffleReader.enabled=false`, a BHJ whose build side
is a BroadcastExchangeExec over a skewed SMJ with materialized stages yields
two candidates, inner first. `reOptimize` runs this rule before the broadcast
is staged, so this shape can occur. Once the inner SMJ is replaced,
`transformUpWithPruning` (TreeNode.scala:633-634) passes the rule a new BHJ
instance, `eq` misses, and the outer split is silently dropped. The fallback
also counts it as accepted. Even if the identity matched, the prebuilt
`replacement` would carry the original build subtree and discard the inner
split. The previous `optimizeSkew
Join` pattern-matched the rebuilt node and split both. Results stay correct,
but the streamed side's skewed partitions now run as stragglers in this shape.
**Verification:**
- **Regression:** An AQE plan with local shuffle reader disabled, in which a
broadcast-hash join's streamed side is a skewed materialized shuffle stage and
whose build side holds a skewed sort-merge join over materialized stages, ends
with both joins marked as skew joins, and the query results match the
non-adaptive run.
- **Behavior:** The existing and new per-join fallback cases still split
exactly the joins they expect, and the fast and force paths produce the same
plans as before for non-nested shapes.
##########
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:
**Nit (P3):** 'Base' only means something while this PR is under review.
After merge, readers cannot tell what it contrasts with. Consider wording it as
the all-at-once attempt (and the previous all-or-nothing rule) also rejecting
this plan, with this arm pinning that the per-join fallback keeps no split the
union's claim cannot take.
**Verification:**
- **Inspection:** The union pass-through layout test comment no longer
refers to a 'Base' implementation and still describes why neither join is split
when the window reads the union's layout.
--
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]
