[SPARK-35455][SQL] Unify empty relation optimization between normal and AQE optimizer

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

@@ -173,7 +173,8 @@ abstract class Optimizer(catalogManager: CatalogManager)
     // LocalRelation and does not trigger many rules.
     Batch("LocalRelation early", fixedPoint,
       ConvertToLocalRelation,
-      PropagateEmptyRelation,
+      PropagateEmptyRelationBasic,
+      PropagateEmptyRelationAdvanced(),
       // PropagateEmptyRelation can change the nullability of an attribute from nullable to
       // non-nullable when an empty relation child of a Union is removed
       UpdateAttributeNullability) ::
@@ -221,7 +222,8 @@ abstract class Optimizer(catalogManager: CatalogManager)
       ReassignLambdaVariableID) :+
     Batch("LocalRelation", fixedPoint,
       ConvertToLocalRelation,
-      PropagateEmptyRelation,
+      PropagateEmptyRelationBasic,
+      PropagateEmptyRelationAdvanced(),
       // PropagateEmptyRelation can change the nullability of an attribute from nullable to
       // non-nullable when an empty relation child of a Union is removed
       UpdateAttributeNullability) :+

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

@@ -20,23 +20,20 @@ package org.apache.spark.sql.catalyst.optimizer
 import org.apache.spark.sql.catalyst.analysis.CastSupport
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.Literal.FalseLiteral
+import org.apache.spark.sql.catalyst.planning.ExtractSingleColumnNullAwareAntiJoin
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
-import org.apache.spark.sql.catalyst.trees.TreePattern.{LOCAL_RELATION, TRUE_OR_FALSE_LITERAL}
+import org.apache.spark.sql.catalyst.trees.TreePattern.LOCAL_RELATION
 
 /**
  * Collapse plans consisting empty local relations generated by [[PruneFilters]].
- * 1. Binary(or Higher)-node Logical Plans
+ * 1. Higher-node Logical Plans
  *    - Union with all empty children.
- *    - Join with one or two empty children (including Intersect/Except).
  * 2. Unary-node Logical Plans
- *    - Project/Filter/Sample/Join/Limit/Repartition with all empty children.
- *    - Join with false condition.
- *    - Aggregate with all empty children and at least one grouping expression.
- *    - Generate(Explode) with all empty children. Others like Hive UDTF may return results.
+ *    - Project/Filter/Sample with all empty children.
  */
-object PropagateEmptyRelation extends Rule[LogicalPlan] with PredicateHelper with CastSupport {
+object PropagateEmptyRelationBasic extends Rule[LogicalPlan] {
   private def isEmptyLocalRelation(plan: LogicalPlan): Boolean = plan match {
     case p: LocalRelation => p.data.isEmpty
     case _ => false
@@ -45,12 +42,8 @@ object PropagateEmptyRelation extends Rule[LogicalPlan] with PredicateHelper wit
   private def empty(plan: LogicalPlan) =
     LocalRelation(plan.output, data = Seq.empty, isStreaming = plan.isStreaming)
 
-  // Construct a project list from plan's output, while the value is always NULL.
-  private def nullValueProjectList(plan: LogicalPlan): Seq[NamedExpression] =
-    plan.output.map{ a => Alias(cast(Literal(null), a.dataType), a.name)(a.exprId) }
-
-  def apply(plan: LogicalPlan): LogicalPlan = plan.transformUpWithPruning(
-    _.containsAnyPattern(LOCAL_RELATION, TRUE_OR_FALSE_LITERAL), ruleId) {
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.transformUpWithPruning(
+    _.containsAnyPattern(LOCAL_RELATION), ruleId) {
     case p: Union if p.children.exists(isEmptyLocalRelation) =>
       val newChildren = p.children.filterNot(isEmptyLocalRelation)
       if (newChildren.isEmpty) {
@@ -72,11 +65,72 @@ object PropagateEmptyRelation extends Rule[LogicalPlan] with PredicateHelper wit
         }
       }
 
+    case p: UnaryNode if p.children.nonEmpty && p.children.forall(isEmptyLocalRelation) => p match {
+      case _: Project => empty(p)
+      case _: Filter => empty(p)
+      case _: Sample => empty(p)
+      case _ => p
+    }
+  }
+}
+
+/**
+ * The rule used by both normal Optimizer and AQE Optimizer for:
+ *  1. Binary-node Logical Plans
+ *     - Join with one or two empty children (including Intersect/Except).
+ *     - Join is single column NULL-aware anti join (NAAJ)
+ *       Broadcasted [[HashedRelation]] is [[HashedRelationWithAllNullKeys]]. Eliminate join to an
+ *       empty [[LocalRelation]].
+ *     - Left semi Join
+ *       Right side is non-empty and condition is empty. Eliminate join to its left side.
+ *     - Left anti join
+ *       Right side is non-empty and condition is empty. Eliminate join to an empty
+ *       [[LocalRelation]].
+ *  2. Unary-node Logical Plans
+ *     - Limit/Repartition with all empty children.
+ *     - Aggregate with all empty children and at least one grouping expression.
+ *     - Generate(Explode) with all empty children. Others like Hive UDTF may return results.
+ *
+ * @param checkRowCount At AQE side, we use the query stage stats to check the check.
+ * @param isRelationWithAllNullKeys At AQE side, we use the broadcast query stage to do the check.
+ */
+case class PropagateEmptyRelationAdvanced(
+    checkRowCount: Option[(LogicalPlan, Boolean) => Boolean] = None,
+    isRelationWithAllNullKeys: Option[LogicalPlan => Boolean] = None)
+  extends Rule[LogicalPlan] with CastSupport {
+
+  private def isEmptyLocalRelation(plan: LogicalPlan): Boolean = {
+    val defaultEmptyRelation: Boolean = plan match {
+      case p: LocalRelation => p.data.isEmpty
+      case _ => false
+    }
+
+    if (checkRowCount.isDefined) {
+      checkRowCount.get.apply(plan, false) || defaultEmptyRelation
+    } else {
+      defaultEmptyRelation
+    }
+  }
+
+  private def empty(plan: LogicalPlan) =
+    LocalRelation(plan.output, data = Seq.empty, isStreaming = plan.isStreaming)
+
+  // Construct a project list from plan's output, while the value is always NULL.
+  private def nullValueProjectList(plan: LogicalPlan): Seq[NamedExpression] =
+    plan.output.map{ a => Alias(cast(Literal(null), a.dataType), a.name)(a.exprId) }
+
+  // We can not use transformUpWithPruning here since this rule is used by both normal Optimizer
+  // and AQE Optimizer. And this may only effective at AQE side.
+  def apply(plan: LogicalPlan): LogicalPlan = plan.transformUp {
+    case j @ ExtractSingleColumnNullAwareAntiJoin(_, _)
+      if isRelationWithAllNullKeys.isDefined && isRelationWithAllNullKeys.get(j.right) =>
+      empty(j)
+
     // Joins on empty LocalRelations generated from streaming sources are not eliminated
     // as stateful streaming joins need to perform other state management operations other than
     // just processing the input data.
     case p @ Join(_, _, joinType, conditionOpt, _)
-        if !p.children.exists(_.isStreaming) =>
+      if !p.children.exists(_.isStreaming) =>
       val isLeftEmpty = isEmptyLocalRelation(p.left)
       val isRightEmpty = isEmptyLocalRelation(p.right)
       val isFalseCondition = conditionOpt match {
@@ -103,14 +157,17 @@ object PropagateEmptyRelation extends Rule[LogicalPlan] with PredicateHelper wit
             Project(nullValueProjectList(p.left) ++ p.right.output, p.right)
           case _ => p
         }
+      } else if (joinType == LeftSemi && conditionOpt.isEmpty &&
+        checkRowCount.isDefined && checkRowCount.get.apply(p.right, true)) {
+        p.left
+      } else if (joinType == LeftAnti && conditionOpt.isEmpty &&
+        checkRowCount.isDefined && checkRowCount.get.apply(p.right, true)) {
+        empty(p)
       } else {
         p
       }
 
     case p: UnaryNode if p.children.nonEmpty && p.children.forall(isEmptyLocalRelation) => p match {
-      case _: Project => empty(p)
-      case _: Filter => empty(p)
-      case _: Sample => empty(p)
       case _: Sort => empty(p)
       case _: GlobalLimit if !p.isStreaming => empty(p)
       case _: LocalLimit if !p.isStreaming => empty(p)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala

@@ -117,7 +117,7 @@ object RuleIdCollection {
       "org.apache.spark.sql.catalyst.optimizer.OptimizeRepartition" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeWindowFunctions" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeUpdateFields"::
-      "org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelation" ::
+      "org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelationBasic" ::
       "org.apache.spark.sql.catalyst.optimizer.PruneFilters" ::
       "org.apache.spark.sql.catalyst.optimizer.PushDownLeftSemiAntiJoin" ::
       "org.apache.spark.sql.catalyst.optimizer.PushExtraPredicateThroughJoin" ::

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeLimitZeroSuite.scala

@@ -33,7 +33,8 @@ class OptimizeLimitZeroSuite extends PlanTest {
       Batch("OptimizeLimitZero", Once,
         ReplaceIntersectWithSemiJoin,
         OptimizeLimitZero,
-        PropagateEmptyRelation) :: Nil
+        PropagateEmptyRelationBasic,
+        PropagateEmptyRelationAdvanced()) :: Nil
   }
 
   val testRelation1 = LocalRelation.fromExternalRows(Seq('a.int), data = Seq(Row(1)))

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizerRuleExclusionSuite.scala

@@ -118,7 +118,8 @@ class OptimizerRuleExclusionSuite extends PlanTest {
   test("Verify optimized plan after excluding CombineUnions rule") {
     val excludedRules = Seq(
       ConvertToLocalRelation.ruleName,
-      PropagateEmptyRelation.ruleName,
+      PropagateEmptyRelationBasic.ruleName,
+      PropagateEmptyRelationAdvanced().ruleName,
       CombineUnions.ruleName)
 
     val testRelation1 = LocalRelation('a.int, 'b.int, 'c.int)

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelationSuite.scala

@@ -38,7 +38,8 @@ class PropagateEmptyRelationSuite extends PlanTest {
         ReplaceIntersectWithSemiJoin,
         PushPredicateThroughNonJoin,
         PruneFilters,
-        PropagateEmptyRelation,
+        PropagateEmptyRelationBasic,
+        PropagateEmptyRelationAdvanced(),
         CollapseProject) :: Nil
   }
 

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

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.execution.adaptive
 
+import org.apache.spark.sql.catalyst.analysis.UpdateAttributeNullability
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, LogicalPlanIntegrity, PlanHelper}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.internal.SQLConf
@@ -27,7 +28,9 @@ import org.apache.spark.util.Utils
  */
 class AQEOptimizer(conf: SQLConf) extends RuleExecutor[LogicalPlan] {
   private val defaultBatches = Seq(
-    Batch("Eliminate Unnecessary Join", Once, EliminateUnnecessaryJoin),
+    Batch("Propagate Empty LocalRelation", Once,
+      PropagateEmptyRelationAdvancedHelper.propagateEmptyRelationAdvanced,
+      UpdateAttributeNullability),
     Batch("Demote BroadcastHashJoin", Once, DemoteBroadcastHashJoin)
   )
 

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

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.adaptive
+
+import org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelationAdvanced
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.execution.joins.HashedRelationWithAllNullKeys
+
+/**
+ * A helper class to provide a AQE side `PropagateEmptyRelationAdvanced` rule.
+ */
+object PropagateEmptyRelationAdvancedHelper {
+
+  private def isRelationWithAllNullKeys(plan: LogicalPlan) = plan match {
+    case LogicalQueryStage(_, stage: BroadcastQueryStageExec)
+      if stage.resultOption.get().isDefined =>
+      stage.broadcast.relationFuture.get().value == HashedRelationWithAllNullKeys
+    case _ => false
+  }
+
+  private def checkRowCount(plan: LogicalPlan, hasRow: Boolean): Boolean = plan match {
+    case LogicalQueryStage(_, stage: QueryStageExec) if stage.resultOption.get().isDefined =>
+      stage.getRuntimeStatistics.rowCount match {
+        case Some(count) => hasRow == (count > 0)
+        case _ => false
+      }
+    case _ => false
+  }
+
+  lazy val propagateEmptyRelationAdvanced: PropagateEmptyRelationAdvanced = {
+    PropagateEmptyRelationAdvanced(Some(checkRowCount), Some(isRelationWithAllNullKeys))
+  }
+}

sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala

@@ -21,6 +21,7 @@ import org.scalatest.GivenWhenThen
 
 import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression}
 import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
+import org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelationAdvanced
 import org.apache.spark.sql.catalyst.plans.ExistenceJoin
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive._
@@ -1382,7 +1383,7 @@ abstract class DynamicPartitionPruningSuiteBase
     withSQLConf(
       SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
       SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
-      SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key -> EliminateUnnecessaryJoin.ruleName) {
+      SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key -> PropagateEmptyRelationAdvanced().ruleName) {
       val df = sql(
         """
           |SELECT * FROM fact_sk f