You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/09/04 06:25:55 UTC

(spark) branch master updated: [SPARK-36082][SQL][FOLLOWUP] Preserve NAAJ broadcast fallback behavior

This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new f8b673c5a530 [SPARK-36082][SQL][FOLLOWUP] Preserve NAAJ broadcast fallback behavior
f8b673c5a530 is described below

commit f8b673c5a5302c97d3b4fc129b494e891abd6a2f
Author: Wenchen Fan <we...@databricks.com>
AuthorDate: Fri Sep 4 14:25:36 2026 +0800

    [SPARK-36082][SQL][FOLLOWUP] Preserve NAAJ broadcast fallback behavior
    
    ### What changes were proposed in this pull request?
    
    This is a follow-up to [#55678](https://github.com/apache/spark/pull/55678). It preserves the
    optimized single-column null-aware anti join (NAAJ) hash path without making the fallback worse
    than it was before that change.
    
    The patch recognizes the structural single-column NAAJ form independently of whether the hash
    optimization is enabled. `SparkStrategies.JoinSelection` then owns the complete physical decision
    for that form. A shared `NullAwareAntiJoinPlanning` decision selects:
    
    1. the specialized null-aware `BroadcastHashJoinExec`, which only supports `BuildRight`, when the
       optimization is enabled, the right side is within `spark.sql.autoBroadcastJoinThreshold`, and
       the generic broadcast nested-loop fallback would also build the right side; or
    2. `BroadcastNestedLoopJoinExec` with the exact build side selected by the generic fallback.
    
    The size boundary is checked before fallback build-side selection. Therefore, an over-threshold
    right side rejects the hash path without computing irrelevant statistics for an uncached left
    subtree. For non-inner/full joins, the nested-loop helper also checks the fixed preferred side
    before consulting the opposite side's statistics.
    
    Logical eligibility checks delegate to the same decision used by physical planning. This avoids
    reimplementing physical hint precedence and build-side rules in the optimizer while still letting
    the optimizer ask whether a plan will use the specialized hash operator.
    
    Aggregate pushdown keeps separate stability guards. Structural NAAJs require both the original and
    the pushed form to select the specialized hash outcome. Ordinary semi and anti joins preserve the
    SPARK-34081 broadcast-eligibility gate: when the original join cannot use broadcast hash join, the
    join remains above the Aggregate so the Aggregate can reduce its input first.
    
    ### Why are the changes needed?
    
    The null-aware hash operator always builds the right side, while the generic broadcast nested-loop
    join can build either side. If the fallback chooses `BuildLeft`, forcing the specialized hash path
    can require broadcasting a much larger right side and regress a query that previously had a valid
    plan.
    
    An over-threshold `BuildRight` input also used the nested-loop fallback before #55678. Planner
    estimates do not provide a reliable proof that every such input fits the runtime hash
    representation, so this follow-up keeps the hash optimization inside the existing broadcast-size
    boundary. This also preserves the nested-loop equality semantics for threshold-disabled
    floating-point NAAJs, including `-0.0` versus `0.0`.
    
    Making the final choice in the physical planner gives the structural NAAJ exactly one owner. It
    also ensures that disabling the hash optimization takes the explicit nested-loop fallback instead
    of depending on the ordinary join-selection path to rediscover the right behavior.
    
    The optimizer must also preserve the existing SPARK-34081 behavior for ordinary semi and anti
    joins. Unconditionally moving such a join below an Aggregate when broadcasting is disabled makes a
    shuffle join consume the ungrouped source rows and reintroduces the performance regression that
    SPARK-34081 avoided. The ordinary branch therefore retains its broadcast-eligibility guard without
    changing the structural NAAJ decision.
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes. Relative to #55678 on unreleased branches, a recognized single-column NAAJ uses the null-aware
    broadcast hash join only when the generic fallback would build the right side and that side is
    within the normal broadcast-size threshold.
    
    If the optimization is disabled, the fallback would build the left side, or the right side exceeds
    the threshold, Spark retains the broadcast nested-loop join. Ordinary semi and anti Aggregate
    pushdown also retains its released broadcast-eligibility behavior. Released Spark behavior is not
    changed.
    
    ### How was this patch tested?
    
    The NAAJ changes were tested with these focused checks:
    
    ```
    build/sbt \
      "catalyst/testOnly org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelperSuite" \
      "catalyst/testOnly org.apache.spark.sql.catalyst.optimizer.LeftSemiAntiJoinPushDownSuite" \
      "sql/testOnly org.apache.spark.sql.JoinSuite -- -z SPARK-36082"
    ```
    
    After restoring the ordinary-join gate, the final head also passed:
    
    ```
    build/sbt 'catalyst/testOnly *LeftSemiAntiJoinPushDownSuite -- -z "SPARK-34081"'
    ```
    
    The tests cover the in-threshold `BuildRight` hash path, disabled optimization, the `BuildLeft`
    fallback, the over-threshold `BuildRight` fallback, uncached-left statistics short-circuiting,
    floating-point signed-zero correctness, structural-NAAJ Aggregate stability, and broadcast-eligible
    versus disabled ordinary semi/anti Aggregate pushdown.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5)
    
    Closes #58404 from cloud-fan/fix-naaj-broadcast-fallback.
    
    Authored-by: Wenchen Fan <we...@databricks.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../optimizer/PushDownLeftSemiAntiJoin.scala       |  63 ++++++++----
 .../spark/sql/catalyst/optimizer/joins.scala       |  86 ++++++++++++++++-
 .../spark/sql/catalyst/planning/patterns.scala     |   9 +-
 .../optimizer/JoinSelectionHelperSuite.scala       |  67 ++++++++++++-
 .../optimizer/LeftSemiAntiJoinPushDownSuite.scala  |  63 +++++++++++-
 .../spark/sql/execution/SparkStrategies.scala      |  54 ++++-------
 .../scala/org/apache/spark/sql/JoinSuite.scala     | 106 ++++++++++++++++++++-
 7 files changed, 375 insertions(+), 73 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownLeftSemiAntiJoin.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownLeftSemiAntiJoin.scala
index a5fcbe6f16b3..d6612bca4e9e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownLeftSemiAntiJoin.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownLeftSemiAntiJoin.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.expressions._
+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.Rule
@@ -60,12 +61,11 @@ object PushDownLeftSemiAntiJoin extends Rule[LogicalPlan]
         }
       }
 
-    // LeftSemi/LeftAnti over Aggregate, only push down if join can be planned as broadcast join.
+    // LeftSemi/LeftAnti over Aggregate
     case join @ Join(agg: Aggregate, rightOp, LeftSemiOrAnti(_), joinCond, _)
         if agg.aggregateExpressions.forall(_.deterministic) && agg.groupingExpressions.nonEmpty &&
           !agg.aggregateExpressions.exists(ScalarSubquery.hasCorrelatedScalarSubquery) &&
-          canPushThroughCondition(agg.children, joinCond, rightOp) &&
-          canPlanAsBroadcastHashJoin(join, conf) =>
+          canPushThroughCondition(agg.children, joinCond, rightOp) =>
       val aliasMap = getAliasMap(agg)
       val canPushDownPredicate = (predicate: Expression) => {
         val replaced = replaceAlias(predicate, aliasMap)
@@ -75,7 +75,20 @@ object PushDownLeftSemiAntiJoin extends Rule[LogicalPlan]
       val makeJoinCondition = (predicates: Seq[Expression]) => {
         replaceAlias(predicates.reduce(And), aliasMap)
       }
-      pushDownJoin(join, canPushDownPredicate, makeJoinCondition)
+      val canPushDownJoin = if (ExtractSingleColumnNullAwareAntiJoin.extract(join).isDefined) {
+        val originalIsBroadcastHash =
+          NullAwareAntiJoinPlanning.decide(join, conf) == NullAwareAntiJoinPlanning.BroadcastHash
+        (pushedJoin: Join) => originalIsBroadcastHash &&
+          NullAwareAntiJoinPlanning.decide(pushedJoin, conf) ==
+            NullAwareAntiJoinPlanning.BroadcastHash
+      } else {
+        (_: Join) => canPlanAsBroadcastHashJoin(join, conf)
+      }
+      pushDownJoin(
+        join,
+        canPushDownPredicate,
+        makeJoinCondition,
+        canPushDownJoin)
 
     // LeftSemi/LeftAnti over Window
     case join @ Join(w: Window, rightOp, LeftSemiOrAnti(_), _, _)
@@ -133,11 +146,17 @@ object PushDownLeftSemiAntiJoin extends Rule[LogicalPlan]
   private def pushDownJoin(
       join: Join,
       canPushDownPredicate: Expression => Boolean,
-      makeJoinCondition: Seq[Expression] => Expression): LogicalPlan = {
+      makeJoinCondition: Seq[Expression] => Expression,
+      canPushDownJoin: Join => Boolean = _ => true): LogicalPlan = {
     assert(join.left.children.length == 1)
 
     if (join.condition.isEmpty) {
-      join.left.withNewChildren(Seq(join.copy(left = join.left.children.head)))
+      val pushedJoin = join.copy(left = join.left.children.head)
+      if (canPushDownJoin(pushedJoin)) {
+        join.left.withNewChildren(Seq(pushedJoin))
+      } else {
+        join
+      }
     } else {
       val (pushDown, stayUp) = splitConjunctivePredicates(join.condition.get)
         .partition(canPushDownPredicate)
@@ -151,19 +170,25 @@ object PushDownLeftSemiAntiJoin extends Rule[LogicalPlan]
       if (pushDown.isEmpty || referRightSideCols)  {
         join
       } else {
-        val newPlan = join.left.withNewChildren(Seq(join.copy(
-          left = join.left.children.head, condition = Some(makeJoinCondition(pushDown)))))
-        // If there is no more filter to stay up, return the new plan that has join pushed down.
-        if (stayUp.isEmpty) {
-          newPlan
+        val pushedJoin = join.copy(
+          left = join.left.children.head, condition = Some(makeJoinCondition(pushDown)))
+        if (!canPushDownJoin(pushedJoin)) {
+          join
         } else {
-          join.joinType match {
-            // In case of Left semi join, the part of the join condition which does not refer to
-            // to attributes of the grandchild are kept as a Filter above.
-            case LeftSemi => Filter(stayUp.reduce(And), newPlan)
-            // In case of left-anti join, the join is pushed down only when the entire join
-            // condition is eligible to be pushed down to preserve the semantics of left-anti join.
-            case _ => join
+          val newPlan = join.left.withNewChildren(Seq(pushedJoin))
+          // If no predicates remain above the join, return the plan with the join pushed down.
+          if (stayUp.isEmpty) {
+            newPlan
+          } else {
+            join.joinType match {
+              // For a left semi join, the non-pushable part of the condition is kept as a Filter
+              // above the join.
+              case LeftSemi => Filter(stayUp.reduce(And), newPlan)
+              // In the case of a left anti join, the join is pushed down only when the entire join
+              // condition is eligible to be pushed down to preserve the semantics of the left anti
+              // join.
+              case _ => join
+            }
           }
         }
       }
@@ -273,5 +298,3 @@ object PushLeftSemiLeftAntiThroughJoin extends Rule[LogicalPlan] with PredicateH
       }
   }
 }
-
-
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
index 8988183fa3c9..6a4b332b35cd 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
@@ -353,6 +353,66 @@ trait JoinSelectionHelper extends Logging {
     }
   }
 
+  def getBroadcastNestedLoopJoinDesiredBuildSide(join: Join): BuildSide = {
+    if (join.joinType.isInstanceOf[InnerLike] || join.joinType == FullOuter) {
+      getSmallerSide(join.left, join.right)
+    } else {
+      // For perf reasons, BroadcastNestedLoopJoinExec prefers to broadcast the left side for a
+      // right join and the right side for a left join. If one side is much smaller, revisiting
+      // that preference may be worthwhile.
+      if (canBuildBroadcastLeft(join.joinType)) BuildLeft else BuildRight
+    }
+  }
+
+  def getBroadcastNestedLoopJoinBuildSide(
+      join: Join,
+      hintOnly: Boolean,
+      conf: SQLConf): Option[BuildSide] = {
+    lazy val buildLeft = if (hintOnly) {
+      hintToBroadcastLeft(join.hint)
+    } else {
+      canBroadcastBySize(join.left, conf) &&
+        !hintToNotBroadcastAndReplicateLeft(join.hint)
+    }
+    lazy val buildRight = if (hintOnly) {
+      hintToBroadcastRight(join.hint)
+    } else {
+      canBroadcastBySize(join.right, conf) &&
+        !hintToNotBroadcastAndReplicateRight(join.hint)
+    }
+
+    if (join.joinType.isInstanceOf[InnerLike] || join.joinType == FullOuter) {
+      if (buildLeft && buildRight) {
+        Some(getBroadcastNestedLoopJoinDesiredBuildSide(join))
+      } else if (buildLeft) {
+        Some(BuildLeft)
+      } else if (buildRight) {
+        Some(BuildRight)
+      } else {
+        None
+      }
+    } else {
+      getBroadcastNestedLoopJoinDesiredBuildSide(join) match {
+        case BuildLeft =>
+          if (buildLeft) Some(BuildLeft) else if (buildRight) Some(BuildRight) else None
+        case BuildRight =>
+          if (buildRight) Some(BuildRight) else if (buildLeft) Some(BuildLeft) else None
+      }
+    }
+  }
+
+  def getBroadcastNestedLoopJoinBuildSide(join: Join, conf: SQLConf): BuildSide = {
+    val hintedBuildSide = if (join.hint.isEmpty) {
+      None
+    } else {
+      getBroadcastNestedLoopJoinBuildSide(join, hintOnly = true, conf)
+    }
+    hintedBuildSide
+      .orElse(getBroadcastNestedLoopJoinBuildSide(join, hintOnly = false, conf))
+      .orElse(getBroadcastNestedLoopJoinBuildSide(join.hint, join.joinType))
+      .getOrElse(getBroadcastNestedLoopJoinDesiredBuildSide(join))
+  }
+
   def getSmallerSide(left: LogicalPlan, right: LogicalPlan): BuildSide = {
     if (right.stats.sizeInBytes <= left.stats.sizeInBytes) BuildRight else BuildLeft
   }
@@ -438,9 +498,13 @@ trait JoinSelectionHelper extends Logging {
       getBroadcastBuildSide(join, hintOnly = true, conf).orElse {
         if (noShufflePlannedBefore) getBroadcastBuildSide(join, hintOnly = false, conf) else None
       }
-    // `JoinSelection` always builds from the right for this shape.
-    case j @ ExtractSingleColumnNullAwareAntiJoin(_, _) =>
-      if (canBroadcastBySize(j.right, conf)) Some(BuildRight) else None
+    case j if ExtractSingleColumnNullAwareAntiJoin.extract(j).isDefined =>
+      if (NullAwareAntiJoinPlanning.decide(j, conf) ==
+          NullAwareAntiJoinPlanning.BroadcastHash) {
+        Some(BuildRight)
+      } else {
+        None
+      }
     case _ => None
   }
 
@@ -564,3 +628,19 @@ trait JoinSelectionHelper extends Logging {
       conf.getConfString("spark.sql.join.forceApplyShuffledHashJoin", "false") == "true"
   }
 }
+
+private[sql] object NullAwareAntiJoinPlanning extends JoinSelectionHelper {
+  sealed trait Decision
+  case object BroadcastHash extends Decision
+  case object BroadcastNestedLoop extends Decision
+
+  def decide(join: Join, conf: SQLConf): Decision = {
+    if (conf.optimizeNullAwareAntiJoin &&
+        canBroadcastBySize(join.right, conf) &&
+        getBroadcastNestedLoopJoinBuildSide(join, conf) == BuildRight) {
+      BroadcastHash
+    } else {
+      BroadcastNestedLoop
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
index c4af18fc8739..5bc5280450d5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
@@ -404,12 +404,11 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre
    * But if it's a single column case O(M*N) calculation could be optimized into O(M)
    * using hash lookup instead of loop lookup.
    */
-  def unapply(join: Join): Option[ReturnType] = join match {
+  private[sql] def extract(join: Join): Option[ReturnType] = join match {
     case Join(left, right, LeftAnti,
       Some(Or(e @ EqualTo(leftAttr: Expression, rightAttr: Expression),
         IsNull(e2 @ EqualTo(_, _)))), _)
-        if SQLConf.get.optimizeNullAwareAntiJoin &&
-          e.semanticEquals(e2) =>
+        if e.semanticEquals(e2) =>
       if (canEvaluate(leftAttr, left) && canEvaluate(rightAttr, right)) {
         Some(Seq(leftAttr), Seq(rightAttr))
       } else if (canEvaluate(leftAttr, right) && canEvaluate(rightAttr, left)) {
@@ -419,6 +418,10 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre
       }
     case _ => None
   }
+
+  def unapply(join: Join): Option[ReturnType] = {
+    if (SQLConf.get.optimizeNullAwareAntiJoin) extract(join) else None
+  }
 }
 
 /**
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala
index 61ef8be59a88..7a593e1a4aaf 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala
@@ -17,15 +17,26 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import java.util.concurrent.atomic.AtomicBoolean
+
 import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{AttributeMap, EqualTo, IsNull, Or}
-import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti, PlanTest}
-import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, HintInfo, Join, JoinHint, NO_BROADCAST_HASH, SHUFFLE_HASH}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, EqualTo, IsNull, Or}
+import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti, PlanTest, RightOuter}
+import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, HintInfo, Join, JoinHint, LeafNode, NO_BROADCAST_HASH, SHUFFLE_HASH, Statistics}
 import org.apache.spark.sql.catalyst.statsEstimation.StatsTestPlan
 import org.apache.spark.sql.internal.SQLConf
 
 class JoinSelectionHelperSuite extends PlanTest with JoinSelectionHelper {
 
+  private case class TrackingStatsTestPlan(
+      override val output: Seq[Attribute],
+      statsAccessed: AtomicBoolean) extends LeafNode {
+    override def computeStats(): Statistics = {
+      statsAccessed.set(true)
+      Statistics(sizeInBytes = 20000000)
+    }
+  }
+
   private val left = StatsTestPlan(
     outputList = Seq($"a".int, $"b".int, $"c".int),
     rowCount = 20000000,
@@ -149,6 +160,22 @@ class JoinSelectionHelperSuite extends PlanTest with JoinSelectionHelper {
     assert(getSmallerSide(left, right) === BuildRight)
   }
 
+  test("getBroadcastNestedLoopJoinBuildSide checks the fixed desired side first") {
+    val leftStatsAccessed = new AtomicBoolean(false)
+    val uncachedLeft = TrackingStatsTestPlan(Seq($"uncachedLeft".int), leftStatsAccessed)
+    val leftAntiJoin = Join(uncachedLeft, right, LeftAnti, None, JoinHint.NONE)
+    val rightStatsAccessed = new AtomicBoolean(false)
+    val uncachedRight = TrackingStatsTestPlan(Seq($"uncachedRight".int), rightStatsAccessed)
+    val rightOuterJoin = Join(right, uncachedRight, RightOuter, None, JoinHint.NONE)
+
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(getBroadcastNestedLoopJoinBuildSide(leftAntiJoin, SQLConf.get) === BuildRight)
+      assert(!leftStatsAccessed.get())
+      assert(getBroadcastNestedLoopJoinBuildSide(rightOuterJoin, SQLConf.get) === BuildLeft)
+      assert(!rightStatsAccessed.get())
+    }
+  }
+
   test("canBroadcastBySize should return true if the plan size is less than 10MB") {
     withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
       assert(canBroadcastBySize(left, SQLConf.get) === false)
@@ -156,18 +183,48 @@ class JoinSelectionHelperSuite extends PlanTest with JoinSelectionHelper {
     }
   }
 
-  test("canPlanAsBroadcastHashJoin should respect size for single-column null-aware anti join") {
+  test("canPlanAsBroadcastHashJoin should respect NAAJ size and nested-loop build side") {
     val leftKey = left.output.head
     val rightKey = right.output.head
     val condition = Or(EqualTo(leftKey, rightKey), IsNull(EqualTo(leftKey, rightKey)))
     val nullAwareAntiJoin = Join(left, right, LeftAnti, Some(condition), JoinHint.NONE)
+    val smallLeft = left.copy(rowCount = 1000, size = Some(1000))
     val largeRight = right.copy(rowCount = 20000000, size = Some(20000000))
 
     withSQLConf(
       SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
       assert(canPlanAsBroadcastHashJoin(nullAwareAntiJoin, SQLConf.get))
-      assert(!canPlanAsBroadcastHashJoin(nullAwareAntiJoin.copy(right = largeRight), SQLConf.get))
+      assert(!canPlanAsBroadcastHashJoin(
+        nullAwareAntiJoin.copy(right = largeRight.copy(rowCount = 1)), SQLConf.get))
+      assert(!canPlanAsBroadcastHashJoin(
+        nullAwareAntiJoin.copy(left = smallLeft, right = largeRight), SQLConf.get))
+      assert(!canPlanAsBroadcastHashJoin(
+        nullAwareAntiJoin.copy(hint = JoinHint(hintBroadcast, None)), SQLConf.get))
+    }
+
+    withSQLConf(
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "false",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(!canPlanAsBroadcastHashJoin(nullAwareAntiJoin, SQLConf.get))
+    }
+  }
+
+  test("canPlanAsBroadcastHashJoin checks the NAAJ right-size short circuit") {
+    val leftStatsAccessed = new AtomicBoolean(false)
+    val uncachedLeft = TrackingStatsTestPlan(Seq($"uncachedLeft".int), leftStatsAccessed)
+    val largeRight = right.copy(rowCount = 20000000, size = Some(20000000))
+    val leftKey = uncachedLeft.output.head
+    val rightKey = largeRight.output.head
+    val condition = Or(EqualTo(leftKey, rightKey), IsNull(EqualTo(leftKey, rightKey)))
+    val nullAwareAntiJoin = Join(
+      uncachedLeft, largeRight, LeftAnti, Some(condition), JoinHint.NONE)
+
+    withSQLConf(
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(!canPlanAsBroadcastHashJoin(nullAwareAntiJoin, SQLConf.get))
+      assert(!leftStatsAccessed.get())
     }
   }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LeftSemiAntiJoinPushDownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LeftSemiAntiJoinPushDownSuite.scala
index ba7386f8c9b5..2b89f752a997 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LeftSemiAntiJoinPushDownSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LeftSemiAntiJoinPushDownSuite.scala
@@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.expressions._
 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.statsEstimation.StatsTestPlan
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.IntegerType
 
@@ -109,6 +110,66 @@ class LeftSemiAntiJoinPushDownSuite extends PlanTest {
     comparePlans(optimized, correctAnswer)
   }
 
+  test("Aggregate: NAAJ pushdown when original and rewritten joins can build right") {
+    val condition = Or($"b" === $"d", IsNull($"b" === $"d"))
+    val originalQuery = testRelation
+      .groupBy($"b")($"b", sum($"c"))
+      .join(testRelation1, joinType = LeftAnti, condition = Some(condition))
+    val correctAnswer = testRelation
+      .join(testRelation1, joinType = LeftAnti, condition = Some(condition))
+      .groupBy($"b")($"b", sum($"c"))
+
+    withSQLConf(SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true") {
+      val analyzedOriginal = originalQuery.analyze
+      val analyzedCorrectAnswer = correctAnswer.analyze
+      val originalJoin = analyzedOriginal.asInstanceOf[Join]
+      val pushedJoin = analyzedCorrectAnswer.asInstanceOf[Aggregate].child.asInstanceOf[Join]
+      assert(PushDownLeftSemiAntiJoin.canPlanAsBroadcastHashJoin(originalJoin, SQLConf.get))
+      assert(PushDownLeftSemiAntiJoin.canPlanAsBroadcastHashJoin(pushedJoin, SQLConf.get))
+      comparePlans(Optimize.execute(analyzedOriginal), analyzedCorrectAnswer)
+    }
+  }
+
+  test("Aggregate: NAAJ no pushdown when the original join would build left") {
+    val leftKey = $"leftKey".int
+    val leftKeyStats = ColumnStat(
+      distinctCount = Some(1),
+      min = Some(0),
+      max = Some(0),
+      nullCount = Some(0),
+      avgLen = Some(4),
+      maxLen = Some(4))
+    val child = StatsTestPlan(
+      Seq(leftKey), 1000, AttributeMap(Seq(leftKey -> leftKeyStats)), Some(1000))
+    val aggregate = Aggregate(Seq(leftKey), Seq(leftKey), child)
+    val right = StatsTestPlan(
+      Seq($"rightKey".int), 1000, AttributeMap(Seq()), Some(1000))
+    val condition = Or(
+      EqualTo(aggregate.output.head, right.output.head),
+      IsNull(EqualTo(aggregate.output.head, right.output.head)))
+    val originalQuery = Join(
+      aggregate, right, LeftAnti, Some(condition), JoinHint.NONE)
+
+    withSQLConf(
+      SQLConf.CBO_ENABLED.key -> "true",
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "100") {
+      assert(aggregate.stats.sizeInBytes <= SQLConf.get.autoBroadcastJoinThreshold)
+      assert(child.stats.sizeInBytes > SQLConf.get.autoBroadcastJoinThreshold)
+      assert(right.stats.sizeInBytes > SQLConf.get.autoBroadcastJoinThreshold)
+      comparePlans(Optimize.execute(originalQuery), originalQuery)
+    }
+  }
+
+  test("Aggregate: ordinary LeftSemi join no pushdown - empty join condition") {
+    val originalQuery = testRelation
+      .groupBy($"b")($"b", sum($"c"))
+      .join(testRelation1, joinType = LeftSemi, condition = None)
+    val correctAnswer = originalQuery.analyze
+
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer)
+  }
+
   test("Aggregate: LeftSemi join no pushdown - non-deterministic aggr expressions") {
     val originalQuery = testRelation
       .groupBy($"b")($"b", Rand(10).as("c"))
@@ -460,7 +521,7 @@ class LeftSemiAntiJoinPushDownSuite extends PlanTest {
   }
 
   Seq(LeftSemi, LeftAnti).foreach { jt =>
-    test(s"SPARK-34081: $jt only push down if join can be planned as broadcast join") {
+    test(s"SPARK-34081: ordinary $jt only pushes down when broadcast-eligible") {
       Seq(-1, 100000).foreach { threshold =>
         withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> threshold.toString) {
           val originalQuery = testRelation
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index e6250dd88736..85011386656c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.{EliminateSubqueryAliases, NamedRelation}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper, NormalizeFloatingNumbers}
+import org.apache.spark.sql.catalyst.optimizer.{BuildRight, BuildSide, JoinSelectionHelper, NormalizeFloatingNumbers, NullAwareAntiJoinPlanning}
 import org.apache.spark.sql.catalyst.planning._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
@@ -340,10 +340,18 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
             .getOrElse(createJoinWithoutHint())
         }
 
-      case j @ ExtractSingleColumnNullAwareAntiJoin(leftKeys, rightKeys)
-          if canBroadcastBySize(j.right, conf) =>
-        Seq(joins.BroadcastHashJoinExec(leftKeys, rightKeys, LeftAnti, BuildRight,
-          None, planLater(j.left), planLater(j.right), isNullAwareAntiJoin = true))
+      case j: logical.Join if ExtractSingleColumnNullAwareAntiJoin.extract(j).isDefined =>
+        val (leftKeys, rightKeys) = ExtractSingleColumnNullAwareAntiJoin.extract(j).get
+        NullAwareAntiJoinPlanning.decide(j, conf) match {
+          case NullAwareAntiJoinPlanning.BroadcastHash =>
+            Seq(joins.BroadcastHashJoinExec(leftKeys, rightKeys, LeftAnti, BuildRight,
+              None, planLater(j.left), planLater(j.right), isNullAwareAntiJoin = true))
+          case NullAwareAntiJoinPlanning.BroadcastNestedLoop =>
+            checkHintNonEquiJoin(j.hint)
+            val buildSide = getBroadcastNestedLoopJoinBuildSide(j, conf)
+            Seq(joins.BroadcastNestedLoopJoinExec(
+              planLater(j.left), planLater(j.right), buildSide, LeftAnti, j.condition))
+        }
 
       // If it is not an equi-join, we first look at the join hints w.r.t. the following order:
       //   1. broadcast hint: pick broadcast nested loop join. If both sides have the broadcast
@@ -361,42 +369,12 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       //   3. Pick broadcast nested loop join as the final solution. It may OOM but we don't have
       //      other choice. It broadcasts the smaller side for inner and full joins, broadcasts the
       //      left side for right join, and broadcasts right side for left join.
-      case logical.Join(left, right, joinType, condition, hint) =>
+      case j @ logical.Join(left, right, joinType, condition, hint) =>
         checkHintNonEquiJoin(hint)
-        val desiredBuildSide = if (joinType.isInstanceOf[InnerLike] || joinType == FullOuter) {
-          getSmallerSide(left, right)
-        } else {
-          // For perf reasons, `BroadcastNestedLoopJoinExec` prefers to broadcast left side if
-          // it's a right join, and broadcast right side if it's a left join.
-          // TODO: revisit it. If left side is much smaller than the right side, it may be better
-          // to broadcast the left side even if it's a left join.
-          if (canBuildBroadcastLeft(joinType)) BuildLeft else BuildRight
-        }
+        val desiredBuildSide = getBroadcastNestedLoopJoinDesiredBuildSide(j)
 
         def createBroadcastNLJoin(onlyLookingAtHint: Boolean) = {
-          val buildLeft = if (onlyLookingAtHint) {
-            hintToBroadcastLeft(hint)
-          } else {
-            canBroadcastBySize(left, conf) && !hintToNotBroadcastAndReplicateLeft(hint)
-          }
-
-          val buildRight = if (onlyLookingAtHint) {
-            hintToBroadcastRight(hint)
-          } else {
-            canBroadcastBySize(right, conf) && !hintToNotBroadcastAndReplicateRight(hint)
-          }
-
-          val maybeBuildSide = if (buildLeft && buildRight) {
-            Some(desiredBuildSide)
-          } else if (buildLeft) {
-            Some(BuildLeft)
-          } else if (buildRight) {
-            Some(BuildRight)
-          } else {
-            None
-          }
-
-          maybeBuildSide.map { buildSide =>
+          getBroadcastNestedLoopJoinBuildSide(j, onlyLookingAtHint, conf).map { buildSide =>
             Seq(joins.BroadcastNestedLoopJoinExec(
               planLater(left), planLater(right), buildSide, joinType, condition))
           }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
index 3c0b00793ca5..d950357ec2dd 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
@@ -1299,14 +1299,114 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
     }
   }
 
-  test("SPARK-36082: only use SingleColumn Null Aware Anti Join when right side " +
-      "can broadcast") {
+  test("SPARK-36082: keep in-threshold NAAJ hash join when the fallback builds right") {
+    withSQLConf(SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString) {
+      val joinExec = assertJoin((
+        "select * from testData where key not in (select b from testData3)",
+        classOf[BroadcastHashJoinExec]))
+      assert(joinExec.asInstanceOf[BroadcastHashJoinExec].isNullAwareAntiJoin)
+    }
+  }
+
+  test("SPARK-36082: disabled NAAJ hash optimization uses nested-loop fallback") {
+    withSQLConf(
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "false",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString) {
+      val joinExec = assertJoin((
+        "select * from testData where key not in (select b from testData3)",
+        classOf[BroadcastNestedLoopJoinExec]))
+      assert(joinExec.asInstanceOf[BroadcastNestedLoopJoinExec].buildSide === BuildRight)
+    }
+  }
+
+  test("SPARK-36082: keep over-threshold NAAJ nested-loop join when fallback builds right") {
     withSQLConf(SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "0") {
       val joinExec = assertJoin((
         "select * from testData where key not in (select b from testData3)",
         classOf[BroadcastNestedLoopJoinExec]))
-      assert(!joinExec.isInstanceOf[BroadcastHashJoinExec])
+      assert(joinExec.asInstanceOf[BroadcastNestedLoopJoinExec].buildSide === BuildRight)
+    }
+  }
+
+  test("SPARK-36082: keep NAAJ nested-loop join when the fallback builds left") {
+    val smallLeft = spark.range(1).selectExpr("id AS key")
+    val largeRight = spark.range(100)
+      .selectExpr("IF(id = 0, CAST(NULL AS BIGINT), id) AS b")
+    val threshold = statisticSizeInByte(smallLeft)
+    assert(threshold < statisticSizeInByte(largeRight))
+
+    withTempView("naajSmallLeft", "naajLargeRight") {
+      smallLeft.createOrReplaceTempView("naajSmallLeft")
+      largeRight.createOrReplaceTempView("naajLargeRight")
+      withSQLConf(
+        SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> threshold.toString) {
+        val result = sql(
+          "select * from naajSmallLeft where key not in (select b from naajLargeRight)")
+        val joinExec = result.queryExecution.sparkPlan.collect {
+          case j: BroadcastNestedLoopJoinExec => j
+        }
+        assert(joinExec.size === 1)
+        assert(joinExec.head.buildSide === BuildLeft)
+        checkAnswer(result, Seq.empty)
+      }
+    }
+  }
+
+  test("SPARK-36082: left-broadcast NAAJ fallback uses nested-loop join") {
+    withSQLConf(
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString) {
+      withTempView("naajHintedLeft", "naajHintedRight") {
+        Seq[java.lang.Double](-0.0d, 2.0d, null).toDF("key")
+          .createOrReplaceTempView("naajHintedLeft")
+        Seq[java.lang.Double](0.0d, 1.0d).toDF("key")
+          .createOrReplaceTempView("naajHintedRight")
+
+        val result = sql(
+          "select /*+ BROADCAST(naajHintedLeft) */ naajHintedLeft.* " +
+            "from naajHintedLeft left anti join naajHintedRight on " +
+            "naajHintedLeft.key = naajHintedRight.key or " +
+            "isnull(naajHintedLeft.key = naajHintedRight.key)")
+        val plan = result.queryExecution.sparkPlan
+        val nestedLoopJoins = plan.collect {
+          case join: BroadcastNestedLoopJoinExec => join
+        }
+        val nullAwareHashJoins = plan.collect {
+          case join: BroadcastHashJoinExec if join.isNullAwareAntiJoin => join
+        }
+        assert(nestedLoopJoins.size === 1)
+        assert(nestedLoopJoins.head.buildSide === BuildLeft)
+        assert(nullAwareHashJoins.isEmpty)
+        checkAnswer(result, Row(2.0d))
+      }
+    }
+  }
+
+  test("SPARK-36082: threshold-disabled NAAJ preserves floating-point equality") {
+    withSQLConf(
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "0") {
+      withTempView("naajFloatingLeft", "naajFloatingRight") {
+        Seq[java.lang.Double](-0.0d, 2.0d, null).toDF("key")
+          .createOrReplaceTempView("naajFloatingLeft")
+        Seq[java.lang.Double](0.0d, 1.0d).toDF("key")
+          .createOrReplaceTempView("naajFloatingRight")
+
+        val result = sql(
+          "select * from naajFloatingLeft " +
+            "where key not in (select key from naajFloatingRight)")
+        val joinExec = result.queryExecution.sparkPlan.collect {
+          case join: BroadcastNestedLoopJoinExec => join
+        }
+        assert(joinExec.size === 1)
+        assert(joinExec.head.buildSide === BuildRight)
+        checkAnswer(result, Row(2.0d))
+      }
     }
   }
 


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