[SPARK-36082][SQL][FOLLOWUP] Preserve NAAJ broadcast fallback behavior

### What changes were proposed in this pull request?

This is a follow-up to [#55678](#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 <wenchen@databricks.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit f8b673c)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: cloud-fan

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
       }
   }
 }
-
-

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
+    }
+  }
+}

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
+  }
 }
 
 /**

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,25 +160,71 @@ 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)
       assert(canBroadcastBySize(right, SQLConf.get) === true)
     }
   }
 
-  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())
     }
   }