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 65c2f63c463e..5454c6d88b00 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,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) } } 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 edd63829a711..8988183fa3c9 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 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 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 bac20c6ed353..61ef8be59a88 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 7f695e90df88..3c0b00793ca5 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 }