Repository navigation
fix: choose the forced shuffled hash join at planning time so Spark keeps its sorts #6785
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,175 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| package org.apache.comet.rules | ||
|
|
||
| import org.apache.spark.sql.SparkSession | ||
| import org.apache.spark.sql.catalyst.expressions.RowOrdering | ||
| import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} | ||
| import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys | ||
| import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} | ||
| import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, JoinStrategyHint, LogicalPlan} | ||
| import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy} | ||
| import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, LogicalQueryStage} | ||
| import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} | ||
| import org.apache.spark.sql.internal.SQLConf | ||
|
|
||
| import org.apache.comet.CometConf | ||
| import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, withFallbackReason, withInfo} | ||
| import org.apache.comet.shims.ShimJoinSelection | ||
|
|
||
| /** | ||
| * Plans a ShuffledHashJoinExec in place of a SortMergeJoinExec, before EnsureRequirements and | ||
| * RemoveRedundantSorts, so the sorts above the join survive. Adapted from Apache Gluten. | ||
| */ | ||
| case class CometShuffledHashJoinStrategy(session: SparkSession) | ||
| extends SparkStrategy | ||
| with ShimJoinSelection { | ||
|
|
||
| override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { | ||
| case join: Join if isEnabled(join) => planJoin(join) | ||
| case _ => Nil | ||
| } | ||
|
|
||
| private def sqlConf: SQLConf = session.sessionState.conf | ||
|
|
||
| // Planner strategies run with Comet off and before CometRule, which applies plan-only mode. | ||
| private def isEnabled(join: Join): Boolean = | ||
| CometConf.COMET_FORCE_SHJ.get(sqlConf) && | ||
| CometConf.COMET_EXEC_ENABLED.get(sqlConf) && | ||
| !CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(sqlConf) && | ||
| !join.isStreaming && | ||
| isCometLoaded(sqlConf) | ||
|
|
||
| private def planJoin(join: Join): Seq[SparkPlan] = join match { | ||
| case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, _, left, right, hint) | ||
| if !isLeftToSpark(join, hint) && canSortMergeJoin(joinType) && | ||
| RowOrdering.isOrderable(leftKeys) && hashJoinSupportedShim(leftKeys, rightKeys) => | ||
| buildSide(join) match { | ||
| case Some(BuildRight) if joinType == LeftSemi || joinType.isInstanceOf[ExistenceJoin] => | ||
| // LeftSemi https://github.com/apache/datafusion-comet/issues/2667 | ||
| // ExistenceJoin https://github.com/apache/datafusion-comet/issues/2697 | ||
| planBySpark(join) { smj => | ||
| withFallbackReason(smj, declineReason(s"BuildRight with $joinType is not supported")) | ||
| } | ||
| case Some(side) => | ||
| buildSideOverLimit(join, side) match { | ||
| case Some(reason) => | ||
| // Not a fallback: the kept sort-merge join still runs natively. | ||
| planBySpark(join)(smj => withInfo(smj, declineReason(reason))) | ||
| case None => | ||
| Seq( | ||
| ShuffledHashJoinExec( | ||
|
Member
There was a problem hiding this comment. [P1] Preserve the final ordering when a checkpointed probe cannot receive a native ordering declaration. In an enabled Comet session, disable AQE and broadcasts, set two shuffle partitions, and enable Evidence: The exact-source Spark 4.1.3 differential harness confirms that the base retains one final Sort and the head retains none for this checkpoint query. The checkpoint is an RDDScanExec with two ordering expressions and no children. Source tracing through CometScanWrapper and CometSparkToColumnarExec reaches operators.scala:3036, causing the declaration at lines 2951–2969 to be empty. Replaying the exported Int64 checkpoint partitions through DataFusion 55.2.0’s right outer hash join preserves all 2,000 rows but produces ordering transitions (60,1960) -> (2,2) and (90,1990) -> (1,1). Using the current SortedInputExec removes both violations. Reproduction sources and logs: /tmp/review6785-153a-validation-qhry1dr7/{CachedOrderProbe.scala,run-checkpoint.log,CheckpointProbe.scala,checkpoint-run.log,checkpoint_native.rs,checkpoint-native-run.log}. This is a planning-plus-kernel reproduction, not an integrated Comet execution.
Contributor
Author
There was a problem hiding this comment. Confirmed. It isn't specific to the forced join: on main, a The fix goes to main first. When the streamed order can't be checked and reaches the join through an input from the JVM (a checkpoint or a cache) or through a nested loop join, the outer hash join falls back to Spark. Joins over native operators in the same plan stay native, since DataFusion sees the order there. Once that lands I'll merge it here and add the forced checkpoint case. |
||
| leftKeys, | ||
| rightKeys, | ||
| joinType, | ||
| side, | ||
| condition, | ||
| planLater(left), | ||
| planLater(right))) | ||
| } | ||
| case None => Nil | ||
| } | ||
| case _ => Nil | ||
| } | ||
|
|
||
| private def isLeftToSpark(join: Join, hint: JoinHint): Boolean = | ||
| isBroadcastStage(join.left) || isBroadcastStage(join.right) || | ||
| canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) | ||
|
|
||
| // canPlanAsBroadcastHashJoin misses a broadcast stage when the adaptive threshold is lower or | ||
| // off, and AQE discards a re-plan that hash joins one. | ||
| private def isBroadcastStage(plan: LogicalPlan): Boolean = plan match { | ||
| case LogicalQueryStage(_, _: BroadcastQueryStageExec) => true | ||
| case _ => false | ||
| } | ||
|
|
||
| // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH are outside JoinStrategyHint.strategies. | ||
| private def hasStrategyHint(hint: JoinHint): Boolean = | ||
| Seq(hint.leftHint, hint.rightHint).exists( | ||
| _.exists(_.strategy.exists(JoinStrategyHint.strategies.contains))) | ||
|
|
||
| private def declineReason(reason: String): String = | ||
| s"Cannot rewrite SortMergeJoin to HashJoin: $reason" | ||
|
|
||
| private def planBySpark(join: Join)( | ||
| record: SortMergeJoinExec => SortMergeJoinExec): Seq[SparkPlan] = { | ||
| val planned = session.sessionState.planner.JoinSelection(join) | ||
| planned.flatMap(_.collect { case smj: SortMergeJoinExec => smj }).foreach(record) | ||
| planned | ||
| } | ||
|
|
||
| private def buildSide(join: Join): Option[BuildSide] = { | ||
| val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) | ||
| val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) | ||
| if (leftBuildable && rightBuildable) { | ||
| Some(optimalBuildSide(join)) | ||
| } else if (leftBuildable) { | ||
| Some(BuildLeft) | ||
| } else if (rightBuildable) { | ||
| Some(BuildRight) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
|
|
||
| private def optimalBuildSide(join: Join): BuildSide = { | ||
| val leftSize = join.left.stats.sizeInBytes | ||
| val rightSize = join.right.stats.sizeInBytes | ||
| val leftRowCount = join.left.stats.rowCount | ||
| val rightRowCount = join.right.stats.rowCount | ||
| if (leftSize == rightSize && rightRowCount.isDefined && leftRowCount.isDefined) { | ||
| if (rightRowCount.get <= leftRowCount.get) BuildRight else BuildLeft | ||
| } else if (rightSize <= leftSize) { | ||
| BuildRight | ||
| } else { | ||
| BuildLeft | ||
| } | ||
| } | ||
|
|
||
| /** The largest build side to hash join, or None when the limit is off. */ | ||
| private def maxBuildSize: Option[BigInt] = | ||
| CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(sqlConf) match { | ||
| case Some(limit) if limit <= 0 => None | ||
| case Some(limit) => Some(BigInt(limit)) | ||
| case None => | ||
| val threshold = sqlConf.autoBroadcastJoinThreshold | ||
| val perPartition = | ||
| if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get | ||
| Some(BigInt(perPartition) * sqlConf.numShufflePartitions) | ||
| } | ||
|
|
||
| /** | ||
| * Why the build side is not strictly under the limit, as in canBuildLocalHashMapBySize. An | ||
| * unknown size is spark.sql.defaultSizeInBytes, so it never fits. | ||
| */ | ||
| private def buildSideOverLimit(join: Join, side: BuildSide): Option[String] = | ||
| maxBuildSize.flatMap { limit => | ||
| val buildSize = side match { | ||
| case BuildLeft => join.left.stats.sizeInBytes | ||
| case BuildRight => join.right.stats.sizeInBytes | ||
| } | ||
| if (buildSize >= limit) { | ||
| Some( | ||
| s"build side size estimate of $buildSize bytes is not under the limit of $limit bytes") | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[P1] Preserve cached probe ordering before exposing the hash join’s ordering to Spark. With forcing enabled and AQE/broadcasts disabled, materialize a probe using
.repartition(2, col("k")).sortWithinPartitions("k", "v").cache(), left-join a smaller table containing only some probe keys, then request.sortWithinPartitions("k", "v"). SelectingShuffledHashJoinExechere lets Spark remove the final sort because it advertises the probe’s entire ordering. Comet reads the cache through nativeScanExec, which carries no ordering metadata. The build-right left outer join becomes DataFusion’sRightjoin, whose unordered-input path appends unmatched probe rows after matched rows. The result therefore violates the requested ordering. This is introduced by the earlier selection: the exact base retains the final(k, v)sort because the original sort-merge join advertises onlyk. Please preserve ordering through the native input/outer join, or retain the sort/fallback until that guarantee is available, and add a cached-probe regression test.Evidence: Spark 4.1.3 differential harness: probe
range(2000).selectExpr("id % 100 as k", "id as v").repartition(2, col("k")).sortWithinPartitions("k", "v").cache()followed by materialization; buildrange(10).selectExpr("id * 10 as k"); left outer join followed bysortWithinPartitions("k", "v"). Applying the exact-base rewrite leaves one finalSort [k, v]; the exact-head strategy leaves none. Separately, native DataFusion HashJoinExec with build keys[1,3], ordered probe keys[0,1,2,3,4], and no ordering metadata returned[1,3,0,2,4]. Its reproduced 55.1.0hash_join/exec.rs,hash_join/stream.rs, andjoins/utils.rsare byte-identical to pinned 55.2.0. Current CometScanExec::newinitializesEquivalencePropertieswithout ordering. Reproduction files and logs are in/tmp/review6785-fdde-probe-v6vo_b5b/. Full Comet native execution was not available.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Confirmed. With
forceShuffledHashJointhe planner now picksShuffledHashJoinExec, which reports the cached probe's(k, v)ordering, so Spark drops the final sort. Comet reads the cache through a scan with no ordering, and DataFusion'sRightjoin puts the unmatched probe rows after the matched ones in each batch. I reproduced it with a cachedrepartition(2, k).sortWithinPartitions(k, v)probe left-joined to a 5-key table. Every partition comes back out of order, with the right rows.The same native gap is already reachable on main without this conf, through a
SHUFFLE_HASHhint or a broadcast hash join. Both fail that test on main. #6835 fixes it at the source: it declares the streamed order on the probe for these two outer shapes and follows the cache scan into the cached plan. With its non-test changes on top of this branch, the forced, hinted and broadcast cases all pass.So I'd like #6835 to land first. I'll then merge it into this branch and add the cached-probe test for the forced and hinted joins. cc @andygrove, since this puts #6785 after #6835.