From cdb511408fd0eed58cec2119debb176abd7896bf Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 5 Aug 2026 17:26:39 +0800 Subject: [PATCH 1/2] [SPARK-58588][SQL] Return the broadcast hash join build side from JoinSelectionHelper instead of a Boolean ### What changes were proposed in this pull request? `canPlanAsBroadcastHashJoin` already computed an `Option[BuildSide]` internally and then discarded the direction with `.isDefined`. A caller that needs the direction therefore had to reimplement the same composition: `PushDownJoinThroughUnion` (SPARK-58449) did exactly that, so one planner decision was modelled twice, with each model incomplete and the coupling invisible in the code. This PR lifts the composition into `JoinSelectionHelper.getBroadcastHashJoinBuildSide`, which returns `Option[BuildSide]`, and redefines `canPlanAsBroadcastHashJoin` as `.isDefined` on it. The truth table of `canPlanAsBroadcastHashJoin` is unchanged, including its two deliberate over-approximations of `JoinSelection`: join keys no hash join supports still reach the size branch, and the method does not model `SHUFFLE_MERGE` or `SHUFFLE_REPLICATE_NL` hints. Both are now stated in the scaladoc instead of being implicit. `PushDownJoinThroughUnion` then drops its own `canPlanAsBroadcastHashJoin` conjunct, since `None` from the new method already covers the equi-key and hash-joinable-key requirements, and shares a `branchJoin` helper between the rewrite and the guard's probe plan. The probe now rewrites the join condition to the branch output, which the previous guard did not need to do because it never extracted equi-join keys. ### Why are the changes needed? Two partial models of the same planner decision drift apart. Strengthening one of them, for example teaching `canPlanAsBroadcastHashJoin` about a hint it currently ignores, silently leaves the other stale, and nothing in the code says the two have to agree. ### Does this PR introduce _any_ user-facing change? No. `canPlanAsBroadcastHashJoin` keeps its truth table, and `spark.sql.optimizer.pushDownJoinThroughUnion.enabled` remains false by default. ### How was this patch tested? `JoinSuite.assertJoin` now also asserts that `getBroadcastHashJoinBuildSide` agrees with the build side the planner picked, which covers every case in that suite that reaches a broadcast hash join. `JoinSelectionHelperSuite` gains five cases asserting the direction itself: a hint on either or both sides, the fall back to the smaller side, a shuffle hash hint, no equi-join keys, and a null-aware anti join. The direction had no coverage before this PR, so both were checked by mutation: inverting `getSmallerSide` fails four cases, and returning `BuildLeft` for the null-aware anti join fails exactly the new case for it. Existing suites pass: `JoinSelectionHelperSuite` (19), `PushDownJoinThroughUnionSuite` in catalyst (24) and core (9), and `JoinSuite` (58). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code --- .../optimizer/PushDownJoinThroughUnion.scala | 86 +++++++++++-------- .../spark/sql/catalyst/optimizer/joins.scala | 38 ++++++-- .../optimizer/JoinSelectionHelperSuite.scala | 52 +++++++++++ .../org/apache/spark/sql/JoinSuite.scala | 8 ++ 4 files changed, 137 insertions(+), 47 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala index 65c2f63c463eb..54773a5275902 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala +++ b/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,67 @@ 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` returning `None` also covers the joins where no broadcast hash + * join can be planned at all, so this one check subsumes the equi-key and hash-joinable-key + * requirements. 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) } } 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 edd63829a7113..a6d2af6d1415a 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 @@ -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 `JoinSelection` could not plan + * one at all. + * + * `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 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 bac20c6ed3533..61ef8be59a883 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 @@ -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)) + } + } + } 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 7f695e90df884..3c0b00793ca5e 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 @@ -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 } From 529d10a273b775a2e915b1c196f7ca5f329ce029 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 6 Aug 2026 16:37:35 +0800 Subject: [PATCH 2/2] [SPARK-58588][SQL][FOLLOWUP] Drop the hash-joinable-key claim from the guard's scaladoc `getBroadcastHashJoinBuildSide` does not carry the hash-joinable-key requirement: with equi keys no hash join supports, `noShufflePlannedBefore` becomes true and the size branch still answers `Some(...)`, while `JoinSelection` falls through to a sort merge join. Only the equi-key requirement of the removed `canPlanAsBroadcastHashJoin` conjunct is carried over. Tighten the method's own summary for the same reason: it returns `None` when a broadcast hash join is ruled out by the join shape or by a hint, which is narrower than "`JoinSelection` could not plan one". --- .../optimizer/PushDownJoinThroughUnion.scala | 13 +++++++------ .../apache/spark/sql/catalyst/optimizer/joins.scala | 4 ++-- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala index 54773a5275902..5454c6d88b001 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala @@ -109,12 +109,13 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf) * 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` returning `None` also covers the joins where no broadcast hash - * join can be planned at all, so this one check subsumes the equi-key and hash-joinable-key - * requirements. 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. + * `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 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 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 a6d2af6d1415a..8988183fa3c94 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 @@ -416,8 +416,8 @@ trait JoinSelectionHelper extends Logging { } /** - * The build side a broadcast hash join would use, or `None` when `JoinSelection` could not plan - * one at all. + * 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