[SPARK-58588][SQL] Return the broadcast hash join build side from JoinSelectionHelper instead of a Boolean

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

@@ -57,12 +57,8 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf)
     plan.transformUpWithPruning(
       _.containsAllPatterns(JOIN, UNION), ruleId) {
 
-    case join @ Join(u: Union, right, joinType, joinCond, hint)
+    case join @ Join(u: Union, right, joinType, _, _)
       if (joinType == Inner || joinType == LeftOuter) &&
-        // Requires equi-join keys, and rejects the join when a shuffle hash hint already decided
-        // the strategy. The check below adds the build side direction, which this one does not
-        // constrain.
-        canPlanAsBroadcastHashJoin(join, conf) &&
         broadcastsRightForEveryBranch(u, join) &&
         // Exclude right subtrees containing subqueries, as DeduplicateRelations
         // may not correctly handle correlated references when cloning.
@@ -79,53 +75,68 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf)
           val deduped = dedupRight(right)
           (deduped, AttributeMap(right.output.zip(deduped.output)))
         }
-        val leftRewrites = AttributeMap(unionHeadOutput.zip(child.output))
-        val newCond = joinCond.map(_.transform {
-          case a: Attribute if leftRewrites.contains(a) => leftRewrites(a)
-          case a: Attribute if rightRewrites.contains(a) => rightRewrites(a)
-        })
-        Join(child, newRight, joinType, newCond, hint)
+        branchJoin(join, unionHeadOutput, child, newRight, rightRewrites)
       }
       u.withNewChildren(newChildren)
   }
   }
 
+  /**
+   * The join for one `Union` branch: the branch on the left, `newRight` on the right, and the
+   * condition rewritten from the `Union` output to the outputs of both new children.
+   */
+  private def branchJoin(
+      join: Join,
+      unionHeadOutput: Seq[Attribute],
+      child: LogicalPlan,
+      newRight: LogicalPlan,
+      rightRewrites: AttributeMap[Attribute]): Join = {
+    val leftRewrites = AttributeMap(unionHeadOutput.zip(child.output))
+    val newCond = join.condition.map(_.transform {
+      case a: Attribute if leftRewrites.contains(a) => leftRewrites(a)
+      case a: Attribute if rightRewrites.contains(a) => rightRewrites(a)
+    })
+    Join(child, newRight, join.joinType, newCond, join.hint)
+  }
+
   /**
    * Whether every join produced by the rewrite is expected to broadcast its right side.
    *
-   * `canPlanAsBroadcastHashJoin` is not enough. It holds when either side is broadcastable, and
-   * for an inner join the planner may build from either side, choosing the smaller one when both
-   * qualify. The rewrite replaces the `Union` on the left with one of its children, so the build
-   * side is decided per branch against a smaller left. Any branch that ends up building from the
-   * left leaves its copy of the right side as a plain probe input, which is not reused, so the
-   * right side is read once per such branch instead of once in total.
+   * Asking whether a broadcast hash join is possible is not enough: for an inner join the planner
+   * may build from either side, choosing the smaller one when both qualify. The rewrite replaces
+   * the `Union` on the left with one of its children, so the build side is decided per branch
+   * against a smaller left. Any branch that ends up building from the left leaves its copy of the
+   * right side as a plain probe input, which is not reused, so the right side is read once per such
+   * branch instead of once in total.
+   *
+   * `getBroadcastHashJoinBuildSide` returns `None` when the join has no equi-join keys, so this one
+   * check also carries the requirement the removed `canPlanAsBroadcastHashJoin` conjunct used to.
+   * It is not a statement about what the planner ends up choosing: with keys no hash join supports
+   * it still answers from the sizes, while `JoinSelection` falls through to a sort merge join. A
+   * `SHUFFLE_MERGE` or `SHUFFLE_REPLICATE_NL` hint is handled here rather than there: the planner
+   * tries those between a hinted broadcast and a size-based one, and falls back to the sizes only
+   * when the hinted strategy does not apply, so declining to rewrite errs on the safe side.
    *
-   * The check follows the planner's broadcast precedence, hints before sizes. It predicts rather
-   * than guarantees the final build side, because later rules and AQE re-estimate the sizes it
-   * reads. It is also all or nothing: one branch small enough to be the build side blocks the
-   * rewrite for the others, trading a missed optimization for never duplicating a probe side. Only
-   * an inner join can build from the left, so for a left outer join this reduces to asking whether
-   * the right side is broadcastable at all.
+   * The result predicts rather than guarantees the final build side, because later rules and AQE
+   * re-estimate the sizes it reads. It is also all or nothing: one branch small enough to be the
+   * build side blocks the rewrite for the others, trading a missed optimization for never
+   * duplicating a probe side. Only an inner join can build from the left, so for a left outer join
+   * this reduces to asking whether the right side is broadcastable at all.
    */
   private def broadcastsRightForEveryBranch(u: Union, join: Join): Boolean = {
-    // The planner tries these hints between a hinted broadcast and a size-based one, so the sizes
-    // below are usually not what decides the strategy. When the hinted strategy turns out not to
-    // apply the planner does fall back to a size-based broadcast, so this errs towards not
-    // rewriting. A shuffle hash hint needs no such check: given hash-joinable keys it makes
-    // `canPlanAsBroadcastHashJoin` reject the join, and a broadcast hint that outranks it is
-    // already answered by the hint-only query below.
     val hintPicksOtherStrategy =
       hintToSortMergeJoin(join.hint) || hintToShuffleReplicateNL(join.hint)
+    val unionHeadOutput = u.children.head.output
     u.children.forall { child =>
-      // Only the join type, the hints and each side's stats are read below, so the condition is
-      // left as is, still referencing the Union output rather than the branch output.
-      val branchJoin = join.copy(left = child)
-      getBroadcastBuildSide(branchJoin, hintOnly = true, conf)
-        .orElse {
-          if (hintPicksOtherStrategy) None
-          else getBroadcastBuildSide(branchJoin, hintOnly = false, conf)
-        }
-        .contains(BuildRight)
+      // The condition has to be rewritten to the branch output: `getBroadcastHashJoinBuildSide`
+      // extracts the equi-join keys, which do not resolve against a branch other than the first.
+      val probe =
+        branchJoin(join, unionHeadOutput, child, join.right, AttributeMap.empty[Attribute])
+      // A hinted broadcast outranks the hints above, so only a size-based answer is vetoed.
+      val hinted = getBroadcastBuildSide(probe, hintOnly = true, conf).isDefined
+      val planned =
+        if (hinted || !hintPicksOtherStrategy) getBroadcastHashJoinBuildSide(probe, conf) else None
+      planned.contains(BuildRight)
     }
   }
 

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

@@ -415,18 +415,38 @@ trait JoinSelectionHelper extends Logging {
     result
   }
 
-  def canPlanAsBroadcastHashJoin(join: Join, conf: SQLConf): Boolean = join match {
+  /**
+   * The build side a broadcast hash join would use, or `None` when one is ruled out by the join
+   * shape or by a hint.
+   *
+   * `Some` does not promise the planner picks a broadcast hash join: a `SHUFFLE_MERGE` or
+   * `SHUFFLE_REPLICATE_NL` hint is tried before the sizes are consulted, join keys no hash join
+   * supports send it to a sort merge join, and AQE re-estimates the sizes at runtime. Within the
+   * broadcast decision itself this does follow the planner's precedence: a hinted broadcast first,
+   * a hinted shuffle hash join as a veto, then the sizes. Callers that only need to know whether a
+   * broadcast hash join is possible should use `canPlanAsBroadcastHashJoin`.
+   */
+  def getBroadcastHashJoinBuildSide(join: Join, conf: SQLConf): Option[BuildSide] = join match {
     case ExtractEquiJoinKeys(_, leftKeys, rightKeys, _, _, _, _, _) =>
-      val hashJoinSupport = hashJoinSupported(leftKeys, rightKeys)
-      val noShufflePlannedBefore =
-        !hashJoinSupport || getShuffleHashJoinBuildSide(join, hintOnly = true, conf).isEmpty
-      getBroadcastBuildSide(join, hintOnly = true, conf).isDefined ||
-        (noShufflePlannedBefore &&
-          getBroadcastBuildSide(join, hintOnly = false, conf).isDefined)
-    case j @ ExtractSingleColumnNullAwareAntiJoin(_, _) => canBroadcastBySize(j.right, conf)
-    case _ => false
+      // A shuffle hash hint outranks a size-based broadcast, so it vetoes one. Keys no hash join
+      // supports cannot honor that hint either, so it does not veto here; the sizes then still
+      // produce an answer, which over-approximates `JoinSelection` (it falls through to a sort
+      // merge join). That over-approximation predates this method and is kept deliberately, so
+      // that `canPlanAsBroadcastHashJoin` keeps its truth table.
+      val noShufflePlannedBefore = !hashJoinSupported(leftKeys, rightKeys) ||
+        getShuffleHashJoinBuildSide(join, hintOnly = true, conf).isEmpty
+      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 _ => None
   }
 
+  def canPlanAsBroadcastHashJoin(join: Join, conf: SQLConf): Boolean =
+    getBroadcastHashJoinBuildSide(join, conf).isDefined
+
   def canPruneLeft(joinType: JoinType): Boolean = joinType match {
     case Inner | LeftSemi | RightOuter => true
     case _ => false

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

@@ -171,4 +171,56 @@ class JoinSelectionHelperSuite extends PlanTest with JoinSelectionHelper {
     }
   }
 
+  test("getBroadcastHashJoinBuildSide returns the hinted side") {
+    val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head)))
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+      assert(getBroadcastHashJoinBuildSide(
+        equiJoin.copy(hint = JoinHint(hintBroadcast, None)), SQLConf.get) === Some(BuildLeft))
+      assert(getBroadcastHashJoinBuildSide(
+        equiJoin.copy(hint = JoinHint(None, hintBroadcast)), SQLConf.get) === Some(BuildRight))
+      // Both sides hinted: the smaller one wins, as in `getBroadcastBuildSide`.
+      assert(getBroadcastHashJoinBuildSide(
+        equiJoin.copy(hint = JoinHint(hintBroadcast, hintBroadcast)), SQLConf.get) ===
+        Some(BuildRight))
+    }
+  }
+
+  test("getBroadcastHashJoinBuildSide falls back to the smaller side") {
+    val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head)))
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      // Only the right side is under the threshold, so it is the only candidate.
+      assert(getBroadcastHashJoinBuildSide(equiJoin, SQLConf.get) === Some(BuildRight))
+    }
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+      assert(getBroadcastHashJoinBuildSide(equiJoin, SQLConf.get).isEmpty)
+    }
+  }
+
+  test("getBroadcastHashJoinBuildSide returns None when a shuffle hash hint applies") {
+    val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head)))
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(getBroadcastHashJoinBuildSide(
+        equiJoin.copy(hint = JoinHint(None, hintShuffleHash)), SQLConf.get).isEmpty)
+    }
+  }
+
+  test("getBroadcastHashJoinBuildSide returns None without equi-join keys") {
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(getBroadcastHashJoinBuildSide(join, SQLConf.get).isEmpty)
+    }
+  }
+
+  test("getBroadcastHashJoinBuildSide builds from the right for a null-aware anti join") {
+    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)
+
+    withSQLConf(
+      SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
+      assert(getBroadcastHashJoinBuildSide(nullAwareAntiJoin, SQLConf.get) === Some(BuildRight))
+    }
+  }
+
 }

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

@@ -83,6 +83,14 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
       canPlanAsBroadcastHashJoin(optimized.asInstanceOf[Join], conf) ===
         operators.head.isInstanceOf[BroadcastHashJoinExec],
       "canPlanAsBroadcastHashJoin not in sync with join selection codepath!")
+    operators.head match {
+      case bhj: BroadcastHashJoinExec =>
+        assert(
+          getBroadcastHashJoinBuildSide(optimized.asInstanceOf[Join], conf)
+            .contains(bhj.buildSide),
+          "getBroadcastHashJoinBuildSide not in sync with join selection codepath!")
+      case _ =>
+    }
     operators.head
   }