You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "sunchao (via GitHub)" <gi...@apache.org> on 2026/03/20 18:11:15 UTC
[PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
sunchao opened a new pull request, #54925:
URL: https://github.com/apache/spark/pull/54925
### What changes were proposed in this pull request?
This PR adds an opt-in AQE fallback path for broadcast join failures caused by broadcast table size/row limits.
When `spark.sql.adaptive.broadcastJoin.fallbackToShuffle.enabled` is enabled, AQE catches qualifying broadcast stage failures and retries adaptive replanning with broadcast joins disabled, so the query can proceed with shuffle-based joins (for example SMJ/SHJ) instead of failing immediately.
The change also includes safety checks to avoid repeatedly retrying the same failed broadcast relation, and test coverage in `AdaptiveQueryExecSuite`.
### Why are the changes needed?
Today, a query can fail if a planned broadcast side exceeds runtime broadcast limits, even though a shuffle join strategy could still complete successfully.
This PR provides a controlled fallback path to improve robustness for those cases while keeping existing behavior unchanged by default.
### Does this PR introduce _any_ user-facing change?
Yes.
A new SQL config is added:
- `spark.sql.adaptive.broadcastJoin.fallbackToShuffle.enabled` (default: `false`)
When enabled, queries that would otherwise fail due to broadcast table row/size limits may instead continue with shuffle-based joins. When disabled, existing behavior remains unchanged.
### How was this patch tested?
- Added/updated unit tests in `sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala`, including:
- fallback enable/disable behavior
- preserving required input broadcast exchange behavior
- rejecting fallback replans that still contain the failed broadcast relation
- Ran AQE adaptive query test coverage relevant to the new fallback path.
- Manually validated with real jobs that the plan can transition from initial BHJ to final SMJ under AQE when fallback is enabled.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
Yes, Codex 5.3 High with a lot of harnessing
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sigmod (via GitHub)" <gi...@apache.org>.
sigmod commented on PR #54925:
URL: https://github.com/apache/spark/pull/54925#issuecomment-4164744775
@maryannxue @liuzqt
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2978409757
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -790,30 +840,63 @@ case class AdaptiveSparkPlanExec(
/**
* Re-optimize and run physical planning on the current logical plan based on the latest stats.
*/
- private def reOptimize(logicalPlan: LogicalPlan): Option[(SparkPlan, LogicalPlan)] = {
+ private def reOptimize(
+ logicalPlan: LogicalPlan,
+ disableBroadcastJoin: Boolean): Option[(SparkPlan, LogicalPlan)] = {
try {
- logicalPlan.invalidateStatsCache()
- val optimized = optimizer.execute(logicalPlan)
- val sparkPlan = context.session.sessionState.planner.plan(ReturnAnswer(optimized)).next()
- val newPlan = applyPhysicalRules(
- applyQueryPostPlannerStrategyRules(sparkPlan),
- preprocessingRules ++ queryStagePreparationRules,
- Some((planChangeLogger, "AQE Replanning")))
-
- // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
- // add the `BroadcastExchangeExec` node manually in the DPP subquery,
- // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is complicated
- // and need to be re-optimized, AQE also need to manually insert the `BroadcastExchangeExec`
- // node to prevent the loss of the `BroadcastExchangeExec` node in DPP subquery.
- // Here, we also need to avoid to insert the `BroadcastExchangeExec` node when the newPlan is
- // already the `BroadcastExchangeExec` plan after apply the `LogicalQueryStageStrategy` rule.
- val finalPlan = inputPlan match {
- case b: BroadcastExchangeLike
- if (!newPlan.isInstanceOf[BroadcastExchangeLike]) => b.withNewChildren(Seq(newPlan))
- case _ => newPlan
+ val optimizedLogicalPlan = if (disableBroadcastJoin) {
+ removeBroadcastHints(logicalPlan)
+ } else {
+ logicalPlan
+ }
+ def planFn(
+ optimizerToUse: AQEOptimizer,
+ plannerToUse: SparkPlanner): (SparkPlan, LogicalPlan) = {
+ optimizedLogicalPlan.invalidateStatsCache()
+ val optimized = optimizerToUse.execute(optimizedLogicalPlan)
+ val sparkPlan = plannerToUse.plan(ReturnAnswer(optimized)).next()
+ val newPlan = applyPhysicalRules(
+ applyQueryPostPlannerStrategyRules(sparkPlan),
+ preprocessingRules ++ queryStagePreparationRules,
+ Some((planChangeLogger, "AQE Replanning")))
+
+ // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
+ // add the `BroadcastExchangeExec` node manually in the DPP subquery,
+ // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is
+ // complicated and need to be re-optimized, AQE also need to manually insert the
+ // `BroadcastExchangeExec` node to prevent the loss of the `BroadcastExchangeExec`
+ // node in DPP subquery.
+ // If this adaptive plan must still produce a broadcast result, preserve the root
+ // broadcast exchange even when join fallback disables broadcast joins inside the plan.
+ // Otherwise, avoid reinserting the exchange when the new plan already contains one
+ // after applying the `LogicalQueryStageStrategy` rule.
+ val finalPlan = inputPlan match {
+ case b: BroadcastExchangeLike
+ if !newPlan.isInstanceOf[BroadcastExchangeLike] =>
+ b.withNewChildren(Seq(newPlan))
+ case _ => newPlan
+ }
+ (finalPlan, optimized)
}
- Some((finalPlan, optimized))
+ val result = if (disableBroadcastJoin) {
+ // Use cloneSession to preserve current session state (e.g. runtime SQLConf and planner
+ // customizations) while keeping fallback conf changes isolated from the original session.
+ val fallbackSession = context.session.cloneSession()
+ val fallbackConf = fallbackSession.sessionState.conf
+ fallbackConf.setConfString(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
Review Comment:
OK I've updated the PR to use `NO_BROADCAST_HASH` - please take a look!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2978405575
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +943,179 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Check error-classes.json for details
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) {
+ stagesToReplace
+ } else {
+ stagesToReplace.filterNot { stage =>
+ failedStages.exists(_.eq(stage)) || (stage match {
+ case broadcastStage: BroadcastQueryStageExec =>
+ failedStages.exists {
+ case failedBroadcast: BroadcastQueryStageExec =>
+ sameBroadcastRelation(
+ broadcastStage.broadcast.child, failedBroadcast.broadcast.child)
+ case _ => false
+ }
+ case _ => false
+ })
+ }
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && plan.exists { p =>
+ broadcastChildPlan(p).exists { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(child, failed.broadcast.child)
+ }
+ }
+ }
+ }
+
+ // Track failed broadcast relations using semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed relation.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ stage: QueryStageExec): Unit = {
+ stage match {
+ case failedStage: BroadcastQueryStageExec =>
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastRelation(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ case _ =>
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ Seq(exchangeStage.plan.canonicalized, exchangeStage.canonicalized).distinct.foreach { key =>
+ context.stageCache.remove(key)
+ }
+ val keysToRemove = context.stageCache.collect {
Review Comment:
Updated.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #54925:
URL: https://github.com/apache/spark/pull/54925#issuecomment-4216254367
Also cc @viirya @cloud-fan @dbtsai to hear your opinions on this PR.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3017708774
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +916,210 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Check error-classes.json for details
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) return stagesToReplace
+
+ val failedBroadcastStages = failedStages.collect {
+ case b: BroadcastQueryStageExec => b
+ }
+
+ stagesToReplace.filterNot {
+ case stage if failedStages.exists(_ eq stage) => true
+ case stage: BroadcastQueryStageExec =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(stage.broadcast.child, failed.broadcast.child)
+ }
+ case _ => false
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ // Check if a Spark physical plan contains any of the previous failed broadcast stages
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && plan.exists { p =>
+ broadcastChildPlan(p).exists { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(child, failed.broadcast.child)
+ }
+ }
+ }
+ }
+
+ // Track failed broadcast relations using semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed relation.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ stage: QueryStageExec): Unit = {
+ stage match {
+ case failedStage: BroadcastQueryStageExec =>
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastRelation(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ case _ =>
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ val planKey = exchangeStage.plan.canonicalized
+ val stageKey = exchangeStage.canonicalized
+
+ context.stageCache.remove(planKey)
+ if (stageKey != planKey) {
+ context.stageCache.remove(stageKey)
+ }
+ case _ =>
+ }
+
+ // Normalize different wrappers to the underlying broadcast child plan when present.
+ private def broadcastChildPlan(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case stage: BroadcastQueryStageExec => Some(stage.broadcast.child)
+ case exchange: BroadcastExchangeLike => Some(exchange.child)
+ case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => Some(exchange.child)
+ case _ => None
+ }
+
+ // Compare broadcast inputs by semantic plan equivalence.
+ // Do not fall back to output-attribute equality: different plans can share
+ // output schemas while producing different relations.
+ private def sameBroadcastRelation(left: SparkPlan, right: SparkPlan): Boolean = {
+ left.sameResult(right)
+ }
+
+ // Root broadcast stages must be preserved (e.g. broadcast subquery output contract).
+ // Only exempt the actual root broadcast stage instance, not semantically equivalent relations.
+ private def isRequiredRootBroadcastStage(stage: QueryStageExec): Boolean = stage match {
+ case broadcastStage: BroadcastQueryStageExec if inputPlan.isInstanceOf[BroadcastExchangeLike] =>
+ @tailrec
+ def findRootBroadcast(current: SparkPlan): Option[BroadcastQueryStageExec] = current match {
+ case root: BroadcastQueryStageExec => Some(root)
+ case resultStage: ResultQueryStageExec => findRootBroadcast(resultStage.plan)
+ case _ => None
+ }
+ findRootBroadcast(currentPhysicalPlan).exists(_ eq broadcastStage)
+ case _ => false
+ }
+
+ // Apply side-specific NO_BROADCAST_HASH hints for relations that previously failed
+ // broadcast stage materialization, so targeted fallback can keep unrelated BHJs.
+ private def addNoBroadcastHashHintsForFailedRelations(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ val failedLogicalPlans = extractFailedBroadcastLogicalPlans(failedBroadcastStages)
+ if (failedLogicalPlans.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case join: Join =>
+ val disableLeft = shouldDisableBroadcastForJoinSide(join.left, failedLogicalPlans)
+ val disableRight = shouldDisableBroadcastForJoinSide(join.right, failedLogicalPlans)
+ if (!disableLeft && !disableRight) {
+ join
+ } else {
+ val newHint = JoinHint(
+ if (disableLeft) {
+ toNoBroadcastHashHint(join.hint.leftHint)
+ } else {
+ join.hint.leftHint
+ },
+ if (disableRight) {
+ toNoBroadcastHashHint(join.hint.rightHint)
+ } else {
+ join.hint.rightHint
+ })
+ if (newHint != join.hint) join.copy(hint = newHint) else join
+ }
+ }
+ }
+ }
+
+ private def extractFailedBroadcastLogicalPlans(
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Seq[LogicalPlan] = {
+ val failedLogicalPlans = new mutable.ArrayBuffer[LogicalPlan]()
+ failedBroadcastStages.foreach { stage =>
+ Seq(stage.broadcast.child.logicalLink, stage.broadcast.logicalLink, stage.logicalLink)
+ .flatten
+ .foreach { logicalPlan =>
+ if (!failedLogicalPlans.exists(_.sameResult(logicalPlan))) {
+ failedLogicalPlans.append(logicalPlan)
+ }
+ }
+ }
+ failedLogicalPlans.toSeq
+ }
+
+ private def shouldDisableBroadcastForJoinSide(
Review Comment:
Good point. Fixed by normalizing both the failed logical plans and the candidate join side to strip nested `LogicalQueryStage` wrappers before doing the comparison. Also added some tests to cover this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3066172178
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -328,14 +344,43 @@ case class AdaptiveSparkPlanExec(
val nextMsg = events.take()
val rem = new util.ArrayList[StageMaterializationEvent]()
events.drainTo(rem)
- (Seq(nextMsg) ++ rem.asScala).foreach {
+ val stageEvents = Seq(nextMsg) ++ rem.asScala
+ val fallbackStages = new mutable.ArrayBuffer[QueryStageExec]()
Review Comment:
Extracted some code into a separate method. Let me know what you think.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3002118799
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +916,210 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Check error-classes.json for details
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) return stagesToReplace
+
+ val failedBroadcastStages = failedStages.collect {
+ case b: BroadcastQueryStageExec => b
+ }
+
+ stagesToReplace.filterNot {
+ case stage if failedStages.exists(_ eq stage) => true
+ case stage: BroadcastQueryStageExec =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(stage.broadcast.child, failed.broadcast.child)
+ }
+ case _ => false
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ // Check if a Spark physical plan contains any of the previous failed broadcast stages
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && plan.exists { p =>
+ broadcastChildPlan(p).exists { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(child, failed.broadcast.child)
+ }
+ }
+ }
+ }
+
+ // Track failed broadcast relations using semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed relation.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ stage: QueryStageExec): Unit = {
+ stage match {
+ case failedStage: BroadcastQueryStageExec =>
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastRelation(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ case _ =>
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ val planKey = exchangeStage.plan.canonicalized
+ val stageKey = exchangeStage.canonicalized
+
+ context.stageCache.remove(planKey)
+ if (stageKey != planKey) {
+ context.stageCache.remove(stageKey)
+ }
+ case _ =>
+ }
+
+ // Normalize different wrappers to the underlying broadcast child plan when present.
+ private def broadcastChildPlan(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case stage: BroadcastQueryStageExec => Some(stage.broadcast.child)
+ case exchange: BroadcastExchangeLike => Some(exchange.child)
+ case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => Some(exchange.child)
+ case _ => None
+ }
+
+ // Compare broadcast inputs by semantic plan equivalence.
+ // Do not fall back to output-attribute equality: different plans can share
+ // output schemas while producing different relations.
+ private def sameBroadcastRelation(left: SparkPlan, right: SparkPlan): Boolean = {
+ left.sameResult(right)
+ }
+
+ // Root broadcast stages must be preserved (e.g. broadcast subquery output contract).
+ // Only exempt the actual root broadcast stage instance, not semantically equivalent relations.
+ private def isRequiredRootBroadcastStage(stage: QueryStageExec): Boolean = stage match {
+ case broadcastStage: BroadcastQueryStageExec if inputPlan.isInstanceOf[BroadcastExchangeLike] =>
+ @tailrec
+ def findRootBroadcast(current: SparkPlan): Option[BroadcastQueryStageExec] = current match {
+ case root: BroadcastQueryStageExec => Some(root)
+ case resultStage: ResultQueryStageExec => findRootBroadcast(resultStage.plan)
+ case _ => None
+ }
+ findRootBroadcast(currentPhysicalPlan).exists(_ eq broadcastStage)
+ case _ => false
+ }
+
+ // Apply side-specific NO_BROADCAST_HASH hints for relations that previously failed
+ // broadcast stage materialization, so targeted fallback can keep unrelated BHJs.
+ private def addNoBroadcastHashHintsForFailedRelations(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ val failedLogicalPlans = extractFailedBroadcastLogicalPlans(failedBroadcastStages)
+ if (failedLogicalPlans.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case join: Join =>
+ val disableLeft = shouldDisableBroadcastForJoinSide(join.left, failedLogicalPlans)
+ val disableRight = shouldDisableBroadcastForJoinSide(join.right, failedLogicalPlans)
+ if (!disableLeft && !disableRight) {
+ join
+ } else {
+ val newHint = JoinHint(
+ if (disableLeft) {
+ toNoBroadcastHashHint(join.hint.leftHint)
+ } else {
+ join.hint.leftHint
+ },
+ if (disableRight) {
+ toNoBroadcastHashHint(join.hint.rightHint)
+ } else {
+ join.hint.rightHint
+ })
+ if (newHint != join.hint) join.copy(hint = newHint) else join
+ }
+ }
+ }
+ }
+
+ private def extractFailedBroadcastLogicalPlans(
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Seq[LogicalPlan] = {
+ val failedLogicalPlans = new mutable.ArrayBuffer[LogicalPlan]()
+ failedBroadcastStages.foreach { stage =>
+ Seq(stage.broadcast.child.logicalLink, stage.broadcast.logicalLink, stage.logicalLink)
+ .flatten
+ .foreach { logicalPlan =>
+ if (!failedLogicalPlans.exists(_.sameResult(logicalPlan))) {
+ failedLogicalPlans.append(logicalPlan)
+ }
+ }
+ }
+ failedLogicalPlans.toSeq
+ }
+
+ private def shouldDisableBroadcastForJoinSide(
Review Comment:
The tests for `addNoBroadcastHashHintsForFailedRelations` invoke it directly on the original logical plan, bypassing the AQE loop. I think this is fine for the common case (broadcast side is a simple scan), but wondering if it could miss when the broadcast side is a complex subplan that already has materialized `LogicalQueryStage` nodes mixed in.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3059775685
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -328,14 +344,43 @@ case class AdaptiveSparkPlanExec(
val nextMsg = events.take()
val rem = new util.ArrayList[StageMaterializationEvent]()
events.drainTo(rem)
- (Seq(nextMsg) ++ rem.asScala).foreach {
+ val stageEvents = Seq(nextMsg) ++ rem.asScala
+ val fallbackStages = new mutable.ArrayBuffer[QueryStageExec]()
Review Comment:
Is it possible to extract fallback logic into a method? The while-loop is very complicated now.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3059861222
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -624,6 +637,814 @@ class AdaptiveQueryExecSuite
}
}
+ test("Fallback to shuffled join should apply NO_BROADCAST_HASH only to failed relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+
+ val joinsWithNoBroadcastHash = targetedLogicalPlan.collect {
+ case join: Join
+ if join.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) ||
+ join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) =>
+ join
+ }
+ assert(joinsWithNoBroadcastHash.nonEmpty)
+ assert(targetedLogicalPlan.collect {
+ case join: Join
+ if !join.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) &&
+ !join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) =>
+ join
+ }.nonEmpty)
+ }
+ }
+
+ test("Fallback to shuffled join should preserve non-broadcast side hint strategies") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val toNoBroadcastHashHint =
+ PrivateMethod[Option[HintInfo]](Symbol("toNoBroadcastHashHint"))
+
+ val noBroadcastAndReplicationHint =
+ Some(HintInfo(strategy = Some(NO_BROADCAST_AND_REPLICATION)))
+ val preservedNoBroadcastAndReplication =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastAndReplicationHint))
+ assert(preservedNoBroadcastAndReplication == noBroadcastAndReplicationHint)
+
+ val noBroadcastHashHint = Some(HintInfo(strategy = Some(NO_BROADCAST_HASH)))
+ val preservedNoBroadcastHash =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastHashHint))
+ assert(preservedNoBroadcastHash == noBroadcastHashHint)
+
+ val broadcastHint = Some(HintInfo(strategy = Some(BROADCAST)))
+ val rewrittenBroadcastHint = adaptivePlan.invokePrivate(toNoBroadcastHashHint(broadcastHint))
+ assert(rewrittenBroadcastHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+
+ Seq(PREFER_SHUFFLE_HASH, SHUFFLE_HASH, SHUFFLE_MERGE).foreach { strategy =>
+ val preservedHint =
+ Some(HintInfo(strategy = Some(strategy)))
+ assert(adaptivePlan.invokePrivate(toNoBroadcastHashHint(preservedHint)) == preservedHint)
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJs after targeted fallback replan") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val baselinePlan = adaptivePlan.invokePrivate(reOptimize(logicalPlan)).get._1
+ assert(adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(baselinePlan, Seq(failedStage))))
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+
+ assert(findTopLevelBroadcastHashJoin(targetedPlan).nonEmpty)
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should not disable other broadcasts from same base relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(failed_t2), BROADCAST(other_t2) */ failed_t2.failed_b, " +
+ "other_t2.other_a " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1) failed_t2 " +
+ "ON t1.key = failed_t2.failed_a " +
+ "JOIN (SELECT a AS other_a FROM testData2 WHERE a <= 2) other_t2 " +
+ "ON t1.key = other_t2.other_a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1")
+ .queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+ val broadcastBuildOutput = finalBhjs.headOption.map { bhj =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left.output.map(_.name)
+ case BuildRight => bhj.right.output.map(_.name)
+ }
+ }.getOrElse(Nil)
+
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(broadcastBuildOutput.contains("other_a"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should match failed relation with nested logical query stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val joinDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a, b FROM testData2 WHERE a > 1) t2 " +
+ "ON t1.key = t2.a")
+ val adaptivePlan = joinDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val relationLogicalPlan = sql("SELECT a, b FROM testData2").queryExecution.optimizedPlan
+ val relationPhysicalPlan = sql("SELECT a, b FROM testData2").queryExecution.sparkPlan
+ val materializedRelation = LogicalQueryStage(
+ relationLogicalPlan,
+ ResultQueryStageExec(0, relationPhysicalPlan, _ => ()))
+
+ val logicalPlanWithStage = joinDf.queryExecution.optimizedPlan.transformDown {
+ case plan if plan.sameResult(relationLogicalPlan) => materializedRelation
+ }
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a, b FROM testData2 WHERE a > 1").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlanWithStage, Seq(failedStage)))
+ val join = targetedLogicalPlan.collectFirst { case j: Join => j }.get
+
+ assert(join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJ for nested failed branch replan") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 5)
+ .selectExpr("id AS k", "concat('small_', cast(id AS string)) AS tag")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t3) */ failed_branch.k, failed_branch.payload, " +
+ "fallback_t3.tag " +
+ "FROM (" +
+ " SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, fallback_t2.payload " +
+ " FROM fallback_t1 " +
+ " JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k" +
+ ") failed_branch " +
+ "JOIN fallback_t3 ON failed_branch.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+
+ val failedBroadcastChild = sql("SELECT k FROM fallback_t2").queryExecution.sparkPlan
+ val failedExchange = adaptivePlan.executedPlan.collectFirst {
+ case exchange: BroadcastExchangeExec if exchange.child.sameResult(failedBroadcastChild) =>
+ exchange
+ }.getOrElse {
+ fail(s"Expected broadcast exchange for fallback_t2 branch:\n${adaptivePlan.executedPlan}")
+ }
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))))
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))))
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(finalBhjs.head.output.exists(_.name == "tag"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should only apply to exclusively broadcast hash join stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true") {
+ def buildBroadcastStage(join: BaseJoinExec): BroadcastQueryStageExec = {
+ val buildPlan = join match {
+ case bhj: BroadcastHashJoinExec =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left
+ case BuildRight => bhj.right
+ }
+ case bnlj: BroadcastNestedLoopJoinExec =>
+ bnlj.buildSide match {
+ case BuildLeft => bnlj.left
+ case BuildRight => bnlj.right
+ }
+ }
+ buildPlan.collectFirst { case stage: BroadcastQueryStageExec => stage }
+ .orElse {
+ buildPlan.collectFirst { case exchange: BroadcastExchangeExec =>
+ BroadcastQueryStageExec(0, exchange, exchange.canonicalized)
+ }
+ }.getOrElse {
+ fail(s"Expected broadcast query stage:\n$join")
+ }
+ }
+
+ val bhjDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a")
+ val bhjAdaptivePlan = bhjDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bhjStage = buildBroadcastStage(findTopLevelBroadcastHashJoin(
+ bhjAdaptivePlan.executedPlan).head)
+
+ val bnljDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key > t2.a")
+ val bnljAdaptivePlan = bnljDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bnljStage = buildBroadcastStage(findTopLevelBroadcastNestedLoopJoin(
+ bnljAdaptivePlan.executedPlan).head)
+
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val mixedDf = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ t1.key " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val mixedAdaptivePlan =
+ mixedDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val mixedBhjs = findTopLevelBroadcastHashJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBnljs = findTopLevelBroadcastNestedLoopJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBhjStage = buildBroadcastStage(mixedBhjs.head)
+ val mixedBnljStage = buildBroadcastStage(mixedBnljs.head)
+
+ assert(mixedBhjs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBnljs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast ne mixedBnljStage.broadcast,
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast.child.sameResult(mixedBnljStage.broadcast.child),
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val subqueryQe = sql("SELECT key FROM testData").queryExecution
+ val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, subqueryQe.sparkPlan)
+ inputPlan.setLogicalLink(subqueryQe.optimizedPlan)
+ val subqueryAdaptivePlan = AdaptiveSparkPlanExec(
+ inputPlan,
+ AdaptiveExecutionContext(spark, subqueryQe),
+ Seq.empty,
+ isSubquery = true)
+ val rootStage = subqueryAdaptivePlan.finalPhysicalPlan match {
+ case stage: BroadcastQueryStageExec => stage
+ case plan =>
+ fail(s"Expected root broadcast query stage:\n$plan")
+ }
+
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should be disabled when conf is false") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val broadcastInputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val stage =
+ BroadcastQueryStageExec(0, broadcastInputPlan, broadcastInputPlan.canonicalized)
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure " +
Review Comment:
These e2e tests look good. Can we add more e2e tests like,
1. self-join
```
SELECT /*+ BROADCAST(t2) */ t1.k, t2.k
FROM large_table t1
JOIN large_table t2 ON t1.k = t2.k
```
2. Two broadcasts fail
```
SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ *
FROM t1
JOIN large_t2 t2 ON t1.k = t2.k
JOIN large_t3 t3 ON t1.k = t3.k
```
large_t2 and large_t3 both are larger than size limit.
3. Check original behavior unchanged when fallback is disabled
```
withSQLConf(
ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false",
MAX_BROADCAST_TABLE_SIZE.key -> "1") {
val error = intercept[SparkException] {
sql("SELECT /*+ BROADCAST(t2) */ ... ").collect()
}
assert(containsBroadcastLimitError(error))
}
```
4. Broadcast failure in subquery
We only have private method test for isRequiredRootBroadcastStage, no e2e test for the case.
```
SELECT *
FROM t1
WHERE t1.k IN (
SELECT /*+ BROADCAST(t2) */ t2.k
FROM large_t2 t2
)
```
5. AQE can re-optimize normally after fallback
Check if AQE optimization (e.g., coalesce partitions) can work after fallback
```
withSQLConf(
COALESCE_PARTITIONS_ENABLED.key -> "true",
ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
MAX_BROADCAST_TABLE_SIZE.key -> "1") {
val df = sql("SELECT /*+ BROADCAST(t2) */ t1.k, count(*) FROM t1 JOIN t2 ... GROUP BY t1.k")
checkAnswer(df, expected)
// check final plan has coalesced partitions
}
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2972083366
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -367,7 +413,11 @@ case class AdaptiveSparkPlanExec(
currentPhysicalPlan = newPhysicalPlan
currentLogicalPlan = newLogicalPlan
stagesToReplace = Seq.empty[QueryStageExec]
+ } else if (forceAdoptBroadcastFallbackPlan) {
Review Comment:
this seems unreachable
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2976401564
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +943,179 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Check error-classes.json for details
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) {
+ stagesToReplace
+ } else {
+ stagesToReplace.filterNot { stage =>
+ failedStages.exists(_.eq(stage)) || (stage match {
+ case broadcastStage: BroadcastQueryStageExec =>
+ failedStages.exists {
+ case failedBroadcast: BroadcastQueryStageExec =>
+ sameBroadcastRelation(
+ broadcastStage.broadcast.child, failedBroadcast.broadcast.child)
+ case _ => false
+ }
+ case _ => false
+ })
+ }
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && plan.exists { p =>
+ broadcastChildPlan(p).exists { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(child, failed.broadcast.child)
+ }
+ }
+ }
+ }
+
+ // Track failed broadcast relations using semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed relation.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ stage: QueryStageExec): Unit = {
+ stage match {
+ case failedStage: BroadcastQueryStageExec =>
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastRelation(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ case _ =>
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ Seq(exchangeStage.plan.canonicalized, exchangeStage.canonicalized).distinct.foreach { key =>
+ context.stageCache.remove(key)
+ }
+ val keysToRemove = context.stageCache.collect {
Review Comment:
It is purely defensive, but I agree that it seems unlikely to happen as the cache is always key'd by canonical plan, so there shouldn't be a case where a materialized stage exec is cached under some extra keys. I can remove it to simplify this code path.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -790,30 +840,63 @@ case class AdaptiveSparkPlanExec(
/**
* Re-optimize and run physical planning on the current logical plan based on the latest stats.
*/
- private def reOptimize(logicalPlan: LogicalPlan): Option[(SparkPlan, LogicalPlan)] = {
+ private def reOptimize(
+ logicalPlan: LogicalPlan,
+ disableBroadcastJoin: Boolean): Option[(SparkPlan, LogicalPlan)] = {
try {
- logicalPlan.invalidateStatsCache()
- val optimized = optimizer.execute(logicalPlan)
- val sparkPlan = context.session.sessionState.planner.plan(ReturnAnswer(optimized)).next()
- val newPlan = applyPhysicalRules(
- applyQueryPostPlannerStrategyRules(sparkPlan),
- preprocessingRules ++ queryStagePreparationRules,
- Some((planChangeLogger, "AQE Replanning")))
-
- // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
- // add the `BroadcastExchangeExec` node manually in the DPP subquery,
- // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is complicated
- // and need to be re-optimized, AQE also need to manually insert the `BroadcastExchangeExec`
- // node to prevent the loss of the `BroadcastExchangeExec` node in DPP subquery.
- // Here, we also need to avoid to insert the `BroadcastExchangeExec` node when the newPlan is
- // already the `BroadcastExchangeExec` plan after apply the `LogicalQueryStageStrategy` rule.
- val finalPlan = inputPlan match {
- case b: BroadcastExchangeLike
- if (!newPlan.isInstanceOf[BroadcastExchangeLike]) => b.withNewChildren(Seq(newPlan))
- case _ => newPlan
+ val optimizedLogicalPlan = if (disableBroadcastJoin) {
+ removeBroadcastHints(logicalPlan)
+ } else {
+ logicalPlan
+ }
+ def planFn(
+ optimizerToUse: AQEOptimizer,
+ plannerToUse: SparkPlanner): (SparkPlan, LogicalPlan) = {
+ optimizedLogicalPlan.invalidateStatsCache()
+ val optimized = optimizerToUse.execute(optimizedLogicalPlan)
+ val sparkPlan = plannerToUse.plan(ReturnAnswer(optimized)).next()
+ val newPlan = applyPhysicalRules(
+ applyQueryPostPlannerStrategyRules(sparkPlan),
+ preprocessingRules ++ queryStagePreparationRules,
+ Some((planChangeLogger, "AQE Replanning")))
+
+ // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
+ // add the `BroadcastExchangeExec` node manually in the DPP subquery,
+ // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is
+ // complicated and need to be re-optimized, AQE also need to manually insert the
+ // `BroadcastExchangeExec` node to prevent the loss of the `BroadcastExchangeExec`
+ // node in DPP subquery.
+ // If this adaptive plan must still produce a broadcast result, preserve the root
+ // broadcast exchange even when join fallback disables broadcast joins inside the plan.
+ // Otherwise, avoid reinserting the exchange when the new plan already contains one
+ // after applying the `LogicalQueryStageStrategy` rule.
+ val finalPlan = inputPlan match {
+ case b: BroadcastExchangeLike
+ if !newPlan.isInstanceOf[BroadcastExchangeLike] =>
+ b.withNewChildren(Seq(newPlan))
+ case _ => newPlan
+ }
+ (finalPlan, optimized)
}
- Some((finalPlan, optimized))
+ val result = if (disableBroadcastJoin) {
+ // Use cloneSession to preserve current session state (e.g. runtime SQLConf and planner
+ // customizations) while keeping fallback conf changes isolated from the original session.
+ val fallbackSession = context.session.cloneSession()
+ val fallbackConf = fallbackSession.sessionState.conf
+ fallbackConf.setConfString(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
Review Comment:
This is a good suggestion. Currently this PR only does a coarse full plan fallback, so if we have multiple joins using BHJ in a query plan, even if only one BHJ failed due to the limits, other BHJs will also be disabled when we enable the fallback, even if they are not on the same input table.
I initially considered this but decided to defer it because it will add extra complexities to this PR, but let me take a look again see if we can still add this into the current PR.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3134123905
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -624,6 +644,959 @@ class AdaptiveQueryExecSuite
}
}
+ test("Fallback to shuffled join should apply NO_BROADCAST_HASH only to failed relation") {
Review Comment:
The suite checks the row-limit branch only through private helper calls, but it never runs a real query that trips MAX_BROADCAST_TABLE_ROWS and exercises the full stage-failure -> AQE fallback -> replan path. Since the new config/documentation promise recovery for both row and size limits, this leaves one of the two advertised triggers validated only at the unit-helper level.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2967533351
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1229,6 +1229,14 @@ object SQLConf {
.bytesConf(ByteUnit.BYTE)
.createOptional
+ val ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED =
+ buildConf("spark.sql.adaptive.broadcastJoin.fallbackToShuffle.enabled")
+ .doc("When true, adaptive execution retries with broadcast joins disabled if a broadcast " +
+ "query stage fails because it exceeds broadcast table row or size limits.")
+ .version("3.5.0")
Review Comment:
Ha thanks @dongjoon-hyun . I was able to update this. Our internal implementation is on branch-3.5 and I ported it for master
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2972075792
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +943,179 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Check error-classes.json for details
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) {
+ stagesToReplace
+ } else {
+ stagesToReplace.filterNot { stage =>
+ failedStages.exists(_.eq(stage)) || (stage match {
+ case broadcastStage: BroadcastQueryStageExec =>
+ failedStages.exists {
+ case failedBroadcast: BroadcastQueryStageExec =>
+ sameBroadcastRelation(
+ broadcastStage.broadcast.child, failedBroadcast.broadcast.child)
+ case _ => false
+ }
+ case _ => false
+ })
+ }
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && plan.exists { p =>
+ broadcastChildPlan(p).exists { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastRelation(child, failed.broadcast.child)
+ }
+ }
+ }
+ }
+
+ // Track failed broadcast relations using semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed relation.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ stage: QueryStageExec): Unit = {
+ stage match {
+ case failedStage: BroadcastQueryStageExec =>
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastRelation(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ case _ =>
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ Seq(exchangeStage.plan.canonicalized, exchangeStage.canonicalized).distinct.foreach { key =>
+ context.stageCache.remove(key)
+ }
+ val keysToRemove = context.stageCache.collect {
Review Comment:
Is this full scan remove guarding against a specific scenario, or is it purely defensive?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #54925:
URL: https://github.com/apache/spark/pull/54925#issuecomment-5154256459
We're closing this PR because it hasn't been updated in a while. This isn't a judgement on the merit of the PR in any way. It's just a way of keeping the PR queue manageable.
If you'd like to revive this PR, please reopen it and ask a committer to remove the Stale tag!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3066175332
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -624,6 +637,814 @@ class AdaptiveQueryExecSuite
}
}
+ test("Fallback to shuffled join should apply NO_BROADCAST_HASH only to failed relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+
+ val joinsWithNoBroadcastHash = targetedLogicalPlan.collect {
+ case join: Join
+ if join.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) ||
+ join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) =>
+ join
+ }
+ assert(joinsWithNoBroadcastHash.nonEmpty)
+ assert(targetedLogicalPlan.collect {
+ case join: Join
+ if !join.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) &&
+ !join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)) =>
+ join
+ }.nonEmpty)
+ }
+ }
+
+ test("Fallback to shuffled join should preserve non-broadcast side hint strategies") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val toNoBroadcastHashHint =
+ PrivateMethod[Option[HintInfo]](Symbol("toNoBroadcastHashHint"))
+
+ val noBroadcastAndReplicationHint =
+ Some(HintInfo(strategy = Some(NO_BROADCAST_AND_REPLICATION)))
+ val preservedNoBroadcastAndReplication =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastAndReplicationHint))
+ assert(preservedNoBroadcastAndReplication == noBroadcastAndReplicationHint)
+
+ val noBroadcastHashHint = Some(HintInfo(strategy = Some(NO_BROADCAST_HASH)))
+ val preservedNoBroadcastHash =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastHashHint))
+ assert(preservedNoBroadcastHash == noBroadcastHashHint)
+
+ val broadcastHint = Some(HintInfo(strategy = Some(BROADCAST)))
+ val rewrittenBroadcastHint = adaptivePlan.invokePrivate(toNoBroadcastHashHint(broadcastHint))
+ assert(rewrittenBroadcastHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+
+ Seq(PREFER_SHUFFLE_HASH, SHUFFLE_HASH, SHUFFLE_MERGE).foreach { strategy =>
+ val preservedHint =
+ Some(HintInfo(strategy = Some(strategy)))
+ assert(adaptivePlan.invokePrivate(toNoBroadcastHashHint(preservedHint)) == preservedHint)
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJs after targeted fallback replan") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val baselinePlan = adaptivePlan.invokePrivate(reOptimize(logicalPlan)).get._1
+ assert(adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(baselinePlan, Seq(failedStage))))
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+
+ assert(findTopLevelBroadcastHashJoin(targetedPlan).nonEmpty)
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should not disable other broadcasts from same base relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(failed_t2), BROADCAST(other_t2) */ failed_t2.failed_b, " +
+ "other_t2.other_a " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1) failed_t2 " +
+ "ON t1.key = failed_t2.failed_a " +
+ "JOIN (SELECT a AS other_a FROM testData2 WHERE a <= 2) other_t2 " +
+ "ON t1.key = other_t2.other_a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1")
+ .queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+ val broadcastBuildOutput = finalBhjs.headOption.map { bhj =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left.output.map(_.name)
+ case BuildRight => bhj.right.output.map(_.name)
+ }
+ }.getOrElse(Nil)
+
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(broadcastBuildOutput.contains("other_a"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should match failed relation with nested logical query stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val joinDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a, b FROM testData2 WHERE a > 1) t2 " +
+ "ON t1.key = t2.a")
+ val adaptivePlan = joinDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val relationLogicalPlan = sql("SELECT a, b FROM testData2").queryExecution.optimizedPlan
+ val relationPhysicalPlan = sql("SELECT a, b FROM testData2").queryExecution.sparkPlan
+ val materializedRelation = LogicalQueryStage(
+ relationLogicalPlan,
+ ResultQueryStageExec(0, relationPhysicalPlan, _ => ()))
+
+ val logicalPlanWithStage = joinDf.queryExecution.optimizedPlan.transformDown {
+ case plan if plan.sameResult(relationLogicalPlan) => materializedRelation
+ }
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a, b FROM testData2 WHERE a > 1").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlanWithStage, Seq(failedStage)))
+ val join = targetedLogicalPlan.collectFirst { case j: Join => j }.get
+
+ assert(join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJ for nested failed branch replan") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 5)
+ .selectExpr("id AS k", "concat('small_', cast(id AS string)) AS tag")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t3) */ failed_branch.k, failed_branch.payload, " +
+ "fallback_t3.tag " +
+ "FROM (" +
+ " SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, fallback_t2.payload " +
+ " FROM fallback_t1 " +
+ " JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k" +
+ ") failed_branch " +
+ "JOIN fallback_t3 ON failed_branch.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+
+ val failedBroadcastChild = sql("SELECT k FROM fallback_t2").queryExecution.sparkPlan
+ val failedExchange = adaptivePlan.executedPlan.collectFirst {
+ case exchange: BroadcastExchangeExec if exchange.child.sameResult(failedBroadcastChild) =>
+ exchange
+ }.getOrElse {
+ fail(s"Expected broadcast exchange for fallback_t2 branch:\n${adaptivePlan.executedPlan}")
+ }
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))))
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))))
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(finalBhjs.head.output.exists(_.name == "tag"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should only apply to exclusively broadcast hash join stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true") {
+ def buildBroadcastStage(join: BaseJoinExec): BroadcastQueryStageExec = {
+ val buildPlan = join match {
+ case bhj: BroadcastHashJoinExec =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left
+ case BuildRight => bhj.right
+ }
+ case bnlj: BroadcastNestedLoopJoinExec =>
+ bnlj.buildSide match {
+ case BuildLeft => bnlj.left
+ case BuildRight => bnlj.right
+ }
+ }
+ buildPlan.collectFirst { case stage: BroadcastQueryStageExec => stage }
+ .orElse {
+ buildPlan.collectFirst { case exchange: BroadcastExchangeExec =>
+ BroadcastQueryStageExec(0, exchange, exchange.canonicalized)
+ }
+ }.getOrElse {
+ fail(s"Expected broadcast query stage:\n$join")
+ }
+ }
+
+ val bhjDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a")
+ val bhjAdaptivePlan = bhjDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bhjStage = buildBroadcastStage(findTopLevelBroadcastHashJoin(
+ bhjAdaptivePlan.executedPlan).head)
+
+ val bnljDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key > t2.a")
+ val bnljAdaptivePlan = bnljDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bnljStage = buildBroadcastStage(findTopLevelBroadcastNestedLoopJoin(
+ bnljAdaptivePlan.executedPlan).head)
+
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val mixedDf = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ t1.key " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val mixedAdaptivePlan =
+ mixedDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val mixedBhjs = findTopLevelBroadcastHashJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBnljs = findTopLevelBroadcastNestedLoopJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBhjStage = buildBroadcastStage(mixedBhjs.head)
+ val mixedBnljStage = buildBroadcastStage(mixedBnljs.head)
+
+ assert(mixedBhjs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBnljs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast ne mixedBnljStage.broadcast,
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast.child.sameResult(mixedBnljStage.broadcast.child),
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val subqueryQe = sql("SELECT key FROM testData").queryExecution
+ val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, subqueryQe.sparkPlan)
+ inputPlan.setLogicalLink(subqueryQe.optimizedPlan)
+ val subqueryAdaptivePlan = AdaptiveSparkPlanExec(
+ inputPlan,
+ AdaptiveExecutionContext(spark, subqueryQe),
+ Seq.empty,
+ isSubquery = true)
+ val rootStage = subqueryAdaptivePlan.finalPhysicalPlan match {
+ case stage: BroadcastQueryStageExec => stage
+ case plan =>
+ fail(s"Expected root broadcast query stage:\n$plan")
+ }
+
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should be disabled when conf is false") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val broadcastInputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val stage =
+ BroadcastQueryStageExec(0, broadcastInputPlan, broadcastInputPlan.canonicalized)
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure " +
Review Comment:
Thanks! Added new tests
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2972129014
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -790,30 +840,63 @@ case class AdaptiveSparkPlanExec(
/**
* Re-optimize and run physical planning on the current logical plan based on the latest stats.
*/
- private def reOptimize(logicalPlan: LogicalPlan): Option[(SparkPlan, LogicalPlan)] = {
+ private def reOptimize(
+ logicalPlan: LogicalPlan,
+ disableBroadcastJoin: Boolean): Option[(SparkPlan, LogicalPlan)] = {
try {
- logicalPlan.invalidateStatsCache()
- val optimized = optimizer.execute(logicalPlan)
- val sparkPlan = context.session.sessionState.planner.plan(ReturnAnswer(optimized)).next()
- val newPlan = applyPhysicalRules(
- applyQueryPostPlannerStrategyRules(sparkPlan),
- preprocessingRules ++ queryStagePreparationRules,
- Some((planChangeLogger, "AQE Replanning")))
-
- // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
- // add the `BroadcastExchangeExec` node manually in the DPP subquery,
- // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is complicated
- // and need to be re-optimized, AQE also need to manually insert the `BroadcastExchangeExec`
- // node to prevent the loss of the `BroadcastExchangeExec` node in DPP subquery.
- // Here, we also need to avoid to insert the `BroadcastExchangeExec` node when the newPlan is
- // already the `BroadcastExchangeExec` plan after apply the `LogicalQueryStageStrategy` rule.
- val finalPlan = inputPlan match {
- case b: BroadcastExchangeLike
- if (!newPlan.isInstanceOf[BroadcastExchangeLike]) => b.withNewChildren(Seq(newPlan))
- case _ => newPlan
+ val optimizedLogicalPlan = if (disableBroadcastJoin) {
+ removeBroadcastHints(logicalPlan)
+ } else {
+ logicalPlan
+ }
+ def planFn(
+ optimizerToUse: AQEOptimizer,
+ plannerToUse: SparkPlanner): (SparkPlan, LogicalPlan) = {
+ optimizedLogicalPlan.invalidateStatsCache()
+ val optimized = optimizerToUse.execute(optimizedLogicalPlan)
+ val sparkPlan = plannerToUse.plan(ReturnAnswer(optimized)).next()
+ val newPlan = applyPhysicalRules(
+ applyQueryPostPlannerStrategyRules(sparkPlan),
+ preprocessingRules ++ queryStagePreparationRules,
+ Some((planChangeLogger, "AQE Replanning")))
+
+ // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
+ // add the `BroadcastExchangeExec` node manually in the DPP subquery,
+ // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is
+ // complicated and need to be re-optimized, AQE also need to manually insert the
+ // `BroadcastExchangeExec` node to prevent the loss of the `BroadcastExchangeExec`
+ // node in DPP subquery.
+ // If this adaptive plan must still produce a broadcast result, preserve the root
+ // broadcast exchange even when join fallback disables broadcast joins inside the plan.
+ // Otherwise, avoid reinserting the exchange when the new plan already contains one
+ // after applying the `LogicalQueryStageStrategy` rule.
+ val finalPlan = inputPlan match {
+ case b: BroadcastExchangeLike
+ if !newPlan.isInstanceOf[BroadcastExchangeLike] =>
+ b.withNewChildren(Seq(newPlan))
+ case _ => newPlan
+ }
+ (finalPlan, optimized)
}
- Some((finalPlan, optimized))
+ val result = if (disableBroadcastJoin) {
+ // Use cloneSession to preserve current session state (e.g. runtime SQLConf and planner
+ // customizations) while keeping fallback conf changes isolated from the original session.
+ val fallbackSession = context.session.cloneSession()
+ val fallbackConf = fallbackSession.sessionState.conf
+ fallbackConf.setConfString(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
Review Comment:
Thanks for adding this very helpful feature!
Instead of globally disabling broadcast joins during fallback, would it be possible to inject NO_BROADCAST_HASH only on the failed join side, while preserving broadcast for other joins that haven’t exceeded the limit?
I understand this could add complexity, just curious whether this approach was considered and intentionally deferred.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3059752743
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +926,275 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Only BHJ fallback is supported here.
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ usedOnlyByBroadcastHashJoin(stage.asInstanceOf[BroadcastQueryStageExec]) &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ private def shouldIgnoreStaleBroadcastStageFailure(
+ stage: BroadcastQueryStageExec,
+ error: Throwable,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249") &&
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(failed.broadcast.child, stage.broadcast.child)
+ }
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) return stagesToReplace
+
+ val failedBroadcastStages = failedStages.collect {
+ case b: BroadcastQueryStageExec => b
+ }
+
+ stagesToReplace.filterNot {
+ case stage if failedStages.exists(_.eq(stage)) => true
+ case stage: BroadcastQueryStageExec =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(stage.broadcast.child, failed.broadcast.child)
+ }
+ case _ => false
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ // Check if a Spark physical plan contains any of the previous failed broadcast stages
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && hasMatchingBroadcastInputPlan(plan) { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(child, failed.broadcast.child)
+ }
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ val planKey = exchangeStage.plan.canonicalized
+ val stageKey = exchangeStage.canonicalized
+
+ context.stageCache.remove(planKey)
+ if (stageKey != planKey) {
+ context.stageCache.remove(stageKey)
+ }
+ case _ =>
+ }
+
+ // Normalize different wrappers to the underlying broadcast child plan when present.
+ private def broadcastChildPlan(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case stage: BroadcastQueryStageExec => Some(stage.broadcast.child)
+ case exchange: BroadcastExchangeLike => Some(exchange.child)
+ case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => Some(exchange.child)
+ case _ => None
+ }
+
+ // Compare full broadcast input plans (including filters/projects under the exchange)
+ // by semantic plan equivalence.
+ // Do not fall back to output-attribute equality: different plans can share
+ // output schemas while producing different relations.
+ private def sameBroadcastInputPlan(left: SparkPlan, right: SparkPlan): Boolean = {
+ left.sameResult(right)
+ }
+
+ // Track failed broadcast input plans by semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed input.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ failedStage: BroadcastQueryStageExec): Unit = {
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastInputPlan(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ }
+
+ private def usedOnlyByBroadcastHashJoin(stage: BroadcastQueryStageExec): Boolean = {
+ val consumers = currentPhysicalPlan.collect {
+ case join: BroadcastHashJoinExec
+ if usesBroadcastInputPlan(
+ broadcastBuildPlan(join.buildSide, join.left, join.right), stage) =>
+ join
+ case join: BroadcastNestedLoopJoinExec
+ if usesBroadcastInputPlan(
+ broadcastBuildPlan(join.buildSide, join.left, join.right), stage) =>
+ join
+ }
+ // We do not support fallback for `BroadcastNestedLoopJoinExec` (BNLJ), as there is no good
+ // "hash join" replacement of it that preserves the same semantics. Under the simpler
+ // relation-based fallback scope, if the same broadcast input plan is consumed by both BHJ
+ // and BNLJ, conservatively abandon the fallback and preserve the original behavior.
+ consumers.nonEmpty && consumers.forall(_.isInstanceOf[BroadcastHashJoinExec])
+ }
+
+ private def broadcastBuildPlan(
+ buildSide: BuildSide,
+ left: SparkPlan,
+ right: SparkPlan): SparkPlan = {
+ buildSide match {
+ case BuildLeft => left
+ case BuildRight => right
+ }
+ }
+
+ // Check whether a join build side consumes the same broadcast input plan as a failed stage.
+ private def usesBroadcastInputPlan(
+ plan: SparkPlan,
+ stage: BroadcastQueryStageExec): Boolean = {
+ hasMatchingBroadcastInputPlan(plan) { child =>
+ sameBroadcastInputPlan(child, stage.broadcast.child)
+ }
+ }
+
+ private def hasMatchingBroadcastInputPlan(
+ plan: SparkPlan)(matches: SparkPlan => Boolean): Boolean = {
+ plan.exists { p =>
+ broadcastChildPlan(p).exists(matches)
+ }
+ }
+
+ // Root broadcast stages must be preserved (e.g. broadcast subquery output contract).
+ // Only exempt the actual root broadcast stage instance, not semantically equivalent relations.
+ private def isRequiredRootBroadcastStage(stage: QueryStageExec): Boolean = stage match {
+ case broadcastStage: BroadcastQueryStageExec if inputPlan.isInstanceOf[BroadcastExchangeLike] =>
+ @tailrec
+ def findRootBroadcast(current: SparkPlan): Option[BroadcastQueryStageExec] = current match {
+ case root: BroadcastQueryStageExec => Some(root)
+ case resultStage: ResultQueryStageExec => findRootBroadcast(resultStage.plan)
+ case _ => None
+ }
+ findRootBroadcast(currentPhysicalPlan).exists(_ eq broadcastStage)
+ case _ => false
+ }
+
+ // Apply NO_BROADCAST_HASH hints for joins whose input matches a previously failed
+ // broadcast relation, so targeted fallback can keep unrelated BHJs.
+ private def addNoBroadcastHashHintsForFailedRelations(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ val failedLogicalPlans = extractFailedBroadcastLogicalPlans(failedBroadcastStages)
+ if (failedLogicalPlans.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case join: Join =>
+ val disableLeft = shouldDisableBroadcastForJoinSide(join.left, failedLogicalPlans)
+ val disableRight = shouldDisableBroadcastForJoinSide(join.right, failedLogicalPlans)
+ if (!disableLeft && !disableRight) {
+ join
+ } else {
+ // If either join side is the failed broadcast relation, add NO_BROADCAST_HASH
+ // where the join does not already carry an explicit non-broadcast strategy.
+ // This preserves existing join-strategy intent such as PREFER_SHUFFLE_HASH.
+ val newHint = JoinHint(
+ toNoBroadcastHashHint(join.hint.leftHint),
+ toNoBroadcastHashHint(join.hint.rightHint))
Review Comment:
Should we add hint to only the side failing broadcast?
```
val newHint = JoinHint(
if (disableLeft) toNoBroadcastHashHint(join.hint.leftHint) else join.hint.leftHint,
if (disableRight) toNoBroadcastHashHint(join.hint.rightHint) else join.hint.rightHint)
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3059767128
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -353,12 +398,36 @@ case class AdaptiveSparkPlanExec(
// plans are updated, we can clear the query stage list because at this point the two
// plans are semantically and physically in sync again.
val logicalPlan = replaceWithQueryStagesInLogicalPlan(currentLogicalPlan, stagesToReplace)
- val afterReOptimize = reOptimize(logicalPlan)
+ val shouldBroadcastFallback = fallbackStages.nonEmpty
+ val failedBroadcastStages = failedBroadcastStagesBuffer.toSeq
+
+ val afterReOptimize = if (shouldBroadcastFallback) {
+ val fallbackLogicalPlan = removeFailedBroadcastStagesFromLogicalPlan(
Review Comment:
Do we still need to run removeFailedBroadcastStagesFromLogicalPlan again here? Looks like redundant.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3134120060
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +907,317 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def processStageMaterializationEvents(
+ stageEvents: Seq[StageMaterializationEvent],
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ errors: mutable.ArrayBuffer[Throwable]): BroadcastFallbackResult = {
+ val fallbackStages = new mutable.ArrayBuffer[BroadcastQueryStageExec]()
+ val fallbackErrors = new mutable.ArrayBuffer[Throwable]()
+
+ stageEvents.foreach {
+ case StageSuccess(stage, res) =>
+ stage.resultOption.set(Some(res))
+ case StageFailure(stage: BroadcastQueryStageExec, ex)
+ if shouldFallbackToShuffleJoin(stage, ex) =>
+ removeStageFromCache(stage)
+ logInfo("Broadcast query stage failed on table size/row limit; retrying " +
+ s"adaptive replanning with matching broadcast hash joins disabled. " +
+ s"Stage ${stage.id}")
+ registerFailedBroadcastStage(failedBroadcastStages, stage)
+ if (!fallbackStages.exists(_.eq(stage))) {
+ fallbackStages.append(stage)
+ }
+ fallbackErrors.append(ex)
+ case StageFailure(stage: BroadcastQueryStageExec, ex)
+ if shouldIgnoreStaleBroadcastStageFailure(stage, ex, failedBroadcastStages.toSeq) =>
+ logDebug("Ignoring stale broadcast query stage failure for a previously failed " +
+ s"broadcast input. Stage ${stage.id}")
+ removeStageFromCache(stage)
+ case StageFailure(stage, ex) =>
+ stage.error.set(Some(ex))
+ errors.append(ex)
+ }
+
+ BroadcastFallbackResult(fallbackStages.toSeq, fallbackErrors.toSeq)
+ }
+
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Only BHJ fallback is supported here.
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ usedOnlyByBroadcastHashJoin(stage.asInstanceOf[BroadcastQueryStageExec]) &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ private def shouldIgnoreStaleBroadcastStageFailure(
+ stage: BroadcastQueryStageExec,
+ error: Throwable,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249") &&
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(failed.broadcast.child, stage.broadcast.child)
+ }
Review Comment:
The new fallback path hard-codes _LEGACY_ERROR_TEMP_2248 and _LEGACY_ERROR_TEMP_2249 to recognize row/size-limit broadcast failures. Those IDs are explicitly temporary legacy error classes, so a future cleanup or renumbering will silently turn this feature off without any compiler signal. This would be safer if the matching were centralized behind a stable predicate/constant, or at least backed by an end-to-end regression test that would fail loudly when the error mapping changes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #54925: [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins
URL: https://github.com/apache/spark/pull/54925
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r2967456492
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1229,6 +1229,14 @@ object SQLConf {
.bytesConf(ByteUnit.BYTE)
.createOptional
+ val ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED =
+ buildConf("spark.sql.adaptive.broadcastJoin.fallbackToShuffle.enabled")
+ .doc("When true, adaptive execution retries with broadcast joins disabled if a broadcast " +
+ "query stage fails because it exceeds broadcast table row or size limits.")
+ .version("3.5.0")
Review Comment:
? If you want to backport this to branch-3.5, it should be `3.5.9` instead of `3.5.0`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3066167808
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +926,275 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Only BHJ fallback is supported here.
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ usedOnlyByBroadcastHashJoin(stage.asInstanceOf[BroadcastQueryStageExec]) &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ private def shouldIgnoreStaleBroadcastStageFailure(
+ stage: BroadcastQueryStageExec,
+ error: Throwable,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249") &&
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(failed.broadcast.child, stage.broadcast.child)
+ }
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) return stagesToReplace
+
+ val failedBroadcastStages = failedStages.collect {
+ case b: BroadcastQueryStageExec => b
+ }
+
+ stagesToReplace.filterNot {
+ case stage if failedStages.exists(_.eq(stage)) => true
+ case stage: BroadcastQueryStageExec =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(stage.broadcast.child, failed.broadcast.child)
+ }
+ case _ => false
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ // Check if a Spark physical plan contains any of the previous failed broadcast stages
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && hasMatchingBroadcastInputPlan(plan) { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(child, failed.broadcast.child)
+ }
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ val planKey = exchangeStage.plan.canonicalized
+ val stageKey = exchangeStage.canonicalized
+
+ context.stageCache.remove(planKey)
+ if (stageKey != planKey) {
+ context.stageCache.remove(stageKey)
+ }
+ case _ =>
+ }
+
+ // Normalize different wrappers to the underlying broadcast child plan when present.
+ private def broadcastChildPlan(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case stage: BroadcastQueryStageExec => Some(stage.broadcast.child)
+ case exchange: BroadcastExchangeLike => Some(exchange.child)
+ case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => Some(exchange.child)
+ case _ => None
+ }
+
+ // Compare full broadcast input plans (including filters/projects under the exchange)
+ // by semantic plan equivalence.
+ // Do not fall back to output-attribute equality: different plans can share
+ // output schemas while producing different relations.
+ private def sameBroadcastInputPlan(left: SparkPlan, right: SparkPlan): Boolean = {
+ left.sameResult(right)
+ }
+
+ // Track failed broadcast input plans by semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed input.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ failedStage: BroadcastQueryStageExec): Unit = {
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastInputPlan(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ }
+
+ private def usedOnlyByBroadcastHashJoin(stage: BroadcastQueryStageExec): Boolean = {
+ val consumers = currentPhysicalPlan.collect {
+ case join: BroadcastHashJoinExec
+ if usesBroadcastInputPlan(
+ broadcastBuildPlan(join.buildSide, join.left, join.right), stage) =>
+ join
+ case join: BroadcastNestedLoopJoinExec
+ if usesBroadcastInputPlan(
+ broadcastBuildPlan(join.buildSide, join.left, join.right), stage) =>
+ join
+ }
+ // We do not support fallback for `BroadcastNestedLoopJoinExec` (BNLJ), as there is no good
+ // "hash join" replacement of it that preserves the same semantics. Under the simpler
+ // relation-based fallback scope, if the same broadcast input plan is consumed by both BHJ
+ // and BNLJ, conservatively abandon the fallback and preserve the original behavior.
+ consumers.nonEmpty && consumers.forall(_.isInstanceOf[BroadcastHashJoinExec])
+ }
+
+ private def broadcastBuildPlan(
+ buildSide: BuildSide,
+ left: SparkPlan,
+ right: SparkPlan): SparkPlan = {
+ buildSide match {
+ case BuildLeft => left
+ case BuildRight => right
+ }
+ }
+
+ // Check whether a join build side consumes the same broadcast input plan as a failed stage.
+ private def usesBroadcastInputPlan(
+ plan: SparkPlan,
+ stage: BroadcastQueryStageExec): Boolean = {
+ hasMatchingBroadcastInputPlan(plan) { child =>
+ sameBroadcastInputPlan(child, stage.broadcast.child)
+ }
+ }
+
+ private def hasMatchingBroadcastInputPlan(
+ plan: SparkPlan)(matches: SparkPlan => Boolean): Boolean = {
+ plan.exists { p =>
+ broadcastChildPlan(p).exists(matches)
+ }
+ }
+
+ // Root broadcast stages must be preserved (e.g. broadcast subquery output contract).
+ // Only exempt the actual root broadcast stage instance, not semantically equivalent relations.
+ private def isRequiredRootBroadcastStage(stage: QueryStageExec): Boolean = stage match {
+ case broadcastStage: BroadcastQueryStageExec if inputPlan.isInstanceOf[BroadcastExchangeLike] =>
+ @tailrec
+ def findRootBroadcast(current: SparkPlan): Option[BroadcastQueryStageExec] = current match {
+ case root: BroadcastQueryStageExec => Some(root)
+ case resultStage: ResultQueryStageExec => findRootBroadcast(resultStage.plan)
+ case _ => None
+ }
+ findRootBroadcast(currentPhysicalPlan).exists(_ eq broadcastStage)
+ case _ => false
+ }
+
+ // Apply NO_BROADCAST_HASH hints for joins whose input matches a previously failed
+ // broadcast relation, so targeted fallback can keep unrelated BHJs.
+ private def addNoBroadcastHashHintsForFailedRelations(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ val failedLogicalPlans = extractFailedBroadcastLogicalPlans(failedBroadcastStages)
+ if (failedLogicalPlans.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case join: Join =>
+ val disableLeft = shouldDisableBroadcastForJoinSide(join.left, failedLogicalPlans)
+ val disableRight = shouldDisableBroadcastForJoinSide(join.right, failedLogicalPlans)
+ if (!disableLeft && !disableRight) {
+ join
+ } else {
+ // If either join side is the failed broadcast relation, add NO_BROADCAST_HASH
+ // where the join does not already carry an explicit non-broadcast strategy.
+ // This preserves existing join-strategy intent such as PREFER_SHUFFLE_HASH.
+ val newHint = JoinHint(
+ toNoBroadcastHashHint(join.hint.leftHint),
+ toNoBroadcastHashHint(join.hint.rightHint))
Review Comment:
Good catch. We should just update one side. Updated.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -353,12 +398,36 @@ case class AdaptiveSparkPlanExec(
// plans are updated, we can clear the query stage list because at this point the two
// plans are semantically and physically in sync again.
val logicalPlan = replaceWithQueryStagesInLogicalPlan(currentLogicalPlan, stagesToReplace)
- val afterReOptimize = reOptimize(logicalPlan)
+ val shouldBroadcastFallback = fallbackStages.nonEmpty
+ val failedBroadcastStages = failedBroadcastStagesBuffer.toSeq
+
+ val afterReOptimize = if (shouldBroadcastFallback) {
+ val fallbackLogicalPlan = removeFailedBroadcastStagesFromLogicalPlan(
Review Comment:
Yea it looks unnecessary here. Removed
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3134186621
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -860,6 +907,317 @@ case class AdaptiveSparkPlanExec(
}
}
+ private def processStageMaterializationEvents(
+ stageEvents: Seq[StageMaterializationEvent],
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ errors: mutable.ArrayBuffer[Throwable]): BroadcastFallbackResult = {
+ val fallbackStages = new mutable.ArrayBuffer[BroadcastQueryStageExec]()
+ val fallbackErrors = new mutable.ArrayBuffer[Throwable]()
+
+ stageEvents.foreach {
+ case StageSuccess(stage, res) =>
+ stage.resultOption.set(Some(res))
+ case StageFailure(stage: BroadcastQueryStageExec, ex)
+ if shouldFallbackToShuffleJoin(stage, ex) =>
+ removeStageFromCache(stage)
+ logInfo("Broadcast query stage failed on table size/row limit; retrying " +
+ s"adaptive replanning with matching broadcast hash joins disabled. " +
+ s"Stage ${stage.id}")
+ registerFailedBroadcastStage(failedBroadcastStages, stage)
+ if (!fallbackStages.exists(_.eq(stage))) {
+ fallbackStages.append(stage)
+ }
+ fallbackErrors.append(ex)
+ case StageFailure(stage: BroadcastQueryStageExec, ex)
+ if shouldIgnoreStaleBroadcastStageFailure(stage, ex, failedBroadcastStages.toSeq) =>
+ logDebug("Ignoring stale broadcast query stage failure for a previously failed " +
+ s"broadcast input. Stage ${stage.id}")
+ removeStageFromCache(stage)
+ case StageFailure(stage, ex) =>
+ stage.error.set(Some(ex))
+ errors.append(ex)
+ }
+
+ BroadcastFallbackResult(fallbackStages.toSeq, fallbackErrors.toSeq)
+ }
+
+ private def shouldFallbackToShuffleJoin(stage: QueryStageExec, error: Throwable): Boolean = {
+ // Detect errors such as:
+ // - Cannot broadcast the table over <maxBroadcastTableRows> rows: <numRows> rows.
+ // - Cannot broadcast the table that is larger than <maxBroadcastTableBytes>: <dataSize>.
+ // Only BHJ fallback is supported here.
+ conf.adaptiveBroadcastJoinFallbackToShuffleEnabled &&
+ !isRequiredRootBroadcastStage(stage) &&
+ stage.isInstanceOf[BroadcastQueryStageExec] &&
+ usedOnlyByBroadcastHashJoin(stage.asInstanceOf[BroadcastQueryStageExec]) &&
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249")
+ }
+
+ private def shouldIgnoreStaleBroadcastStageFailure(
+ stage: BroadcastQueryStageExec,
+ error: Throwable,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ hasErrorClass(error, "_LEGACY_ERROR_TEMP_2248", "_LEGACY_ERROR_TEMP_2249") &&
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(failed.broadcast.child, stage.broadcast.child)
+ }
+ }
+
+ // Remove failed fallback stages from the replacement set before replanning.
+ // For broadcast stages, remove semantically equivalent siblings as well, so
+ // reused/duplicated broadcast stages for the same relation don't leak through.
+ private def filterStagesToReplaceForFallback(
+ stagesToReplace: Seq[QueryStageExec],
+ failedStages: Seq[QueryStageExec]): Seq[QueryStageExec] = {
+ if (failedStages.isEmpty) return stagesToReplace
+
+ val failedBroadcastStages = failedStages.collect {
+ case b: BroadcastQueryStageExec => b
+ }
+
+ stagesToReplace.filterNot {
+ case stage if failedStages.exists(_.eq(stage)) => true
+ case stage: BroadcastQueryStageExec =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(stage.broadcast.child, failed.broadcast.child)
+ }
+ case _ => false
+ }
+ }
+
+ // Remove failed broadcast query-stage wrappers that may have been embedded in prior adopted
+ // logical plans, so fallback replanning does not keep retrying the same failed relation.
+ private def removeFailedBroadcastStagesFromLogicalPlan(
+ plan: LogicalPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): LogicalPlan = {
+ if (failedBroadcastStages.isEmpty) {
+ plan
+ } else {
+ plan.transformDown {
+ case stage: LogicalQueryStage
+ if hasFailedBroadcastRelation(stage.physicalPlan, failedBroadcastStages) =>
+ stage.logicalPlan
+ }
+ }
+ }
+
+ // Check if a Spark physical plan contains any of the previous failed broadcast stages
+ private def hasFailedBroadcastRelation(
+ plan: SparkPlan,
+ failedBroadcastStages: Seq[BroadcastQueryStageExec]): Boolean = {
+ failedBroadcastStages.nonEmpty && hasMatchingBroadcastInputPlan(plan) { child =>
+ failedBroadcastStages.exists { failed =>
+ sameBroadcastInputPlan(child, failed.broadcast.child)
+ }
+ }
+ }
+
+ // Drop stale cache entries for a failed/replaced exchange stage so AQE can rebuild it.
+ private def removeStageFromCache(stage: QueryStageExec): Unit = stage match {
+ case exchangeStage: ExchangeQueryStageExec =>
+ val planKey = exchangeStage.plan.canonicalized
+ val stageKey = exchangeStage.canonicalized
+
+ context.stageCache.remove(planKey)
+ if (stageKey != planKey) {
+ context.stageCache.remove(stageKey)
+ }
+ case _ =>
+ }
+
+ // Normalize different wrappers to the underlying broadcast child plan when present.
+ private def broadcastChildPlan(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case stage: BroadcastQueryStageExec => Some(stage.broadcast.child)
+ case exchange: BroadcastExchangeLike => Some(exchange.child)
+ case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => Some(exchange.child)
+ case _ => None
+ }
+
+ // Compare full broadcast input plans (including filters/projects under the exchange)
+ // by semantic plan equivalence.
+ // Do not fall back to output-attribute equality: different plans can share
+ // output schemas while producing different relations.
+ private def sameBroadcastInputPlan(left: SparkPlan, right: SparkPlan): Boolean = {
+ left.sameResult(right)
+ }
+
+ // Track failed broadcast input plans by semantic equality so equivalent stages
+ // in later AQE replans are treated as the same failed input.
+ private def registerFailedBroadcastStage(
+ failedBroadcastStages: mutable.ArrayBuffer[BroadcastQueryStageExec],
+ failedStage: BroadcastQueryStageExec): Unit = {
+ if (!failedBroadcastStages.exists { existing =>
+ sameBroadcastInputPlan(existing.broadcast.child, failedStage.broadcast.child)
+ }) {
+ failedBroadcastStages.append(failedStage)
+ }
+ }
+
+ private def usedOnlyByBroadcastHashJoin(stage: BroadcastQueryStageExec): Boolean = {
Review Comment:
A few of the new helpers use TreeNode.collect / .exists / .transformDown, which don't descend into the plans held by ExecSubqueryExpression. The affected call sites are:
- usedOnlyByBroadcastHashJoin — currentPhysicalPlan.collect { ... }
- hasMatchingBroadcastInputPlan — plan.exists { ... } (used by both usesBroadcastInputPlan and hasFailedBroadcastRelation)
- removeFailedBroadcastStagesFromLogicalPlan and addNoBroadcastHashHintsForFailedRelations — plan.transformDown
In practice, subqueries are wrapped in their own AdaptiveSparkPlanExec by InsertAdaptiveSparkPlan, so a broadcast failure inside a subquery is handled by that nested AQE and never flows into the outer AQE's event loop. That's why the existing "recover from broadcast failure in subquery" test passes. But the main-plan traversals still have two visible gaps:
1. usedOnlyByBroadcastHashJoin can under-count BNLJ consumers. If the same broadcast relation (by sameResult) is consumed by a BHJ in the main plan and by a BNLJ inside a nested subquery, the scan only sees the BHJ and concludes "BHJ-only", triggering the fallback. The outer fallback only rewrites the main plan's BHJ, so this doesn't produce wrong results in today's layout — but it violates the conservative invariant the helper's comment advertises ("if the same broadcast input plan is consumed by both BHJ and BNLJ, conservatively abandon the fallback").
2. hasFailedBroadcastRelation has a blind spot at the subquery boundary. Combined with addNoBroadcastHashHintsForFailedRelations (which also won't descend through subqueries), a replan can reintroduce the same failed broadcast relation below a subquery boundary without being caught by the rejection check. The outer AQE would then accept a plan that still contains the previously-failed relation under a subquery.
If the intent is full-plan safety, these scans should switch to the subquery-aware helpers (collectWithSubqueries, foreachWithSubqueries, etc.). If the intent is main-plan only (relying on nested AQE instances to handle their own broadcast failures independently), then:
- the comments on these helpers should state that boundary explicitly, and
- it would be worth adding an end-to-end test that exercises "outer AQE with a subquery" to pin down the chosen semantics — the current subquery
test only validates that the nested AQE can fall back on its own, not that the outer AQE's guards behave correctly when a subquery is present.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3134362909
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -624,6 +644,959 @@ class AdaptiveQueryExecSuite
}
}
+ test("Fallback to shuffled join should apply NO_BROADCAST_HASH only to failed relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+
+ assert(joins.count(_.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 0)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ }
+ }
+
+ test("Fallback to shuffled join should preserve non-broadcast side hint strategies") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val toNoBroadcastHashHint =
+ PrivateMethod[Option[HintInfo]](Symbol("toNoBroadcastHashHint"))
+
+ val noBroadcastAndReplicationHint =
+ Some(HintInfo(strategy = Some(NO_BROADCAST_AND_REPLICATION)))
+ val preservedNoBroadcastAndReplication =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastAndReplicationHint))
+ assert(preservedNoBroadcastAndReplication == noBroadcastAndReplicationHint)
+
+ val noBroadcastHashHint = Some(HintInfo(strategy = Some(NO_BROADCAST_HASH)))
+ val preservedNoBroadcastHash =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastHashHint))
+ assert(preservedNoBroadcastHash == noBroadcastHashHint)
+
+ val broadcastHint = Some(HintInfo(strategy = Some(BROADCAST)))
+ val rewrittenBroadcastHint = adaptivePlan.invokePrivate(toNoBroadcastHashHint(broadcastHint))
+ assert(rewrittenBroadcastHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+
+ Seq(PREFER_SHUFFLE_HASH, SHUFFLE_HASH, SHUFFLE_MERGE).foreach { strategy =>
+ val preservedHint =
+ Some(HintInfo(strategy = Some(strategy)))
+ assert(adaptivePlan.invokePrivate(toNoBroadcastHashHint(preservedHint)) == preservedHint)
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJs after targeted fallback replan") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val baselinePlan = adaptivePlan.invokePrivate(reOptimize(logicalPlan)).get._1
+ assert(adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(baselinePlan, Seq(failedStage))))
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+
+ assert(findTopLevelBroadcastHashJoin(targetedPlan).nonEmpty)
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should not disable other broadcasts from same base relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(failed_t2), BROADCAST(other_t2) */ failed_t2.failed_b, " +
+ "other_t2.other_a " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1) failed_t2 " +
+ "ON t1.key = failed_t2.failed_a " +
+ "JOIN (SELECT a AS other_a FROM testData2 WHERE a <= 2) other_t2 " +
+ "ON t1.key = other_t2.other_a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1")
+ .queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+ val broadcastBuildOutput = finalBhjs.headOption.map { bhj =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left.output.map(_.name)
+ case BuildRight => bhj.right.output.map(_.name)
+ }
+ }.getOrElse(Nil)
+
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(broadcastBuildOutput.contains("other_a"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should match failed relation with nested logical query stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val joinDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a, b FROM testData2 WHERE a > 1) t2 " +
+ "ON t1.key = t2.a")
+ val adaptivePlan = joinDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val relationLogicalPlan = sql("SELECT a, b FROM testData2").queryExecution.optimizedPlan
+ val relationPhysicalPlan = sql("SELECT a, b FROM testData2").queryExecution.sparkPlan
+ val materializedRelation = LogicalQueryStage(
+ relationLogicalPlan,
+ ResultQueryStageExec(0, relationPhysicalPlan, _ => ()))
+
+ val logicalPlanWithStage = joinDf.queryExecution.optimizedPlan.transformDown {
+ case plan if plan.sameResult(relationLogicalPlan) => materializedRelation
+ }
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a, b FROM testData2 WHERE a > 1").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlanWithStage, Seq(failedStage)))
+ val join = targetedLogicalPlan.collectFirst { case j: Join => j }.get
+
+ assert(join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJ for nested failed branch replan") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 5)
+ .selectExpr("id AS k", "concat('small_', cast(id AS string)) AS tag")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t3) */ failed_branch.k, failed_branch.payload, " +
+ "fallback_t3.tag " +
+ "FROM (" +
+ " SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, fallback_t2.payload " +
+ " FROM fallback_t1 " +
+ " JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k" +
+ ") failed_branch " +
+ "JOIN fallback_t3 ON failed_branch.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+
+ val failedBroadcastChild = sql("SELECT k FROM fallback_t2").queryExecution.sparkPlan
+ val failedExchange = adaptivePlan.executedPlan.collectFirst {
+ case exchange: BroadcastExchangeExec if exchange.child.sameResult(failedBroadcastChild) =>
+ exchange
+ }.getOrElse {
+ fail(s"Expected broadcast exchange for fallback_t2 branch:\n${adaptivePlan.executedPlan}")
+ }
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val bhjBuildOutputs = finalBhjs.map(broadcastHashJoinBuildOutput)
+
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))))
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))))
+ assert(bhjBuildOutputs.exists(_.contains("tag")), s"$targetedPlan")
+ assert(!bhjBuildOutputs.exists(_.contains("payload")), s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should only apply to exclusively broadcast hash join stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true") {
+ def buildBroadcastStage(join: BaseJoinExec): BroadcastQueryStageExec = {
+ val buildPlan = join match {
+ case bhj: BroadcastHashJoinExec =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left
+ case BuildRight => bhj.right
+ }
+ case bnlj: BroadcastNestedLoopJoinExec =>
+ bnlj.buildSide match {
+ case BuildLeft => bnlj.left
+ case BuildRight => bnlj.right
+ }
+ }
+ buildPlan.collectFirst { case stage: BroadcastQueryStageExec => stage }
+ .orElse {
+ buildPlan.collectFirst { case exchange: BroadcastExchangeExec =>
+ BroadcastQueryStageExec(0, exchange, exchange.canonicalized)
+ }
+ }.getOrElse {
+ fail(s"Expected broadcast query stage:\n$join")
+ }
+ }
+
+ val bhjDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a")
+ val bhjAdaptivePlan = bhjDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bhjStage = buildBroadcastStage(findTopLevelBroadcastHashJoin(
+ bhjAdaptivePlan.executedPlan).head)
+
+ val bnljDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key > t2.a")
+ val bnljAdaptivePlan = bnljDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bnljStage = buildBroadcastStage(findTopLevelBroadcastNestedLoopJoin(
+ bnljAdaptivePlan.executedPlan).head)
+
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val mixedDf = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ t1.key " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val mixedAdaptivePlan =
+ mixedDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val mixedBhjs = findTopLevelBroadcastHashJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBnljs = findTopLevelBroadcastNestedLoopJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBhjStage = buildBroadcastStage(mixedBhjs.head)
+ val mixedBnljStage = buildBroadcastStage(mixedBnljs.head)
+
+ assert(mixedBhjs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBnljs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast ne mixedBnljStage.broadcast,
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast.child.sameResult(mixedBnljStage.broadcast.child),
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val subqueryQe = sql("SELECT key FROM testData").queryExecution
+ val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, subqueryQe.sparkPlan)
+ inputPlan.setLogicalLink(subqueryQe.optimizedPlan)
+ val subqueryAdaptivePlan = AdaptiveSparkPlanExec(
+ inputPlan,
+ AdaptiveExecutionContext(spark, subqueryQe),
+ Seq.empty,
+ isSubquery = true)
+ val rootStage = subqueryAdaptivePlan.finalPhysicalPlan match {
+ case stage: BroadcastQueryStageExec => stage
+ case plan =>
+ fail(s"Expected root broadcast query stage:\n$plan")
+ }
+
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should be disabled when conf is false") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val broadcastInputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val stage =
+ BroadcastQueryStageExec(0, broadcastInputPlan, broadcastInputPlan.canonicalized)
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure " +
+ "without explicit broadcast hint") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should preserve broadcast failure when disabled") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ val error = intercept[SparkException] {
+ df.collect()
+ }
+ assert(containsBroadcastLimitError(error), error)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from broadcast size failure in self-join") {
+ withTempView("fallback_self") {
+ spark.range(0, 100)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_self")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2) */ COUNT(*) " +
+ "FROM fallback_self t1 JOIN fallback_self t2 ON t1.k = t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(100))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover when two broadcasts fail") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('y', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2), BROADCAST(fallback_t3) */ COUNT(*) " +
+ "FROM fallback_t1 " +
+ "JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k " +
+ "JOIN fallback_t3 ON fallback_t1.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).size == 2)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBaseJoin(finalPlan).size == 2, s"$finalPlan")
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).size == 2, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from broadcast failure in subquery") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT k FROM fallback_t1 WHERE k IN (" +
+ "SELECT /*+ BROADCAST(fallback_t2) */ fallback_t2.k " +
+ "FROM fallback_t3 JOIN fallback_t2 ON fallback_t3.k = fallback_t2.k)")
+
+ checkAnswer(df, (0L until 50L).map(Row(_)))
+
+ val finalPlan = df.queryExecution.executedPlan
+ assert(collectWithSubqueries(finalPlan) {
+ case j: BroadcastHashJoinExec => j
+ }.isEmpty, s"$finalPlan")
+ assert(collectWithSubqueries(finalPlan) {
+ case j: SortMergeJoinExec => j
+ }.nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should keep AQE coalesce after fallback replan") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 1000).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 1000)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1MB",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "20") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k " +
+ "GROUP BY fallback_t1.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, (0L until 1000L).map(i => Row(i, 1L)))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty, s"$finalPlan")
+ assert(collect(finalPlan) {
+ case read: AQEShuffleReadExec if read.hasCoalescedPartition => read
+ }.nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should not retry mixed BHJ and BNLJ consumers at runtime") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ COUNT(*) " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logAppender = new LogAppender("mixed BHJ BNLJ fallback")
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).size == 1)
+ assert(findTopLevelBroadcastNestedLoopJoin(adaptivePlan.executedPlan).size == 1)
+
+ withLogAppender(
+ logAppender,
+ loggerNames = Seq(AdaptiveSparkPlanExec.getClass.getName.dropRight(1)),
+ level = Some(Level.INFO)) {
+ val error = intercept[SparkException] {
+ df.collect()
+ }
+ assert(containsBroadcastLimitError(error), error)
+ }
+
+ assert(!logAppender.loggingEvents.exists { event =>
+ event.getMessage.getFormattedMessage.contains(
+ "retrying adaptive replanning with matching broadcast hash joins disabled")
+ })
+ }
+ }
+
+ test("Fallback to shuffled join should recover from nested failed broadcast branch") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ val widePayload = "x" * 2000
Review Comment:
Oh I saw it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56065][SQL] Add AQE fallback from failed broadcast joins to shuffle joins [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #54925:
URL: https://github.com/apache/spark/pull/54925#discussion_r3134360958
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -624,6 +644,959 @@ class AdaptiveQueryExecSuite
}
}
+ test("Fallback to shuffled join should apply NO_BROADCAST_HASH only to failed relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+
+ assert(joins.count(_.hint.leftHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 0)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ }
+ }
+
+ test("Fallback to shuffled join should preserve non-broadcast side hint strategies") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val toNoBroadcastHashHint =
+ PrivateMethod[Option[HintInfo]](Symbol("toNoBroadcastHashHint"))
+
+ val noBroadcastAndReplicationHint =
+ Some(HintInfo(strategy = Some(NO_BROADCAST_AND_REPLICATION)))
+ val preservedNoBroadcastAndReplication =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastAndReplicationHint))
+ assert(preservedNoBroadcastAndReplication == noBroadcastAndReplicationHint)
+
+ val noBroadcastHashHint = Some(HintInfo(strategy = Some(NO_BROADCAST_HASH)))
+ val preservedNoBroadcastHash =
+ adaptivePlan.invokePrivate(toNoBroadcastHashHint(noBroadcastHashHint))
+ assert(preservedNoBroadcastHash == noBroadcastHashHint)
+
+ val broadcastHint = Some(HintInfo(strategy = Some(BROADCAST)))
+ val rewrittenBroadcastHint = adaptivePlan.invokePrivate(toNoBroadcastHashHint(broadcastHint))
+ assert(rewrittenBroadcastHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+
+ Seq(PREFER_SHUFFLE_HASH, SHUFFLE_HASH, SHUFFLE_MERGE).foreach { strategy =>
+ val preservedHint =
+ Some(HintInfo(strategy = Some(strategy)))
+ assert(adaptivePlan.invokePrivate(toNoBroadcastHashHint(preservedHint)) == preservedHint)
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJs after targeted fallback replan") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2), BROADCAST(t3) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a " +
+ "JOIN testData3 t3 ON t1.key = t3.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode, sql("SELECT * FROM testData2").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val baselinePlan = adaptivePlan.invokePrivate(reOptimize(logicalPlan)).get._1
+ assert(adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(baselinePlan, Seq(failedStage))))
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+
+ assert(findTopLevelBroadcastHashJoin(targetedPlan).nonEmpty)
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should not disable other broadcasts from same base relation") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(failed_t2), BROADCAST(other_t2) */ failed_t2.failed_b, " +
+ "other_t2.other_a " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1) failed_t2 " +
+ "ON t1.key = failed_t2.failed_a " +
+ "JOIN (SELECT a AS other_a FROM testData2 WHERE a <= 2) other_t2 " +
+ "ON t1.key = other_t2.other_a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a AS failed_a, b AS failed_b FROM testData2 WHERE a > 1")
+ .queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val finalSmjs = findTopLevelSortMergeJoin(targetedPlan)
+ val broadcastBuildOutput = finalBhjs.headOption.map { bhj =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left.output.map(_.name)
+ case BuildRight => bhj.right.output.map(_.name)
+ }
+ }.getOrElse(Nil)
+
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))) == 1)
+ assert(joins.count(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))) == 1)
+ assert(finalBhjs.size == 1, s"$targetedPlan")
+ assert(broadcastBuildOutput.contains("other_a"), s"$targetedPlan")
+ assert(finalSmjs.size == 1, s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+
+ test("Fallback to shuffled join should match failed relation with nested logical query stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ val joinDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a, b FROM testData2 WHERE a > 1) t2 " +
+ "ON t1.key = t2.a")
+ val adaptivePlan = joinDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+
+ val relationLogicalPlan = sql("SELECT a, b FROM testData2").queryExecution.optimizedPlan
+ val relationPhysicalPlan = sql("SELECT a, b FROM testData2").queryExecution.sparkPlan
+ val materializedRelation = LogicalQueryStage(
+ relationLogicalPlan,
+ ResultQueryStageExec(0, relationPhysicalPlan, _ => ()))
+
+ val logicalPlanWithStage = joinDf.queryExecution.optimizedPlan.transformDown {
+ case plan if plan.sameResult(relationLogicalPlan) => materializedRelation
+ }
+
+ val failedExchange = BroadcastExchangeExec(
+ IdentityBroadcastMode,
+ sql("SELECT a, b FROM testData2 WHERE a > 1").queryExecution.sparkPlan)
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlanWithStage, Seq(failedStage)))
+ val join = targetedLogicalPlan.collectFirst { case j: Join => j }.get
+
+ assert(join.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH)))
+ }
+ }
+
+ test("Fallback to shuffled join should keep unrelated BHJ for nested failed branch replan") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 5)
+ .selectExpr("id AS k", "concat('small_', cast(id AS string)) AS tag")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t3) */ failed_branch.k, failed_branch.payload, " +
+ "fallback_t3.tag " +
+ "FROM (" +
+ " SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, fallback_t2.payload " +
+ " FROM fallback_t1 " +
+ " JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k" +
+ ") failed_branch " +
+ "JOIN fallback_t3 ON failed_branch.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
+
+ val reOptimize =
+ PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
+ val addNoBroadcastHashHintsForFailedRelations =
+ PrivateMethod[LogicalPlan](Symbol("addNoBroadcastHashHintsForFailedRelations"))
+ val hasFailedBroadcastRelation =
+ PrivateMethod[Boolean](Symbol("hasFailedBroadcastRelation"))
+
+ val failedBroadcastChild = sql("SELECT k FROM fallback_t2").queryExecution.sparkPlan
+ val failedExchange = adaptivePlan.executedPlan.collectFirst {
+ case exchange: BroadcastExchangeExec if exchange.child.sameResult(failedBroadcastChild) =>
+ exchange
+ }.getOrElse {
+ fail(s"Expected broadcast exchange for fallback_t2 branch:\n${adaptivePlan.executedPlan}")
+ }
+ val failedStage = BroadcastQueryStageExec(0, failedExchange, failedExchange.canonicalized)
+
+ val targetedLogicalPlan = adaptivePlan.invokePrivate(
+ addNoBroadcastHashHintsForFailedRelations(logicalPlan, Seq(failedStage)))
+ val joins = targetedLogicalPlan.collect { case join: Join => join }
+ val targetedPlan = adaptivePlan.invokePrivate(reOptimize(targetedLogicalPlan)).get._1
+ val finalBhjs = findTopLevelBroadcastHashJoin(targetedPlan)
+ val bhjBuildOutputs = finalBhjs.map(broadcastHashJoinBuildOutput)
+
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(NO_BROADCAST_HASH))))
+ assert(joins.exists(_.hint.rightHint.exists(_.strategy.contains(BROADCAST))))
+ assert(bhjBuildOutputs.exists(_.contains("tag")), s"$targetedPlan")
+ assert(!bhjBuildOutputs.exists(_.contains("payload")), s"$targetedPlan")
+ assert(!adaptivePlan.invokePrivate(
+ hasFailedBroadcastRelation(targetedPlan, Seq(failedStage))))
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should only apply to exclusively broadcast hash join stages") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true") {
+ def buildBroadcastStage(join: BaseJoinExec): BroadcastQueryStageExec = {
+ val buildPlan = join match {
+ case bhj: BroadcastHashJoinExec =>
+ bhj.buildSide match {
+ case BuildLeft => bhj.left
+ case BuildRight => bhj.right
+ }
+ case bnlj: BroadcastNestedLoopJoinExec =>
+ bnlj.buildSide match {
+ case BuildLeft => bnlj.left
+ case BuildRight => bnlj.right
+ }
+ }
+ buildPlan.collectFirst { case stage: BroadcastQueryStageExec => stage }
+ .orElse {
+ buildPlan.collectFirst { case exchange: BroadcastExchangeExec =>
+ BroadcastQueryStageExec(0, exchange, exchange.canonicalized)
+ }
+ }.getOrElse {
+ fail(s"Expected broadcast query stage:\n$join")
+ }
+ }
+
+ val bhjDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a")
+ val bhjAdaptivePlan = bhjDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bhjStage = buildBroadcastStage(findTopLevelBroadcastHashJoin(
+ bhjAdaptivePlan.executedPlan).head)
+
+ val bnljDf = sql(
+ "SELECT /*+ BROADCAST(t2) */ * " +
+ "FROM testData t1 JOIN testData2 t2 ON t1.key > t2.a")
+ val bnljAdaptivePlan = bnljDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val bnljStage = buildBroadcastStage(findTopLevelBroadcastNestedLoopJoin(
+ bnljAdaptivePlan.executedPlan).head)
+
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(bhjAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!bnljAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ bnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val mixedDf = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ t1.key " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val mixedAdaptivePlan =
+ mixedDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val mixedBhjs = findTopLevelBroadcastHashJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBnljs = findTopLevelBroadcastNestedLoopJoin(mixedAdaptivePlan.executedPlan)
+ val mixedBhjStage = buildBroadcastStage(mixedBhjs.head)
+ val mixedBnljStage = buildBroadcastStage(mixedBnljs.head)
+
+ assert(mixedBhjs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBnljs.size == 1, s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast ne mixedBnljStage.broadcast,
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(mixedBhjStage.broadcast.child.sameResult(mixedBnljStage.broadcast.child),
+ s"${mixedAdaptivePlan.executedPlan}")
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBhjStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!mixedAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ mixedBnljStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+
+ val subqueryQe = sql("SELECT key FROM testData").queryExecution
+ val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, subqueryQe.sparkPlan)
+ inputPlan.setLogicalLink(subqueryQe.optimizedPlan)
+ val subqueryAdaptivePlan = AdaptiveSparkPlanExec(
+ inputPlan,
+ AdaptiveExecutionContext(spark, subqueryQe),
+ Seq.empty,
+ isSubquery = true)
+ val rootStage = subqueryAdaptivePlan.finalPhysicalPlan match {
+ case stage: BroadcastQueryStageExec => stage
+ case plan =>
+ fail(s"Expected root broadcast query stage:\n$plan")
+ }
+
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!subqueryAdaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ rootStage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should be disabled when conf is false") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val broadcastInputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ qe.sparkPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val stage =
+ BroadcastQueryStageExec(0, broadcastInputPlan, broadcastInputPlan.canonicalized)
+ val shouldFallbackToShuffleJoin =
+ PrivateMethod[Boolean](Symbol("shouldFallbackToShuffleJoin"))
+
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableBytesError(1, 2))))
+ assert(!adaptivePlan.invokePrivate(shouldFallbackToShuffleJoin(
+ stage,
+ QueryExecutionErrors.cannotBroadcastTableOverMaxTableRowsError(1, 2))))
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from runtime broadcast size failure " +
+ "without explicit broadcast hint") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty)
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should preserve broadcast failure when disabled") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50).selectExpr("id AS k").createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "false",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ val error = intercept[SparkException] {
+ df.collect()
+ }
+ assert(containsBroadcastLimitError(error), error)
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from broadcast size failure in self-join") {
+ withTempView("fallback_self") {
+ spark.range(0, 100)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_self")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(t2) */ COUNT(*) " +
+ "FROM fallback_self t1 JOIN fallback_self t2 ON t1.k = t2.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, Row(100))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover when two broadcasts fail") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('y', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2), BROADCAST(fallback_t3) */ COUNT(*) " +
+ "FROM fallback_t1 " +
+ "JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k " +
+ "JOIN fallback_t3 ON fallback_t1.k = fallback_t3.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).size == 2)
+
+ checkAnswer(df, Row(50))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBaseJoin(finalPlan).size == 2, s"$finalPlan")
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).size == 2, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should recover from broadcast failure in subquery") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 50)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ spark.range(0, 100).selectExpr("id AS k").createOrReplaceTempView("fallback_t3")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+ val df = sql(
+ "SELECT k FROM fallback_t1 WHERE k IN (" +
+ "SELECT /*+ BROADCAST(fallback_t2) */ fallback_t2.k " +
+ "FROM fallback_t3 JOIN fallback_t2 ON fallback_t3.k = fallback_t2.k)")
+
+ checkAnswer(df, (0L until 50L).map(Row(_)))
+
+ val finalPlan = df.queryExecution.executedPlan
+ assert(collectWithSubqueries(finalPlan) {
+ case j: BroadcastHashJoinExec => j
+ }.isEmpty, s"$finalPlan")
+ assert(collectWithSubqueries(finalPlan) {
+ case j: SortMergeJoinExec => j
+ }.nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should keep AQE coalesce after fallback replan") {
+ withTempView("fallback_t1", "fallback_t2") {
+ spark.range(0, 1000).selectExpr("id AS k").createOrReplaceTempView("fallback_t1")
+ spark.range(0, 1000)
+ .selectExpr("id AS k", "repeat('x', 2000) AS payload")
+ .createOrReplaceTempView("fallback_t2")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "16KB",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1MB",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "20") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(fallback_t2) */ fallback_t1.k, COUNT(*) " +
+ "FROM fallback_t1 JOIN fallback_t2 ON fallback_t1.k = fallback_t2.k " +
+ "GROUP BY fallback_t1.k")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).nonEmpty)
+
+ checkAnswer(df, (0L until 1000L).map(i => Row(i, 1L)))
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBroadcastHashJoin(finalPlan).isEmpty, s"$finalPlan")
+ assert(findTopLevelSortMergeJoin(finalPlan).nonEmpty, s"$finalPlan")
+ assert(collect(finalPlan) {
+ case read: AQEShuffleReadExec if read.hasCoalescedPartition => read
+ }.nonEmpty, s"$finalPlan")
+ }
+ }
+ }
+
+ test("Fallback to shuffled join should not retry mixed BHJ and BNLJ consumers at runtime") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB",
+ SQLConf.ADAPTIVE_BROADCAST_JOIN_FALLBACK_TO_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.MAX_BROADCAST_TABLE_SIZE.key -> "1") {
+ val df = sql(
+ "SELECT /*+ BROADCAST(eq_t2), BROADCAST(range_t2) */ COUNT(*) " +
+ "FROM testData t1 " +
+ "JOIN (SELECT a FROM testData2) eq_t2 ON t1.key = eq_t2.a " +
+ "JOIN (SELECT a FROM testData2) range_t2 ON t1.key > range_t2.a")
+ val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+ val logAppender = new LogAppender("mixed BHJ BNLJ fallback")
+
+ assert(findTopLevelBroadcastHashJoin(adaptivePlan.executedPlan).size == 1)
+ assert(findTopLevelBroadcastNestedLoopJoin(adaptivePlan.executedPlan).size == 1)
+
+ withLogAppender(
+ logAppender,
+ loggerNames = Seq(AdaptiveSparkPlanExec.getClass.getName.dropRight(1)),
+ level = Some(Level.INFO)) {
+ val error = intercept[SparkException] {
+ df.collect()
+ }
+ assert(containsBroadcastLimitError(error), error)
+ }
+
+ assert(!logAppender.loggingEvents.exists { event =>
+ event.getMessage.getFormattedMessage.contains(
+ "retrying adaptive replanning with matching broadcast hash joins disabled")
+ })
+ }
+ }
+
+ test("Fallback to shuffled join should recover from nested failed broadcast branch") {
+ withTempView("fallback_t1", "fallback_t2", "fallback_t3") {
+ val widePayload = "x" * 2000
Review Comment:
Is this used in the test?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org