diff --git a/.github/labeler.yml b/.github/labeler.yml index 83db8b4c6ac..1f7fec4b70a 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -87,7 +87,7 @@ area:joins: - changed-files: - any-glob-to-any-file: - 'native/core/src/execution/operators/dynamic_filter/**' - - 'spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala' + - 'spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala' - 'spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala' - 'spark/src/main/**/*Join*' diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 52096160fe5..a8c87b2fadb 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -30,9 +30,12 @@ to configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`. Comet does to test with both for your specific workloads. To configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`, set `spark.comet.exec.forceShuffledHashJoin=true`. -The conversion only happens when the build side is under a size limit, and a join whose build side has no statistics -is left as a `SortMergeJoin`. The size is Spark's planning estimate of the build side, or under AQE the materialized -shuffle size of the build side. By default the limit is Spark's own rule for choosing a `ShuffledHashJoin`: +Comet makes the choice while Spark plans the query, so a join Spark would broadcast, or one with a join strategy hint +such as `MERGE`, is planned as Spark would plan it. The choice also applies to joins of static tables inside a +streaming query, which then run as Spark's own `ShuffledHashJoin`. The conversion only happens when the build side is +under a size limit, and a join whose build side has no statistics is left as a `SortMergeJoin`. The size is Spark's +planning estimate of the build side, or under AQE the materialized shuffle size of the build side. By default the +limit is Spark's own rule for choosing a `ShuffledHashJoin`: `spark.sql.autoBroadcastJoinThreshold` times the initial shuffle partition count (`spark.sql.adaptive.coalescePartitions.initialPartitionNum` when AQE and partition coalescing are both on and it is set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive threshold, Spark's default diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 7789e8eacc2..aa8b6ba2a62 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -224,6 +224,8 @@ Keep the following in mind when reading the reports: native plan still counts as accelerated, so treat the percentage as an upper bound. - Comet's split Iceberg V2 write (`spark.comet.write.iceberg.enabled`) is declined in plan-only mode, so such writes run on, and are reported as, Spark. +- `spark.comet.exec.forceShuffledHashJoin` is not applied in plan-only mode, so the report shows + the sort-merge joins Spark plans rather than the hash joins a real run would use. - Under AQE the report describes the plan before any adaptive re-planning, so coverage of the plan that finally executes can differ. In particular, AQE plans subqueries into the outer query only after the report is produced, so the outer report counts their operators as diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 11472de837d..63823a40e01 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometCoalesceShufflePartitions, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule} +import org.apache.comet.rules.{CometCoalesceShufflePartitions, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometShuffledHashJoinStrategy, CometSpark34AqeDppFallbackRule} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -113,6 +113,7 @@ class CometSparkSessionExtensions injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) injectQueryStageOptimizerRuleShim(extensions, CometCoalesceShufflePartitions) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } + extensions.injectPlannerStrategy { session => CometShuffledHashJoinStrategy(session) } } case class CometColumnar(session: SparkSession) extends ColumnarRule { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index a9e3f442f7c..7d325939523 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -971,34 +971,23 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { - // The rewrite removes the sorts of each join it converts, so put back any sort a kept - // operator above still needs. This runs before transform(), which converts the sorts it - // adds like any other SortExec. - val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - RewriteJoin.restoreRequiredOrdering(plan.transformUp { case p => - RewriteJoin.rewrite(p, conf) - }) - } else { - plan - } - // Tag Partial aggregates that must not be converted to Comet because a // corresponding Final or PartialMerge cannot be converted and the intermediate buffer // formats are incompatible. This runs before transform() so the tags are checked // during the bottom-up conversion. Tags persist through AQE stage creation. - tagUnsafePartialAggregates(planWithJoinRewritten) + tagUnsafePartialAggregates(plan) // Tag leaves on broadcast join build sides so shouldApplySparkToColumnar can bridge an // unsupported build-side scan (e.g. Text) to Arrow and let the join run natively (#6008). - tagBroadcastBuildSideLeaves(planWithJoinRewritten) + tagBroadcastBuildSideLeaves(plan) - var transformed = transform(planWithJoinRewritten) + var transformed = transform(plan) // A build-side bridge is only worth keeping when its broadcast join went native. Where the // join stayed on Spark, clearWastedBuildSideBridges clears the tag from those scans and we // transform once more so the build side rebuilds exactly as it would without the feature. // See issue #6008. if (clearWastedBuildSideBridges(transformed)) { - transformed = transform(planWithJoinRewritten) + transformed = transform(plan) } var newPlan = revertUnsafePartialAggregates(transformed) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala new file mode 100644 index 00000000000..39b0b086283 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -0,0 +1,175 @@ +/* + * 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.comet.rules + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.RowOrdering +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, JoinStrategyHint, LogicalPlan} +import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy} +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, LogicalQueryStage} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, withFallbackReason, withInfo} +import org.apache.comet.shims.ShimJoinSelection + +/** + * Plans a ShuffledHashJoinExec in place of a SortMergeJoinExec, before EnsureRequirements and + * RemoveRedundantSorts, so the sorts above the join survive. Adapted from Apache Gluten. + */ +case class CometShuffledHashJoinStrategy(session: SparkSession) + extends SparkStrategy + with ShimJoinSelection { + + override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { + case join: Join if isEnabled(join) => planJoin(join) + case _ => Nil + } + + private def sqlConf: SQLConf = session.sessionState.conf + + // Planner strategies run with Comet off and before CometRule, which applies plan-only mode. + private def isEnabled(join: Join): Boolean = + CometConf.COMET_FORCE_SHJ.get(sqlConf) && + CometConf.COMET_EXEC_ENABLED.get(sqlConf) && + !CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(sqlConf) && + !join.isStreaming && + isCometLoaded(sqlConf) + + private def planJoin(join: Join): Seq[SparkPlan] = join match { + case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, _, left, right, hint) + if !isLeftToSpark(join, hint) && canSortMergeJoin(joinType) && + RowOrdering.isOrderable(leftKeys) && hashJoinSupportedShim(leftKeys, rightKeys) => + buildSide(join) match { + case Some(BuildRight) if joinType == LeftSemi || joinType.isInstanceOf[ExistenceJoin] => + // LeftSemi https://github.com/apache/datafusion-comet/issues/2667 + // ExistenceJoin https://github.com/apache/datafusion-comet/issues/2697 + planBySpark(join) { smj => + withFallbackReason(smj, declineReason(s"BuildRight with $joinType is not supported")) + } + case Some(side) => + buildSideOverLimit(join, side) match { + case Some(reason) => + // Not a fallback: the kept sort-merge join still runs natively. + planBySpark(join)(smj => withInfo(smj, declineReason(reason))) + case None => + Seq( + ShuffledHashJoinExec( + leftKeys, + rightKeys, + joinType, + side, + condition, + planLater(left), + planLater(right))) + } + case None => Nil + } + case _ => Nil + } + + private def isLeftToSpark(join: Join, hint: JoinHint): Boolean = + isBroadcastStage(join.left) || isBroadcastStage(join.right) || + canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) + + // canPlanAsBroadcastHashJoin misses a broadcast stage when the adaptive threshold is lower or + // off, and AQE discards a re-plan that hash joins one. + private def isBroadcastStage(plan: LogicalPlan): Boolean = plan match { + case LogicalQueryStage(_, _: BroadcastQueryStageExec) => true + case _ => false + } + + // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH are outside JoinStrategyHint.strategies. + private def hasStrategyHint(hint: JoinHint): Boolean = + Seq(hint.leftHint, hint.rightHint).exists( + _.exists(_.strategy.exists(JoinStrategyHint.strategies.contains))) + + private def declineReason(reason: String): String = + s"Cannot rewrite SortMergeJoin to HashJoin: $reason" + + private def planBySpark(join: Join)( + record: SortMergeJoinExec => SortMergeJoinExec): Seq[SparkPlan] = { + val planned = session.sessionState.planner.JoinSelection(join) + planned.flatMap(_.collect { case smj: SortMergeJoinExec => smj }).foreach(record) + planned + } + + private def buildSide(join: Join): Option[BuildSide] = { + val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) + val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) + if (leftBuildable && rightBuildable) { + Some(optimalBuildSide(join)) + } else if (leftBuildable) { + Some(BuildLeft) + } else if (rightBuildable) { + Some(BuildRight) + } else { + None + } + } + + private def optimalBuildSide(join: Join): BuildSide = { + val leftSize = join.left.stats.sizeInBytes + val rightSize = join.right.stats.sizeInBytes + val leftRowCount = join.left.stats.rowCount + val rightRowCount = join.right.stats.rowCount + if (leftSize == rightSize && rightRowCount.isDefined && leftRowCount.isDefined) { + if (rightRowCount.get <= leftRowCount.get) BuildRight else BuildLeft + } else if (rightSize <= leftSize) { + BuildRight + } else { + BuildLeft + } + } + + /** The largest build side to hash join, or None when the limit is off. */ + private def maxBuildSize: Option[BigInt] = + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(sqlConf) match { + case Some(limit) if limit <= 0 => None + case Some(limit) => Some(BigInt(limit)) + case None => + val threshold = sqlConf.autoBroadcastJoinThreshold + val perPartition = + if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + Some(BigInt(perPartition) * sqlConf.numShufflePartitions) + } + + /** + * Why the build side is not strictly under the limit, as in canBuildLocalHashMapBySize. An + * unknown size is spark.sql.defaultSizeInBytes, so it never fits. + */ + private def buildSideOverLimit(join: Join, side: BuildSide): Option[String] = + maxBuildSize.flatMap { limit => + val buildSize = side match { + case BuildLeft => join.left.stats.sizeInBytes + case BuildRight => join.right.stats.sizeInBytes + } + if (buildSize >= limit) { + Some( + s"build side size estimate of $buildSize bytes is not under the limit of $limit bytes") + } else { + None + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala deleted file mode 100644 index 5f543cc84aa..00000000000 --- a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala +++ /dev/null @@ -1,183 +0,0 @@ -/* - * 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.comet.rules - -import org.apache.spark.sql.catalyst.expressions.SortOrder -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} -import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} -import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} - -/** - * Adapted from equivalent rule in Apache Gluten. - * - * This rule replaces [[SortMergeJoinExec]] with [[ShuffledHashJoinExec]]. - */ -object RewriteJoin extends JoinSelectionHelper { - - private def getSmjBuildSide(join: SortMergeJoinExec): Option[BuildSide] = { - val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) - val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) - if (!leftBuildable && !rightBuildable) { - return None - } - if (!leftBuildable) { - return Some(BuildRight) - } - if (!rightBuildable) { - return Some(BuildLeft) - } - val side = join.logicalLink - .flatMap { - case join: Join => Some(getOptimalBuildSide(join)) - case _ => None - } - .getOrElse { - // If smj has no logical link, or its logical link is not a join, - // then we always choose left as build side. - BuildLeft - } - Some(side) - } - - private def removeSort(plan: SparkPlan) = plan match { - case _: SortExec => plan.children.head - case _ => plan - } - - /** - * Adds a local sort under any operator whose child no longer has the ordering the operator - * requires, using the same check as EnsureRequirements. EnsureRequirements puts no sort above a - * sort-merge join whose output ordering already satisfies its parent, so when rewrite turns - * that join into a hash join and removes its sorts, a kept parent such as another sort-merge - * join on the same key would read unsorted input. Plans that already satisfy their required - * orderings are returned unchanged. - */ - def restoreRequiredOrdering(plan: SparkPlan): SparkPlan = plan.transformUp { case p => - p.withNewChildren(p.children.zip(p.requiredChildOrdering).map { - case (child, required) if !SortOrder.orderingSatisfies(child.outputOrdering, required) => - SortExec(required, global = false, child = child) - case (child, _) => child - }) - } - - /** - * The largest build side the rewrite converts, or None when the limit is switched off. Unset - * means Spark's own rule for choosing a shuffled hash join by size: the broadcast threshold - * times the initial shuffle partition count. A non-positive threshold only disables broadcast - * joins and says nothing about the hash table an executor can hold, so Spark's default - * threshold stands in for it. - */ - private def maxBuildSize(conf: SQLConf): Option[BigInt] = - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(conf) match { - case Some(limit) if limit <= 0 => None - case Some(limit) => Some(BigInt(limit)) - case None => - val threshold = conf.autoBroadcastJoinThreshold - val perPartition = - if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get - Some(BigInt(perPartition) * conf.numShufflePartitions) - } - - /** - * Why the build side is too large to rewrite, or None when it fits. Like Spark's - * canBuildLocalHashMapBySize, the size must be strictly under the limit. The hash join cannot - * spill, so a build side without statistics is refused as well. Spark reports an unknown size - * as Long.MaxValue, which fails the comparison on its own. - */ - private def buildSideOverLimit( - smj: SortMergeJoinExec, - buildSide: BuildSide, - conf: SQLConf): Option[String] = - maxBuildSize(conf).flatMap { limit => - smj.logicalLink match { - case Some(join: Join) => - val buildSize = buildSide match { - case BuildLeft => join.left.stats.sizeInBytes - case BuildRight => join.right.stats.sizeInBytes - } - if (buildSize >= limit) { - Some( - s"build side size estimate of $buildSize bytes is not under the limit of " + - s"$limit bytes") - } else { - None - } - case _ => Some("no statistics are available for the build side") - } - } - - def rewrite(plan: SparkPlan, conf: SQLConf): SparkPlan = plan match { - case smj: SortMergeJoinExec => - getSmjBuildSide(smj) match { - case Some(BuildRight) - if smj.joinType == LeftSemi || smj.joinType.isInstanceOf[ExistenceJoin] => - // LeftSemi https://github.com/apache/datafusion-comet/issues/2667 - // ExistenceJoin https://github.com/apache/datafusion-comet/issues/2697 - withFallbackReason( - smj, - "Cannot rewrite SortMergeJoin to HashJoin: " + - s"BuildRight with ${smj.joinType} is not supported") - plan - case Some(buildSide) => - buildSideOverLimit(smj, buildSide, conf) match { - case Some(reason) => - // The sort-merge join still runs natively, so this is information for extended - // explain rather than a fallback reason. - withInfo(smj, s"Cannot rewrite SortMergeJoin to HashJoin: $reason") - plan - case None => - ShuffledHashJoinExec( - smj.leftKeys, - smj.rightKeys, - smj.joinType, - buildSide, - smj.condition, - removeSort(smj.left), - removeSort(smj.right), - smj.isSkewJoin) - } - case _ => plan - } - case _ => plan - } - - def getOptimalBuildSide(join: Join): BuildSide = { - val leftSize = join.left.stats.sizeInBytes - val rightSize = join.right.stats.sizeInBytes - val leftRowCount = join.left.stats.rowCount - val rightRowCount = join.right.stats.rowCount - if (leftSize == rightSize && rightRowCount.isDefined && leftRowCount.isDefined) { - if (rightRowCount.get <= leftRowCount.get) { - return BuildRight - } - return BuildLeft - } - if (rightSize <= leftSize) { - return BuildRight - } - BuildLeft - } -} diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala new file mode 100644 index 00000000000..c26ffe295f0 --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala @@ -0,0 +1,35 @@ +/* + * 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.comet.shims + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper +import org.apache.spark.sql.catalyst.plans.JoinType + +trait ShimJoinSelection extends JoinSelectionHelper { + + /** Spark 3.x has no collations. */ + protected def hashJoinSupportedShim( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Boolean = true + + /** Spark 3.x has no LeftSingle. */ + protected def canSortMergeJoin(joinType: JoinType): Boolean = true +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala new file mode 100644 index 00000000000..f975aa5ddbb --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala @@ -0,0 +1,35 @@ +/* + * 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.comet.shims + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper +import org.apache.spark.sql.catalyst.plans.{JoinType, LeftSingle} + +trait ShimJoinSelection extends JoinSelectionHelper { + + /** False for keys whose collation is not binary stable, which Spark never hash joins. */ + protected def hashJoinSupportedShim( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Boolean = hashJoinSupported(leftKeys, rightKeys) + + /** Spark plans a LeftSingle join as a hash or nested loop join, never a sort-merge join. */ + protected def canSortMergeJoin(joinType: JoinType): Boolean = joinType != LeftSingle +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 2d5a5228259..6d7b673b484 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -25,6 +25,8 @@ import org.scalactic.source.Position import org.scalatest.Tag import org.apache.hadoop.fs.Path +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.LoggerContext import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.SparkException @@ -33,10 +35,10 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, IsNotNull, SortOrder} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} -import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} -import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.{LeftOuter, LeftSemi} +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, LogicalPlan, NO_BROADCAST_HASH} import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, CommandResultExec, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SparkStrategy, SQLExecution, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException @@ -47,8 +49,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} -import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} -import org.apache.comet.rules.RewriteJoin +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus, isSpark40Plus} class CometJoinSuite extends CometTestBase { @@ -1076,11 +1077,8 @@ class CometJoinSuite extends CometTestBase { } } - // The MERGE hint makes Spark plan a SortMergeJoin, which forceShuffledHashJoin then rewrites. - private val forcedSmjQuery = - "SELECT /*+ MERGE(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + private val forcedSmjQuery = "SELECT * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" - // tbl_b is the smaller side, so the rewrite builds on the right. private def withForcedSmjTables(f: => Unit): Unit = { withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") { withParquetTable((0 until 10).map(i => (i % 10, i + 2)), "tbl_b") { @@ -1092,9 +1090,7 @@ class CometJoinSuite extends CometTestBase { private def rewriteReason(reason: String): String = s"Cannot rewrite SortMergeJoin to HashJoin: $reason" - // Checks that the SortMergeJoin and its sorts survived and that the kept join, which still - // runs natively, is not reported as a fallback. Returns Spark's size estimate of the build - // side, read through the same logical link the rewrite used. + // Returns Spark's size estimate of the build side. private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = { val joins = collect(plan) { case j: CometSortMergeJoinExec => j } assert(joins.size == 1, plan) @@ -1138,6 +1134,8 @@ class CometJoinSuite extends CometTestBase { test(s"forceShuffledHashJoin keeps SortMergeJoin over maxBuildSize, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { withForcedSmjTables { @@ -1150,6 +1148,8 @@ class CometJoinSuite extends CometTestBase { test(s"forceShuffledHashJoin rewrites SortMergeJoin under maxBuildSize, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1g") { withForcedSmjTables { @@ -1166,6 +1166,7 @@ class CometJoinSuite extends CometTestBase { var estimate = BigInt(0) withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1174,6 +1175,7 @@ class CometJoinSuite extends CometTestBase { // A build side exactly at the limit is not under it, as in canBuildLocalHashMapBySize. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> estimate.toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1181,6 +1183,7 @@ class CometJoinSuite extends CometTestBase { } withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> (estimate + 1).toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1193,6 +1196,7 @@ class CometJoinSuite extends CometTestBase { withSQLConf( CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "0", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { withForcedSmjTables { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1202,21 +1206,22 @@ class CometJoinSuite extends CometTestBase { } test("forceShuffledHashJoin without maxBuildSize uses Spark's shuffled hash join size rule") { - // The limit is autoBroadcastJoinThreshold times the shuffle partition count. - withSQLConf( - CometConf.COMET_FORCE_SHJ.key -> "true", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { - withForcedSmjTables { + withForcedSmjTables { + var estimate = BigInt(0) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + estimate = assertSortMergeJoinKept(cometPlan) assertKeptWithReason(cometPlan, 2) } - } - withSQLConf( - CometConf.COMET_FORCE_SHJ.key -> "true", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1g") { - withForcedSmjTables { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> (estimate - 1).toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) assertRewrittenToHashJoin(cometPlan) } @@ -1248,30 +1253,47 @@ class CometJoinSuite extends CometTestBase { } } - test("forceShuffledHashJoin keeps a SortMergeJoin that has no statistics") { - val left = spark.range(1).toDF("a").queryExecution.sparkPlan - val right = spark.range(1).toDF("b").queryExecution.sparkPlan - def smj: SortMergeJoinExec = - SortMergeJoinExec(left.output, right.output, Inner, None, left, right) - assert(smj.logicalLink.isEmpty) - - val kept = smj - assert(RewriteJoin.rewrite(kept, spark.sessionState.conf) eq kept) - assert(!hasFallbackReason(kept)) - val info = kept.getTagValue(CometExplainInfo.EXTENSION_INFO) - assert( - info.exists(_.contains(rewriteReason("no statistics are available for the build side"))), - info) - - withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - val rewritten = RewriteJoin.rewrite(smj, spark.sessionState.conf) - assert(rewritten.isInstanceOf[ShuffledHashJoinExec], rewritten) + test("forceShuffledHashJoin keeps a SortMergeJoin whose build side has no size estimate") { + val schema = StructType(Seq(StructField("k", IntegerType), StructField("v", IntegerType))) + def unsized(name: String): Unit = spark + .createDataFrame(spark.sparkContext.parallelize((0 until 10).map(i => Row(i, i))), schema) + .createOrReplaceTempView(name) + val query = "SELECT * FROM rdd_a JOIN rdd_b ON rdd_a.k = rdd_b.k" + withTempView("rdd_a", "rdd_b") { + unsized("rdd_a") + unsized("rdd_b") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_FORCE_SHJ.key -> "true") { + val (_, cometPlan) = checkSparkAnswer(query) + val kept = collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j.originalPlan + } + assert(kept.size == 1, cometPlan) + val estimate = kept.head.logicalLink match { + case Some(join: Join) => join.right.stats.sizeInBytes + case other => fail(s"Expected a Join logical link, got $other") + } + assert(estimate == spark.sessionState.conf.defaultSizeInBytes, estimate) + val info = kept.head.getTagValue(CometExplainInfo.EXTENSION_INFO).toSeq.flatten + val expected = rewriteReason(s"build side size estimate of $estimate bytes") + assert(info.exists(_.contains(expected)), info) + + withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + val (_, cometPlan) = checkSparkAnswer(query) + val hashJoins = collect(cometPlan) { + case j: ShuffledHashJoinExec => j + case j: CometHashJoinExec => j + } + assert(hashJoins.size == 1, cometPlan) + } + } } } - // Both joins are on the same key, so EnsureRequirements puts no sort between them: the inner - // sort-merge join's output ordering satisfies the outer one. When the inner join is rewritten - // to a hash join and the outer one is kept, the outer join needs that ordering restored. + // Both joins are on the same key, so EnsureRequirements puts no sort between them. private def withChainedJoinTables(midRows: Int, midKeys: Int)(f: => Unit): Unit = { withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { @@ -1282,6 +1304,16 @@ class CometJoinSuite extends CometTestBase { } } + private def withForcedShjTables(adaptive: Boolean)(f: => Unit): Unit = withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100)(f) + } + // Whether a hash join is reachable without leaving the stage. private def reachesHashJoinInStage(plan: SparkPlan): Boolean = plan match { case _: CometHashJoinExec => true @@ -1300,8 +1332,6 @@ class CometJoinSuite extends CometTestBase { case _ => false } - // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join - // through a sort. private def checkSortRestoredOverHashJoin(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) val inputs = collect(cometPlan) { @@ -1343,23 +1373,10 @@ class CometJoinSuite extends CometTestBase { } } - // The rewrite never converts a LeftSemi or ExistenceJoin that builds on the right, whatever - // the build size, so these joins are kept above the rewritten inner join with no size limit. - // Each predicate also refers to small, so neither the filter nor the semi join is pushed below - // the inner join. + // Each predicate refers to small, so neither the filter nor the semi join is pushed down. for (adaptive <- Seq(false, true)) { - def withKeptSemiJoinConf(f: => Unit): Unit = withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_FORCE_SHJ.key -> "true", - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - withChainedJoinTables(midRows = 3000, midKeys = 100)(f) - } - test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") { - withKeptSemiJoinConf { + withForcedShjTables(adaptive) { checkSortRestoredOverHashJoin( "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + "WHERE small._2 < 0 OR EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1)") @@ -1370,7 +1387,7 @@ class CometJoinSuite extends CometTestBase { withSQLConf( SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> "org.apache.spark.sql.catalyst.optimizer.PushLeftSemiLeftAntiThroughJoin") { - withKeptSemiJoinConf { + withForcedShjTables(adaptive) { checkSortRestoredOverHashJoin( "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + "WHERE EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1 AND mid._2 > small._2)") @@ -1383,19 +1400,340 @@ class CometJoinSuite extends CometTestBase { // output ordering satisfies it, so the window needs that ordering restored too. for (adaptive <- Seq(false, true)) { test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { + withForcedShjTables(adaptive) { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + + "FROM big JOIN small ON big._1 = small._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + } + } + } + + test("forceShuffledHashJoin chain keeps early join dynamic filtering") { + withForcedShjTables(adaptive = false) { + val query = "SELECT big._1, big._2, mid._2, small._2 FROM big " + + "JOIN mid ON big._1 = mid._1 JOIN small ON big._1 = small._1" + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> enabled.toString) { + val (_, plan) = checkSparkAnswerAndOperator(sql(query)) + val joins = collect(plan) { case join: CometHashJoinExec => join } + assert(joins.map(_.buildSide) == Seq(BuildRight, BuildRight), plan) + // big has 100 rows per key, mid 30, and small one row for keys 0 to 90 by 10. + val outputs = joins.map(_.metrics("output_rows").value).sorted + assert(outputs == (if (enabled) Seq(30000L, 30000L) else Seq(30000L, 300000L))) + if (enabled) { + assert( + joins.map(_.metrics("dynamic_filter_early_rows_evaluated").value).sum == 10000L) + assert(joins.map(_.metrics("dynamic_filter_early_rows_pruned").value).sum == 9000L) + } else { + assert(joins.forall(!_.metrics.contains("dynamic_filter_early_rows_pruned"))) + } + } + } + } + } + + private def checkPartitionsSortedOverHashJoin(df: => DataFrame): Unit = { + // withSQLConf returns Unit on Spark 3.x, so the expected rows are assigned inside it. + var expected = Seq.empty[String] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = df.collect().map(_.toString).sorted.toSeq + } + val cometDf = df + assert(cometDf.collect().map(_.toString).sorted.toSeq == expected) + val cometPlan = cometDf.queryExecution.executedPlan + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + val sorted = cometDf.queryExecution.toRdd + .mapPartitions { it => + val keys = it.map(_.getInt(0)).toArray + Iterator(keys.sameElements(keys.sorted)) + } + .collect() + assert(sorted.forall(identity), sorted.mkString(", ")) + } + + private val sortLossJoin = + "SELECT big._1 AS k, big._2 AS v FROM big JOIN small ON big._1 = small._1" + + // Spark drops a sort that the sort-merge join's ordering satisfies (#6770). + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin keeps sortWithinPartitions on the join key, AQE=$adaptive") { + withForcedShjTables(adaptive) { + checkPartitionsSortedOverHashJoin(sql(sortLossJoin).sortWithinPartitions("k")) + } + } + + // Unsorted partition values make the single dynamic partition writer reopen a closed file. + test(s"forceShuffledHashJoin keeps the sort of a partitioned insert, AQE=$adaptive") { + withForcedShjTables(adaptive) { + withTable("out") { + sql("CREATE TABLE out (v INT, w INT, k INT) USING parquet PARTITIONED BY (k)") + val insert = sql( + "INSERT OVERWRITE TABLE out " + + "SELECT big._2, small._2, big._1 FROM big JOIN small ON big._1 = small._1") + val written = insert.queryExecution.executedPlan match { + case c: CommandResultExec => c.commandPhysicalPlan + case p => p + } + assert(collect(written) { case j: CometHashJoinExec => j }.nonEmpty, written) + assert(spark.table("out").count() == 1000) + val files = spark.table("out").inputFiles + val perPartition = files.groupBy(_.split("/").find(_.startsWith("k="))) + assert(perPartition.size == 10, files.mkString("\n")) + assert(perPartition.values.forall(_.length == 1), files.mkString("\n")) + } + } + } + + for ((joinType, key) <- Seq("RIGHT" -> "big._1", "FULL" -> "coalesce(big._1, small._1)")) { + test( + s"forceShuffledHashJoin keeps sortWithinPartitions over a $joinType join, " + + s"AQE=$adaptive") { + withForcedShjTables(adaptive) { + checkPartitionsSortedOverHashJoin( + sql(s"SELECT $key AS k, big._2 AS v, small._2 AS w " + + s"FROM small $joinType JOIN big ON small._1 = big._1") + .sortWithinPartitions("k")) + } + } + } + } + + test("forceShuffledHashJoin keeps probe side order when no exchange is added") { + withForcedShjTables(adaptive = false) { + val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") + val build = spark.table("small") + checkPartitionsSortedOverHashJoin( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2"), build("_2")) + .sortWithinPartitions("k")) + } + } + + // The cache reaches the native join with no declared order, so the join must declare it. + test("forceShuffledHashJoin keeps the order of a sorted cache through an outer join") { + withForcedShjTables(adaptive = false) { withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_FORCE_SHJ.key -> "true", - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - withChainedJoinTables(midRows = 3000, midKeys = 100) { + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + val streamed = spark + .table("big") + .repartition(2, $"_1") + .sortWithinPartitions("_1", "_2") + .cache() + try { + streamed.count() + val build = spark.table("small") + val valueOrder = SortOrder(BoundReference(1, IntegerType, nullable = true), Ascending) + for { + hint <- Seq(None, Some("shuffle_hash")) + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + } { + val hinted = hint.fold(build)(build.hint(_)) + val joined = if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + def df: DataFrame = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k", "v") + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.map(_.buildSide) == Seq(buildSide), plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) + assert(collect(plan) { case s @ (_: SortExec | _: CometSortExec) => s }.isEmpty, plan) + sameRowsAsSpark(df) + assert( + partitionsInOrder(df, keyOrder(Ascending), valueOrder), + s"$joinType with hint $hint is out of order\n$plan") + } + } finally { + streamed.unpersist() + } + } + } + } + + test("forceShuffledHashJoin respects join strategy hints") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + def hinted(hint: String): DataFrame = + sql(s"SELECT /*+ $hint */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1") + val (_, merge) = checkSparkAnswerAndOperator(hinted("MERGE(tbl_a)")) + assert(collect(merge) { case j: CometSortMergeJoinExec => j }.size == 1, merge) + assert(collect(merge) { case j: CometHashJoinExec => j }.isEmpty, merge) + // BuildLeft proves Spark planned it: the forced conversion would build the smaller tbl_b. + val (_, shuffleHash) = checkSparkAnswer(hinted("SHUFFLE_HASH(tbl_a)")) + assert( + collect(shuffleHash) { case j: CometHashJoinExec => j.buildSide } == Seq(BuildLeft), + shuffleHash) + val (_, broadcast) = checkSparkAnswer(hinted("BROADCAST(tbl_b)")) + assert(collect(broadcast) { case j: CometBroadcastHashJoinExec => j }.nonEmpty, broadcast) + } + } + } + + test("forceShuffledHashJoin ignores AQE's NO_BROADCAST_HASH hint") { + var hints = Seq.empty[JoinHint] + val recordJoinHints = new SparkStrategy { + override def apply(plan: LogicalPlan): Seq[SparkPlan] = { + plan match { + case j: Join if CometConf.COMET_ENABLED.get(SQLConf.get) => hints :+= j.hint + case _ => + } + Nil + } + } + withForcedShjTables(adaptive = true) { + // oneKey fills 1 of 10 shuffle partitions, so AQE adds NO_BROADCAST_HASH on re-plan. + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") { + withParquetTable((0 until 100).map(i => (10, i)), "oneKey") { + val strategies = spark.experimental.extraStrategies + spark.experimental.extraStrategies = strategies :+ recordJoinHints + try { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._2, oneKey._2 FROM big JOIN oneKey ON big._1 = oneKey._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + val sortMergeJoins = collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j + } + assert(sortMergeJoins.isEmpty, cometPlan) + val strategyHints = + hints.flatMap(h => Seq(h.leftHint, h.rightHint)).flatten.flatMap(_.strategy) + assert(strategyHints.contains(NO_BROADCAST_HASH), hints) + } finally { + spark.experimental.extraStrategies = strategies + } + } + } + } + } + + test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { + withForcedShjTables(adaptive = true) { + withSQLConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) + val broadcasts = collect(cometPlan) { + case j: CometBroadcastHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(broadcasts.nonEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + } + } + } + + test("forceShuffledHashJoin handles a skewed join under AQE") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10", + SQLConf.SKEW_JOIN_ENABLED.key -> "true", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100", + CometConf.COMET_FORCE_SHJ.key -> "true") { + withParquetTable((0 until 10000).map(i => (if (i < 9000) 0 else i % 100, i)), "skewed") { + withParquetTable((0 until 100).map(i => (i, i)), "keys") { val (_, cometPlan) = checkSparkAnswer( - "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + - "FROM big JOIN small ON big._1 = small._1") - assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) - assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + "SELECT skewed._2, keys._2 FROM skewed JOIN keys ON skewed._1 = keys._1") + val hashJoins = collect(cometPlan) { + case j: ShuffledHashJoinExec => j + case j: CometHashJoinExec => j.originalPlan + } + val skewJoins = hashJoins.collect { case j: ShuffledHashJoinExec if j.isSkewJoin => j } + assert(skewJoins.nonEmpty, cometPlan) + } + } + } + } + + // A computed key stops Spark from rewriting the join to compare collation keys. + test("forceShuffledHashJoin keeps a SortMergeJoin on collated keys") { + assume(isSpark40Plus) + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTable("c1", "c2") { + sql("CREATE TABLE c1 (s STRING COLLATE UTF8_LCASE, v INT) USING parquet") + sql("INSERT INTO c1 VALUES ('a', 1), ('A', 2), ('b', 3)") + sql("CREATE TABLE c2 (s STRING COLLATE UTF8_LCASE, w INT) USING parquet") + sql("INSERT INTO c2 VALUES ('a', 10), ('B', 20)") + val (_, cometPlan) = + checkSparkAnswer("SELECT v, w FROM c1 JOIN c2 ON trim(c1.s) = c2.s") + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + } + } + } + + // AQE discards a re-plan that hash joins the broadcast stage, so only its warning shows one. + test("forceShuffledHashJoin leaves a planned broadcast stage to AQE") { + val appender = new LogAppender("AQE re-plans") + val logger = classOf[AdaptiveSparkPlanExec].getName + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val isUnconfigured = !context.getConfiguration.getLoggers.containsKey(logger) + val threshold = 10L * 1024 + withForcedShjTables(adaptive = true) { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> threshold.toString, + SQLConf.ADAPTIVE_EXECUTION_LOG_LEVEL.key -> "WARN") { + val bigSize = spark.table("big").queryExecution.optimizedPlan.stats.sizeInBytes + val smallSize = spark.table("small").queryExecution.optimizedPlan.stats.sizeInBytes + assert(bigSize > threshold, s"big is $bigSize bytes, not over $threshold") + assert(smallSize <= threshold, s"small is $smallSize bytes, over $threshold") + try { + withLogAppender(appender, Seq(logger), Some(Level.WARN)) { + val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) + val broadcasts = collect(cometPlan) { + case j: CometBroadcastHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(broadcasts.nonEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + } + } finally { + // withLogAppender never removes the config it adds for an unconfigured logger. + if (isUnconfigured) { + context.getConfiguration.removeLogger(logger) + context.updateLoggers() + } + } + } + } + val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage) + assert(messages.exists(_.contains("Final plan")), messages.mkString("\n")) + val invalidStages = messages.filter(_.contains("Invalid broadcast query stage")) + assert(invalidStages.isEmpty, invalidStages.mkString("\n")) + } + + for ((key, value) <- Seq( + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true")) { + test(s"forceShuffledHashJoin does nothing with $key=$value") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + key -> value, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + val df = sql(forcedSmjQuery) + df.collect() + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: ShuffledHashJoinExec => j }.isEmpty, plan) + assert(collect(plan) { case j: SortMergeJoinExec => j }.nonEmpty, plan) } } } @@ -2441,10 +2779,7 @@ class CometJoinSuite extends CometTestBase { } test("ExistenceJoin via SortMergeJoin stays on Spark even with forceShuffledHashJoin") { - // RewriteJoin refuses to rewrite an existence SortMergeJoin into a BuildRight ShuffledHashJoin - // (https://github.com/apache/datafusion-comet/issues/2697), so even with forced SHJ the join - // must stay on Spark rather than run as a native hash join. Assert on the plan (no native join - // form) rather than the fallback reason, which is the robust signal for this rewrite guard. + // Never a BuildRight hash join: https://github.com/apache/datafusion-comet/issues/2697 withSQLConf( CometConf.COMET_EXEC_EXISTENCE_JOIN_ENABLED.key -> "true", CometConf.COMET_FORCE_SHJ.key -> "true",