ulysses-you commented on code in PR #57181: URL: https://github.com/apache/spark/pull/57181#discussion_r3568250246
########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ReplaceSortMergeJoinToShuffledHashJoin.scala: ########## @@ -0,0 +1,132 @@ +/* + * 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.spark.sql.execution.adaptive + +import scala.annotation.tailrec + +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} +import org.apache.spark.sql.catalyst.plans.LeftExistence +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SortExec, SparkPlan} +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, EnsureRequirements} +import org.apache.spark.sql.execution.joins.{BaseJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.window.{WindowExecBase, WindowGroupLimitExec} + +/** + * Converts a [[SortMergeJoinExec]] into a [[ShuffledHashJoinExec]] during adaptive execution when + * a build side's materialized shuffle statistics show it is small enough for a local hash map. + * Unlike [[DynamicJoinSelection]], this runs on the physical plan, so it can reach the input + * shuffle through operators (aggregate, project, filter, window, etc...) sitting above it. + * + * The swap is shuffle-free since both joins are `ShuffledJoin`s with the same distribution and + * partitioning; only the child sorts become unnecessary. As a shuffled hash join loses the sort + * merge join's output ordering, [[EnsureRequirements]] is re-run to restore any ordering an + * ancestor still needs, and AQE's [[CostEvaluator]] decides whether to adopt the converted plan. + */ +case class ReplaceSortMergeJoinToShuffledHashJoin(ensureRequirements: EnsureRequirements) + extends Rule[SparkPlan] with JoinSelectionHelper { + + /** + * Chooses the build side for the shuffled hash join. A side is eligible only if it is allowed + * as a build side for this join type and its input shuffle is small enough to build a local + * hash map. When both sides are eligible, the smaller one (by total shuffle bytes) is chosen. + */ + private def selectBuildSide( + smj: SortMergeJoinExec, + left: ShuffleQueryStageExec, + right: ShuffleQueryStageExec): Option[BuildSide] = { + val canBuildLeft = canBuildShuffledHashJoinLeft(smj.joinType) && + preferShuffledHashJoin(left.mapStats.get) + val canBuildRight = canBuildShuffledHashJoinRight(smj.joinType) && + preferShuffledHashJoin(right.mapStats.get) + if (canBuildLeft && canBuildRight) { + if (left.mapStats.get.bytesByPartitionId.sum < right.mapStats.get.bytesByPartitionId.sum) { + Some(BuildLeft) + } else { + Some(BuildRight) + } + } else if (canBuildLeft) { + Some(BuildLeft) + } else if (canBuildRight) { + Some(BuildRight) + } else { + None + } + } + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.convertSortMergeJoinToShuffledHashJoinEnabled) { + return plan + } + val optimizedPlan = plan.transformUp { + case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, Review Comment: Addressed, and generalized: `hasJoinStrategyHint` reads the SMJ's `logicalLink` and skips conversion when either side carries any strategy hint (not just `MERGE`), matching `DynamicJoinSelection`, which never overrides an existing strategy hint. Added an enabled-config regression test that `/*+ MERGE */` keeps the join as a sort merge join. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ReplaceSortMergeJoinToShuffledHashJoin.scala: ########## @@ -0,0 +1,132 @@ +/* + * 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.spark.sql.execution.adaptive + +import scala.annotation.tailrec + +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} +import org.apache.spark.sql.catalyst.plans.LeftExistence +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SortExec, SparkPlan} +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, EnsureRequirements} +import org.apache.spark.sql.execution.joins.{BaseJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.window.{WindowExecBase, WindowGroupLimitExec} + +/** + * Converts a [[SortMergeJoinExec]] into a [[ShuffledHashJoinExec]] during adaptive execution when + * a build side's materialized shuffle statistics show it is small enough for a local hash map. + * Unlike [[DynamicJoinSelection]], this runs on the physical plan, so it can reach the input + * shuffle through operators (aggregate, project, filter, window, etc...) sitting above it. + * + * The swap is shuffle-free since both joins are `ShuffledJoin`s with the same distribution and + * partitioning; only the child sorts become unnecessary. As a shuffled hash join loses the sort + * merge join's output ordering, [[EnsureRequirements]] is re-run to restore any ordering an + * ancestor still needs, and AQE's [[CostEvaluator]] decides whether to adopt the converted plan. + */ +case class ReplaceSortMergeJoinToShuffledHashJoin(ensureRequirements: EnsureRequirements) + extends Rule[SparkPlan] with JoinSelectionHelper { + + /** + * Chooses the build side for the shuffled hash join. A side is eligible only if it is allowed + * as a build side for this join type and its input shuffle is small enough to build a local + * hash map. When both sides are eligible, the smaller one (by total shuffle bytes) is chosen. + */ + private def selectBuildSide( + smj: SortMergeJoinExec, + left: ShuffleQueryStageExec, + right: ShuffleQueryStageExec): Option[BuildSide] = { + val canBuildLeft = canBuildShuffledHashJoinLeft(smj.joinType) && + preferShuffledHashJoin(left.mapStats.get) + val canBuildRight = canBuildShuffledHashJoinRight(smj.joinType) && + preferShuffledHashJoin(right.mapStats.get) + if (canBuildLeft && canBuildRight) { + if (left.mapStats.get.bytesByPartitionId.sum < right.mapStats.get.bytesByPartitionId.sum) { + Some(BuildLeft) + } else { + Some(BuildRight) + } + } else if (canBuildLeft) { + Some(BuildLeft) + } else if (canBuildRight) { + Some(BuildRight) + } else { + None + } + } + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.convertSortMergeJoinToShuffledHashJoinEnabled) { + return plan + } + val optimizedPlan = plan.transformUp { + case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, + ExtractShuffleStage(left), ExtractShuffleStage(right), false) => + selectBuildSide(smj, left, right) match { + case Some(buildSide) => + ShuffledHashJoinExec(leftKeys, rightKeys, joinType, buildSide, condition, Review Comment: Fixed. Conversion is now gated on `hashJoinSupported(leftKeys, rightKeys)`, so non-binary-stable keys fall through to SMJ, same as `JoinSelection.createShuffleHashJoin` and the `DynamicJoinSelection` hint path. Added a `UTF8_LCASE` regression test. Note: for a bare `attr = attr` collated equi-key, `RewriteCollationJoin` injects a binary-stable `CollationKey` before physical planning, so those are legitimately convertible; the guard catches the cases the rewrite does not cover (e.g. a collated key wrapped in an expression), where the physical keys stay non-binary-stable. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala: ########## @@ -126,6 +128,7 @@ case class AdaptiveSparkPlanExec( AdjustShuffleExchangePosition, ValidateSparkPlan, ReplaceHashWithSortAgg, + ReplaceSortMergeJoinToShuffledHashJoin(ensureRequirements), Review Comment: Fixed - the rule now runs **before** `ReplaceHashWithSortAgg`, so the ordering is already gone when `ReplaceHashWithSortAgg` makes its decision. This avoids both the strictly-worse sort-aggregate-plus-reinserted-sort plan and the `countLocalSort`-rejects-the-candidate case you described. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala: ########## @@ -37,24 +37,52 @@ case class SimpleCost(value: Long) extends Cost { /** * A skew join aware implementation of [[CostEvaluator]], which counts the number of - * [[ShuffleExchangeLike]] nodes and skew join nodes in the plan. + * [[ShuffleExchangeLike]] nodes, skew join nodes and (optionally) local [[SortExec]] nodes in the + * plan. + * + * The cost is packed into a single [[Long]] so that the components are compared in priority order. + * From the most significant bits to the least significant: + * - `-numSkewJoins` (only when `forceOptimizeSkewedJoin` is true), so that more skew joins means + * lower cost and is compared first; + * - `numShuffles`, so that fewer shuffles means lower cost; + * - `numLocalSorts` (only when `countLocalSort` is true), the lowest-priority tiebreaker, so that + * among plans with the same number of skew joins and shuffles the one with fewer local sorts is + * preferred (e.g. a shuffled hash join over a sort merge join when the conversion does not push + * extra sorts elsewhere). */ -case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean) extends CostEvaluator { +case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean, countLocalSort: Boolean) + extends CostEvaluator { + + import SimpleCostEvaluator._ + override def evaluateCost(plan: SparkPlan): Cost = { - val numShuffles = plan.collect { - case s: ShuffleExchangeLike => s - }.size + var numShuffles = 0 + var numLocalSorts = 0 + plan.foreach { + case _: ShuffleExchangeLike => numShuffles += 1 + case s: SortExec if !s.global => numLocalSorts += 1 + case _ => + } + val sortCost = if (countLocalSort) numLocalSorts.toLong else 0L + val shuffleAndSortCost = (numShuffles.toLong << SHUFFLE_SHIFT) | sortCost Review Comment: Done. Replaced the packed `Long` with `SimpleCost(numSkewJoins, numShuffles, numLocalSorts)` compared lexicographically (skew joins descending, then shuffles, then local sorts), so there is no shift/overflow concern and no shift-comment to keep in sync. -- 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]
