sunchao commented on code in PR #6785: URL: https://github.com/apache/datafusion-comet/pull/6785#discussion_r4268916352
########## 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 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")`. Selecting `ShuffledHashJoinExec` here lets Spark remove the final sort because it advertises the probe’s entire ordering. Comet reads the cache through native `ScanExec`, which carries no ordering metadata. The build-right left outer join becomes DataFusion’s `Right` join, 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 only `k`. Please preserve ordering through the native input/outer join, or retain the sort/fallback until t hat 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; build `range(10).selectExpr("id * 10 as k")`; left outer join followed by `sortWithinPartitions("k", "v")`. Applying the exact-base rewrite leaves one final `Sort [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.0 `hash_join/exec.rs`, `hash_join/stream.rs`, and `joins/utils.rs` are byte-identical to pinned 55.2.0. Current Comet `ScanExec::new` initializes `EquivalenceProperties` without ordering. Reproduction files and logs are in `/tmp/review6785-fdde-probe-v6vo_b5b/`. Full Comet native execution was not available. -- 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]
