[SPARK-48155][SQL] AQEPropagateEmptyRelation for join should check if remain child is just BroadcastQueryStageExec

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelation.scala

@@ -65,6 +65,8 @@ abstract class PropagateEmptyRelationBase extends Rule[LogicalPlan] with CastSup
  private def nullValueProjectList(plan: LogicalPlan): Seq[NamedExpression] =
  plan.output.map{ a => Alias(cast(Literal(null), a.dataType), a.name)(a.exprId) }
 
+ protected def canPropagate(plan: LogicalPlan): Boolean = true
+
  protected def commonApplyFunc: PartialFunction[LogicalPlan, LogicalPlan] = {
  case p: Union if p.children.exists(isEmpty) =>
  val newChildren = p.children.filterNot(isEmpty)
@@ -111,18 +113,19 @@ abstract class PropagateEmptyRelationBase extends Rule[LogicalPlan] with CastSup
  case LeftSemi if isRightEmpty | isFalseCondition => empty(p)
  case LeftAnti if isRightEmpty | isFalseCondition => p.left
  case FullOuter if isLeftEmpty && isRightEmpty => empty(p)
- case LeftOuter | FullOuter if isRightEmpty =>
+ case LeftOuter | FullOuter if isRightEmpty && canPropagate(p.left) =>
  Project(p.left.output ++ nullValueProjectList(p.right), p.left)
  case RightOuter if isRightEmpty => empty(p)
- case RightOuter | FullOuter if isLeftEmpty =>
+ case RightOuter | FullOuter if isLeftEmpty && canPropagate(p.right) =>
  Project(nullValueProjectList(p.left) ++ p.right.output, p.right)
- case LeftOuter if isFalseCondition =>
+ case LeftOuter if isFalseCondition && canPropagate(p.left) =>
  Project(p.left.output ++ nullValueProjectList(p.right), p.left)
- case RightOuter if isFalseCondition =>
+ case RightOuter if isFalseCondition && canPropagate(p.right) =>
  Project(nullValueProjectList(p.left) ++ p.right.output, p.right)
  case _ => p
  }
- } else if (joinType == LeftSemi && conditionOpt.isEmpty && nonEmpty(p.right)) {
+ } else if (joinType == LeftSemi && conditionOpt.isEmpty &&
+ nonEmpty(p.right) && canPropagate(p.left)) {
  p.left
  } else if (joinType == LeftAnti && conditionOpt.isEmpty && nonEmpty(p.right)) {
  empty(p)

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala

@@ -82,6 +82,11 @@ object AQEPropagateEmptyRelation extends PropagateEmptyRelationBase {
  case _ => false
  }
 
+ override protected def canPropagate(plan: LogicalPlan): Boolean = plan match {
+ case LogicalQueryStage(_, _: BroadcastQueryStageExec) => false
+ case _ => true
+ }
+
  override protected def applyInternal(p: LogicalPlan): LogicalPlan = p.transformUpWithPruning(
  // LOCAL_RELATION and TRUE_OR_FALSE_LITERAL pattern are matched at
  // `PropagateEmptyRelationBase.commonApplyFunc`

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala

@@ -164,6 +164,12 @@ class AdaptiveQueryExecSuite
  }
  }
 
+ private def findTopLevelUnion(plan: SparkPlan): Seq[UnionExec] = {
+ collect(plan) {
+ case l: UnionExec => l
+ }
+ }
+
  private def findReusedExchange(plan: SparkPlan): Seq[ReusedExchangeExec] = {
  collectWithSubqueries(plan) {
  case ShuffleQueryStageExec(_, e: ReusedExchangeExec, _) => e
@@ -2772,6 +2778,30 @@ class AdaptiveQueryExecSuite
  }
  }
 
+ test("SPARK-48155: AQEPropagateEmptyRelation check remained child for join") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+ val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
+ """
+ |SELECT /*+ BROADCAST(t3) */ t3.b, count(t3.a) FROM testData2 t1
+ |INNER JOIN (
+ | SELECT * FROM testData2
+ | WHERE b = 0
+ | UNION ALL
+ | SELECT * FROM testData2
+ | WHErE b != 0
+ |) t2
+ |ON t1.b = t2.b AND t1.a = 0
+ |RIGHT OUTER JOIN testData2 t3
+ |ON t1.a > t3.a
+ |GROUP BY t3.b
+ """.stripMargin
+ )
+ assert(findTopLevelBroadcastNestedLoopJoin(adaptivePlan).size == 1)
+ assert(findTopLevelUnion(adaptivePlan).size == 0)
+ }
+ }
+
  test("SPARK-39915: Dataset.repartition(N) may not create N partitions") {
  withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "6") {
  // partitioning: HashPartitioning