From 8e7b798c5e28c0724c090335b8904bbe0ec02a3c Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 00:23:24 +0700 Subject: [PATCH 1/7] fix: choose the forced shuffled hash join at planning time so Spark keeps its sorts With spark.comet.exec.forceShuffledHashJoin on, RewriteJoin turned sort-merge joins into hash joins inside CometExecRule, after EnsureRequirements and RemoveRedundantSorts had run. A sort Spark had dropped because the sort-merge join already produced that order was then gone for good, so sortWithinPartitions and SORT BY returned unsorted partitions and partitioned inserts failed with FileAlreadyExistsException. The choice now happens in a planner strategy injected ahead of Spark's JoinSelection, so Spark places and keeps every sort the plan needs. It also runs on each AQE re-plan with the materialized shuffle sizes. The build side size limit and the LeftSemi and ExistenceJoin exclusions carry over. Joins Spark would broadcast, joins with a strategy hint, keys that cannot be sorted, and on Spark 4 collated keys and LeftSingle joins are left to Spark. Closes #6770. --- .github/labeler.yml | 2 +- .../user-guide/latest/tuning/operators.md | 9 +- .../latest/understanding-comet-plans.md | 2 + .../comet/CometSparkSessionExtensions.scala | 3 +- .../apache/comet/rules/CometExecRule.scala | 15 +- .../rules/CometShuffledHashJoinStrategy.scala | 194 ++++++++++ .../org/apache/comet/rules/RewriteJoin.scala | 183 --------- .../comet/shims/ShimJoinSelection.scala | 35 ++ .../comet/shims/ShimJoinSelection.scala | 35 ++ .../apache/comet/exec/CometJoinSuite.scala | 359 +++++++++++++++--- 10 files changed, 588 insertions(+), 249 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala delete mode 100644 spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala create mode 100644 spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala diff --git a/.github/labeler.yml b/.github/labeler.yml index 5307106f723..ea419934978 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 97baabfccad..8b3e123e4c5 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 13733f0b510..228bfd878e6 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.splitOperator.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 a47f1c2a6b4..35731cc8aac 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -994,24 +994,13 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } else { val normalizedPlan = normalizeDivisors(plan) - // 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(normalizedPlan.transformUp { case p => - RewriteJoin.rewrite(p, conf) - }) - } else { - normalizedPlan - } - // 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(normalizedPlan) - var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten)) + var newPlan = revertUnsafePartialAggregates(transform(normalizedPlan)) // if the plan cannot be run fully natively then explain why (when appropriate // config is enabled) 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..cd8ec4c9ecc --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -0,0 +1,194 @@ +/* + * 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.{BROADCAST, Join, JoinHint, LogicalPlan, SHUFFLE_HASH, SHUFFLE_MERGE, SHUFFLE_REPLICATE_NL} +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 an equi-join as a ShuffledHashJoinExec where Spark would plan a SortMergeJoinExec, when + * spark.comet.exec.forceShuffledHashJoin is on. Choosing the join here, before EnsureRequirements + * and RemoveRedundantSorts run, lets Spark place and keep every sort the plan needs above the + * join. Planner strategies also run again on every AQE re-plan, where the join's children carry + * the materialized shuffle sizes. + * + * Adapted from the equivalent rule in 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 whether or not Comet is enabled, and before CometRule, so plan-only + // mode needs its own check here. + 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) => + // The sort-merge join still runs natively, so this is information for extended + // explain rather than a fallback reason. + 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 + } + + // Broadcasts, including the ones AQE plans once a stage turns out small, and joins with a + // strategy hint are planned by Spark as usual. + private def isLeftToSpark(join: Join, hint: JoinHint): Boolean = + isBroadcastStage(join.left) || isBroadcastStage(join.right) || + canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) + + private def isBroadcastStage(plan: LogicalPlan): Boolean = plan match { + case LogicalQueryStage(_, _: BroadcastQueryStageExec) => true + case _ => false + } + + // AQE's own hints, such as NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH, do not count. + private def hasStrategyHint(hint: JoinHint): Boolean = + Seq(hint.leftHint, hint.rightHint).flatten.flatMap(_.strategy).exists { + case BROADCAST | SHUFFLE_MERGE | SHUFFLE_HASH | SHUFFLE_REPLICATE_NL => true + case _ => false + } + + private def declineReason(reason: String): String = + s"Cannot rewrite SortMergeJoin to HashJoin: $reason" + + // Lets Spark plan a join this strategy declines and records why on each sort-merge join in + // the result, so the reason reaches extended explain. + 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 this strategy plans as a hash join, 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: 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 too large for a hash join, 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 of unknown size, which Spark reports as spark.sql.defaultSizeInBytes + * (Long.MaxValue by default), is refused by the same comparison. + */ + 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..68791fe54d7 --- /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, so every key it can hash join is hashed correctly. */ + protected def hashJoinSupportedShim( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Boolean = true + + /** Every join type Spark 3.x can hash join can also be planned as a sort-merge join. */ + 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 369c2517636..0a4a9d5d7ae 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,20 +31,18 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, CommandResultExec, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, 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 { @@ -894,11 +892,10 @@ 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" + // Without forceShuffledHashJoin, and with broadcasts off, Spark plans a SortMergeJoin. + 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. + // tbl_b is the smaller side, so the hash join 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") { @@ -912,7 +909,7 @@ class CometJoinSuite extends CometTestBase { // 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. + // side, read through the join's logical link. private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = { val joins = collect(plan) { case j: CometSortMergeJoinExec => j } assert(joins.size == 1, plan) @@ -956,6 +953,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 { @@ -968,6 +967,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 { @@ -984,6 +985,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)) @@ -992,6 +994,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)) @@ -999,6 +1002,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)) @@ -1011,6 +1015,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)) @@ -1020,21 +1025,26 @@ 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 { + // The limit is autoBroadcastJoinThreshold times the shuffle partition count. AQE is off so + // every run compares the same planning estimate. + 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 { + // A threshold just under the build side keeps it from being broadcast, while twice the + // threshold is over it. + 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) } @@ -1066,30 +1076,51 @@ 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") { + // A relation read from an RDD has no statistics, so Spark reports + // spark.sql.defaultSizeInBytes for it, which no limit admits. + 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. + // sort-merge join's output ordering satisfies the outer one. When the inner join becomes a + // hash join and the outer one is kept, the outer join needs a sort of its own. 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") { @@ -1118,7 +1149,7 @@ class CometJoinSuite extends CometTestBase { case _ => false } - // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join + // Checks the answer first, then that a kept sort-merge join reads the forced hash join // through a sort. private def checkSortRestoredOverHashJoin(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) @@ -1161,8 +1192,9 @@ 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. + // forceShuffledHashJoin never converts a LeftSemi or ExistenceJoin that builds on the right, + // whatever the build size, so these joins are kept above the forced 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. for (adaptive <- Seq(false, true)) { @@ -1198,7 +1230,7 @@ class CometJoinSuite extends CometTestBase { } // A window partitioned by the join key gets no sort of its own, since the sort-merge join's - // output ordering satisfies it, so the window needs that ordering restored too. + // output ordering satisfies it, so the window needs a sort once the join is a hash join. for (adaptive <- Seq(false, true)) { test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { withSQLConf( @@ -1219,6 +1251,237 @@ class CometJoinSuite extends CometTestBase { } } + // Spark removes a sort above a sort-merge join whose output ordering satisfies it, so the + // forced hash join must be chosen before that happens or the sort is lost (#6770). + private def withSortLossConf(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") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + f + } + } + } + + // Checks that a hash join ran, that the rows match Spark's, and that every partition is sorted + // on the first column. A local sort leaves rows with equal keys in any order, so the rows are + // compared without their order. + 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" + + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin keeps sortWithinPartitions on the join key, AQE=$adaptive") { + withSortLossConf(adaptive) { + checkPartitionsSortedOverHashJoin(sql(sortLossJoin).sortWithinPartitions("k")) + } + } + + test(s"forceShuffledHashJoin keeps SORT BY on the join key, AQE=$adaptive") { + withSortLossConf(adaptive) { + checkPartitionsSortedOverHashJoin(sql(s"$sortLossJoin SORT BY k")) + } + } + + // The planned write sorts on the partition column, and the single dynamic partition writer + // fails when a partition value comes back after its file was closed. + test(s"forceShuffledHashJoin keeps the sort of a partitioned insert, AQE=$adaptive") { + withSortLossConf(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")) + } + } + } + + // RightOuter builds on the smaller left side. A full outer join reports no output ordering, + // so Spark keeps that sort on its own; the case guards the BuildLeft full outer join. + 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") { + withSortLossConf(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")) + } + } + } + } + + // The streamed side is already partitioned and sorted on the key, so no exchange is added + // under it and Spark trusts the hash join to keep its order instead of sorting again. The sort + // runs natively in the same plan as the join, so the native join sees the ordered input. + test("forceShuffledHashJoin keeps probe side order when no exchange is added") { + withSortLossConf(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")) + } + } + + 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) + // Spark builds the hinted side, where 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) + } + } + } + + // AQE plans a broadcast once the build side's shuffle turns out to be small. + test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { + withSortLossConf(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) + } + } + } + + // OptimizeSkewedJoin splits the skewed partitions of the forced hash join. + 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 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) + } + } + } + } + + // Spark does not hash join keys whose collation is not binary stable, so the forced conversion + // leaves them to Spark too. Spark rewrites a comparison of two collated columns to compare + // their collation keys, so the key here is computed. + 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) + } + } + } + + // A join Spark broadcasts from the static estimate is left to Spark, and stays a broadcast + // after AQE re-plans with the adaptive threshold off. + test("forceShuffledHashJoin leaves a planned broadcast stage to AQE") { + withSortLossConf(adaptive = true) { + withSQLConf(SQLConf.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) + } + } + } + + 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) + } + } + } + } + test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") { withSQLConf( SQLConf.PREFER_SORTMERGEJOIN.key -> "false", @@ -2058,10 +2321,10 @@ 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 + // forceShuffledHashJoin never plans an existence join as 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. + // form) rather than the fallback reason, which is the robust signal for this guard. withSQLConf( CometConf.COMET_EXEC_EXISTENCE_JOIN_ENABLED.key -> "true", CometConf.COMET_FORCE_SHJ.key -> "true", From 2ca7eb93022fcda4f92dc420ffb512f9942f9c85 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sat, 10 Oct 2026 01:13:17 +0700 Subject: [PATCH 2/7] fix: match strategy hints by JoinStrategyHint.strategies and pin the broadcast stage check hasStrategyHint now uses Spark's own set of strategy hints. The forced hash join tests share one helper, and the SORT BY case is dropped since it plans the same local sort as sortWithinPartitions. New coverage: a side AQE marks NO_BROADCAST_HASH keeps the hash join, and the planned broadcast test now fails without the broadcast stage check, which keeps AQE from discarding a re-plan that hash joins the stage. --- .../rules/CometShuffledHashJoinStrategy.scala | 15 +- .../apache/comet/exec/CometJoinSuite.scala | 181 +++++++++++------- 2 files changed, 125 insertions(+), 71 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala index cd8ec4c9ecc..45fa317f135 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -24,7 +24,7 @@ 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.{BROADCAST, Join, JoinHint, LogicalPlan, SHUFFLE_HASH, SHUFFLE_MERGE, SHUFFLE_REPLICATE_NL} +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} @@ -102,17 +102,20 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) isBroadcastStage(join.left) || isBroadcastStage(join.right) || canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) + // Spark's LogicalQueryStageStrategy plans a join over a broadcast stage as a broadcast join, + // but runs after this strategy. canPlanAsBroadcastHashJoin can miss such a join, since it + // checks the stage's runtime size against the adaptive threshold, which can be lower or off. + // AQE then throws away a re-plan that hash joins the stage, with everything else it changed. private def isBroadcastStage(plan: LogicalPlan): Boolean = plan match { case LogicalQueryStage(_, _: BroadcastQueryStageExec) => true case _ => false } - // AQE's own hints, such as NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH, do not count. + // The hints in JoinStrategyHint.strategies count, including a SHUFFLE_HASH that AQE sets; + // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH do not. private def hasStrategyHint(hint: JoinHint): Boolean = - Seq(hint.leftHint, hint.rightHint).flatten.flatMap(_.strategy).exists { - case BROADCAST | SHUFFLE_MERGE | SHUFFLE_HASH | SHUFFLE_REPLICATE_NL => true - case _ => false - } + 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" 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 0a4a9d5d7ae..466bd4a91b5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -23,6 +23,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 @@ -31,9 +33,9 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.catalyst.plans.logical.Join +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, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, CommandResultExec, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, CommandResultExec, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SparkStrategy, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} @@ -1131,6 +1133,18 @@ class CometJoinSuite extends CometTestBase { } } + // forceShuffledHashJoin with no build size limit and broadcasts off, over the chained join + // tables. + 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 @@ -1198,18 +1212,8 @@ class CometJoinSuite extends CometTestBase { // Each predicate also refers to small, so neither the filter nor the semi join is pushed below // the inner join. 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)") @@ -1220,7 +1224,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)") @@ -1233,35 +1237,12 @@ class CometJoinSuite extends CometTestBase { // output ordering satisfies it, so the window needs a sort once the join is a hash join. for (adaptive <- Seq(false, true)) { test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { - 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) { - 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) - } - } - } - } - - // Spark removes a sort above a sort-merge join whose output ordering satisfies it, so the - // forced hash join must be chosen before that happens or the sort is lost (#6770). - private def withSortLossConf(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") { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - f + 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) } } } @@ -1291,23 +1272,19 @@ class CometJoinSuite extends CometTestBase { private val sortLossJoin = "SELECT big._1 AS k, big._2 AS v FROM big JOIN small ON big._1 = small._1" + // Spark removes a sort above a sort-merge join whose output ordering satisfies it, so the + // forced hash join must be chosen before that happens or the sort is lost (#6770). for (adaptive <- Seq(false, true)) { test(s"forceShuffledHashJoin keeps sortWithinPartitions on the join key, AQE=$adaptive") { - withSortLossConf(adaptive) { + withForcedShjTables(adaptive) { checkPartitionsSortedOverHashJoin(sql(sortLossJoin).sortWithinPartitions("k")) } } - test(s"forceShuffledHashJoin keeps SORT BY on the join key, AQE=$adaptive") { - withSortLossConf(adaptive) { - checkPartitionsSortedOverHashJoin(sql(s"$sortLossJoin SORT BY k")) - } - } - // The planned write sorts on the partition column, and the single dynamic partition writer // fails when a partition value comes back after its file was closed. test(s"forceShuffledHashJoin keeps the sort of a partitioned insert, AQE=$adaptive") { - withSortLossConf(adaptive) { + withForcedShjTables(adaptive) { withTable("out") { sql("CREATE TABLE out (v INT, w INT, k INT) USING parquet PARTITIONED BY (k)") val insert = sql( @@ -1333,7 +1310,7 @@ class CometJoinSuite extends CometTestBase { test( s"forceShuffledHashJoin keeps sortWithinPartitions over a $joinType join, " + s"AQE=$adaptive") { - withSortLossConf(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") @@ -1347,7 +1324,7 @@ class CometJoinSuite extends CometTestBase { // under it and Spark trusts the hash join to keep its order instead of sorting again. The sort // runs natively in the same plan as the join, so the native join sees the ordered input. test("forceShuffledHashJoin keeps probe side order when no exchange is added") { - withSortLossConf(adaptive = false) { + withForcedShjTables(adaptive = false) { val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") val build = spark.table("small") checkPartitionsSortedOverHashJoin( @@ -1380,9 +1357,51 @@ class CometJoinSuite extends CometTestBase { } } + // AQE marks the side of a join whose shuffle has few non-empty partitions NO_BROADCAST_HASH + // when it re-plans. That hint is not in JoinStrategyHint.strategies, so the join stays a hash + // join instead of going back to a sort-merge join with the same shuffles, which AQE would + // adopt. + test("forceShuffledHashJoin ignores AQE's NO_BROADCAST_HASH hint") { + var hints = Seq.empty[JoinHint] + // Records only the Comet run; checkSparkAnswer's Spark run gets the same hint. + 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 has a single key, so 1 of its 10 shuffle partitions is non-empty. + 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 + } + } + } + } + } + // AQE plans a broadcast once the build side's shuffle turns out to be small. test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { - withSortLossConf(adaptive = true) { + withForcedShjTables(adaptive = true) { withSQLConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) val broadcasts = collect(cometPlan) { @@ -1444,21 +1463,53 @@ class CometJoinSuite extends CometTestBase { } } - // A join Spark broadcasts from the static estimate is left to Spark, and stays a broadcast - // after AQE re-plans with the adaptive threshold off. + // A join Spark broadcasts from the static estimate stays a broadcast after AQE re-plans with + // the adaptive threshold off, and no re-plan hash joins the broadcast stage. The 10 KB + // threshold lets Spark broadcast small (about 3 KB) from its static estimate but not big + // (about 46 KB). On the re-plan, small is a broadcast stage measured against the adaptive + // threshold, which is off, so only the broadcast stage check keeps the strategy away from the + // join. AQE discards a re-plan that hash joins the stage, so the final plan cannot show one; + // the warning AQE logs for it can. test("forceShuffledHashJoin leaves a planned broadcast stage to AQE") { - withSortLossConf(adaptive = true) { - withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { - val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) - val broadcasts = collect(cometPlan) { - case j: CometBroadcastHashJoinExec => j - case j: BroadcastHashJoinExec => j + 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 { + // For a logger with no config of its own, withLogAppender adds one and never removes + // it. That config does not pass events up, so it would keep them out of the test log. + if (isUnconfigured) { + context.getConfiguration.removeLogger(logger) + context.updateLoggers() + } } - assert(broadcasts.nonEmpty, cometPlan) - assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) - assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) } } + val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage) + // AQE logs the final plan at the same level, so this shows the appender saw its events. + 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( From 5467694447be7e4bee71a6cea36ef9173a7d3ff2 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sun, 11 Oct 2026 21:41:50 +0700 Subject: [PATCH 3/7] test: keep early join dynamic filtering for a forced shuffled hash join chain Chain two forced shuffled hash joins on the same key with the join dynamic filter off and on. With it on, the outer join's build keys filter the big side below the intermediate join, so both joins see the same row count and the early filter prunes the keys the outer build lacks. --- .../apache/comet/exec/CometJoinSuite.scala | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) 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 baa17c8652d..30db7d0ee4f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1349,6 +1349,33 @@ class CometJoinSuite extends CometTestBase { } } + // Both joins use big._1, so no exchange separates them and the outer join's filter on small's + // keys reaches big below the intermediate join. Each build side is the smaller input. + 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 for each of its 100 keys, mid 30, and small one for keys 0 to 90 + // in steps of 10. Each of the two shuffle partitions holds some of small's keys. + 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"))) + } + } + } + } + } + // Checks that a hash join ran, that the rows match Spark's, and that every partition is sorted // on the first column. A local sort leaves rows with equal keys in any order, so the rows are // compared without their order. From 60bcf09c2dbf75f6ef325320fdc66fe1ea55f9f6 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 21 Oct 2026 23:49:28 +0700 Subject: [PATCH 4/7] docs: keep only the comments that explain the hash join strategy's choices Drop comments that restate the code or a test's name, and shorten the rest to the Spark behavior and limits they depend on. --- .../rules/CometShuffledHashJoinStrategy.scala | 42 +++------ .../comet/shims/ShimJoinSelection.scala | 4 +- .../apache/comet/exec/CometJoinSuite.scala | 86 +++---------------- 3 files changed, 26 insertions(+), 106 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala index 45fa317f135..39b0b086283 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -35,13 +35,8 @@ import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, withFallback import org.apache.comet.shims.ShimJoinSelection /** - * Plans an equi-join as a ShuffledHashJoinExec where Spark would plan a SortMergeJoinExec, when - * spark.comet.exec.forceShuffledHashJoin is on. Choosing the join here, before EnsureRequirements - * and RemoveRedundantSorts run, lets Spark place and keep every sort the plan needs above the - * join. Planner strategies also run again on every AQE re-plan, where the join's children carry - * the materialized shuffle sizes. - * - * Adapted from the equivalent rule in Apache Gluten. + * 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 @@ -54,8 +49,7 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) private def sqlConf: SQLConf = session.sessionState.conf - // Planner strategies run whether or not Comet is enabled, and before CometRule, so plan-only - // mode needs its own check here. + // 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) && @@ -77,8 +71,7 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) case Some(side) => buildSideOverLimit(join, side) match { case Some(reason) => - // The sort-merge join still runs natively, so this is information for extended - // explain rather than a fallback reason. + // Not a fallback: the kept sort-merge join still runs natively. planBySpark(join)(smj => withInfo(smj, declineReason(reason))) case None => Seq( @@ -96,23 +89,18 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) case _ => Nil } - // Broadcasts, including the ones AQE plans once a stage turns out small, and joins with a - // strategy hint are planned by Spark as usual. private def isLeftToSpark(join: Join, hint: JoinHint): Boolean = isBroadcastStage(join.left) || isBroadcastStage(join.right) || canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) - // Spark's LogicalQueryStageStrategy plans a join over a broadcast stage as a broadcast join, - // but runs after this strategy. canPlanAsBroadcastHashJoin can miss such a join, since it - // checks the stage's runtime size against the adaptive threshold, which can be lower or off. - // AQE then throws away a re-plan that hash joins the stage, with everything else it changed. + // 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 } - // The hints in JoinStrategyHint.strategies count, including a SHUFFLE_HASH that AQE sets; - // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH do not. + // 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))) @@ -120,8 +108,6 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) private def declineReason(reason: String): String = s"Cannot rewrite SortMergeJoin to HashJoin: $reason" - // Lets Spark plan a join this strategy declines and records why on each sort-merge join in - // the result, so the reason reaches extended explain. private def planBySpark(join: Join)( record: SortMergeJoinExec => SortMergeJoinExec): Seq[SparkPlan] = { val planned = session.sessionState.planner.JoinSelection(join) @@ -157,13 +143,7 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) } } - /** - * The largest build side this strategy plans as a hash join, 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. - */ + /** 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 @@ -176,10 +156,8 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) } /** - * Why the build side is too large for a hash join, 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 of unknown size, which Spark reports as spark.sql.defaultSizeInBytes - * (Long.MaxValue by default), is refused by the same comparison. + * 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 => 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 index 68791fe54d7..c26ffe295f0 100644 --- 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 @@ -25,11 +25,11 @@ import org.apache.spark.sql.catalyst.plans.JoinType trait ShimJoinSelection extends JoinSelectionHelper { - /** Spark 3.x has no collations, so every key it can hash join is hashed correctly. */ + /** Spark 3.x has no collations. */ protected def hashJoinSupportedShim( leftKeys: Seq[Expression], rightKeys: Seq[Expression]): Boolean = true - /** Every join type Spark 3.x can hash join can also be planned as a sort-merge join. */ + /** Spark 3.x has no LeftSingle. */ protected def canSortMergeJoin(joinType: JoinType): Boolean = true } 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 5496835db8f..6d7b673b484 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1077,10 +1077,8 @@ class CometJoinSuite extends CometTestBase { } } - // Without forceShuffledHashJoin, and with broadcasts off, Spark plans a SortMergeJoin. 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 hash join 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 join's logical link. + // 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) @@ -1210,8 +1206,6 @@ 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. AQE is off so - // every run compares the same planning estimate. withForcedSmjTables { var estimate = BigInt(0) withSQLConf( @@ -1223,8 +1217,6 @@ class CometJoinSuite extends CometTestBase { estimate = assertSortMergeJoinKept(cometPlan) assertKeptWithReason(cometPlan, 2) } - // A threshold just under the build side keeps it from being broadcast, while twice the - // threshold is over it. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_FORCE_SHJ.key -> "true", @@ -1262,8 +1254,6 @@ class CometJoinSuite extends CometTestBase { } test("forceShuffledHashJoin keeps a SortMergeJoin whose build side has no size estimate") { - // A relation read from an RDD has no statistics, so Spark reports - // spark.sql.defaultSizeInBytes for it, which no limit admits. 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) @@ -1303,9 +1293,7 @@ class CometJoinSuite extends CometTestBase { } } - // 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 becomes a - // hash join and the outer one is kept, the outer join needs a sort of its own. + // 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") { @@ -1316,8 +1304,6 @@ class CometJoinSuite extends CometTestBase { } } - // forceShuffledHashJoin with no build size limit and broadcasts off, over the chained join - // tables. private def withForcedShjTables(adaptive: Boolean)(f: => Unit): Unit = withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -1346,8 +1332,6 @@ class CometJoinSuite extends CometTestBase { case _ => false } - // Checks the answer first, then that a kept sort-merge join reads the forced hash join - // through a sort. private def checkSortRestoredOverHashJoin(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) val inputs = collect(cometPlan) { @@ -1389,11 +1373,7 @@ class CometJoinSuite extends CometTestBase { } } - // forceShuffledHashJoin never converts a LeftSemi or ExistenceJoin that builds on the right, - // whatever the build size, so these joins are kept above the forced 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)) { test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") { withForcedShjTables(adaptive) { @@ -1417,7 +1397,7 @@ class CometJoinSuite extends CometTestBase { } // A window partitioned by the join key gets no sort of its own, since the sort-merge join's - // output ordering satisfies it, so the window needs a sort once the join is a hash join. + // 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) { @@ -1430,8 +1410,6 @@ class CometJoinSuite extends CometTestBase { } } - // Both joins use big._1, so no exchange separates them and the outer join's filter on small's - // keys reaches big below the intermediate join. Each build side is the smaller input. test("forceShuffledHashJoin chain keeps early join dynamic filtering") { withForcedShjTables(adaptive = false) { val query = "SELECT big._1, big._2, mid._2, small._2 FROM big " + @@ -1441,8 +1419,7 @@ class CometJoinSuite extends CometTestBase { 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 for each of its 100 keys, mid 30, and small one for keys 0 to 90 - // in steps of 10. Each of the two shuffle partitions holds some of small's keys. + // 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) { @@ -1457,9 +1434,6 @@ class CometJoinSuite extends CometTestBase { } } - // Checks that a hash join ran, that the rows match Spark's, and that every partition is sorted - // on the first column. A local sort leaves rows with equal keys in any order, so the rows are - // compared without their order. 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] @@ -1482,8 +1456,7 @@ class CometJoinSuite extends CometTestBase { private val sortLossJoin = "SELECT big._1 AS k, big._2 AS v FROM big JOIN small ON big._1 = small._1" - // Spark removes a sort above a sort-merge join whose output ordering satisfies it, so the - // forced hash join must be chosen before that happens or the sort is lost (#6770). + // 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) { @@ -1491,8 +1464,7 @@ class CometJoinSuite extends CometTestBase { } } - // The planned write sorts on the partition column, and the single dynamic partition writer - // fails when a partition value comes back after its file was closed. + // 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") { @@ -1514,8 +1486,6 @@ class CometJoinSuite extends CometTestBase { } } - // RightOuter builds on the smaller left side. A full outer join reports no output ordering, - // so Spark keeps that sort on its own; the case guards the BuildLeft full outer join. for ((joinType, key) <- Seq("RIGHT" -> "big._1", "FULL" -> "coalesce(big._1, small._1)")) { test( s"forceShuffledHashJoin keeps sortWithinPartitions over a $joinType join, " + @@ -1530,9 +1500,6 @@ class CometJoinSuite extends CometTestBase { } } - // The streamed side is already partitioned and sorted on the key, so no exchange is added - // under it and Spark trusts the hash join to keep its order instead of sorting again. The sort - // runs natively in the same plan as the join, so the native join sees the ordered input. test("forceShuffledHashJoin keeps probe side order when no exchange is added") { withForcedShjTables(adaptive = false) { val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") @@ -1545,13 +1512,8 @@ class CometJoinSuite extends CometTestBase { } } - // The streamed side is a cache sorted within its partitions, so Spark drops the sort above - // the join, and the cache reaches the native join as a JVM input that declares no order. The - // join keeps the unmatched streamed rows in place only because it declares the cache's order - // on its probe input. + // 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") { - // Under AQE the cached plan, which has an exchange, is adaptive and reports no order, so - // Spark keeps the sort above the join. withForcedShjTables(adaptive = false) { withSQLConf( CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", @@ -1607,7 +1569,7 @@ class CometJoinSuite extends CometTestBase { 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) - // Spark builds the hinted side, where the forced conversion would build the smaller tbl_b. + // 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), @@ -1618,13 +1580,8 @@ class CometJoinSuite extends CometTestBase { } } - // AQE marks the side of a join whose shuffle has few non-empty partitions NO_BROADCAST_HASH - // when it re-plans. That hint is not in JoinStrategyHint.strategies, so the join stays a hash - // join instead of going back to a sort-merge join with the same shuffles, which AQE would - // adopt. test("forceShuffledHashJoin ignores AQE's NO_BROADCAST_HASH hint") { var hints = Seq.empty[JoinHint] - // Records only the Comet run; checkSparkAnswer's Spark run gets the same hint. val recordJoinHints = new SparkStrategy { override def apply(plan: LogicalPlan): Seq[SparkPlan] = { plan match { @@ -1635,7 +1592,7 @@ class CometJoinSuite extends CometTestBase { } } withForcedShjTables(adaptive = true) { - // oneKey has a single key, so 1 of its 10 shuffle partitions is non-empty. + // 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 @@ -1660,7 +1617,6 @@ class CometJoinSuite extends CometTestBase { } } - // AQE plans a broadcast once the build side's shuffle turns out to be small. test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { withForcedShjTables(adaptive = true) { withSQLConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { @@ -1676,7 +1632,6 @@ class CometJoinSuite extends CometTestBase { } } - // OptimizeSkewedJoin splits the skewed partitions of the forced hash join. test("forceShuffledHashJoin handles a skewed join under AQE") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", @@ -1702,9 +1657,7 @@ class CometJoinSuite extends CometTestBase { } } - // Spark does not hash join keys whose collation is not binary stable, so the forced conversion - // leaves them to Spark too. Spark rewrites a comparison of two collated columns to compare - // their collation keys, so the key here is computed. + // A computed key stops Spark from rewriting the join to compare collation keys. test("forceShuffledHashJoin keeps a SortMergeJoin on collated keys") { assume(isSpark40Plus) withSQLConf( @@ -1724,13 +1677,7 @@ class CometJoinSuite extends CometTestBase { } } - // A join Spark broadcasts from the static estimate stays a broadcast after AQE re-plans with - // the adaptive threshold off, and no re-plan hash joins the broadcast stage. The 10 KB - // threshold lets Spark broadcast small (about 3 KB) from its static estimate but not big - // (about 46 KB). On the re-plan, small is a broadcast stage measured against the adaptive - // threshold, which is off, so only the broadcast stage check keeps the strategy away from the - // join. AQE discards a re-plan that hash joins the stage, so the final plan cannot show one; - // the warning AQE logs for it can. + // 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 @@ -1757,8 +1704,7 @@ class CometJoinSuite extends CometTestBase { assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) } } finally { - // For a logger with no config of its own, withLogAppender adds one and never removes - // it. That config does not pass events up, so it would keep them out of the test log. + // withLogAppender never removes the config it adds for an unconfigured logger. if (isUnconfigured) { context.getConfiguration.removeLogger(logger) context.updateLoggers() @@ -1767,7 +1713,6 @@ class CometJoinSuite extends CometTestBase { } } val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage) - // AQE logs the final plan at the same level, so this shows the appender saw its events. assert(messages.exists(_.contains("Final plan")), messages.mkString("\n")) val invalidStages = messages.filter(_.contains("Invalid broadcast query stage")) assert(invalidStages.isEmpty, invalidStages.mkString("\n")) @@ -2834,10 +2779,7 @@ class CometJoinSuite extends CometTestBase { } test("ExistenceJoin via SortMergeJoin stays on Spark even with forceShuffledHashJoin") { - // forceShuffledHashJoin never plans an existence join as 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 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", From 09c28d2f4e5798a1cb0adc2a85134cf0c51fa442 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 23 Oct 2026 22:13:26 +0700 Subject: [PATCH 5/7] test: keep the order of a sorted checkpoint through a forced outer hash join With forceShuffledHashJoin on, Spark drops the final sort above an outer join over a sorted localCheckpoint, whose order Comet cannot check. The join now stays in Spark's ShuffledHashJoinExec, so both outer shapes keep the requested (k, v) order. --- .../apache/comet/exec/CometJoinSuite.scala | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) 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 1f281f997bc..d2aaca3c5df 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1558,6 +1558,40 @@ class CometJoinSuite extends CometTestBase { } } + // Spark drops the final sort for the checkpoint's order, which Comet cannot check. + test("forceShuffledHashJoin keeps the order of a sorted checkpoint through an outer join") { + withOrderedCheckpoint { streamed => + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> "true") { + val build = spark.range(10).selectExpr("id * 10 AS k") + val order = + (0 to 1).map(i => SortOrder(BoundReference(i, LongType, nullable = true), Ascending)) + for { + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + } { + def df: DataFrame = { + val on = streamed("k") === build("k") + val joined = if (joinType == "left_outer") { + streamed.join(build, on, joinType) + } else { + build.join(streamed, on, joinType) + } + joined.select(streamed("k"), streamed("v")).sortWithinPartitions("k", "v") + } + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + assert(partitionsInOrder(df, order: _*), context) + checkSparkAnswerAndFallbackReason(df, CometHashJoin.uncheckedStreamedOrderReason) + val joins = collect(plan) { case j: ShuffledHashJoinExec => j.buildSide } + assert(joins == Seq(buildSide), context) + assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + } + } + } + } + test("forceShuffledHashJoin respects join strategy hints") { withSQLConf( CometConf.COMET_FORCE_SHJ.key -> "true", From dfe7f902b99d4e25cba225f6bfee851c7f4103b9 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sun, 25 Oct 2026 23:35:42 +0700 Subject: [PATCH 6/7] fix: keep join order claims true under forced hash joins - Build the forced hash join on the side AQE's PREFER_SHUFFLE_HASH hint names when only one buildable side carries it. - Run a broadcast nested loop join in Spark when Spark reports an output order for it: the native join emits rows build row by build row, so a sort Spark dropped on the strength of that order was lost (#6898). - Declare the streamed order on hash joins that keep their probe order when Comet can check it, so a native window above sees it, and let the join dynamic filter walks pass through SortedInputExec so reader pruning still applies below the declared order. --- .../latest/compatibility/operators.md | 15 +- docs/source/user-guide/latest/operators.md | 12 +- .../operators/dynamic_filter/early.rs | 16 +- .../operators/dynamic_filter/join/tests.rs | 95 ++++++ .../dynamic_filter/parquet_reader.rs | 15 +- .../src/execution/operators/sorted_input.rs | 19 +- native/core/src/execution/planner.rs | 45 ++- .../rules/CometShuffledHashJoinStrategy.scala | 11 +- .../apache/spark/sql/comet/operators.scala | 57 ++-- .../apache/comet/exec/CometJoinSuite.scala | 281 ++++++++++++++++-- 10 files changed, 486 insertions(+), 80 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 8785874c6f6..78f9312fb17 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -133,19 +133,26 @@ express falls back to Spark and records a fallback reason (see an operator above the join relies on the order. An outer join falls back to Spark when Comet cannot check the order Spark reports and that order -reaches the join through an input from the JVM or through a nested loop join. Spark keeps the rows -in the order they arrive. One such input is an RDD checkpoint. It reaches a native join only with +reaches the join through an input from the JVM. Spark keeps the rows in the order they arrive. One +such input is an RDD checkpoint. It reaches a native join only with `spark.comet.convert.rdd.enabled=true`, which defaults to `false`. Without that setting the checkpoint stays a Spark input and the join stays in Spark anyway. Another is a cache written by an outer hash join whose own probe input was not a scan -([#6922](https://github.com/apache/datafusion-comet/issues/6922)). The output of a nested loop -join can already be out of order when it reaches the join +([#6922](https://github.com/apache/datafusion-comet/issues/6922)). + +The native nested loop join falls back to Spark when Spark reports an output order for it, since +it emits its rows build row by build row ([#6898](https://github.com/apache/datafusion-comet/issues/6898)). When the order Comet cannot check comes only from another native hash join in the same plan, the join does not fall back. It runs natively and declares no order, since DataFusion keeps the order it can see there. +Hash joins of the other types that report an order, such as inner joins, keep their streamed rows +in place either way. They declare the order when Comet can check and express it, so a native +operator above them, a window for instance, sees it. Otherwise they declare nothing and still run +natively. + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 334d34673f2..232bce582af 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -96,12 +96,12 @@ omitted from the tables below and may be reconsidered based on demand: ## Joins -| Operator | Status | Notes | -| ----------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `BroadcastHashJoinExec` | ✅ | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering. | -| `ShuffledHashJoinExec` | ✅ | | -| `SortMergeJoinExec` | ✅ | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark. | -| `BroadcastNestedLoopJoinExec` | ✅ | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)). | +| Operator | Status | Notes | +| ----------------------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `BroadcastHashJoinExec` | ✅ | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering. | +| `ShuffledHashJoinExec` | ✅ | | +| `SortMergeJoinExec` | ✅ | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark. | +| `BroadcastNestedLoopJoinExec` | ✅ | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)), and when Spark reports an output ordering for the join, which the native join does not keep ([#6898](https://github.com/apache/datafusion-comet/issues/6898)). | `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual diff --git a/native/core/src/execution/operators/dynamic_filter/early.rs b/native/core/src/execution/operators/dynamic_filter/early.rs index 9b101044bd9..2a51c652f87 100644 --- a/native/core/src/execution/operators/dynamic_filter/early.rs +++ b/native/core/src/execution/operators/dynamic_filter/early.rs @@ -28,11 +28,12 @@ use datafusion::common::{internal_err, Result}; use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use datafusion_comet_operators::CometFilterExec; use super::parquet_reader::is_direct_column_null_checks; use super::{DynamicFilterExec, DynamicFilterJoinExec}; +use crate::execution::operators::SortedInputExec; use datafusion_comet_operators::CometProjectionExec; pub(super) fn place_early_filter( @@ -126,6 +127,19 @@ fn place( { return Ok(Some(filter.with_execution_input(child))); } + } else if let Some(sorted) = input.downcast_ref::() { + // It only declares its input's order, which dropping rows below it keeps. + if let Some(child) = place( + sorted.input(), + Arc::clone(&predicate), + metrics, + crossed_join, + )? { + return Ok(Some(Arc::clone(input).replace_children( + vec![child], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?)); + } } } // A terminal consumer is useful only after crossing a join. Above other diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index 54ae6737edf..e55ab5b38af 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs @@ -2010,3 +2010,98 @@ fn early_filter_stops_at_intermediate_join_fetch() { &place_early_filter(&wrapped, predicate, &ExecutionPlanMetricsSet::new()).unwrap() )); } + +/// An order that `key@index` declares, as the planner declares a sorted probe's. +fn declared_order(input: Arc, index: usize) -> Arc { + use crate::execution::operators::SortedInputExec; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + let key = Arc::new(Column::new(input.schema().field(index).name(), index)); + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(key)]).unwrap(); + Arc::new(SortedInputExec::new(input, ordering)) +} + +#[test] +fn reader_filter_crosses_a_declared_probe_order() { + // The order wrapper only passes rows through, so the reader below it can drop rows too. + let session = Arc::new(SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + )); + let (_file, scan) = parquet_probe((0..400).collect(), &session, 100); + let probe = declared_order(filtered_probe(&scan), 0); + let ordering = probe.properties().output_ordering().unwrap().to_string(); + let join = single_key_join_plans(two_batch_build(), probe, PartitionMode::Partitioned); + let wrapper = + DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone()) + .unwrap(); + let runtime = wrapper.build_runtime_join().unwrap(); + assert!(runtime.reader_filter_attached); + let consumer = runtime + .join + .right() + .downcast_ref::() + .unwrap(); + let sorted = consumer.children()[0]; + assert_eq!("SortedInputExec", sorted.name()); + assert_eq!( + ordering, + sorted.properties().output_ordering().unwrap().to_string() + ); + let filter = sorted.children()[0] + .downcast_ref::() + .unwrap(); + let reader = filter.input().downcast_ref::().unwrap(); + let (_, source) = reader.downcast_to_file_source::().unwrap(); + let reader_filter = source.filter().unwrap(); + let reader_filter = find_dynamic_filter(&reader_filter).unwrap(); + assert_eq!( + reader_filter.expression_id().unwrap(), + produced_join_filter(&runtime.join).expression_id().unwrap() + ); + // The rebuilt join still keeps the declared order in its own output. + assert!(runtime.join.properties().output_ordering().is_some()); +} + +#[tokio::test] +async fn ancestor_filter_crosses_a_declared_probe_order() { + // As in the test above without the wrapper: the ancestor's filter still reaches the + // intermediate join's probe. + let ctx = SessionContext::new(); + let fact = input( + vec![Some(0), Some(1), Some(2), Some(2), Some(3), None], + &DataType::Int32, + 0, + ); + let dimension = input( + vec![Some(0), Some(1), Some(2), Some(2), Some(3)], + &DataType::Int32, + 0, + ); + let ancestor_build = input(vec![Some(2), Some(2)], &DataType::Int32, 0); + let inner = chain_join(dimension, fact, 0, 0); + let inner = Arc::new( + DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ) as Arc; + let probe = declared_order(inner, 2); + let outer = chain_join(ancestor_build, probe, 0, 2); + let expected = collect( + Arc::new(outer.builder().reset_state().build().unwrap()), + ctx.task_ctx(), + ) + .await + .unwrap(); + assert_eq!(row_count(&expected), 8); + let outer = Arc::new( + DynamicFilterJoinExec::try_new(&outer, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ) as Arc; + let actual = collect(Arc::clone(&outer), ctx.task_ctx()).await.unwrap(); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_evaluated"), 6); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_pruned"), 4); +} diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs index 59009f58fa9..d0d84d05436 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs @@ -29,10 +29,12 @@ use datafusion::physical_expr::expressions::{ }; use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use datafusion_comet_operators::CometFilterExec; +use crate::execution::operators::SortedInputExec; + mod schema_adapter; use schema_adapter::RuntimeFilterSchemaAdapterFactory; @@ -70,6 +72,17 @@ pub(super) fn try_attach_parquet_reader_filter( .map(|reader| filter.with_execution_input(reader)), ); } + // It only declares its input's order, which dropping rows below it keeps. + if let Some(sorted) = input.downcast_ref::() { + return try_attach_parquet_reader_filter(sorted.input(), predicate, config)? + .map(|reader| { + Arc::clone(input).replace_children( + vec![reader], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + }) + .transpose(); + } // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND // chains of inferred null checks. A reader predicate can cross those direct // checks because both operations only discard rows. Broader deterministic diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index a31053bd6e8..94acefc7c32 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -37,13 +37,14 @@ use datafusion::physical_plan::{ /// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none, /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only /// when the probe input declares an order. Otherwise it moves them after the matched rows of each -/// batch, which loses an order Spark has already dropped a sort for. The planner puts this only -/// on that join's probe input, with the order Spark reports for it, and only when the rows reach -/// that input from the JVM through operators that keep their order. The join keeps the order in -/// its own output, so the operators above it in the native plan see it too: an aggregate grouped -/// by the key runs in sorted, streaming mode, a sort aggregate takes the order as already -/// satisfied, and a top-k sort whose keys start with it stops reading early. A wrong order would -/// therefore give wrong groups or wrong rows, not only rows out of order. +/// batch, which loses an order Spark has already dropped a sort for. The planner puts this on a +/// hash join's probe input, with the order Spark reports for it, when the rows reach that input +/// from the JVM or a file scan through operators that keep their order, and on the native range, +/// whose order is real, so that the join and the operators above it in the native plan see the +/// order: an aggregate grouped by the key runs in sorted, streaming mode, a window finds the order +/// it requires, a sort aggregate takes the order as already satisfied, and a top-k sort whose keys +/// start with it stops reading early. A wrong order would therefore give wrong groups or wrong +/// rows, not only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, @@ -70,6 +71,10 @@ impl SortedInputExec { cache, } } + + pub fn input(&self) -> &Arc { + &self.input + } } impl DisplayAs for SortedInputExec { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a47c19d1a23..25cdb02c486 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4044,7 +4044,7 @@ mod tests { .expect("a hash join") } - fn sorted_outer_hash_join( + fn sorted_hash_join( join_type: spark_operator::JoinType, build_side: spark_operator::BuildSide, children: Vec, @@ -4078,7 +4078,7 @@ mod tests { ), ] { let operator = - sorted_outer_hash_join(join_type, build_side, vec![create_scan(), create_scan()]); + sorted_hash_join(join_type, build_side, vec![create_scan(), create_scan()]); let (_, _, planned) = PhysicalPlanner::default() .create_plan(&operator, &mut vec![], 1) .unwrap(); @@ -4099,6 +4099,39 @@ mod tests { } } + #[test] + fn inner_hash_join_declares_streamed_order_on_probe_input() { + // An inner join keeps its probe order whether or not the probe declares it, but the + // operators above it, a window for one, see that order only once the probe declares it. + for build_side in [ + spark_operator::BuildSide::BuildLeft, + spark_operator::BuildSide::BuildRight, + ] { + let operator = sorted_hash_join( + spark_operator::JoinType::Inner, + build_side, + vec![create_scan(), create_scan()], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let join = planned_hash_join(&planned.native_plan); + assert_eq!(&datafusion::common::JoinType::Inner, join.join_type()); + assert_eq!("ScanExec", join.left().name(), "{build_side:?}"); + let probe = join.right(); + assert_eq!("SortedInputExec", probe.name(), "{build_side:?}"); + assert_eq!("ScanExec", probe.children()[0].name(), "{build_side:?}"); + assert_eq!( + "col_0@0 ASC", + probe.properties().output_ordering().unwrap().to_string() + ); + assert!( + planned.native_plan.properties().output_ordering().is_some(), + "{build_side:?}" + ); + } + } + #[test] fn hash_join_declares_order_through_operators_that_keep_it() { let projection = Operator { @@ -4108,7 +4141,7 @@ mod tests { })), ..Default::default() }; - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, spark_operator::BuildSide::BuildLeft, vec![create_scan(), projection], @@ -4157,7 +4190,7 @@ mod tests { (LeftAnti, BuildRight, DF::RightAnti), (Existence, BuildRight, DF::RightMark), ] { - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, BuildLeft, vec![ @@ -4208,7 +4241,7 @@ mod tests { ), ] { let operator = - sorted_outer_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); + sorted_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); let (_, _, planned) = PhysicalPlanner::default() .create_plan(&operator, &mut vec![], 1) .unwrap(); @@ -4229,7 +4262,7 @@ mod tests { })), ..Default::default() }; - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, spark_operator::BuildSide::BuildLeft, vec![create_scan(), sort], diff --git a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala index 39b0b086283..3a55bd8e85b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -100,7 +100,8 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) case _ => false } - // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH are outside JoinStrategyHint.strategies. + // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH are outside JoinStrategyHint.strategies, so + // the join stays with this strategy, which builds on the side PREFER_SHUFFLE_HASH names. private def hasStrategyHint(hint: JoinHint): Boolean = Seq(hint.leftHint, hint.rightHint).exists( _.exists(_.strategy.exists(JoinStrategyHint.strategies.contains))) @@ -115,10 +116,16 @@ case class CometShuffledHashJoinStrategy(session: SparkSession) planned } + // When AQE hints only one side to build on, that side wins, much as in + // getShuffleHashJoinBuildSide. private def buildSide(join: Join): Option[BuildSide] = { val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) - if (leftBuildable && rightBuildable) { + val leftHinted = leftBuildable && hintToPreferShuffleHashJoinLeft(join.hint) + val rightHinted = rightBuildable && hintToPreferShuffleHashJoinRight(join.hint) + if (leftHinted != rightHinted) { + Some(if (leftHinted) BuildLeft else BuildRight) + } else if (leftBuildable && rightBuildable) { Some(optimalBuildSide(join)) } else if (leftBuildable) { Some(BuildLeft) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index da26c91f9c7..006a126fd42 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -2939,26 +2939,29 @@ trait CometHashJoin { // and an input sorted outside the native plan, a relation sorted before it was cached for // instance, declares no order to DataFusion. Send the order so the native join declares it // on the probe input; see SortedInputExec. Every other join type keeps the probe order as - // it is. The null-aware anti join is the exception, as Comet builds that one on its streamed - // side, but it runs natively only when that side reports no order, so it reports none. + // it is, but an operator above it in the native plan, a window for one, sees that order only + // if the probe input declares it, so the order is sent for those joins too. The null-aware + // anti join is the exception, as Comet builds that one on its streamed side, but it runs + // natively only when that side reports no order, so it reports none. // - // The order Spark reports can be false to begin with. A Comet join that may move its - // streamed rows reports it all the same: a right outer join like this one whose probe input - // declared no order, and the nested loop join, which keeps the order of neither input. A - // cache reports the order of the plan that wrote it, so the rows of such a join can reach - // this one through a cache, as a JVM input that the native planner cannot look behind. - // Declaring the order there would let an operator above rely on it, an - // aggregate in sorted mode for instance, which then splits groups. So nothing is declared - // when such a join feeds the streamed side with no sort in between, or when a leaf reports - // an order that cannot be checked; see mayReportFalseOrder. When it comes from the JVM or a - // nested loop join, the join stays in Spark, as Spark may have dropped a sort above it. + // The order Spark reports can be false to begin with. A right outer join like this one whose + // probe input declared no order reports it all the same. A cache reports the order of the + // plan that wrote it, so the rows of such a join can reach this one through a cache, as a JVM + // input that the native planner cannot look behind. Declaring the order there would let an + // operator above rely on it, an aggregate in sorted mode for instance, which then splits + // groups. So nothing is declared when such a join feeds the streamed side with no sort in + // between, or when a leaf reports an order that cannot be checked; see mayReportFalseOrder. + // When it comes from the JVM, an outer join of the two shapes stays in Spark, as Spark may + // have dropped a sort above it. // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix // of the order that Comet can express is enough. A key that exprToProto takes counts as // expressible, unless it holds a collated string, at the top level or nested, which Spark // orders by its collation but the native side would declare in raw byte order. That is the // check `supportedSortType` makes for a sort. Nothing here can tell whether an operator above - // relies on the order, so a join whose first key Comet cannot express stays in Spark. + // relies on the order, so an outer join of the two shapes whose first key Comet cannot express + // stays in Spark. Any other join keeps its probe rows in place regardless, so it stays native + // and declares what it can. val (streamed, _) = streamedSide(join.buildSide, join.left, join.right) val keepsStreamedOrder = runsAsRightOuterJoin(join.joinType, join.buildSide) && join.outputOrdering.nonEmpty @@ -2967,14 +2970,14 @@ trait CometHashJoin { return None } val streamedSortOrders = - if (keepsStreamedOrder && !mayReportFalseOrder(streamed, countInPlan = true)) { + if (join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed, countInPlan = true)) { val orders = join.outputOrdering.iterator .takeWhile(order => !hasNonDefaultStringCollation(order.dataType)) .map(exprToProto(_, streamed.output)) .takeWhile(_.isDefined) .flatten .toSeq - if (orders.isEmpty) { + if (orders.isEmpty && keepsStreamedOrder) { withFallbackReason( join, "the streamed side's sort order, which this outer join has to keep, cannot be " + @@ -3008,10 +3011,9 @@ trait CometHashJoin { // Whether the order that `plan` reports may be false, because something that feeds it with no // sort in between cannot be relied on to keep it. Comet's hash joins report the order of their - // streamed side, which Spark's own joins keep, but two Comet joins do not always: an outer join - // that runs as DataFusion's right outer join, which keeps its probe rows in place only if the - // planner declared an order on its probe input; and the nested loop join, which DataFusion runs - // build row by build row. So does a leaf that reports an order from state this walk cannot + // streamed side, which Spark's own joins keep, but an outer join that runs as DataFusion's right + // outer join does not always, as it keeps its probe rows in place only if the planner declared + // an order on its probe input. So does a leaf that reports an order from state this walk cannot // check, an RDD checkpointed from a plan that is gone for one. // // The walk follows the order, so a hash join, which passes on only its streamed side's, is @@ -3028,9 +3030,6 @@ trait CometHashJoin { mayReportFalseOrder(adaptive.executedPlan, countInPlan) case _ if plan.outputOrdering.isEmpty => false case _: SortExec | _: CometSortExec => false - // Its rows stay in order only over a one-row build side, and a native join above would move - // them. Over more build rows they are already out of order (#6898); drop this once fixed. - case _: CometBroadcastNestedLoopJoinExec => true case join: CometBroadcastHashJoinExec => hashJoinMayReportFalseOrder( join.nativeOp, @@ -3174,11 +3173,23 @@ object CometBroadcastNestedLoopJoinExec extends CometOperatorSerde[BroadcastNest " DataFusion's NestedLoopJoinExec does not provide. Affects: LeftOuter+BuildLeft," + " RightOuter+BuildRight, FullOuter, LeftSemi+BuildLeft, LeftAnti+BuildLeft." + // Spark's nested loop join keeps its streamed side's order and reports it, so a sort above it + // may be gone, but DataFusion's emits its rows build row by build row. Remove this fallback once + // the native join keeps that order. + private[apache] val streamedOrderReason: String = + "nested loop join over an ordered streamed side is not supported natively, because the " + + "native join does not keep that order " + + "(https://github.com/apache/datafusion-comet/issues/6898)" + override def getSupportLevel(op: BroadcastNestedLoopJoinExec): SupportLevel = (op.joinType, op.buildSide) match { case (Inner, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) | (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => - Compatible(None) + if (op.outputOrdering.nonEmpty) { + Unsupported(Some(streamedOrderReason)) + } else { + Compatible(None) + } case (LeftOuter | RightOuter | FullOuter | LeftSemi | LeftAnti, _) => Unsupported(Some(broadcastBuildReplicationReason)) case _ => Unsupported(Some(s"Join type ${op.joinType} is not supported")) 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 d2aaca3c5df..631711f1fcb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -36,20 +36,21 @@ 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, JoinHint, LogicalPlan, NO_BROADCAST_HASH} +import org.apache.spark.sql.catalyst.plans.logical.{HintInfo, Join, JoinHint, LogicalPlan, NO_BROADCAST_HASH, PREFER_SHUFFLE_HASH} import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} 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 import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{array, count, lit, struct, sum} 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, isSpark40Plus} +import org.apache.comet.rules.CometShuffledHashJoinStrategy class CometJoinSuite extends CometTestBase { @@ -1651,6 +1652,90 @@ class CometJoinSuite extends CometTestBase { } } + test("forceShuffledHashJoin builds on the side AQE's PREFER_SHUFFLE_HASH hint names") { + // Every shuffle partition of the right side fits a local hash map, so AQE hints that side on + // re-plan. The skewed left side is smaller in all, so the strategy would build on it otherwise. + 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 + } + } + def isPreferred(hint: Option[HintInfo]): Boolean = + hint.exists(_.strategy.contains(PREFER_SHUFFLE_HASH)) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SKEW_JOIN_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "16", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "131072", + SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "1572864", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_FORCE_SHJ.key -> "true") { + val left = spark + .range(90000) + .selectExpr("IF(id < 60000, 0L, id) AS lk", "sha2(cast(id AS string), 256) AS ls") + val right = + spark.range(200000).selectExpr("id AS rk", "sha2(cast(id AS string), 256) AS rs") + val strategies = spark.experimental.extraStrategies + spark.experimental.extraStrategies = strategies :+ recordJoinHints + try { + val (_, cometPlan) = checkSparkAnswer( + left + .join(right, $"lk" === $"rk") + .selectExpr("sum(length(ls))", "sum(length(rs))")) + assert(hints.exists(h => isPreferred(h.rightHint) && !isPreferred(h.leftHint)), hints) + assert( + collect(cometPlan) { case j: CometHashJoinExec => j.buildSide } == Seq(BuildRight), + cometPlan) + } finally { + spark.experimental.extraStrategies = strategies + } + } + } + + test("forceShuffledHashJoin builds on the one buildable side PREFER_SHUFFLE_HASH names") { + // tbl_b is the smaller table, so with no hint, or both sides hinted, the strategy builds on + // it. A LEFT ANTI join cannot build on its left side, so a hint there changes nothing. + val prefer = Some(HintInfo(strategy = Some(PREFER_SHUFFLE_HASH))) + val left = JoinHint(prefer, None) + val right = JoinHint(None, prefer) + val both = JoinHint(prefer, prefer) + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + def buildSide(query: String, hint: JoinHint): Seq[BuildSide] = { + val join = sql(query).queryExecution.optimizedPlan.collectFirst { case j: Join => j } + assert(join.isDefined, query) + CometShuffledHashJoinStrategy(spark)(join.get.copy(hint = hint)).collect { + case j: ShuffledHashJoinExec => j.buildSide + } + } + val bigFirst = "SELECT * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + val smallFirst = "SELECT * FROM tbl_b JOIN tbl_a ON tbl_a._2 = tbl_b._1" + val anti = "SELECT * FROM tbl_a LEFT ANTI JOIN tbl_b ON tbl_a._2 = tbl_b._1" + for ((query, hint, expected) <- Seq( + (bigFirst, JoinHint.NONE, BuildRight), + (bigFirst, left, BuildLeft), + (bigFirst, both, BuildRight), + (smallFirst, JoinHint.NONE, BuildLeft), + (smallFirst, right, BuildRight), + (smallFirst, both, BuildLeft), + (anti, left, BuildRight))) { + assert(buildSide(query, hint) == Seq(expected), s"$query $hint") + } + } + } + } + test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { withForcedShjTables(adaptive = true) { withSQLConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { @@ -2658,6 +2743,46 @@ class CometJoinSuite extends CometTestBase { } } + test("BroadcastNestedLoopJoin over an ordered streamed side falls back and keeps its order") { + // Spark's nested loop join keeps the native range's order, and Spark drops the sort for it, + // either directly above the join or above a hash join over it. The native nested loop join + // emits its rows build row by build row, so it stays in Spark. + val nestedLoop = + "SELECT /*+ BROADCAST(n) */ p.id AS x, n.id AS z FROM range(0, 20, 1, 1) p " + + "JOIN range(0, 2, 1, 1) n ON p.id % 2 <= n.id" + def overHashJoin(hint: String): String = + s"SELECT $hint j.x, j.z, b.id AS w FROM ($nestedLoop) j JOIN range(0, 2, 1, 1) b " + + "ON j.x % 4 = b.id * 2 ORDER BY x" + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + for ((query, force, hashJoins) <- Seq( + (s"SELECT * FROM ($nestedLoop) j ORDER BY x", false, 0), + (overHashJoin(""), true, 1), + (overHashJoin("/*+ SHUFFLE_HASH(b) */"), false, 1), + (overHashJoin("/*+ BROADCAST(b) */"), false, 1))) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + val (_, plan) = checkSparkAnswerAndFallbackReason( + sql(query), + CometBroadcastNestedLoopJoinExec.streamedOrderReason) + val context = s"$query\n$plan" + assert(collect(plan) { case j: BroadcastNestedLoopJoinExec => j }.size == 1, context) + assert(collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j }.isEmpty, context) + // The hash join over the nested loop join stays in Spark as well. + val joins = collect(plan) { + case j: ShuffledHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(joins.size == hashJoins, context) + val xs = sql(query).collect().map(_.getLong(0)).toSeq + assert(xs == xs.sorted, s"$xs\n$context") + } + } + } + } + test("ExistenceJoin via BroadcastHashJoin (EXISTS combined with OR)") { withSQLConf( CometConf.COMET_EXEC_EXISTENCE_JOIN_ENABLED.key -> "true", @@ -3033,6 +3158,85 @@ class CometJoinSuite extends CometTestBase { } } + test("an inner hash join declares a sorted cache's order for a native window above it") { + // Spark drops the window's sort, as the join reports the cache's order, so the native window + // needs that order declared on the join's probe input, which reads the cache from the JVM. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + val cached = spark + .range(2000) + .selectExpr("id % 20 AS k", "id AS v") + .repartition(2, $"k") + .sortWithinPartitions("k", "v") + .cache() + try { + cached.count() + val keys = spark.range(20).toDF("id") + for ((build, force) <- Seq( + keys -> true, + keys.hint("shuffle_hash") -> false, + keys.hint("broadcast") -> false)) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = cached + .join(build, $"k" === $"id") + .selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") + sameRowsAsSpark(df) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val declared = collect(plan) { + case j: CometHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + case j: CometBroadcastHashJoinExec => + j.nativeOp.getHashJoin.getStreamedSortOrdersCount + } + assert(declared == Seq(2), plan) + } + } + } finally { + cached.unpersist() + } + } + } + + test("an inner hash join declares the part of the streamed order Comet can express") { + // An inner join keeps its probe rows in place without a declared order, so one whose first + // key Comet cannot express declares nothing and still runs natively. With Add disabled, Comet + // cannot express `_1 + 1`. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Add") -> "false") { + withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + val build = spark.table("small") + try { + for ((order, declared) <- Seq(Seq($"_1" + 1) -> 0, Seq($"_1", $"_1" + 1) -> 1)) { + val cached = + spark.table("big").repartition(2, $"_1").sortWithinPartitions(order: _*).cache() + cached.count() + val (_, plan) = checkSparkAnswer( + cached.join(build.hint("shuffle_hash"), cached("_1") === build("_1"))) + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.size == 1, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount == declared, plan) + cached.unpersist() + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + test("outer hash joins keep a descending, nulls last streamed order with null keys") { // Null keys match nothing, so all of them take the unmatched path, at the end of each // partition. @@ -3165,6 +3369,32 @@ class CometJoinSuite extends CometTestBase { } } + test("an inner hash join over a checkpointed streamed side declares no order") { + // Comet cannot check the checkpoint's order, so it declares none, and the inner join keeps its + // probe rows in place without one, so it stays native. + withOrderedCheckpoint { streamed => + val build = spark.range(10).selectExpr("id * 10 AS k", "id AS w") + val order = + (0 to 1).map(i => SortOrder(BoundReference(i, LongType, nullable = true), Ascending)) + for (hint <- Seq("shuffle_hash", "broadcast")) { + def df: DataFrame = streamed + .join(build.hint(hint), streamed("k") === build("k")) + .select(streamed("k"), streamed("v"), build("w")) + .sortWithinPartitions("k", "v") + val (_, plan) = checkSparkAnswer(df) + val declared = collect(plan) { + case j: CometHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + case j: CometBroadcastHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + } + assert(declared == Seq(0), s"$hint\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$hint\n$plan") + assert(partitionsInOrder(df, order: _*), s"$hint\n$plan") + } + } + } + test("chained outer hash joins over a native window keep its order natively") { // The window and both joins share one native plan, where DataFusion sees the order itself. withParquetTable((0 until 2000).map(i => (i % 40, i)), "t") { @@ -3702,36 +3932,27 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a native nested loop join falls back") { - // Over two build rows the nested loop join emits rows out of order (#6898), so they are - // compared as a multiset. Over one build row they stay in order, and Spark's join keeps it. + test("an outer hash join over a nested loop join over a sorted cache keeps its order") { + // The nested loop join stays in Spark, which keeps the cache's order over any number of build + // rows, and so does the outer join over it, so the order holds without the sort Spark drops. withSortedCache() { withParquetTable(Seq(Tuple1(0), Tuple1(1)), "two") { - withParquetTable(Seq(Tuple1(1)), "one") { - def query(nlBuild: String): String = - s"""SELECT /*+ BROADCAST(n), BROADCAST(b) */ a.k, a.v, n._1 AS z, b._2 AS w - |FROM cached a JOIN $nlBuild n ON a.v % 2 <= n._1 - |LEFT JOIN small b ON a.k = b._1""".stripMargin - val reason = CometHashJoin.uncheckedStreamedOrderReason - def checkFallback(df: => DataFrame): SparkPlan = { - sameRowsAsSpark(df) - val plan = df.queryExecution.executedPlan - val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) - assert(reasons.exists(_.contains(reason)), plan) - val nestedLoopJoins = collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j } - assert(nestedLoopJoins.size == 1, plan) - val outer = collect(plan) { case j: BroadcastHashJoinExec => j.joinType } - assert(outer == Seq(LeftOuter), plan) - plan - } - - checkFallback(sql(query("two"))) - - def ordered: DataFrame = sql(query("one")).sortWithinPartitions("k") - val plan = checkFallback(ordered) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) - assert(partitionsInOrder(ordered, keyOrder(Ascending)), plan) - } + def df: DataFrame = + sql("""SELECT /*+ BROADCAST(n), BROADCAST(b) */ a.k, a.v, n._1 AS z, b._2 AS w + |FROM cached a JOIN two n ON a.v % 2 <= n._1 + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .sortWithinPartitions("k") + sameRowsAsSpark(df) + val plan = df.queryExecution.executedPlan + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert( + reasons.exists(_.contains(CometBroadcastNestedLoopJoinExec.streamedOrderReason)), + plan) + assert(collect(plan) { case j: BroadcastNestedLoopJoinExec => j }.size == 1, plan) + val outer = collect(plan) { case j: BroadcastHashJoinExec => j.joinType } + assert(outer == Seq(LeftOuter), plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } } } From c20eb20f887231ff875ea0ff7cf79f8938fd49d4 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 28 Oct 2026 02:04:12 +0700 Subject: [PATCH 7/7] fix: sort a native window's input when its partition keys arrive unordered With the forced shuffled hash join planned before Spark places sorts, the join reports its streamed side's order and Spark drops the sort below a window. When that side is a checkpoint, Comet cannot verify the order and declares none, so the native window rejected its input with "Expects PARTITION BY expression to be ordered". Hinted joins and windows directly over a checkpoint failed the same way. The window now sorts its input by the partition keys, then the order keys, whenever its input does not show every partition key as ordered. That is the condition the window rejects, so plans that ran before are unchanged. The sort is reported with the window's metrics. --- .../latest/compatibility/operators.md | 4 + native/core/src/execution/planner.rs | 1 + .../planner/tests/window_input_order.rs | 406 ++++++++++++++++++ native/core/src/execution/planner/window.rs | 50 ++- .../apache/comet/exec/CometJoinSuite.scala | 66 ++- .../comet/exec/CometWindowExecSuite.scala | 26 +- 6 files changed, 544 insertions(+), 9 deletions(-) create mode 100644 native/core/src/execution/planner/tests/window_input_order.rs diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 78f9312fb17..25deaf3ab05 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -210,6 +210,10 @@ supported range of 0001-01-01 to 9999-12-31, such as one produced by `date_add` native frame can treat the bound as the edge of the partition where Spark compares against it ([#5022](https://github.com/apache/datafusion-comet/issues/5022)). +Spark drops the sort below a window when it believes the window's input is already ordered. For a window with +`PARTITION BY` keys, when Comet cannot verify that order, as with rows read from a checkpoint, the native window sorts +its input by the `PARTITION BY` keys and then the `ORDER BY` keys. The results match Spark, at the cost of a sort. + `WindowGroupLimitExec` (window-based limit pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`) runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` (default: true). diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 40829ddae76..79672fd9b51 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3096,6 +3096,7 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool { mod tests { mod empty_native_scan; mod multi_store_scan; + mod window_input_order; use futures::{poll, StreamExt}; use std::{ diff --git a/native/core/src/execution/planner/tests/window_input_order.rs b/native/core/src/execution/planner/tests/window_input_order.rs new file mode 100644 index 00000000000..6dbd3bb8685 --- /dev/null +++ b/native/core/src/execution/planner/tests/window_input_order.rs @@ -0,0 +1,406 @@ +// 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. + +//! A native window sorts its input unless every PARTITION BY key is ordered. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Float64Array, Int32Array}; +use arrow::compute::cast; +use arrow::datatypes::DataType; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_comet_proto::spark_expression::{ + self, data_type::DataTypeId, expr::ExprStruct, Expr, +}; +use datafusion_comet_proto::spark_operator::{ + self, lower_window_frame_bound::LowerFrameBoundStruct, operator::OpStruct, + upper_window_frame_bound::UpperFrameBoundStruct, CurrentRow, LowerWindowFrameBound, Operator, + UnboundedPreceding, UpperWindowFrameBound, WindowFrame, WindowFrameType, WindowSpecDefinition, +}; + +use super::collect_with_inputs; +use crate::execution::metrics::utils::to_native_metric_node; +use crate::execution::operators::InputBatch; +use crate::execution::planner::PhysicalPlanner; + +/// `row_number` runs in `BoundedWindowAggExec` and `percent_rank` in `WindowAggExec`. +const FUNCTIONS: [(&str, &str); 2] = [ + ("row_number", "BoundedWindowAggExec"), + ("percent_rank", "WindowAggExec"), +]; + +fn data_type(type_id: DataTypeId) -> spark_expression::DataType { + spark_expression::DataType { + type_id: type_id as i32, + type_info: None, + } +} + +fn int32() -> spark_expression::DataType { + data_type(DataTypeId::Int32) +} + +fn typed_column(index: i32, type_id: DataTypeId) -> Expr { + Expr { + expr_struct: Some(ExprStruct::Bound(spark_expression::BoundReference { + index, + datatype: Some(data_type(type_id)), + })), + ..Default::default() + } +} + +fn column(index: i32) -> Expr { + typed_column(index, DataTypeId::Int32) +} + +/// Spark's default null ordering: ASC NULLS FIRST, DESC NULLS LAST. +fn sort_order_of(child: Expr, descending: bool) -> Expr { + Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(child)), + direction: i32::from(descending), + null_ordering: i32::from(descending), + }, + ))), + ..Default::default() + } +} + +fn sort_order(index: i32, descending: bool) -> Expr { + sort_order_of(column(index), descending) +} + +/// Columns from the JVM, which declares no order. +fn scan_of(fields: Vec) -> Operator { + Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields, + source: String::new(), + })), + ..Default::default() + } +} + +/// Three Int32 columns, (a, b, v). +fn scan() -> Operator { + scan_of(vec![int32(); 3]) +} + +fn sort_by(child: Operator, sort_orders: Vec) -> Operator { + Operator { + children: vec![child], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders, + ..Default::default() + })), + ..Default::default() + } +} + +fn sort(child: Operator, order: &[(i32, bool)]) -> Operator { + let orders = order.iter().map(|&(i, desc)| sort_order(i, desc)).collect(); + sort_by(child, orders) +} + +fn window(child: Operator, function: &str, partition_by: &[i32], order_by: &[i32]) -> Operator { + let partition_by = partition_by.iter().map(|&i| column(i)).collect(); + let order_by = order_by.iter().map(|&i| sort_order(i, false)).collect(); + window_by(child, function, partition_by, order_by, None) +} + +/// A window with one function. A `result_type` that differs from the native one makes the +/// planner cast the result in a projection above the window. +fn window_by( + child: Operator, + function: &str, + partition_by_list: Vec, + order_by_list: Vec, + result_type: Option, +) -> Operator { + let expr = spark_operator::WindowExpr { + built_in_window_function: Some(Expr { + expr_struct: Some(ExprStruct::ScalarFunc(spark_expression::ScalarFunc { + func: function.to_string(), + ..Default::default() + })), + ..Default::default() + }), + spec: Some(WindowSpecDefinition { + frame_specification: Some(WindowFrame { + frame_type: WindowFrameType::Rows as i32, + lower_bound: Some(LowerWindowFrameBound { + lower_frame_bound_struct: Some(LowerFrameBoundStruct::UnboundedPreceding( + UnboundedPreceding {}, + )), + }), + upper_bound: Some(UpperWindowFrameBound { + upper_frame_bound_struct: Some(UpperFrameBoundStruct::CurrentRow( + CurrentRow {}, + )), + }), + }), + ..Default::default() + }), + result_type, + ..Default::default() + }; + Operator { + children: vec![child], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![expr], + partition_by_list, + order_by_list, + ..Default::default() + }))), + ..Default::default() + } +} + +fn find_window(plan: &Arc) -> Arc { + if matches!(plan.name(), "BoundedWindowAggExec" | "WindowAggExec") { + return Arc::clone(plan); + } + find_window(plan.children().first().expect("no window in the plan")) +} + +fn count_sorts(plan: &Arc) -> usize { + usize::from(plan.downcast_ref::().is_some()) + + plan.children().into_iter().map(count_sorts).sum::() +} + +/// The keys of the sort directly below the window as (column, descending, nulls first). +fn input_sort_keys(window: &Arc) -> Option> { + let sort = window.children()[0].downcast_ref::()?; + Some( + sort.expr() + .iter() + .map(|e| { + let column = e.expr.downcast_ref::().expect("a column sort key"); + (column.index(), e.options.descending, e.options.nulls_first) + }) + .collect(), + ) +} + +struct Case { + name: &'static str, + input_order: Option<&'static [(i32, bool)]>, + partition_by: &'static [i32], + order_by: &'static [i32], + /// The keys of the sort the planner adds, as (column, descending, nulls first). + added_sort: Option>, +} + +#[test] +fn window_sorts_input_without_partition_key_order() { + let asc = |index: usize| (index, false, true); + let cases = [ + Case { + name: "no declared order", + input_order: None, + partition_by: &[0], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(2)]), + }, + // The ORDER BY key is not declared either. Only the PARTITION BY keys are checked. + Case { + name: "ordered on the partition key", + input_order: Some(&[(0, false)]), + partition_by: &[0], + order_by: &[2], + added_sort: None, + }, + Case { + name: "ordered on the first of two partition keys", + input_order: Some(&[(0, false)]), + partition_by: &[0, 1], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(1), asc(2)]), + }, + Case { + name: "ordered on the partition keys in another order", + input_order: Some(&[(1, false), (0, false)]), + partition_by: &[0, 1], + order_by: &[2], + added_sort: None, + }, + Case { + name: "ordered descending on the partition key", + input_order: Some(&[(0, true)]), + partition_by: &[0], + order_by: &[2], + added_sort: None, + }, + Case { + name: "no PARTITION BY", + input_order: None, + partition_by: &[], + order_by: &[2], + added_sort: None, + }, + ]; + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + for (function, exec) in FUNCTIONS { + for case in &cases { + let child = match case.input_order { + Some(order) => sort(scan(), order), + None => scan(), + }; + let operator = window(child, function, case.partition_by, case.order_by); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let window = find_window(&planned.native_plan); + let context = format!("{function}, {}", case.name); + assert_eq!(window.name(), exec, "{context}"); + assert_eq!( + count_sorts(&window), + usize::from(case.input_order.is_some()) + usize::from(case.added_sort.is_some()), + "{context}" + ); + if case.added_sort.is_some() { + assert_eq!(input_sort_keys(&window), case.added_sort, "{context}"); + } + } + } +} + +#[test] +fn window_keeps_input_sorted_on_a_float_partition_key() { + // Both the sort key and the PARTITION BY key go through float normalization, so the window + // sees the input as ordered on its key. + let key = || typed_column(0, DataTypeId::Double); + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + for (function, exec) in FUNCTIONS { + let input = scan_of(vec![data_type(DataTypeId::Double), int32(), int32()]); + let child = sort_by(input, vec![sort_order_of(key(), false)]); + let order_by = vec![sort_order(2, false)]; + let operator = window_by(child, function, vec![key()], order_by, None); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let window = find_window(&planned.native_plan); + assert_eq!(window.name(), exec, "{function}"); + assert_eq!(count_sorts(&window), 1, "{function}"); + } +} + +/// Rows of (a, b, v) for the execution tests. +fn input_columns() -> (Vec, usize) { + let a = vec![1, 0, 1, 0, 1]; + let v = vec![5, 4, 3, 2, 1]; + let rows = a.len(); + let columns: Vec = vec![ + Arc::new(Int32Array::from(a)), + Arc::new(Int32Array::from(vec![0; rows])), + Arc::new(Int32Array::from(v)), + ]; + (columns, rows) +} + +#[tokio::test] +async fn window_reports_its_input_sort_metrics() { + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + for (cast_result, sorted_input) in [(false, false), (false, true), (true, false), (true, true)] + { + let context = format!("cast_result={cast_result} sorted_input={sorted_input}"); + let child = if sorted_input { + sort(scan(), &[(0, false)]) + } else { + scan() + }; + // `row_number` returns UInt64 natively, so an Int32 result type adds a cast. + let result_type = cast_result.then(int32); + let partition_by = vec![column(0)]; + let order_by = vec![sort_order(2, false)]; + let operator = window_by(child, "row_number", partition_by, order_by, result_type); + let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + + let additional: Vec<&str> = planned + .additional_native_plans + .iter() + .map(|p| p.name()) + .collect(); + let expected: &[&str] = match (cast_result, sorted_input) { + (false, false) => &["SortExec"], + (false, true) => &[], + (true, false) => &["BoundedWindowAggExec", "SortExec"], + (true, true) => &["BoundedWindowAggExec"], + }; + assert_eq!(additional, expected, "{context}"); + + let (columns, rows) = input_columns(); + let mut inputs = [vec![InputBatch::Batch(columns, rows), InputBatch::EOF].into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + rows, + "{context}" + ); + + let elapsed = |plan: &Arc| { + plan.metrics() + .and_then(|m| m.elapsed_compute()) + .unwrap_or(0) as i64 + }; + let total_time = std::iter::once(&planned.native_plan) + .chain(&planned.additional_native_plans) + .map(elapsed) + .sum::(); + let metrics = to_native_metric_node(&planned).unwrap(); + // Output rows come from the top operator only, and the time covers all of them. + assert_eq!(metrics.metrics["output_rows"], rows as i64, "{context}"); + assert_eq!(metrics.metrics["elapsed_compute"], total_time, "{context}"); + } +} + +#[tokio::test] +async fn window_over_unordered_input_ranks_each_partition() { + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + for (function, expected) in [ + ("row_number", [1.0, 2.0, 1.0, 2.0, 3.0]), + ("percent_rank", [0.0, 1.0, 0.0, 0.5, 1.0]), + ] { + let operator = window(scan(), function, &[0], &[2]); + let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let (columns, num_rows) = input_columns(); + let mut inputs = [vec![InputBatch::Batch(columns, num_rows), InputBatch::EOF].into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + let mut rows = vec![]; + for batch in &batches { + let ints = |i: usize| { + let column = batch.column(i).as_any().downcast_ref::(); + column.unwrap().clone() + }; + let result = cast(batch.column(3), &DataType::Float64).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let (keys, values) = (ints(0), ints(2)); + for row in 0..batch.num_rows() { + rows.push((keys.value(row), values.value(row), result.value(row))); + } + } + rows.sort_by_key(|&(k, v, _)| (k, v)); + let ranks: Vec = rows.iter().map(|r| r.2).collect(); + assert_eq!(rows.len(), num_rows, "{function}"); + assert_eq!(ranks, expected, "{function}"); + } +} diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 2d15fdc86e1..0cbfe186f7b 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -36,8 +36,9 @@ use datafusion::{ physical_plan::{ limit::LocalLimitExec, projection::ProjectionExec, + sorts::sort::SortExec, windows::{BoundedWindowAggExec, WindowAggExec}, - ExecutionPlan, InputOrderMode, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }, }; use datafusion_comet_operators::{PartitionedRankLimitExec, WindowFnKind}; @@ -98,6 +99,36 @@ impl OperatorBuilder for WindowBuilder { let sort_exprs = &sort_exprs?; let partition_exprs = &partition_exprs?; + // Sort the input unless it is already ordered on every PARTITION BY key; rows from the + // JVM declare no order even when Spark reports them as sorted. The check looks only at + // the PARTITION BY keys; the sort orders by them ascending with nulls first, then by the + // ORDER BY keys. + let (_, ordered) = child + .native_plan + .equivalence_properties() + .find_longest_permutation(partition_exprs)?; + let required = partition_exprs + .iter() + .map(|e| { + PhysicalSortExpr::new( + Arc::clone(e), + SortOptions { + descending: false, + nulls_first: true, + }, + ) + }) + .chain(sort_exprs.iter().cloned()); + let input_sort: Option> = match LexOrdering::new(required) { + Some(ordering) if ordered.len() < partition_exprs.len() => Some(Arc::new( + SortExec::new(ordering, Arc::clone(&child.native_plan)), + )), + _ => None, + }; + let window_input = input_sort + .as_ref() + .map_or_else(|| Arc::clone(&child.native_plan), Arc::clone); + let window_expr: Result>, ExecutionError> = wnd .window_expr .iter() @@ -134,14 +165,14 @@ impl OperatorBuilder for WindowBuilder { let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + window_input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + window_input, !partition_exprs.is_empty(), )?) }; @@ -183,16 +214,21 @@ impl OperatorBuilder for WindowBuilder { proj_exprs, Arc::clone(&window_agg), )?); - // The window is an additional plan so that its time is reported with the - // projection's rather than left out. + // The window and any input sort are additional plans so that their time is + // reported with the projection's rather than left out. SparkPlan::new_with_additional( spark_plan.plan_id, projection, vec![child], - vec![window_agg], + std::iter::once(window_agg).chain(input_sort).collect(), ) } else { - SparkPlan::new(spark_plan.plan_id, window_agg, vec![child]) + SparkPlan::new_with_additional( + spark_plan.plan_id, + window_agg, + vec![child], + input_sort.into_iter().collect(), + ) } }; 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 631711f1fcb..270006d6127 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeRefer 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.{HintInfo, Join, JoinHint, LogicalPlan, NO_BROADCAST_HASH, PREFER_SHUFFLE_HASH} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec, CometWindowGroupLimitExec} 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 @@ -3302,6 +3302,15 @@ class CometJoinSuite extends CometTestBase { .localCheckpoint(eager = true)) } + /** + * Build sides for an inner join with a checkpoint, each with whether to force a shuffled hash + * join: unhinted, then hinted for a shuffled hash join, then for a broadcast. + */ + private def checkpointJoinBuildSides: Seq[(DataFrame, Boolean)] = { + val keys = spark.range(100).toDF("id") + Seq(keys -> true, keys.hint("shuffle_hash") -> false, keys.hint("broadcast") -> false) + } + test("outer hash joins over a checkpointed streamed side fall back and keep its order") { // Spark drops the sort above the join for the checkpoint's order, which Comet cannot check. withOrderedCheckpoint { streamed => @@ -3395,6 +3404,61 @@ class CometJoinSuite extends CometTestBase { } } + test("a native window above an inner hash join over a checkpointed streamed side runs") { + // Spark drops the window's sort for the checkpoint's order, which the join reports, but Comet + // cannot check that order, so the join declares none and the native window sorts its input. + withOrderedCheckpoint { streamed => + for { + (build, force) <- checkpointJoinBuildSides + window <- Seq( + "dense_rank() OVER (PARTITION BY k ORDER BY v)", + "sum(v) OVER (PARTITION BY k)") + } { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = streamed + .join(build, $"k" === $"id") + .selectExpr("k", "v", s"$window AS r") + val plan = df.queryExecution.executedPlan + val context = s"force=$force $window\n$plan" + val joins = collect(plan) { + case j: CometHashJoinExec => j + case j: CometBroadcastHashJoinExec => j + } + assert(joins.size == 1, context) + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + } + } + } + } + + test("a native window group limit above an inner hash join over a checkpoint ranks it") { + assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") + // The native limit declares the order Spark requires of its input without checking it, which + // holds for the checkpoint's rows. + withOrderedCheckpoint { streamed => + for ((build, force) <- checkpointJoinBuildSides) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = streamed + .join(build, $"k" === $"id") + .selectExpr("k", "v", "rank() OVER (PARTITION BY k ORDER BY v) AS rk") + .where("rk <= 2") + val plan = df.queryExecution.executedPlan + val context = s"force=$force\n$plan" + assert(collect(plan) { case l: CometWindowGroupLimitExec => l }.nonEmpty, context) + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + } + } + } + } + test("chained outer hash joins over a native window keep its order natively") { // The window and both joins share one native plan, where DataFusion sees the order itself. withParquetTable((0 until 2000).map(i => (i % 40, i)), "t") { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 642a73e337a..2c7651913c3 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -28,7 +28,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Divide, Expression, MakeDecimal, WindowExpression} import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} -import org.apache.spark.sql.execution.{InputAdapter, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{InputAdapter, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.{count, expr, lead, sum} @@ -543,6 +543,30 @@ class CometWindowExecSuite extends CometTestBase { } } + test("a native window over a sorted checkpoint runs without Spark's sort") { + // Spark drops the window's sort for the checkpoint's order, but those rows enter the native + // plan with no declared order, so the native window sorts them itself. + // AQE is off, as an adaptive plan reports no order. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + val checkpoint = spark + .range(2000) + .selectExpr("id % 20 AS k", "id AS v") + .repartition(2, $"k") + .sortWithinPartitions("k", "v") + .localCheckpoint(eager = true) + for (window <- Seq( + "dense_rank() OVER (PARTITION BY k ORDER BY v)", + "sum(v) OVER (PARTITION BY k)")) { + val (_, plan) = checkSparkAnswer(checkpoint.selectExpr("k", "v", s"$window AS r")) + val context = s"$window\n$plan" + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, context) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, context) + } + } + } + test("Window range frame with long boundary should not fail") { val df = Seq((1L, "1"), (1L, "1"), (2147483650L, "1"), (3L, "2"), (2L, "1"), (2147483650L, "2"))