peter-toth commented on code in PR #59630:
URL: https://github.com/apache/spark/pull/59630#discussion_r4495486184
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala:
##########
@@ -53,9 +53,20 @@ case class SimpleCost(numSkewJoins: Int, numShuffles: Int,
numLocalSorts: Int) e
/**
* A skew join aware implementation of [[CostEvaluator]], which counts the
number of
* [[ShuffleExchangeLike]] nodes, skew join nodes and (optionally) local
[[SortExec]] nodes in the
- * plan. See [[SimpleCost]] for how the components are compared.
+ * plan. With union shuffle push-down enabled, tagged new exchanges are
excluded and each union
+ * identified by [[PushDownShuffleExchangeThroughUnion.hasPushedShuffle]]
counts as one shuffle.
+ * Query-stage contents are not visited. Inspecting the same shape in current
and re-planned unions
+ * keeps an equivalent re-plan from gaining a shuffle-cost credit when a moved
branch finishes.
+ * Ordinary exchanges and recognized unions have equal weight, with no
separate tie-breaker.
+ * This is shape-based compensation, not a record of executed work: pruning
the only qualifying
+ * branch can remove the union cost even while the union survives. That
decrease can offset an
+ * additional exchange elsewhere, so adoption need not match the
push-down-disabled plan.
+ * See [[SimpleCost]] for how the components are compared.
*/
-case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean,
countLocalSort: Boolean)
+case class SimpleCostEvaluator(
+ forceOptimizeSkewedJoin: Boolean,
+ countLocalSort: Boolean,
+ pushDownShuffleThroughUnion: Boolean)
Review Comment:
**Finding 11.** Every two-argument `SimpleCostEvaluator(...)` call stops
compiling. A default keeps them working, and `AdaptiveSparkPlanExec` still
passes the conf:
```suggestion
pushDownShuffleThroughUnion: Boolean = false)
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala:
##########
@@ -46,7 +46,13 @@ case class InsertAdaptiveSparkPlan(
override def conf: SQLConf =
adaptiveExecutionContext.session.sessionState.conf
- override def apply(plan: SparkPlan): SparkPlan = applyInternal(plan, false)
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!adaptiveExecutionContext.queryContainsSubqueries &&
+ plan.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY,
DYNAMIC_PRUNING_SUBQUERY)) {
Review Comment:
**Finding 9.** Any scalar subquery in the physical plan sets this flag. That
includes the Bloom filter subqueries `InjectRuntimeFilter` adds by default. So
with `allowWithSubqueries` at its default, enabling the rule does nothing for a
query with a selective join anywhere, though the user wrote no subquery.
I measured it at this head, with AQE on, the rule and `lookThroughOperators`
enabled, broadcast off and `applicationSideScanSizeThreshold = 0`:
```sql
SELECT a.k, a.c, b.j FROM (
SELECT k, count(*) c FROM (
SELECT k, sum(v) v FROM p1 GROUP BY k UNION ALL SELECT k, v FROM p2)
GROUP BY k) a
JOIN p3 b ON a.k = b.k WHERE b.j = 1
```
- With Bloom filters on (the default), the plan has `might_contain`, and the
rule pushes no shuffle.
- With `spark.sql.optimizer.runtime.bloomFilter.enabled=false`, it pushes
one below the union.
A heuristic Bloom filter is built from a filtered scan or a materialized
cache, so its subquery has no exchange whose reuse the rewrite could break. The
gate could skip a `ScalarSubquery` that is the first argument of a
`BloomFilterMightContain`. Otherwise the `allowWithSubqueries` doc and the
description should name these subqueries, since users don't write them.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -8764,6 +8766,139 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED =
+ buildConf("spark.sql.adaptive.pushDownShuffleThroughUnion.enabled")
+ .internal()
+ .doc("When true, pushes an EnsureRequirements hash shuffle below a union
during AQE " +
+ "query stage preparation. Eligible children already reporting the
selected layout " +
+ "are retained; other children are retargeted or receive a new
exchange. The target " +
+ "partition count must be at least the smaller of the replaced
exchange's count and " +
+ "the effective initial shuffle count, including " +
+ "spark.sql.adaptive.coalescePartitions.initialPartitionNum when set.
Requires " +
Review Comment:
**Finding 12.** Four keys in the new docs are defined earlier in this file,
so they can be `${X.key}` and follow a rename:
- `spark.sql.adaptive.coalescePartitions.initialPartitionNum` (here):
`COALESCE_PARTITIONS_INITIAL_PARTITION_NUM`
- `spark.sql.unionOutputPartitioning` (line 8778):
`UNION_OUTPUT_PARTITIONING`
- `spark.sql.adaptive.coalescePartitions.enabled` (line 8816):
`COALESCE_PARTITIONS_ENABLED`
- `spark.sql.requireAllClusterKeysForCoPartition` (line 8856):
`REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION`
Only the `canChangeCachedPlanOutputPartitioning` doc needs the literal,
since that conf is defined before the push-down conf.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -1456,6 +1456,7 @@ case class EnsureRequirements(
}
}
if (hasSemanticEqualPartitioning(child.outputPartitioning)) {
+ PartitioningOrigin.record(child, upper, shuffleOrigin)
Review Comment:
**Finding 10.** This runs on every `EnsureRequirements` pass, in classic
preparation and in AQE with the rule off, though only
`PushDownShuffleExchangeThroughUnion` reads the tag. So the description's "no
change to default behavior" doesn't quite hold. `TreeNode.copyTagsFrom` copies
only onto a node that has no tags. Returning `child` here used to give a
tagless child the removed repartition's tags, its logical link among them. Once
`record` has tagged the child, it gets none. Planner-produced children usually
carry their own logical link, so this rarely shows. This is derived from
reading.
`AdaptiveSparkPlanExec` builds its own `EnsureRequirements` instance. It
could pass a flag, e.g. `recordElidedRepartitions =
conf.pushDownShuffleThroughUnionEnabled`, so that only a preparation that runs
this rule writes the tag.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PushDownShuffleExchangeThroughUnion.scala:
##########
@@ -0,0 +1,608 @@
+/*
+ * 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 scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap,
Expression, SessionWindow, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Partial =>
PartialAggregate}
+import org.apache.spark.sql.catalyst.optimizer.BuildLeft
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution,
HashPartitioning, HashShuffleSpec, Partitioning, PartitioningCollection}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.execution.{
+ FilterExec, ProjectExec, SortExec, SparkPlan, UnaryExecNode,
UnionConfSnapshot,
+ UnionExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS,
EnsureRequirements, Exchange, PartitioningOrigin, REPARTITION_BY_COL,
REPARTITION_BY_NUM, ShuffleExchangeExec, ShuffleExchangeLike,
ValidateRequirements}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
BroadcastNestedLoopJoinExec, CartesianProductExec, ShuffledJoin}
+import org.apache.spark.sql.execution.window.{Partial =>
PartialWindowGroupLimit, WindowGroupLimitExec}
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Push down a shuffle [[EnsureRequirements]] put above a [[UnionExec]] below
the union. Eligible
+ * children already reporting the target retain their layouts; other children
are retargeted or
+ * receive a new shuffle. Local sorts above the shuffle and operators between
it and the union
+ * are rebuilt above the union:
+ *
+ * {{{
+ * SortMergeJoin SortMergeJoin
+ * :- Sort :- Sort
+ * : +- Exchange => : +- Union
+ * : +- Union : :- HashAggregate
+ * : :- HashAggregate : : +- Exchange
+ * : : +- Exchange : : +- Scan
+ * : : +- Scan : +- Exchange
+ * : +- Project : +- Project
+ * : +- Scan : +- Scan
+ * +- ... +- ...
+ * }}}
+ *
+ * The left child already answers the target, since it shuffles on the join
key for its aggregate,
+ * so it keeps its layout and only the right child takes a shuffle.
+ *
+ * The operators between the shuffle and the union are looked through only
under
+ *
`spark.sql.adaptive.pushDownShuffleThroughUnion.lookThroughOperators.enabled`,
and only when
+ * moving the shuffle past them preserves results: a deterministic project, a
local sort, a
+ * filter, a partial aggregate and a partial window group limit. The reducing
ones are why the
+ * gate exists: the child is then shuffled before the reduction, which can
send more data than the
+ * one shuffle it replaces.
+ *
+ * Targets are hash layouts reported by the union's children, ranked by how
many children offer
+ * them when size estimates are unavailable; otherwise estimated bytes kept
decide first. The
+ * count must be at least the smaller of the replaced exchange's count and the
effective initial
+ * shuffle partition count (`coalescePartitions.initialPartitionNum` when
set). At least one
+ * child keeps its layout; others are retargeted or shuffled without changing
child order.
+ * `spark.sql.adaptive.pushDownShuffleThroughUnion.maxMovedBranches` limits
moved children, and
+ * `maxUnionChildren` limits all branches whose tails run within each consumer
task.
+ * The rebuilt child must satisfy its consumer's distribution and pair with
fixed siblings.
+ *
+ * Removing the union's exchange can reduce AQE join conversion and skew
handling opportunities.
+ * Keeping a shuffled join child's layout can make skew join optimization or
local shuffle reads
+ * invalidate an ancestor's distribution. AQE can discard the optimization for
the round, or add
+ * another shuffle when skew optimization is forced. Join inputs and join
branches are opt-in.
+ * Allowing targets via
+ *
`spark.sql.adaptive.pushDownShuffleThroughUnion.keepNonCoalescibleBranches.enabled`
+ * can prevent coalescing of the consumer's compatible shuffle group. Forced
skew optimization
+ * can restore a union exchange after moved branches have shuffled, shuffling
them a second time.
+ *
+ * Allowing subqueries can make pruning builds differ from main-query
broadcast builds. With
+ * broadcast-only dynamic pruning, failed broadcast reuse removes the pruning
filter and can scan
+ * all fact partitions; that cost can outweigh the saved build-side shuffle.
+ *
+ * The rule runs only in AQE preparation, after `EnsureRequirements` and before
+ * `StampUnionDecisions`. It preserves the root distribution and all user
repartition layouts
+ * recorded when `EnsureRequirements` elides exchanges, including requirements
below the root.
+ * These origin checks apply only to this rule's proposals, not to other
adaptive optimizations.
+ * Cost compensation belongs to SimpleCostEvaluator. Custom evaluators should
account for the
+ * one-to-many exchange rewrite when this rule is enabled, or they may reject
the entire re-plan.
+ *
+ * TODO: support push down of a keyed layout.
+ */
+case class PushDownShuffleExchangeThroughUnion(
+ snapshot: UnionConfSnapshot,
+ ensureRequirements: EnsureRequirements,
+ queryContainsSubqueries: Boolean) extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!snapshot.outputPartitioning ||
+ !conf.pushDownShuffleThroughUnionEnabled ||
+ (queryContainsSubqueries &&
+
!conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_WITH_SUBQUERIES_ENABLED)))
{
+ return plan
+ }
+
+ val lookThrough =
+
conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_LOOK_THROUGH_OPERATORS_ENABLED)
+ val policy = LayoutPolicy(
+ allowNonCoalescible = !conf.getConf(SQLConf.COALESCE_PARTITIONS_ENABLED)
|| conf.getConf(
+
SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_NON_COALESCIBLE_BRANCHES_ENABLED),
+ maxMovedBranches =
+
conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_MAX_MOVED_BRANCHES),
+ maxUnionChildren =
+
conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_MAX_UNION_CHILDREN),
+ keepJoinBranches =
+
conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_JOIN_BRANCHES_ENABLED),
+ allowSubsetTargets = conf.getConf(SQLConf
+
.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_PARTITION_KEYS_SUBSET_OF_CLUSTER_KEYS_ENABLED))
+ val allowJoinInputs =
+
conf.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_JOIN_INPUTS_ENABLED)
+
+ // An exchange the plan reuses elsewhere is worth more than this rewrite:
moving it below the
+ // union loses the stage the others reuse. See
`AQEUtils.duplicatedShuffleForms`. With reuse
+ // off every exchange gets its own stage, so there is nothing to protect.
+ lazy val sharedExchanges: Set[SparkPlan] =
+ if (conf.exchangeReuseEnabled) AQEUtils.duplicatedShuffleForms(plan)
else Set.empty
+
+ val pushed = plan.transformUp {
+ case join: ShuffledJoin if !allowJoinInputs => join
+ case operator: SparkPlan =>
+ val matched = operator.children.map { child =>
+ findPushableShuffle(child, lookThrough).filter { pushable =>
+ pushable.union.children.length <= policy.maxUnionChildren &&
+ (sharedExchanges.isEmpty ||
+ !sharedExchanges.contains(pushable.shuffle.canonicalized))
+ }
+ }
+ if (matched.forall(_.isEmpty)) {
+ operator
+ } else {
+ // Read requirements only after finding a site; some operators log
when asked.
+ val required = operator.requiredChildDistribution
+ if (!required.forall(_.isInstanceOf[ClusteredDistribution])) {
+ operator
+ } else {
+ val requirements =
required.map(_.asInstanceOf[ClusteredDistribution])
+ pushDownThroughUnion(
+ operator,
+ matched.zip(requirements),
+ sharedExchanges = sharedExchanges,
+ policy = policy.copy(allowNonCoalescible =
+ policy.allowNonCoalescible ||
!allExchangeInputs(operator))).getOrElse(operator)
+ }
+ }
+ }
+ if (pushed.fastEquals(plan)) {
+ plan
+ } else {
+ // Local checks do not cover ancestors or elided repartitions below the
root. Reject the
+ // whole proposal if any plan requirement or recorded output layout is
lost.
+ val satisfied = ensureRequirements.requiredDistribution match {
+ // A user repartition elided above was asked for on all its keys, so a
layout on a subset
+ // of them does not stand in for it.
+ case Some(required: ClusteredDistribution) =>
+ ValidateRequirements.validate(pushed,
required.copy(requireAllClusterKeys = true))
+ case Some(required) =>
+ ValidateRequirements.validate(pushed, required)
+ case None =>
+ ValidateRequirements.validate(pushed)
+ }
+ if (satisfied && !pushed.exists(node =>
!PartitioningOrigin.validateNode(node))) {
+ pushed
+ } else {
+ plan
+ }
+ }
+ }
+
+ /**
+ * One of an operator's children this rule can rewrite: the shuffle it moves
and the union that
+ * shuffle is moved past.
+ */
+ private case class PushableChild(shuffle: ShuffleExchangeExec, union:
UnionExec)
+
+ /**
+ * Layout selection settings, read once per pass.
+ */
+ private case class LayoutPolicy(
+ allowNonCoalescible: Boolean,
+ maxMovedBranches: Int,
+ maxUnionChildren: Int,
+ keepJoinBranches: Boolean,
+ allowSubsetTargets: Boolean)
+
+ /**
+ * The shape this rule acts on below one of an operator's children: a
pushable shuffle and the
+ * union it is moved past, with the operators in between looked through to
reach it. Only a hash
+ * shuffle is pushable, since the layout that replaces it comes from the
union's own children.
+ *
+ * What the child owes its operator is not read here: this is asked of every
child of every
+ * operator, and `requiredChildDistribution` is not free of side effects:
`WindowExecBase` logs
+ * from it for an unpartitioned window.
+ */
+ @tailrec
+ private def findPushableShuffle(
+ plan: SparkPlan,
+ lookThrough: Boolean,
+ found: Option[ShuffleExchangeExec] = None): Option[PushableChild] = {
+ plan match {
+ // A local sort above the shuffle is looked through, since the rebuilt
child is sorted back
+ // to the same requirement; one below it is an operator in between and
needs the gate.
+ case SortExec(_, false, child, _) if found.isEmpty || lookThrough =>
+ findPushableShuffle(child, lookThrough, found)
+ case shuffle @ ShuffleExchangeExec(partitioning, child,
ENSURE_REQUIREMENTS, _, _)
+ if found.isEmpty && supportPushDown(partitioning) =>
+ findPushableShuffle(child, lookThrough, Some(shuffle))
+ // AttributeMap keeps the last duplicate exprId, while runtime binding
uses the first.
+ // Positional rebinding is ambiguous for these unions, so retain their
exchange.
+ case union: UnionExec if union.output.map(_.exprId).distinct.size !=
union.output.size => None
+ // Below it: the operators this rule can cross, down to the union the
shuffle is moved past.
+ case union: UnionExec =>
+ found.map(PushableChild(_, union))
+ case operator: UnaryExecNode if found.isDefined && canCross(operator,
lookThrough) =>
+ findPushableShuffle(operator.child, lookThrough, found)
+ case _ => None
+ }
+ }
+
+ /**
+ * Rewrites eligible children in order. If a joint rewrite fails, retry
individual children
+ * against all current siblings, keeping children without a viable target
unchanged.
+ *
+ * `pushable` holds one entry per child, in child order: what can be
rewritten about it, and
+ * what it owes the operator.
+ */
+ private def pushDownThroughUnion(
+ operator: SparkPlan,
+ pushable: Seq[(Option[PushableChild], ClusteredDistribution)],
+ sharedExchanges: Set[SparkPlan],
+ policy: LayoutPolicy): Option[SparkPlan] = {
+ val children = operator.children
+ val requiredOrderings = operator.requiredChildOrdering
+
+ val newChildren = new ArrayBuffer[SparkPlan](children.length)
+ def rebuildChildren(joint: Boolean): Boolean = {
+ newChildren.clear()
+ children.indices.forall { index =>
+ val (pushableChild, required) = pushable(index)
+ val rebuilt = pushableChild.flatMap { child =>
+ val fixed = children.indices.flatMap { position =>
+ val (sibling, distribution) = pushable(position)
+ Option.when(position != index && (!joint || position < index ||
sibling.isEmpty)) {
+ val fixedChild = if (position < index) {
+ newChildren(position)
+ } else {
+ children(position)
+ }
+ (fixedChild, distribution)
+ }
+ }
+ rebuildChild(child, required, fixed, sharedExchanges, policy)
+ }
+ if (joint && pushableChild.isDefined && rebuilt.isEmpty) {
+ false
+ } else {
+ newChildren += rebuilt.getOrElse(children(index))
+ true
+ }
+ }
+ }
+
+ // Joint selection lets later pushable siblings follow; the fallback fixes
every sibling.
+ if (!rebuildChildren(joint = true)) {
+ rebuildChildren(joint = false)
+ }
+
+ // The ordering is put back once the layout is fixed. A local sort leaves
the layout alone,
+ // and a child that already answers keeps its own.
+ val sorted = newChildren.toSeq.lazyZip(requiredOrderings).map {
+ case (child, ordering) => withOrdering(child, ordering)
+ }
+ Some(operator.withNewChildren(sorted))
+ }
+
+ /**
+ * Rebuild on the best eligible union layout satisfying the consumer and
pairing with fixed
+ * siblings. Subset keys require both the rule's permission and the
co-partition policy, even
+ * when a single-input consumer accepts a coarser distribution.
+ */
+ private def rebuildChild(
+ pushable: PushableChild,
+ required: ClusteredDistribution,
+ fixed: Seq[(SparkPlan, ClusteredDistribution)],
+ sharedExchanges: Set[SparkPlan],
+ policy: LayoutPolicy): Option[SparkPlan] = {
+ val layouts = childLayouts(pushable.union, policy)
Review Comment:
**Finding 13.** On every `rebuildChild` call, `childLayouts` walks each
branch with `keepableLayout` and re-derives each child's partitioning, and
`candidateTargets` reads each child's stats. `pushDownThroughUnion` calls it
once per pushable child in the joint pass and again in the fallback. Both
results depend only on the `PushableChild` and the policy, so they can be
computed once per site and passed 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]
