sunchao commented on code in PR #6785: URL: https://github.com/apache/datafusion-comet/pull/6785#discussion_r4321400980
########## spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala: ########## @@ -0,0 +1,197 @@ +/* + * 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 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( Review 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 `forceShuffledHashJoin`, `spark.comet.convert.rdd.enabled`, and `spark.comet.exec.range.enabled`. Build the probe with `spark.range(2000).selectExpr("id % 100 as k", "id as v").repartition(2, col("k")).sortWithinPartitions("k", "v").localCheckpoint(eager = true)`, left-join `spark.range(10).selectExpr("id * 10 as k")`, then request `sortWithinPartitions("k", "v")`. Selecting the hash join here makes Spark remove the final sort because `RDDScanExec` retains the checkpoint’s ordering. Unlike the sorted-cache case already discussed, this input hits `mayReportFalseOrder`’s unknown-leaf branch. `CometHashJoin.doConvert` therefore sends no `streamedSortOrders` but still converts the join. DataFusion then emits unmatched probe rows after matched rows, violating the requested ordering. T he exact base retains the final `(k, v)` sort. Please retain that sort or fall back to Spark when an ordered probe cannot safely receive the declaration. 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
