cloud-fan commented on code in PR #57181: URL: https://github.com/apache/spark/pull/57181#discussion_r3580998483
########## 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: `left` here is bound by the case-class extractor to `smj.left` (same for `right`/`smj.right`), so both arguments to `wideningFactor` are the *same* `Seq[Attribute]`: ```scala val leftFactor = wideningFactor(smj.left.output, left.output) // left eq smj.left val rightFactor = wideningFactor(smj.right.output, right.output) // right eq smj.right ``` `wideningFactor` then computes `getSizePerRow(build) / getSizePerRow(shuffle)` on identical outputs, which is always `1.0`, so it collapses to `math.max(minWideningFactor, 1.0)` — i.e. just the manual `minWideningFactor` floor (default `1.0`). The automatic per-row widening estimate never contributes. Per its own doc ("the estimated per-row byte-size ratio of the build subtree's output to its **input shuffle's** output") and PR-description item 3, the second argument should be the input shuffle stage's output — `leftStage.get.output` / `rightStage.get.output` — not the join child's output. The effect: on the look-through path an operator that widens each row through a size-bounded expression (e.g. `cast(int as string)`, which the comment at :245 tolerates *because* `wideningFactor` was meant to account for it) is no longer scaled, so the build-size bound degrades to raw shuffle bytes — the same non-spillable-hash-map OOM class @sunchao's P1 blocked on. It's inert on the default direct-shuffle path (there the build output legitimately equals the shuffle output), so it only bites when `lookThroughOperators.enabled=true`. A one-line swap to `leftStage.get.output` won't work as-is, since `leftFactor` is computed before the `leftStage.isDefined` check (`.get` would throw when the stage is `None`). Cleanest fix is to fold the `wideningFactor` call into the `canBuildLeft`/`canBuildRight` guards where `leftStage`/`rightStage` are already known-defined, e.g.: ```scala val canBuildLeft = leftStage.isDefined && canBuildShuffledHashJoinLeft(smj.joinType) && preferShuffledHashJoin( leftStage.get.mapStats.get, wideningFactor(smj.left.output, leftStage.get.output)) ``` This is my own miss from the prior round — sorry for not catching it then. The existing `MinWideningFactor` test can't detect it because it drives the `math.max` *floor* (factor `1000.0`), never a computed ratio; a regression test with a size-bounded-but-widening operator (e.g. `cast`) between shuffle and join — where the correct ratio would reject a conversion that ratio=1.0 wrongly accepts — would lock this in. -- 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]
