ulysses-you commented on code in PR #57181: URL: https://github.com/apache/spark/pull/57181#discussion_r3584106562
########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ConvertSortMergeJoinToShuffledHashJoin.scala: ########## @@ -0,0 +1,257 @@ +/* + * 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.MapOutputStatistics +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, CaseWhen, Cast, Coalesce, Expression, If, Literal, String2TrimExpression, Substring, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, JoinSelectionHelper} +import org.apache.spark.sql.catalyst.plans.LeftExistence +import org.apache.spark.sql.catalyst.plans.logical.{Join, SHUFFLE_MERGE} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{CollectMetricsExec, 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} +import org.apache.spark.sql.internal.SQLConf + +/** + * 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. + * + * This runs on the physical plan and owns the shuffled-hash-over-sort-merge selection that AQE + * makes from materialized shuffle statistics. It is gated by + * `spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled` (default true, the master + * switch), and has two modes: + * - Default: it looks through the sort merge join's own required [[SortExec]] to reach a + * *direct* input shuffle. + * - Behind `...convertSortMergeJoinToShuffledHashJoin.lookThroughOperators.enabled`: it + * additionally looks through non-shuffle operators (aggregate, project, filter, window, + * left-existence join) sitting above the shuffle. + * + * 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. + * + * A shuffled hash join builds a non-spillable local hash map, so the traversed operators must not + * blow up the build size that the input shuffle statistics estimate: + * - the traversal only looks through an operator whose output expressions are all size-bounded + * (see [[isSizeBoundedExpr]]), so no operator can widen a row in a way the shuffle statistics + * cannot see; and + * - the build-side estimate is scaled by [[wideningFactor]] to account for the width change the + * traversed operators do introduce. + */ +case class ConvertSortMergeJoinToShuffledHashJoin(ensureRequirements: EnsureRequirements) + extends Rule[SparkPlan] with JoinSelectionHelper { + + private def preferShuffledHashJoin( + mapStats: MapOutputStatistics, + sizeInBytesFactor: Double): Boolean = { + val maxShuffledHashJoinLocalMapThreshold = + conf.getConf(SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD) + val advisoryPartitionSize = conf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) + advisoryPartitionSize <= maxShuffledHashJoinLocalMapThreshold && + mapStats.bytesByPartitionId.forall( + _ * sizeInBytesFactor <= maxShuffledHashJoinLocalMapThreshold) + } + + /** + * The estimated per-row byte-size ratio of the build subtree's output to its input shuffle's + * output, i.e. how much the traversed operators widen each row. The traversed operators never + * increase the row count (`N_build <= N_shuffle`), so scaling the input shuffle bytes by this + * ratio keeps them a valid upper bound on the hash-map build size once row width is accounted + * for: `buildSize = N_build * buildRowWidth <= shuffleBytes * (buildRowWidth / shuffleRowWidth)`. + * + * Floored at `spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.minWideningFactor` + * (default 1.0). Unlike `SizeInBytesOnlyStatsPlanVisitor`, which computes a best-effort size and + * lets a narrowing operator shrink it, the default keeps a conservative bound for a non-spillable + * build: `getSizePerRow` under-estimates a variable-width column (it uses `defaultSize`), so a + * `factor < 1` could push the scaled bytes below the real build size and reintroduce the + * out-of-memory risk, whereas the raw shuffle bytes are always a valid bound when the build side + * is no wider than the shuffle row. Raising the floor above 1.0 is more conservative still. + */ + private def wideningFactor(buildOutput: Seq[Attribute], shuffleOutput: Seq[Attribute]): Double = { + val buildRowSize = EstimationUtils.getSizePerRow(buildOutput).toDouble + val shuffleRowSize = EstimationUtils.getSizePerRow(shuffleOutput).toDouble + math.max(conf.convertSortMergeJoinToShuffledHashJoinMinWideningFactor, + buildRowSize / shuffleRowSize) + } + + private def hasSortMergeJoinHint(smj: SortMergeJoinExec): Boolean = smj.logicalLink.exists { + case j: Join => + j.hint.leftHint.exists(_.strategy.contains(SHUFFLE_MERGE)) || + j.hint.rightHint.exists(_.strategy.contains(SHUFFLE_MERGE)) + case _ => false + } + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.convertSortMergeJoinToShuffledHashJoinEnabled) { + return plan + } + val lookThroughOperatorsEnabled = + conf.convertSortMergeJoinToShuffledHashJoinLookThroughOperatorsEnabled + val optimizedPlan = plan.transformUp { + case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, + left, right, isSkewJoin) + // Do not convert if the join keys are not hash-join-compatible (e.g. collated or other + // non-binary-stable string keys), since a hash join matches keys by binary equality and + // would return wrong results. This mirrors the guard on the other SHJ-planning paths. + if !hasSortMergeJoinHint(smj) && hashJoinSupported(leftKeys, rightKeys) => + val leftStage = findShuffleStage(left, lookThroughOperatorsEnabled) + val rightStage = findShuffleStage(right, lookThroughOperatorsEnabled) + val leftFactor = wideningFactor(smj.left.output, left.output) Review Comment: Good catch, and thanks for the detailed diagnosis. Fixed in a40f1d68a1f: `wideningFactor` now takes the input shuffle stage's output as its second argument (`leftStage.get.output` / `rightStage.get.output`), computed only when the stage is defined -- `leftFactor`/`rightFactor` are now `Option`s derived from `leftStage.map(...)`/`rightStage.map(...)`, so there's no `.get`-on-None hazard. Added a regression test (`Widening factor uses the input shuffle's row width, not the join child's`) with a size-bounded-but-widening `cast(count(*) AS string)` between the shuffle and the join: the shuffle row is 20 bytes (c1:int, count:long), the build row is 32 bytes (c1:int, s:string), factor 1.6; the build side's largest partition (372 bytes) scaled by 1.6 is ~595, above the 500-byte threshold, so the join stays SMJ. I verified the test fails on the old self-comparison (ratio 1.0 -> converts) and passes with the fix. -- 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]
