You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by su...@apache.org on 2026/03/17 02:55:29 UTC
(spark) 03/11: address codex comments
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch bhj-fallback
in repository https://gitbox.apache.org/repos/asf/spark.git
commit d429090302c94844aa89f49e4ea47cc8d7d66643
Author: Chao Sun <ch...@openai.com>
AuthorDate: Mon Mar 9 18:21:04 2026 -0700
address codex comments
---
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 12 +++---
.../adaptive/AdaptiveQueryExecSuite.scala | 48 +++++++++++++++++-----
2 files changed, 43 insertions(+), 17 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index 9fb35bfd7339..60856754c904 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
@@ -51,6 +51,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec._
import org.apache.spark.sql.execution.bucketing.{CoalesceBucketsInJoin, DisableUnnecessaryBucketedScan}
import org.apache.spark.sql.execution.columnar.InMemoryTableScanLike
import org.apache.spark.sql.execution.exchange._
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
import org.apache.spark.sql.execution.ui.{SparkListenerSQLAdaptiveExecutionUpdate, SparkListenerSQLAdaptiveSQLMetricUpdates, SQLPlanMetric}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -897,11 +898,10 @@ case class AdaptiveSparkPlanExec(
Seq.empty
}
- private def containsBroadcastExchangeOrStage(plan: SparkPlan): Boolean = {
+ private def containsBroadcastJoin(plan: SparkPlan): Boolean = {
plan.exists {
- case _: BroadcastQueryStageExec => true
- case _: BroadcastExchangeLike => true
- case ReusedExchangeExec(_, _: BroadcastExchangeLike) => true
+ case _: BroadcastHashJoinExec => true
+ case _: BroadcastNestedLoopJoinExec => true
case _ => false
}
}
@@ -909,13 +909,13 @@ case class AdaptiveSparkPlanExec(
private def shouldRejectBroadcastPlanInFallback(
plan: SparkPlan,
broadcastJoinDisabledForFallback: Boolean): Boolean = {
- broadcastJoinDisabledForFallback && containsBroadcastExchangeOrStage(plan)
+ broadcastJoinDisabledForFallback && containsBroadcastJoin(plan)
}
private def cancelSupersededStages(oldPlan: SparkPlan, newPlan: SparkPlan): Unit = {
val newStages = collectExchangeQueryStages(newPlan)
collectExchangeQueryStages(oldPlan).foreach { stage =>
- if (!newStages.exists(_.eq(stage))) {
+ if (!newStages.exists(_.sameResult(stage))) {
try {
stage.cancel()
} catch {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 4c4eb3551aa7..be37f5de7ea6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -714,19 +714,36 @@ class AdaptiveQueryExecSuite
val logicalPlan = adaptivePlan.inputPlan.logicalLink.get
val reOptimize =
PrivateMethod[Option[(SparkPlan, LogicalPlan)]](Symbol("reOptimize"))
- val containsBroadcastExchangeOrStage =
- PrivateMethod[Boolean](Symbol("containsBroadcastExchangeOrStage"))
+ val containsBroadcastJoin =
+ PrivateMethod[Boolean](Symbol("containsBroadcastJoin"))
val shouldRejectBroadcastPlanInFallback =
PrivateMethod[Boolean](Symbol("shouldRejectBroadcastPlanInFallback"))
val fallbackPlan = adaptivePlan.invokePrivate(reOptimize(logicalPlan, true)).get._1
assert(findTopLevelBroadcastNestedLoopJoin(fallbackPlan).nonEmpty)
- assert(adaptivePlan.invokePrivate(containsBroadcastExchangeOrStage(fallbackPlan)))
+ assert(adaptivePlan.invokePrivate(containsBroadcastJoin(fallbackPlan)))
assert(adaptivePlan.invokePrivate(
shouldRejectBroadcastPlanInFallback(fallbackPlan, true)))
}
}
+ test("Fallback to shuffled join should allow non-join broadcast plans") {
+ val qe = sql("SELECT key FROM testData").queryExecution
+ val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
+ val adaptivePlan = AdaptiveSparkPlanExec(
+ inputPlan,
+ AdaptiveExecutionContext(spark, qe),
+ Seq.empty,
+ isSubquery = false)
+ val containsBroadcastJoin = PrivateMethod[Boolean](Symbol("containsBroadcastJoin"))
+ val shouldRejectBroadcastPlanInFallback =
+ PrivateMethod[Boolean](Symbol("shouldRejectBroadcastPlanInFallback"))
+
+ assert(!adaptivePlan.invokePrivate(containsBroadcastJoin(inputPlan)))
+ assert(!adaptivePlan.invokePrivate(
+ shouldRejectBroadcastPlanInFallback(inputPlan, true)))
+ }
+
test("Fallback to shuffled join should cancel superseded stages") {
val qe = sql("SELECT key FROM testData").queryExecution
val adaptivePlan = AdaptiveSparkPlanExec(
@@ -737,13 +754,18 @@ class AdaptiveQueryExecSuite
val cancelSupersededStages =
PrivateMethod[Unit](Symbol("cancelSupersededStages"))
- val reusedStage = new TestExchangeQueryStageExec(1)
- val staleStage = new TestExchangeQueryStageExec(2)
- val oldPlan = UnionExec(Seq(reusedStage, staleStage))
+ val originalStage = new TestExchangeQueryStageExec(1)
+ val reusedStage = originalStage.newReuseInstance(2, originalStage.output)
+ .asInstanceOf[TestExchangeQueryStageExec]
+ val staleStage = new TestExchangeQueryStageExec(
+ 3,
+ _canonicalized = sql("SELECT a FROM testData2").queryExecution.sparkPlan.canonicalized)
+ val oldPlan = UnionExec(Seq(originalStage, staleStage))
val newPlan = UnionExec(Seq(reusedStage))
adaptivePlan.invokePrivate(cancelSupersededStages(oldPlan, newPlan))
+ assert(originalStage.cancelCounter.get() == 0)
assert(reusedStage.cancelCounter.get() == 0)
assert(staleStage.cancelCounter.get() == 1)
}
@@ -3045,10 +3067,10 @@ class AdaptiveQueryExecSuite
private case class TestExchangeQueryStageExec(
override val id: Int,
- cancelCounter: AtomicInteger = new AtomicInteger(0)) extends ExchangeQueryStageExec {
- private val leafPlan = LocalTableScanExec(Nil, Nil)
- override val plan: SparkPlan = leafPlan
- override val _canonicalized: SparkPlan = leafPlan.canonicalized
+ cancelCounter: AtomicInteger = new AtomicInteger(0),
+ override val _canonicalized: SparkPlan = TestExchangeQueryStageExec.leafPlan.canonicalized)
+ extends ExchangeQueryStageExec {
+ override val plan: SparkPlan = TestExchangeQueryStageExec.leafPlan
override protected def doMaterialize(): Future[Any] = Future.successful(null)
@@ -3056,7 +3078,7 @@ class AdaptiveQueryExecSuite
newStageId: Int,
newOutput: Seq[
org.apache.spark.sql.catalyst.expressions.Attribute]): ExchangeQueryStageExec = {
- new TestExchangeQueryStageExec(newStageId)
+ new TestExchangeQueryStageExec(newStageId, cancelCounter, _canonicalized)
}
override def cancel(): Unit = {
@@ -3065,6 +3087,10 @@ class AdaptiveQueryExecSuite
override def getRuntimeStatistics: Statistics = Statistics(sizeInBytes = 1)
}
+
+ private object TestExchangeQueryStageExec {
+ val leafPlan: SparkPlan = LocalTableScanExec(Nil, Nil)
+ }
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org