(spark) 07/11: address codex comments

[su...@apache.org]

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 b0693e6a689cc438750ba39160d4d282a938a82b
Author: Chao Sun <ch...@openai.com>
AuthorDate: Mon Mar 9 21:36:01 2026 -0700

    address codex comments
---
 .../execution/adaptive/AdaptiveSparkPlanExec.scala | 23 +++++++++-------------
 .../adaptive/AdaptiveQueryExecSuite.scala          | 10 ++++------
 2 files changed, 13 insertions(+), 20 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 733012dc84a8..90e55cbb3af0 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
@@ -802,7 +802,7 @@ case class AdaptiveSparkPlanExec(
           preprocessingRules ++ queryStagePreparationRules,
           Some((planChangeLogger, "AQE Replanning")))
 
-        val finalPlan = wrapInputBroadcastExchangeIfNeeded(newPlan, disableBroadcastJoin)
+        val finalPlan = wrapInputBroadcastExchangeIfNeeded(newPlan)
         (finalPlan, optimized)
       }
 
@@ -967,25 +967,20 @@ case class AdaptiveSparkPlanExec(
     stages.toSeq
   }
 
-  private def wrapInputBroadcastExchangeIfNeeded(
-      newPlan: SparkPlan,
-      disableBroadcastJoin: Boolean): SparkPlan = {
+  private def wrapInputBroadcastExchangeIfNeeded(newPlan: SparkPlan): SparkPlan = {
     // 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 inserting the `BroadcastExchangeExec` node when broadcast
-    // joins are disabled for fallback or when the new plan already contains the broadcast
-    // exchange after applying the `LogicalQueryStageStrategy` rule.
-    if (disableBroadcastJoin) {
-      newPlan
-    } else {
-      inputPlan match {
-        case b: BroadcastExchangeLike
+    // 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.
+    inputPlan match {
+      case b: BroadcastExchangeLike
           if (!newPlan.isInstanceOf[BroadcastExchangeLike]) => b.withNewChildren(Seq(newPlan))
-        case _ => newPlan
-      }
+      case _ => newPlan
     }
   }
 
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 e136e021c4db..e69036c350af 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
@@ -731,7 +731,7 @@ class AdaptiveQueryExecSuite
     }
   }
 
-  test("Fallback to shuffled join should not preserve input broadcast exchange") {
+  test("Fallback to shuffled join should preserve input broadcast exchange") {
     val qe = sql("SELECT key FROM testData").queryExecution
     val inputPlan = BroadcastExchangeExec(IdentityBroadcastMode, qe.sparkPlan)
     val adaptivePlan = AdaptiveSparkPlanExec(
@@ -743,13 +743,11 @@ class AdaptiveQueryExecSuite
     val wrapInputBroadcastExchangeIfNeeded =
       PrivateMethod[SparkPlan](Symbol("wrapInputBroadcastExchangeIfNeeded"))
 
-    val wrapped = adaptivePlan.invokePrivate(
-      wrapInputBroadcastExchangeIfNeeded(newPlan, false))
+    val wrapped = adaptivePlan.invokePrivate(wrapInputBroadcastExchangeIfNeeded(newPlan))
     assert(wrapped.isInstanceOf[BroadcastExchangeExec])
 
-    val unwrapped = adaptivePlan.invokePrivate(
-      wrapInputBroadcastExchangeIfNeeded(newPlan, true))
-    assert(!unwrapped.isInstanceOf[BroadcastExchangeExec])
+    val fallbackWrapped = adaptivePlan.invokePrivate(wrapInputBroadcastExchangeIfNeeded(newPlan))
+    assert(fallbackWrapped.isInstanceOf[BroadcastExchangeExec])
   }
 
   test("Fallback to shuffled join should reject replans that still require broadcast") {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org