From 0430322258a5aeee5b2c3d74467e5e0be29e810c Mon Sep 17 00:00:00 2001
From: Angerszhuuuu <angers.zhu@gmail.com>
Date: Mon, 1 Jul 2024 10:30:37 +0800
Subject: [PATCH] [SPARK-48750][SQL] AQEPropagateEmptyRelation convert
 broadcast query stage plan to empty relation causing error

---
 .../execution/adaptive/AdaptiveSparkPlanExec.scala    | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

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 f21960aeedd64..ef75ca6f532e4 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
@@ -386,7 +386,16 @@ case class AdaptiveSparkPlanExec(
               sideBySide(currentPhysicalPlan.treeString, newPhysicalPlan.treeString).mkString("\n")
             logOnLevel(log"Plan changed:\n${MDC(QUERY_PLAN, plans)}")
             cleanUpTempTags(newPhysicalPlan)
-            currentPhysicalPlan = newPhysicalPlan
+            currentPhysicalPlan = currentPhysicalPlan match {
+              case broadcast: BroadcastQueryStageExec =>
+                broadcast.plan match {
+                  case b: BroadcastExchangeExec =>
+                    broadcast.copy(plan = b.copy(child = newPhysicalPlan))
+                  case ReusedExchangeExec(_, b: BroadcastExchangeExec) =>
+                    broadcast.copy(plan = b.copy(child = newPhysicalPlan))
+                }
+              case p => p
+            }
             currentLogicalPlan = newLogicalPlan
             stagesToReplace = Seq.empty[QueryStageExec]
           }