viirya commented on code in PR #57181:
URL: https://github.com/apache/spark/pull/57181#discussion_r3566948657
##########
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:
The new rule runs after `ReplaceHashWithSortAgg`. That rule replaces a hash
aggregate with a sort aggregate when the child is already sorted (e.g. on top
of a sort merge join), and this rule then destroys exactly that ordering. When
both are enabled: with `countLocalSort` off we adopt a strictly worse plan
(sort aggregate plus the sort that `EnsureRequirements` re-inserts, instead of
the original hash aggregate); with `countLocalSort` on, the extra sort makes
the cost evaluator reject the whole candidate plan, so we also lose an
otherwise beneficial SMJ-to-SHJ conversion. Placing this rule before
`ReplaceHashWithSortAgg` avoids both cases, since the ordering is already gone
when `ReplaceHashWithSortAgg` makes its decision.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1337,6 +1337,30 @@ object SQLConf {
.bytesConf(ByteUnit.BYTE)
.createWithDefault(0L)
+ val ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED =
+
buildConf("spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled")
+ .doc("When true, Spark converts a sort merge join to a shuffled hash
join during adaptive " +
+ "execution when the build side's materialized per-partition sizes are
all within " +
+ s"${ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key}, even when
non-shuffle " +
+ "operators sit between the join and its input shuffle.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val ADAPTIVE_COST_EVALUATOR_COUNT_LOCAL_SORT_ENABLED =
Review Comment:
Enabling only `convertSortMergeJoinToShuffledHashJoin.enabled` without
`costEvaluator.countLocalSort.enabled` means a conversion that pushes extra
sorts elsewhere is still adopted (your own test shows the 5-sorts-vs-4 case).
Is there a reason to keep them independent? If the cost-evaluator change is
kept behind its own flag for safety, the docs should at least recommend
enabling them together - or `countLocalSort` could arguably default to true,
since it is a no-op for cost comparison when nothing else changes the sort
count.
##########
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:
`(numShuffles.toLong << SHUFFLE_SHIFT) | sortCost` is not masked, so the
shuffle count only safely fits in the 16 bits between `SHUFFLE_SHIFT` and
`SKEW_SHIFT`. The previous encoding was safe by construction (`numShuffles` is
an `Int` in the low 32 bits), while the new one silently wraps if a plan ever
has 65536+ shuffles - the overflow bits get absorbed into the sign-extended
skew term. Not a practical concern for any real plan, but since we're now
packing three components, I'd suggest replacing the bit tricks with a small
`Cost` implementation holding `(numSkewJoins, numShuffles, numLocalSorts)` and
comparing lexicographically - it also removes the need to keep the shift
comment in sync (the comment currently says the shuffle band has 32 bits, but
it has 16). Clamping the shuffle count is fine too if you prefer keeping the
packed `Long`.
##########
docs/sql-performance-tuning.md:
##########
@@ -375,6 +375,22 @@ AQE converts sort-merge join to shuffled hash join when
all post shuffle partiti
</td>
<td>3.2.0</td>
</tr>
+ <tr>
+
<td><code>spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled</code></td>
Review Comment:
Nit: this entry only mentions the per-partition size condition, but
`preferShuffledHashJoin` also requires
`spark.sql.adaptive.advisoryPartitionSizeInBytes <=
maxShuffledHashJoinLocalMapThreshold`, which is worth mentioning like the
existing 3.2.0 entry above does.
##########
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)
Review Comment:
Nit: the rule name `ReplaceSortMergeJoinToShuffledHashJoin` mixes the two
naming patterns nearby (`ReplaceHashWithSortAgg` vs the config's
`convert...To...`) - `ReplaceSortMergeJoinWithShuffledHashJoin` or
`ConvertSortMergeJoinToShuffledHashJoin` would be more consistent.
--
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]
