peter-toth commented on code in PR #59630: URL: https://github.com/apache/spark/pull/59630#discussion_r4495483202
########## 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 Review Comment: **Finding 8.** Reproduced on `118d30a` with your confs. Push-down off gives 1,000 first rows, and on gives 4,582, in both your variants. The mis-mapping is older than this PR. Give both branches a repartition of their own, so that no rule has to add one: ```scala val x = spark.range(1000).selectExpr("id x", "id c", "id u").repartition($"c") val r = p1.union(p2).repartition(p1.col("p")) ``` Then the window gets 4,582 first rows with push-down off at this head, and also on the merge base `352ac9bf228`, with AQE on and off. With `spark.sql.unionOutputPartitioning=false` it gets 1,000. `prepareOutputPartitioning` maps `r`'s `hash(p#)` through `AttributeMap(childAttrs.zip(unionOutput))`, which keeps the last `p#` position (`c`), while the exchange hashes the first one. So the union reports `hash(c)`, and the window keeps no exchange of its own. Checking every child, as you suggest, closes this rule's case. The base case needs the fix in `prepareOutputPartitioning`: map each attribute to its first position, as binding does, or report no partitioning for a child whose output repeats an exprId. That fix would close both. I found no ticket for the base case. ulysses-you, if you agree it's a bug, could you open one? -- 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]
