philo-he commented on code in PR #12922: URL: https://github.com/apache/gluten/pull/12922#discussion_r3953006072
########## backends-velox/src/main/scala/org/apache/gluten/extension/VeloxBroadcastNestedLoopJoinRewriteRule.scala: ########## @@ -0,0 +1,244 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, IsNull, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.plans.{FullOuter, LeftOuter, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan, UnionExec} +import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.types.BooleanType + +/** + * Rewrites `BroadcastNestedLoopJoinExec(FullOuter)` into a union of two nested-loop joins that + * Velox already supports natively: + * 1. left outer join to produce matches plus unmatched streamed-side rows + * 2. an outer join with a synthetic non-null marker on the opposite side to identify unmatched + * broadcast-side rows without relying on data columns being non-null + */ +case class VeloxBroadcastNestedLoopJoinRewriteRule() extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan.transformUp { + case bnlj: BroadcastNestedLoopJoinExec + if bnlj.joinType == FullOuter && shouldRewriteFullOuter(bnlj) && conditionOffloadable( + bnlj) && broadcastSideRelocatable(bnlj) => + rewriteFullOuter(bnlj) + } + + /** + * The rewrite reuses the original broadcast side in two roles at once: [[rewriteFullOuter]] keeps + * it as the build (broadcast) side of `branchA`, while [[buildUnmatchedBroadcastSide]] calls + * [[unwrapBroadcast]] on it and consumes the unwrapped subtree as a normal STREAMED input of + * `branchB`. That is only safe when the broadcast side can be cleanly re-materialized as a + * partitioned plan. Reject the rewrite otherwise, e.g. for the MERGE cardinality-check join (`ON + * t.pk > s.pk` with an `autoBroadcastJoinThreshold = -1` broadcast of a reused `Union` source): + * there the broadcast side does not unwrap to a plain partitioned subtree, so after the rewrite a + * `ColumnarBroadcastExchangeExec` ends up in `branchB`'s streamed slot and is executed via + * `ColumnarInputAdapter.doExecuteColumnar -> executeColumnar()`, which the broadcast exchange + * does not support, crashing with `[INTERNAL_ERROR] ... has column support mismatch`. + * + * A broadcast side is considered relocatable only when: + * - it is an exclusively-owned broadcast, i.e. NOT a [[ReusedExchangeExec]] (a reused/shared + * exchange must not be turned into a streamed input); and + * - its unwrapped payload does not itself contain a nested broadcast, which would otherwise + * leak into the streamed position of `branchB`. + */ + private def broadcastSideRelocatable(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val broadcastSide = bnlj.buildSide match { + case BuildLeft => bnlj.left + case BuildRight => bnlj.right + } + isCleanRelocatableBroadcast(broadcastSide) + } + + private def isCleanRelocatableBroadcast(plan: SparkPlan): Boolean = plan match { + case stage: BroadcastQueryStageExec => isCleanRelocatableBroadcast(stage.plan) + case _: ReusedExchangeExec => false + case exchange: BroadcastExchangeLike => !containsBroadcast(exchange.child) + case _ => false + } + + private def containsBroadcast(plan: SparkPlan): Boolean = + plan.exists { + case _: BroadcastExchangeLike => true + case _: BroadcastQueryStageExec => true + case _: ReusedExchangeExec => true + case _ => false + } + + private def shouldRewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val threshold = VeloxConfig.get.broadcastNestedLoopJoinFullOuterRewriteThreshold + bnlj.logicalLink.collect { + case join: Join => + val leftSize = join.left.stats.sizeInBytes + val rightSize = join.right.stats.sizeInBytes + leftSize >= 0 && rightSize >= 0 && leftSize <= threshold && rightSize <= threshold + }.getOrElse(false) + } + + private def extractChildLogicalSizes( + bnlj: BroadcastNestedLoopJoinExec): Option[(BigInt, BigInt)] = + for { + leftLogical <- bnlj.left.logicalLink + rightLogical <- bnlj.right.logicalLink + } yield (leftLogical.stats.sizeInBytes, rightLogical.stats.sizeInBytes) + + private def conditionOffloadable(bnlj: BroadcastNestedLoopJoinExec): Boolean = + bnlj.condition.exists { + cond => + cond.references.exists(bnlj.left.outputSet.contains) && + cond.references.exists(bnlj.right.outputSet.contains) + } + + private def rewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): SparkPlan = { + val matchesAndStreamedUnmatched = bnlj.buildSide match { + case BuildRight => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.left, + bnlj.right, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + case BuildLeft => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.right, + bnlj.left, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + } + + val unmatchedBroadcastRows = bnlj.buildSide match { + case BuildRight => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.right, + otherSide = bnlj.left, + unmatchedSideIsLeft = false, + condition = bnlj.condition, + output = bnlj.output) + case BuildLeft => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.left, + otherSide = bnlj.right, + unmatchedSideIsLeft = true, + condition = bnlj.condition, + output = bnlj.output) + } + + val union = UnionExec(Seq(matchesAndStreamedUnmatched, unmatchedBroadcastRows)) + ProjectExec( + union.output.zip(bnlj.output).map { + case (childAttr, targetAttr) => + Alias(childAttr, targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + }, + union + ) + } + + private def buildUnmatchedBroadcastSide( + unmatchedSide: SparkPlan, + otherSide: SparkPlan, + unmatchedSideIsLeft: Boolean, + condition: Option[Expression], + output: Seq[Attribute]): SparkPlan = { + val unmatchedSideBase = unwrapBroadcast(unmatchedSide) + val otherSideBase = unwrapBroadcast(otherSide) + val markerAttr = + AttributeReference("__gluten_bnlj_matched_build_side", BooleanType, nullable = false)() + val markedOtherSide = ProjectExec( + otherSideBase.output.map(attr => aliasTo(attr, attr)) :+ + Alias(Literal.TrueLiteral, markerAttr.name)(exprId = markerAttr.exprId), + otherSideBase) + val unmatchedJoin = if (unmatchedSideIsLeft) { + BroadcastNestedLoopJoinExec( + unmatchedSideBase, + ensureBroadcast(markedOtherSide), + BuildRight, + LeftOuter, + condition) + } else { + BroadcastNestedLoopJoinExec( + ensureBroadcast(markedOtherSide), + unmatchedSideBase, + BuildLeft, + RightOuter, + condition) + } + val unmatchedOnly = FilterExec(IsNull(markerAttr), unmatchedJoin) + val projected = if (unmatchedSideIsLeft) { + output.zipWithIndex.map { + case (targetAttr, idx) if idx < unmatchedSideBase.output.size => + aliasTo(unmatchedSideBase.output(idx), targetAttr) + case (targetAttr, _) => + Alias(Literal.create(null, targetAttr.dataType), targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + } + } else { + output.zipWithIndex.map { + case (targetAttr, idx) if idx < otherSide.output.size => + Alias(Literal.create(null, targetAttr.dataType), targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) Review Comment: ditto ########## backends-velox/src/main/scala/org/apache/gluten/extension/VeloxBroadcastNestedLoopJoinRewriteRule.scala: ########## @@ -0,0 +1,244 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, IsNull, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.plans.{FullOuter, LeftOuter, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan, UnionExec} +import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.types.BooleanType + +/** + * Rewrites `BroadcastNestedLoopJoinExec(FullOuter)` into a union of two nested-loop joins that + * Velox already supports natively: + * 1. left outer join to produce matches plus unmatched streamed-side rows + * 2. an outer join with a synthetic non-null marker on the opposite side to identify unmatched + * broadcast-side rows without relying on data columns being non-null + */ +case class VeloxBroadcastNestedLoopJoinRewriteRule() extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan.transformUp { + case bnlj: BroadcastNestedLoopJoinExec + if bnlj.joinType == FullOuter && shouldRewriteFullOuter(bnlj) && conditionOffloadable( + bnlj) && broadcastSideRelocatable(bnlj) => + rewriteFullOuter(bnlj) + } + + /** + * The rewrite reuses the original broadcast side in two roles at once: [[rewriteFullOuter]] keeps + * it as the build (broadcast) side of `branchA`, while [[buildUnmatchedBroadcastSide]] calls + * [[unwrapBroadcast]] on it and consumes the unwrapped subtree as a normal STREAMED input of + * `branchB`. That is only safe when the broadcast side can be cleanly re-materialized as a + * partitioned plan. Reject the rewrite otherwise, e.g. for the MERGE cardinality-check join (`ON + * t.pk > s.pk` with an `autoBroadcastJoinThreshold = -1` broadcast of a reused `Union` source): + * there the broadcast side does not unwrap to a plain partitioned subtree, so after the rewrite a + * `ColumnarBroadcastExchangeExec` ends up in `branchB`'s streamed slot and is executed via + * `ColumnarInputAdapter.doExecuteColumnar -> executeColumnar()`, which the broadcast exchange + * does not support, crashing with `[INTERNAL_ERROR] ... has column support mismatch`. + * + * A broadcast side is considered relocatable only when: + * - it is an exclusively-owned broadcast, i.e. NOT a [[ReusedExchangeExec]] (a reused/shared + * exchange must not be turned into a streamed input); and + * - its unwrapped payload does not itself contain a nested broadcast, which would otherwise + * leak into the streamed position of `branchB`. + */ + private def broadcastSideRelocatable(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val broadcastSide = bnlj.buildSide match { + case BuildLeft => bnlj.left + case BuildRight => bnlj.right + } + isCleanRelocatableBroadcast(broadcastSide) + } + + private def isCleanRelocatableBroadcast(plan: SparkPlan): Boolean = plan match { + case stage: BroadcastQueryStageExec => isCleanRelocatableBroadcast(stage.plan) + case _: ReusedExchangeExec => false + case exchange: BroadcastExchangeLike => !containsBroadcast(exchange.child) + case _ => false + } + + private def containsBroadcast(plan: SparkPlan): Boolean = + plan.exists { + case _: BroadcastExchangeLike => true + case _: BroadcastQueryStageExec => true + case _: ReusedExchangeExec => true + case _ => false + } + + private def shouldRewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val threshold = VeloxConfig.get.broadcastNestedLoopJoinFullOuterRewriteThreshold + bnlj.logicalLink.collect { + case join: Join => + val leftSize = join.left.stats.sizeInBytes + val rightSize = join.right.stats.sizeInBytes + leftSize >= 0 && rightSize >= 0 && leftSize <= threshold && rightSize <= threshold + }.getOrElse(false) + } + + private def extractChildLogicalSizes( + bnlj: BroadcastNestedLoopJoinExec): Option[(BigInt, BigInt)] = + for { + leftLogical <- bnlj.left.logicalLink + rightLogical <- bnlj.right.logicalLink + } yield (leftLogical.stats.sizeInBytes, rightLogical.stats.sizeInBytes) + + private def conditionOffloadable(bnlj: BroadcastNestedLoopJoinExec): Boolean = + bnlj.condition.exists { + cond => + cond.references.exists(bnlj.left.outputSet.contains) && + cond.references.exists(bnlj.right.outputSet.contains) + } + + private def rewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): SparkPlan = { + val matchesAndStreamedUnmatched = bnlj.buildSide match { + case BuildRight => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.left, + bnlj.right, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + case BuildLeft => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.right, + bnlj.left, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + } + + val unmatchedBroadcastRows = bnlj.buildSide match { + case BuildRight => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.right, + otherSide = bnlj.left, + unmatchedSideIsLeft = false, + condition = bnlj.condition, + output = bnlj.output) + case BuildLeft => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.left, + otherSide = bnlj.right, + unmatchedSideIsLeft = true, + condition = bnlj.condition, + output = bnlj.output) + } + + val union = UnionExec(Seq(matchesAndStreamedUnmatched, unmatchedBroadcastRows)) + ProjectExec( + union.output.zip(bnlj.output).map { + case (childAttr, targetAttr) => + Alias(childAttr, targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + }, + union + ) + } + + private def buildUnmatchedBroadcastSide( + unmatchedSide: SparkPlan, + otherSide: SparkPlan, + unmatchedSideIsLeft: Boolean, + condition: Option[Expression], + output: Seq[Attribute]): SparkPlan = { + val unmatchedSideBase = unwrapBroadcast(unmatchedSide) + val otherSideBase = unwrapBroadcast(otherSide) + val markerAttr = + AttributeReference("__gluten_bnlj_matched_build_side", BooleanType, nullable = false)() + val markedOtherSide = ProjectExec( + otherSideBase.output.map(attr => aliasTo(attr, attr)) :+ + Alias(Literal.TrueLiteral, markerAttr.name)(exprId = markerAttr.exprId), + otherSideBase) + val unmatchedJoin = if (unmatchedSideIsLeft) { + BroadcastNestedLoopJoinExec( + unmatchedSideBase, + ensureBroadcast(markedOtherSide), + BuildRight, + LeftOuter, + condition) + } else { + BroadcastNestedLoopJoinExec( + ensureBroadcast(markedOtherSide), + unmatchedSideBase, + BuildLeft, + RightOuter, + condition) + } + val unmatchedOnly = FilterExec(IsNull(markerAttr), unmatchedJoin) + val projected = if (unmatchedSideIsLeft) { + output.zipWithIndex.map { + case (targetAttr, idx) if idx < unmatchedSideBase.output.size => + aliasTo(unmatchedSideBase.output(idx), targetAttr) + case (targetAttr, _) => + Alias(Literal.create(null, targetAttr.dataType), targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + } + } else { + output.zipWithIndex.map { + case (targetAttr, idx) if idx < otherSide.output.size => + Alias(Literal.create(null, targetAttr.dataType), targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + case (targetAttr, idx) => + aliasTo(unmatchedSideBase.output(idx - otherSide.output.size), targetAttr) + } + } + ProjectExec(projected, unmatchedOnly) + } + + private def unwrapBroadcast(plan: SparkPlan): SparkPlan = plan match { + case stage: BroadcastQueryStageExec => unwrapBroadcast(stage.plan) + case reused: ReusedExchangeExec => unwrapBroadcast(reused.child) + case exchange: BroadcastExchangeLike => exchange.child + case other => other + } + + private def ensureBroadcast(plan: SparkPlan): SparkPlan = plan match { + case exchange: BroadcastExchangeLike => exchange + case other => BroadcastExchangeExec(IdentityBroadcastMode, other) + } + + private def projectToOutput(child: SparkPlan, output: Seq[Attribute]): ProjectExec = { + val sourceByExprId = child.output.map(attr => attr.exprId -> attr).toMap + ProjectExec( + output.map(targetAttr => aliasTo(sourceByExprId(targetAttr.exprId), targetAttr)), + child) + } + + private def aliasTo(childAttr: Attribute, targetAttr: Attribute): NamedExpression = { + Alias(childAttr, targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + } +} Review Comment: How about adding the following util method also? ``` private def nullAliasFor(targetAttr: Attribute): Alias = { aliasTo(Literal.create(null, targetAttr.dataType), targetAttr) } } ``` ########## backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala: ########## @@ -238,6 +241,16 @@ object VeloxConfig extends ConfigRegistry { .bytesConf(ByteUnit.BYTE) .createWithDefaultString("32MB") + val VELOX_BROADCAST_NESTED_LOOP_JOIN_FULL_OUTER_REWRITE_THRESHOLD = + buildConf( + "spark.gluten.sql.columnar.backend.velox.broadcastNLJ.fullOuterRewriteThreshold") + .doc( + "Maximum per-side plan size in bytes for rewriting a full outer broadcast nested loop " + + "join into two outer joins plus union. The rewrite is applied only when both sides " + + "have known statistics and each side is at or below this threshold.") + .bytesConf(ByteUnit.BYTE) + .createWithDefaultString("10MB") Review Comment: If -1 is set, the rewriting will be disabled, right? If so, could you document this behavior also? ########## backends-velox/src/main/scala/org/apache/gluten/extension/VeloxBroadcastNestedLoopJoinRewriteRule.scala: ########## @@ -0,0 +1,244 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, IsNull, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.plans.{FullOuter, LeftOuter, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan, UnionExec} +import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.types.BooleanType + +/** + * Rewrites `BroadcastNestedLoopJoinExec(FullOuter)` into a union of two nested-loop joins that + * Velox already supports natively: + * 1. left outer join to produce matches plus unmatched streamed-side rows + * 2. an outer join with a synthetic non-null marker on the opposite side to identify unmatched + * broadcast-side rows without relying on data columns being non-null + */ +case class VeloxBroadcastNestedLoopJoinRewriteRule() extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan.transformUp { + case bnlj: BroadcastNestedLoopJoinExec + if bnlj.joinType == FullOuter && shouldRewriteFullOuter(bnlj) && conditionOffloadable( + bnlj) && broadcastSideRelocatable(bnlj) => + rewriteFullOuter(bnlj) + } + + /** + * The rewrite reuses the original broadcast side in two roles at once: [[rewriteFullOuter]] keeps + * it as the build (broadcast) side of `branchA`, while [[buildUnmatchedBroadcastSide]] calls + * [[unwrapBroadcast]] on it and consumes the unwrapped subtree as a normal STREAMED input of + * `branchB`. That is only safe when the broadcast side can be cleanly re-materialized as a + * partitioned plan. Reject the rewrite otherwise, e.g. for the MERGE cardinality-check join (`ON + * t.pk > s.pk` with an `autoBroadcastJoinThreshold = -1` broadcast of a reused `Union` source): + * there the broadcast side does not unwrap to a plain partitioned subtree, so after the rewrite a + * `ColumnarBroadcastExchangeExec` ends up in `branchB`'s streamed slot and is executed via + * `ColumnarInputAdapter.doExecuteColumnar -> executeColumnar()`, which the broadcast exchange + * does not support, crashing with `[INTERNAL_ERROR] ... has column support mismatch`. + * + * A broadcast side is considered relocatable only when: + * - it is an exclusively-owned broadcast, i.e. NOT a [[ReusedExchangeExec]] (a reused/shared + * exchange must not be turned into a streamed input); and + * - its unwrapped payload does not itself contain a nested broadcast, which would otherwise + * leak into the streamed position of `branchB`. + */ + private def broadcastSideRelocatable(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val broadcastSide = bnlj.buildSide match { + case BuildLeft => bnlj.left + case BuildRight => bnlj.right + } + isCleanRelocatableBroadcast(broadcastSide) + } + + private def isCleanRelocatableBroadcast(plan: SparkPlan): Boolean = plan match { + case stage: BroadcastQueryStageExec => isCleanRelocatableBroadcast(stage.plan) + case _: ReusedExchangeExec => false + case exchange: BroadcastExchangeLike => !containsBroadcast(exchange.child) + case _ => false + } + + private def containsBroadcast(plan: SparkPlan): Boolean = + plan.exists { + case _: BroadcastExchangeLike => true + case _: BroadcastQueryStageExec => true + case _: ReusedExchangeExec => true + case _ => false + } + + private def shouldRewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): Boolean = { + val threshold = VeloxConfig.get.broadcastNestedLoopJoinFullOuterRewriteThreshold + bnlj.logicalLink.collect { + case join: Join => + val leftSize = join.left.stats.sizeInBytes + val rightSize = join.right.stats.sizeInBytes + leftSize >= 0 && rightSize >= 0 && leftSize <= threshold && rightSize <= threshold + }.getOrElse(false) + } + + private def extractChildLogicalSizes( + bnlj: BroadcastNestedLoopJoinExec): Option[(BigInt, BigInt)] = + for { + leftLogical <- bnlj.left.logicalLink + rightLogical <- bnlj.right.logicalLink + } yield (leftLogical.stats.sizeInBytes, rightLogical.stats.sizeInBytes) + + private def conditionOffloadable(bnlj: BroadcastNestedLoopJoinExec): Boolean = + bnlj.condition.exists { + cond => + cond.references.exists(bnlj.left.outputSet.contains) && + cond.references.exists(bnlj.right.outputSet.contains) + } + + private def rewriteFullOuter(bnlj: BroadcastNestedLoopJoinExec): SparkPlan = { + val matchesAndStreamedUnmatched = bnlj.buildSide match { + case BuildRight => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.left, + bnlj.right, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + case BuildLeft => + projectToOutput( + BroadcastNestedLoopJoinExec( + bnlj.right, + bnlj.left, + BuildRight, + LeftOuter, + bnlj.condition), + bnlj.output) + } + + val unmatchedBroadcastRows = bnlj.buildSide match { + case BuildRight => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.right, + otherSide = bnlj.left, + unmatchedSideIsLeft = false, + condition = bnlj.condition, + output = bnlj.output) + case BuildLeft => + buildUnmatchedBroadcastSide( + unmatchedSide = bnlj.left, + otherSide = bnlj.right, + unmatchedSideIsLeft = true, + condition = bnlj.condition, + output = bnlj.output) + } + + val union = UnionExec(Seq(matchesAndStreamedUnmatched, unmatchedBroadcastRows)) + ProjectExec( + union.output.zip(bnlj.output).map { + case (childAttr, targetAttr) => + Alias(childAttr, targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) + }, + union + ) + } + + private def buildUnmatchedBroadcastSide( + unmatchedSide: SparkPlan, + otherSide: SparkPlan, + unmatchedSideIsLeft: Boolean, + condition: Option[Expression], + output: Seq[Attribute]): SparkPlan = { + val unmatchedSideBase = unwrapBroadcast(unmatchedSide) + val otherSideBase = unwrapBroadcast(otherSide) + val markerAttr = + AttributeReference("__gluten_bnlj_matched_build_side", BooleanType, nullable = false)() + val markedOtherSide = ProjectExec( + otherSideBase.output.map(attr => aliasTo(attr, attr)) :+ + Alias(Literal.TrueLiteral, markerAttr.name)(exprId = markerAttr.exprId), + otherSideBase) + val unmatchedJoin = if (unmatchedSideIsLeft) { + BroadcastNestedLoopJoinExec( + unmatchedSideBase, + ensureBroadcast(markedOtherSide), + BuildRight, + LeftOuter, + condition) + } else { + BroadcastNestedLoopJoinExec( + ensureBroadcast(markedOtherSide), + unmatchedSideBase, + BuildLeft, + RightOuter, + condition) + } + val unmatchedOnly = FilterExec(IsNull(markerAttr), unmatchedJoin) + val projected = if (unmatchedSideIsLeft) { + output.zipWithIndex.map { + case (targetAttr, idx) if idx < unmatchedSideBase.output.size => + aliasTo(unmatchedSideBase.output(idx), targetAttr) + case (targetAttr, _) => + Alias(Literal.create(null, targetAttr.dataType), targetAttr.name)( + exprId = targetAttr.exprId, + qualifier = targetAttr.qualifier, + explicitMetadata = Some(targetAttr.metadata)) Review Comment: Recommend to use the suggested nullAliasFor. ########## backends-velox/src/main/scala/org/apache/gluten/extension/VeloxBroadcastNestedLoopJoinRewriteRule.scala: ########## @@ -0,0 +1,244 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, IsNull, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.plans.{FullOuter, LeftOuter, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan, UnionExec} +import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.types.BooleanType + +/** + * Rewrites `BroadcastNestedLoopJoinExec(FullOuter)` into a union of two nested-loop joins that + * Velox already supports natively: + * 1. left outer join to produce matches plus unmatched streamed-side rows + * 2. an outer join with a synthetic non-null marker on the opposite side to identify unmatched + * broadcast-side rows without relying on data columns being non-null + */ +case class VeloxBroadcastNestedLoopJoinRewriteRule() extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan.transformUp { + case bnlj: BroadcastNestedLoopJoinExec Review Comment: It would be better to add a fast path to skip if the config is set to -1. ########## gluten-ut/common/src/test/scala/org/apache/spark/sql/execution/GlutenBroadcastNestedLoopJoinFullOuterSuite.scala: ########## @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution + +import org.apache.gluten.config.GlutenConfig +import org.apache.gluten.execution.{BroadcastNestedLoopJoinExecTransformer, SortMergeJoinExecTransformer} +import org.apache.gluten.utils.BackendTestUtils + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{Dataset, GlutenSQLTestsTrait, Row} +import org.apache.spark.sql.catalyst.plans.FullOuter +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper} +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.internal.SQLConf + +import scala.reflect.ClassTag + +/** + * Spark-version-agnostic tests for the full outer `BroadcastNestedLoopJoinExec` rewrite. These + * cases only exercise Gluten/Velox behavior and vanilla Spark SQL APIs, so they live in the shared + * `gluten-ut` common test module and run against every supported Spark version instead of being + * pinned to a single version-specific suite. Concrete suites live in the Spark-version-specific + * `gluten-ut` modules so test discovery only instantiates them when backend components are present + * on the classpath. + * + * The full outer BNLJ rewrite is a Velox backend feature, hence each test is guarded with + * `assumeVeloxBackend()` so the ClickHouse backend skips them. + */ +abstract class GlutenBroadcastNestedLoopJoinFullOuterSuiteBase + extends GlutenSQLTestsTrait + with AdaptiveSparkPlanHelper { + import testImplicits._ + + // Disable the forced shuffled hash join rewrite so explicit join hints retain their semantics. + override def sparkConf: SparkConf = { + super.sparkConf + .set(GlutenConfig.COLUMNAR_FORCE_SHUFFLED_HASH_JOIN_ENABLED.key, "false") + } + + private def assumeVeloxBackend(): Unit = assume(BackendTestUtils.isVeloxBackendLoaded()) + + private def materializePlan(df: Dataset[_]): SparkPlan = { + val materializedDf = df.toDF() + val executedPlan = materializedDf.queryExecution.executedPlan + executedPlan.execute() + stripAQEPlan(executedPlan match { + case adaptivePlan: AdaptiveSparkPlanExec => adaptivePlan.executedPlan + case otherPlan => otherPlan + }) + } + + private def assertPlanCount[T <: SparkPlan: ClassTag]( + df: Dataset[_], + expectedCount: Int): Unit = { + val targetClass = implicitly[ClassTag[T]].runtimeClass + val plan = materializePlan(df) + val matchedNodes = plan.collect { + case node if targetClass.isInstance(node) => node + } + assert( + matchedNodes.size === expectedCount, + s"Expected $expectedCount ${targetClass.getSimpleName} node(s), but found " + + s"${matchedNodes.size}:\n" + plan.treeString + ) + } + + private def assertNoSparkFullOuterBNLJ(df: Dataset[_]): SparkPlan = { + val plan = materializePlan(df) + val rawFullOuterBnljs = plan.collect { + case bnlj: BroadcastNestedLoopJoinExec if bnlj.joinType == FullOuter => bnlj + } + assert( + rawFullOuterBnljs.isEmpty, + s"Expected rewritten/supported final plan without raw Spark FullOuter " + + s"BroadcastNestedLoopJoinExec, but found ${rawFullOuterBnljs.size}:\n" + + plan.treeString + ) + plan + } + + private def assertSupportedFullOuterPlan(df: Dataset[_]): Unit = { + val plan = assertNoSparkFullOuterBNLJ(df) + val nativeBnljCount = plan.collect { case _: BroadcastNestedLoopJoinExecTransformer => 1 }.size + val nativeSmjCount = plan.collect { case _: SortMergeJoinExecTransformer => 1 }.size + assert( + nativeBnljCount + nativeSmjCount > 0, + s"Expected a supported native full outer plan after rewrite/planning, but found neither " + + s"${classOf[BroadcastNestedLoopJoinExecTransformer].getSimpleName} nor " + + s"${classOf[SortMergeJoinExecTransformer].getSimpleName}:\n" + + plan.treeString + ) + } + + testGluten("Full outer BroadcastNestedLoopJoinExec should be rewritten into supported stages") { + assumeVeloxBackend() + val df1 = spark.range(4).select($"id".as("k1")) + val df2 = spark.range(3).select($"id".as("k2")) + + Seq(true, false).foreach { + codegenEnabled => + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString, + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val fullOuterJoin = df1.hint("broadcast").join(df2, $"k1" < $"k2", "full_outer") + assertNoSparkFullOuterBNLJ(fullOuterJoin) + assertPlanCount[BroadcastNestedLoopJoinExecTransformer]( + fullOuterJoin, + expectedCount = 2) + checkAnswer( + fullOuterJoin, + Seq( + Row(0, 1), + Row(0, 2), + Row(1, 2), + Row(2, null), + Row(3, null), + Row(null, 0))) + } + } + } + + testGluten( + "Full outer BroadcastNestedLoopJoin rewrite should preserve null semantics for equals") { + assumeVeloxBackend() + val df1 = Seq[java.lang.Integer](null, 1, 2, null).toDF("k1") + val df2 = Seq[java.lang.Integer](null, 1, 3, null).toDF("k2") + + Seq(true, false).foreach { + codegenEnabled => + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString, + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val fullOuterJoin = df1.hint("broadcast").join(df2, $"k1" === $"k2", "full_outer") + assertSupportedFullOuterPlan(fullOuterJoin) + checkAnswer( + fullOuterJoin, + Seq( + Row(null, null), + Row(null, null), + Row(null, null), + Row(null, null), + Row(1, 1), + Row(2, null), + Row(null, 3))) + } + } + } + + testGluten( + "Full outer BNLJ rewrite should preserve null semantics for null-safe equals") { + assumeVeloxBackend() + val df1 = Seq[java.lang.Integer](null, 1, 2, null).toDF("k1") + val df2 = Seq[java.lang.Integer](null, 1, 3, null).toDF("k2") + + Seq(true, false).foreach { + codegenEnabled => + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> Long.MaxValue.toString, + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val fullOuterJoin = df1.hint("broadcast").join(df2, $"k1" <=> $"k2", "full_outer") + assertSupportedFullOuterPlan(fullOuterJoin) + checkAnswer( + fullOuterJoin, + Seq( + Row(null, null), + Row(null, null), + Row(null, null), + Row(null, null), + Row(1, 1), + Row(2, null), + Row(null, 3))) + } + } + } +} Review Comment: Can we move this test to gluten-ut/test? gluten-ut/sparkxx is only used to maintain the imported Spark tests. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
