dwsmith1983 commented on code in PR #6785: URL: https://github.com/apache/datafusion-comet/pull/6785#discussion_r4278036358
########## 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: > 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. Confirmed. With `forceShuffledHashJoin` the planner now picks `ShuffledHashJoinExec`, 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's `Right` join puts the unmatched probe rows after the matched ones in each batch. I reproduced it with a cached `repartition(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_HASH` hint 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. -- 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]
