cloud-fan commented on code in PR #59630: URL: https://github.com/apache/spark/pull/59630#discussion_r4420462021
########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PushDownShuffleExchangeThroughUnion.scala: ########## @@ -0,0 +1,538 @@ +/* + * 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.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, so that + * only the union's children that are not already laid out on the target take one. A local sort + * above the shuffle, and the 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.optimizer.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.optimizer.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.optimizer.pushDownShuffleThroughUnion.keepNonCoalescibleBranches.enabled` + * can prevent coalescing of the consumer's compatible shuffle group. + * + * 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. + * + * 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.getConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED) || + (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 + } 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)) + // 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) + // Do not reduce parallelism below both the replaced exchange and the configured default. + val minPartitions = math.min( + pushable.shuffle.outputPartitioning.numPartitions, conf.numShufflePartitions) + + // `findPushableShuffle` already checked this chain, including local sorts. + def rebuild(plan: SparkPlan, target: Partitioning with Expression): Option[SparkPlan] = + plan match { + case union: UnionExec => rebuildUnion(union, layouts, target, sharedExchanges) + case operator: UnaryExecNode => + rebuild(operator.child, target).map { child => + operator.withNewChildren(Seq(child)) + } + case other => + throw SparkException.internalError( + s"a shuffle push-down reached ${other.nodeName}, not the union it was found above") + } + + candidateTargets(layouts, pushable.union.children, minPartitions).view + .filter(target => layouts.count(child => !child.contains(target)) <= policy.maxMovedBranches) + .flatMap(target => rebuild(pushable.shuffle.child, target)) + .map(child => PartitioningOrigin.inherit(pushable.shuffle, child)) + .find { child => + val partitioning = child.outputPartitioning + partitioning.satisfies(required) && + PartitioningCollection.specsForPairing(partitioning, required).exists { spec => + spec.canCreatePartitioning && (policy.allowSubsetTargets || (spec match { + case hash: HashShuffleSpec => + required.areAllClusterKeysMatched(hash.partitioning.expressions) + case _ => false + })) + } && + ValidateRequirements.satisfiesForPairing( + child +: fixed.map(_._1), required +: fixed.map(_._2)) + } + } + + private case class CandidateBenefit( + offeringChildren: Int, + keptBytes: BigInt, + firstOccurrence: Int) + + /** + * Rank by estimated bytes kept when every child has a usable estimate. Missing estimates + * trigger count-based ranking rather than treating unknown branches as smaller. Partition + * count and first occurrence break ties deterministically. + */ + private def candidateTargets( + layouts: Seq[Seq[Partitioning with Expression]], + children: Seq[SparkPlan], + minPartitions: Int): Seq[Partitioning with Expression] = { + val sizes = children.map { child => + child.logicalLink.map(_.stats.sizeInBytes).filter(_ < conf.defaultSizeInBytes) + } + val useSizes = sizes.forall(_.isDefined) + val benefits = mutable.HashMap.empty[Partitioning with Expression, CandidateBenefit] + layouts.zip(sizes).foreach { case (reported, size) => + reported.filter(_.numPartitions >= minPartitions).distinct.foreach { target => + val previous = benefits.getOrElse(target, CandidateBenefit(0, BigInt(0), benefits.size)) + benefits.update(target, previous.copy( + offeringChildren = previous.offeringChildren + 1, + keptBytes = previous.keptBytes + size.getOrElse(BigInt(0)))) + } + } + benefits.toSeq.sortBy { case (target, benefit) => + val keptBytes = if (useSizes) { benefit.keptBytes } else { BigInt(0) } + (-keptBytes, -benefit.offeringChildren, -target.numPartitions, benefit.firstOccurrence) + } + .map(_._1) + } + + /** + * The layouts each child of `union` reports, in child order, each expressed in the union's + * output attributes. An eligible child reporting a [[PartitioningCollection]] offers its hash + * members. A child that offers nothing takes a shuffle whichever layout is picked. + * + * They are compared by equality below, which is what `UnionExec` itself compares its children + * with: `prepareOutputPartitioning` puts every child in the union's attributes. Where equality + * errs it errs towards shuffling a child it did not recognise. + */ + private def childLayouts( + union: UnionExec, + policy: LayoutPolicy): Seq[Seq[Partitioning with Expression]] = { + union.prepareOutputPartitioning().zip(union.children).map { case (partitioning, child) => + if (keepableLayout(child, policy)) { + PartitioningCollection.flatten(partitioning).collect { + case hash: HashPartitioning => hash + } + } else { + Nil + } + } + } + + /** + * Whether every input path reaches an exchange boundary. Stop at unstaged exchanges as well + * as query stages: their input scans do not prevent coalescing. A non-exchange leaf already + * prevents the consumer's compatible group from coalescing, independently of target selection. + */ + private def allExchangeInputs(plan: SparkPlan): Boolean = plan match { + case _: Exchange => true + case _: ExchangeQueryStageExec => true + case leaf if leaf.children.isEmpty => false + case other => other.children.forall(allExchangeInputs) + } + + /** + * Whether a child may offer a target. Shuffle boundaries isolate their input layouts; join + * policies apply above those boundaries, and broadcast joins only depend on streamed sources. + */ + private def keepableLayout(plan: SparkPlan, policy: LayoutPolicy): Boolean = plan match { + case shuffle: ShuffleExchangeLike => + shuffle.shuffleOrigin match { + case ENSURE_REQUIREMENTS | REPARTITION_BY_COL => true + case REPARTITION_BY_NUM => policy.allowNonCoalescible + case _ => false + } + case stage: ShuffleQueryStageExec => keepableLayout(stage.shuffle, policy) + case _: ShuffledJoin if !policy.keepJoinBranches => false + case _: BroadcastNestedLoopJoinExec | _: CartesianProductExec + if !policy.allowNonCoalescible => false + case join: BroadcastHashJoinExec => Review Comment: Thanks, verified in e3e43d9656c. A broadcast probe backed by a requirement shuffle, stage or adaptive read now needs keepJoinBranches, and the runtime test checks the converted join's streamed-side local read and the retained union exchange. This review raises two documentation follow-ups on the same guard: the keepJoinBranches conf doc does not mention broadcast probes, and the keepableLayout doc still describes requirement shuffles as rejected. <!-- SPARK_DEV_REVIEW_REPLY {"feedback_id":"inline:4361513953","thread_id":"inline:4361513953","verdict_sha256":"3e8febb23918757ca150c4722b1c9dc886490b9509809c986e163ebe64e247cd"} --> ########## sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/PushDownShuffleExchangeThroughUnionSuite.scala: ########## @@ -0,0 +1,2311 @@ +/* + * 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 org.apache.spark.rdd.RDD +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{ + Alias, Ascending, Attribute, AttributeReference, ExprId, GreaterThan, IsNotNull, Literal, Rand, + RowNumber, SortOrder} +import org.apache.spark.sql.catalyst.optimizer.BuildLeft +import org.apache.spark.sql.catalyst.plans.physical.{ + ClusteredDistribution, CoalescedBoundary, CoalescedHashPartitioning, Distribution, + HashPartitioning, NullAwareHashPartitioning, Partitioning, RoundRobinPartitioning, + UnknownPartitioning} +import org.apache.spark.sql.execution.{ + FilterExec, InputAdapter, LocalLimitExec, ProjectExec, QueryExecution, ScalarSubquery, SortExec, + SparkPlan, StampUnionDecisions, SubqueryExec, UnaryExecNode, UnionConfSnapshot, UnionExec, + WholeStageCodegenExec} +import org.apache.spark.sql.execution.exchange.{ + ENSURE_REQUIREMENTS, EnsureRequirements, REPARTITION_BY_COL, REPARTITION_BY_NUM, + ShuffleExchangeExec, ValidateRequirements} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.window.{Final => FinalWindowGroupLimit, Partial => PartialWindowGroupLimit, WindowGroupLimitExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.IntegerType + +class PushDownShuffleExchangeThroughUnionSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + // The suite exercises the rule itself, so it is on for every test whatever the default is. + override protected def sparkConf = + super.sparkConf + .set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key, "true") + .set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_JOIN_INPUTS_ENABLED.key, "true") + .set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_JOIN_BRANCHES_ENABLED.key, "true") + .set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_WITH_SUBQUERIES_ENABLED.key, "true") + .set(SQLConf + .PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_PARTITION_KEYS_SUBSET_OF_CLUSTER_KEYS_ENABLED.key, + "true") + .set(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key, "false") + + private val query: String = + """ + |SELECT id, v, row_number() OVER (PARTITION BY id ORDER BY v) AS rn + |FROM ( + | SELECT id, sum(v) AS v FROM t1 GROUP BY id + | UNION ALL + | SELECT id, v FROM t2 + |) + |""".stripMargin + + // A shuffle that moves a union as a whole, not one the rewrite puts under one of its children to + // lay that child out on the target. + private def numShufflesOnUnion(plan: SparkPlan): Int = + plan.collect { case s: ShuffleExchangeExec if isUnionPlan(s.child) => s }.size + + // The union may be wrapped in a `WholeStageCodegenExec` by `CollapseCodegenStages`, or in an + // `InputAdapter` when it feeds a codegen'd parent, so look through both when checking whether + // a shuffle's child is (effectively) a union. + private def isUnionPlan(plan: SparkPlan): Boolean = plan match { + case _: UnionExec => true + case w: WholeStageCodegenExec => isUnionPlan(w.child) + case i: InputAdapter => isUnionPlan(i.child) + case _ => false + } + + // A shuffle with a union anywhere below it, through the operators this rule can cross and the + // codegen wrappers `CollapseCodegenStages` may have added. + private def numShufflesOverUnion(plan: SparkPlan): Int = + countShuffles(plan, underUnion = false) { s => reachesUnion(s.child) } + + private def countShuffles( + plan: SparkPlan, + underUnion: Boolean)(matches: ShuffleExchangeExec => Boolean): Int = { + val here = plan match { + case s: ShuffleExchangeExec if !underUnion && matches(s) => 1 + case _ => 0 + } + here + plan.children.map(countShuffles(_, underUnion || plan.isInstanceOf[UnionExec])(matches)) + .sum + } + + private def reachesUnion(plan: SparkPlan): Boolean = plan match { + case _: UnionExec => true + case unary: UnaryExecNode => reachesUnion(unary.child) + case _ => false + } + + private def withLookThrough(f: => Unit): Unit = withSQLConf( + SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_LOOK_THROUGH_OPERATORS_ENABLED.key -> "true")(f) + + private def aqeInitialPlan(df: DataFrame): SparkPlan = df.queryExecution.executedPlan match { + case a: AdaptiveSparkPlanExec => a.initialPlan + case other => fail(s"expected an adaptive plan:\n$other") + } + + private def pushDown(plan: SparkPlan, required: Option[Distribution] = None): SparkPlan = { + val snapshot = UnionConfSnapshot(spark.sessionState.conf) + PushDownShuffleExchangeThroughUnion( + snapshot, EnsureRequirements(requiredDistribution = required), + queryContainsSubqueries = false)(plan) + } + + private def withPushDown(f: => Unit): Unit = withSQLConf( + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", + SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key -> "true", + SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_NON_COALESCIBLE_BRANCHES_ENABLED.key -> + "true")(f) + + private def attr(name: String): AttributeReference = AttributeReference(name, IntegerType)() + + private def leaf( + output: Attribute, + partitioning: Partitioning = UnknownPartitioning(0)): TestPartitioningPlan = + TestPartitioningPlan(outputPartitioning = partitioning, output = Seq(output)) + + private def hash(a: Attribute, n: Int): HashPartitioning = HashPartitioning(Seq(a), n) + + private def clustered(a: Attribute): ClusteredDistribution = ClusteredDistribution(Seq(a)) + + private def ensureShuffle(partitioning: Partitioning, child: SparkPlan): ShuffleExchangeExec = + ShuffleExchangeExec(partitioning, child, ENSURE_REQUIREMENTS) + + private def assertReshuffled(plan: SparkPlan, expected: Partitioning, child: SparkPlan): Unit = { + plan match { + case s: ShuffleExchangeExec => + assert(s.shuffleOrigin === ENSURE_REQUIREMENTS) + assert(s.outputPartitioning === expected, s"unexpected target:\n$plan") + assert(s.child.fastEquals(child), s"the shuffled child should be unchanged:\n$plan") + case other => fail(s"expected a shuffle into $expected:\n$other") + } + } + + private def createTables(): Unit = { + Seq((1, 10), (1, 20), (2, 5), (3, 7)).toDF("id", "v").createOrReplaceTempView("t1") + Seq((1, 40), (2, 8)).toDF("id", "v").createOrReplaceTempView("t2") + } + + private def expectedResult(sql: String): Seq[Row] = + withSQLConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key -> "false") { + spark.sql(sql).collect().toSeq + } + + Seq("MERGE", "SHUFFLE_HASH").foreach { hint => + Seq("INNER", "FULL OUTER").foreach { joinType => + test(s"SPARK-60379: push down through union for $hint $joinType join") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.SHUFFLE_SPREAD_NULL_JOIN_KEYS_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTempView("t1", "t2") { + Seq((Option(1), 10), (Option(1), 20), (None, 30)).toDF("id", "v") + .createOrReplaceTempView("t1") + Seq((Option(2), 40), (None, 50)).toDF("id", "v").createOrReplaceTempView("t2") + val query = s"""SELECT /*+ $hint(u, r) */ u.id, u.v, r.v FROM ( + |SELECT /*+ REPARTITION(4, id) */ id, v FROM t1 + |UNION ALL SELECT id, v FROM t2) u + |$joinType JOIN t1 r ON u.id = r.id""".stripMargin + val expected = expectedResult(query) + Seq(false, true).foreach { enabled => + withSQLConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key -> enabled.toString) { + val df = spark.sql(query) + checkAnswer(df, expected) + val initial = aqeInitialPlan(df) + assert(numShufflesOnUnion(initial) === (if (enabled) 0 else 1), + s"unexpected shuffle over the union:\n$initial") + if (hint == "MERGE") { + assert(initial.exists(_.isInstanceOf[SortMergeJoinExec]), + s"expected a sort-merge join:\n$initial") + } else { + assert(initial.exists(_.isInstanceOf[ShuffledHashJoinExec]), + s"expected a shuffled hash join:\n$initial") + } + assert(ValidateRequirements.validate(initial)) + } + } + } + } + } + } + } + + Seq(2, 8, 64).foreach { partitions => + test(s"SPARK-60379: the rebuilt union carries its own count with a " + + s"$partitions-partition child") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "8") { + withTempView("t1", "t2") { + createTables() + val query = s"""SELECT id, v, + |row_number() OVER (PARTITION BY id ORDER BY v) AS rn FROM ( + |SELECT /*+ REPARTITION($partitions, id) */ id, v FROM t1 + |UNION ALL SELECT id, v FROM t2)""".stripMargin + val expected = expectedResult(query) + Seq(false, true).foreach { enabled => + withSQLConf(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key -> enabled.toString) { + val df = spark.sql(query) + checkAnswer(df, expected) + val initial = aqeInitialPlan(df) + // A target below what the replaced exchange carried is not picked, since the + // operator above would then run with fewer partitions than before. + val pushedDown = enabled && partitions >= 8 + assert(numShufflesOnUnion(initial) === (if (pushedDown) 0 else 1), + s"unexpected shuffle over the union:\n$initial") + // The window runs with the child's count once the layout is pushed down, and with + // `spark.sql.shuffle.partitions` while the exchange above the union is kept. + val expectedPartitions = if (pushedDown) partitions else 8 + assert(initial.outputPartitioning.numPartitions === expectedPartitions, + s"unexpected count:\n$initial") + assert(df.rdd.getNumPartitions === expectedPartitions) + } + } + } + } + } + } + + test("SPARK-60379: do not push down a null-aware shuffle") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.SHUFFLE_SPREAD_NULL_JOIN_KEYS_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTempView("t1", "t2") { + Seq((Option(1), 10), (Option(1), 20), (None, 30)).toDF("id", "v") + .createOrReplaceTempView("t1") + Seq((Option(2), 40), (None, 50)).toDF("id", "v").createOrReplaceTempView("t2") + // A full outer join on nullable keys lets null keys spread, so its exchange is a + // `NullAwareHashPartitioning`, which a child laid out on a plain layout cannot report. + val query = + """SELECT /*+ MERGE(u, r) */ u.id, u.v, r.v FROM ( + |SELECT /*+ REPARTITION(4, id) */ id, v FROM t1 + |UNION ALL SELECT id, v FROM t2) u + |FULL OUTER JOIN t1 r ON u.id = r.id""".stripMargin + val df = spark.sql(query) + checkAnswer(df, expectedResult(query)) + + val initial = aqeInitialPlan(df) + assert(initial.collect { + case s: ShuffleExchangeExec + if s.outputPartitioning.isInstanceOf[NullAwareHashPartitioning] => s + }.nonEmpty, s"the join should spread null keys:\n$initial") + assert(numShufflesOnUnion(initial) === 1, s"the union shuffle must stay:\n$initial") + } + } + } + + test("SPARK-60379: select compatible targets for two union sides") { + withPushDown { + val a = attr("a") + val b = attr("b") + val a1 = attr("a1") + val b1 = attr("b1") + val left = UnionExec(Seq(leaf(a, hash(a, 4)), leaf(a1))) + val right = UnionExec(Seq(leaf(b, hash(b, 4)), leaf(b1))) + val parent = TestPartitioningPlan( + children = Seq(ensureShuffle(hash(a, 4), left), ensureShuffle(hash(b, 4), right)), + required = Seq(clustered(a), clustered(b))) + val result = pushDown(parent) + assert(result.children.forall(_.isInstanceOf[UnionExec])) + assert(ValidateRequirements.validate(result)) + } + } + + test("SPARK-60379: shuffle push-down is not registered in classic preparations") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assert(!QueryExecution.preparations(spark, subquery = false) + .exists(_.isInstanceOf[PushDownShuffleExchangeThroughUnion])) + withTempView("t1", "t2") { + createTables() + val df = spark.sql(query) + checkAnswer(df, expectedResult(query)) + assert(numShufflesOnUnion(df.queryExecution.executedPlan) === 1) + } + } + } + + test("SPARK-60379: push down shuffle through union") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + createTables() + val df = spark.sql(query) + + // The rule only rewrites the physical plan, so the result must be unchanged. + checkAnswer(df, expectedResult(query)) + + val plan = aqeInitialPlan(df) + assert(numShufflesOnUnion(plan) === 0, + s"shuffle should be pushed down below the union:\n$plan") + } + } + } + + test("SPARK-60379: do not push down when no child is co-located with the target") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + Seq((1, 10), (2, 5)).toDF("id", "v").createOrReplaceTempView("t1") + Seq((1, 40), (2, 8)).toDF("id", "v").createOrReplaceTempView("t2") + + val df = spark.sql( + """ + |SELECT id, v, row_number() OVER (PARTITION BY id ORDER BY v) AS rn + |FROM ( + | SELECT id, v FROM t1 + | UNION ALL + | SELECT id, v FROM t2 + |) + |""".stripMargin) + + // Pushing the shuffle down would turn one shuffle into two with no child saved. + val plan = aqeInitialPlan(df) + assert(numShufflesOnUnion(plan) === 1, + s"shuffle on top of union should be kept when no child is co-located:\n$plan") + } + } + } + + test("SPARK-60379: push down shuffle through union with AQE enabled") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + createTables() + val df = spark.sql(query) + checkAnswer(df, expectedResult(query)) + + val initial = aqeInitialPlan(df) + assert(numShufflesOnUnion(initial) === 0, + s"shuffle should be pushed below the union in the initial plan:\n$initial") + val union = initial.collectFirst { case u: UnionExec => u }.get + assert(union.children.count(_.isInstanceOf[ShuffleExchangeExec]) === 1, + s"only the scan child should be re-shuffled:\n$initial") + } + } + } + + Seq(false, true).foreach { adaptive => + test(s"SPARK-60379: retain the user-requested root partition count with AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "5") { + withTempView("t1", "t2") { + createTables() + val df = spark.table("t1").repartition(8, $"id") + .union(spark.table("t2")) + .withColumn("rn", org.apache.spark.sql.functions.row_number().over( + org.apache.spark.sql.expressions.Window.partitionBy("id").orderBy("v"))) + .repartition(5, $"id").select("id", "v", "rn") + checkAnswer(df, Seq( + Row(1, 10, 1), Row(1, 20, 2), Row(1, 40, 3), Row(2, 5, 1), Row(2, 8, 2), + Row(3, 7, 1))) + + val plan = if (adaptive) { + aqeInitialPlan(df) + } else { + df.queryExecution.executedPlan + } + if (adaptive) { + // The 8-partition target passes the floor of 5, but not the elided root repartition's + // explicit count of 5. Only the root check rejects this rewrite. + assert(numShufflesOnUnion(plan) === 1, + s"the root repartition's layout must be kept:\n$plan") + } + assert(df.rdd.getNumPartitions === 5, s"unexpected root:\n$plan") + } + } + } + } + + test("SPARK-60379: keep a union child that is a join on the picked layout") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "5") { + withTempView("t1", "t2", "t3", "t4") { + Seq((1, 10), (2, 20)).toDF("id", "v").createOrReplaceTempView("t1") + Seq((1, 100), (3, 300)).toDF("id", "v").createOrReplaceTempView("t2") + Seq((1, 7), (3, 8)).toDF("id", "v").createOrReplaceTempView("t3") + Seq((1, 1), (3, 3)).toDF("id", "v2").createOrReplaceTempView("t4") + // The union's first child is an inner join whose sides are shuffled on the join key, so + // it reports their layouts as a `PartitioningCollection`. That branch answers the target + // and keeps its exchanges; only the scan branch is shuffled onto it. + val sql = + """SELECT u.id, u.v, r.v2 FROM ( + | SELECT l.id AS id, l.v AS v FROM t1 l JOIN t2 r ON l.id = r.id + | UNION ALL + | SELECT id, v FROM t3 + |) u JOIN t4 r ON u.id = r.id""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, Seq(Row(1, 10, 1), Row(1, 7, 1), Row(3, 8, 3))) + + val plan = aqeInitialPlan(df) + assert(numShufflesOnUnion(plan) === 0, + s"the union exchange should be pushed down:\n$plan") + val union = plan.collectFirst { case u: UnionExec => u }.get + assert(union.children.head.collect { case _: SortMergeJoinExec => true }.nonEmpty, + s"the join branch should keep its layout:\n$plan") + assert(union.children(1).isInstanceOf[ShuffleExchangeExec], + s"only the scan branch should be shuffled:\n$plan") + } + } + } + + test("SPARK-60379: push down through union below a sort-merge join") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTempView("t1", "t2") { + createTables() + // The local sort stays above the rebuilt union; the other side's layout constrains it. + val sql = + """ + |SELECT t1.id, u.v FROM t1 JOIN ( + | SELECT id, v FROM (SELECT id, v FROM t1 DISTRIBUTE BY id) + | UNION ALL + | SELECT id, v FROM t2 + |) u ON t1.id = u.id + |""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, expectedResult(sql)) + val plan = aqeInitialPlan(df) + assert(plan.collect { case j: SortMergeJoinExec => j }.nonEmpty, s"expected SMJ:\n$plan") + assert(numShufflesOnUnion(plan) === 0, + s"the join union shuffle should be pushed down:\n$plan") + } + } + } + + test("SPARK-60379: push down shuffle with multiple partition key columns") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + Seq((1, "a", 10), (1, "a", 20), (2, "b", 5)).toDF("id", "grp", "v") + .createOrReplaceTempView("t1") + Seq((1, "a", 40), (2, "b", 8)).toDF("id", "grp", "v").createOrReplaceTempView("t2") + + val sql = + """ + |SELECT id, grp, v, row_number() OVER (PARTITION BY id, grp ORDER BY v) AS rn + |FROM ( + | SELECT id, grp, sum(v) AS v FROM t1 GROUP BY id, grp + | UNION ALL + | SELECT id, grp, v FROM t2 + |) + |""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, expectedResult(sql)) + assert(numShufflesOnUnion(aqeInitialPlan(df)) === 0, + s"shuffle should be pushed below the union:\n${df.queryExecution.executedPlan}") + } + } + } + + test("SPARK-60379: push down through union with a nullability-adjusted co-located child") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + // `id` is nullable so that `IsNotNull` in the Filter actually tightens its nullability. + Seq((Option(1), 10), (Option(1), 20), (Option(2), 5)).toDF("id", "v") + .createOrReplaceTempView("t1") + Seq((Option(3), 7)).toDF("id", "v").createOrReplaceTempView("t2") + + val sql = + """ + |SELECT id, v, row_number() OVER (PARTITION BY id ORDER BY v) AS rn + |FROM ( + | SELECT id, v FROM (SELECT id, v FROM t1 DISTRIBUTE BY id) WHERE id IS NOT NULL + | UNION ALL + | SELECT id, v FROM t2 + |) + |""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, expectedResult(sql)) + + // The DISTRIBUTE BY child is hash-partitioned on id and co-located with the union's + // output despite the Filter tightening `id`'s nullability, so the shuffle for the window + // should be pushed below the union and only the scan child re-shuffled. + val plan = aqeInitialPlan(df) + assert(numShufflesOnUnion(plan) === 0, s"shuffle should be pushed below the union:\n$plan") + } + } + } + + test("SPARK-60379: push down shuffle through union with whole-stage codegen disabled") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") { + withTempView("t1", "t2") { + createTables() + val df = spark.sql(query) + checkAnswer(df, expectedResult(query)) + val finalPlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + .executedPlan + assert(!finalPlan.exists(_.isInstanceOf[WholeStageCodegenExec])) + assert(numShufflesOnUnion(aqeInitialPlan(df)) === 0, + s"shuffle should be pushed below the union:\n${df.queryExecution.executedPlan}") + } + } + } + + test("SPARK-60379: relax the target to a child's coarser partitioning") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + Seq((1, "a", 10), (1, "a", 20), (2, "b", 5)).toDF("id", "grp", "v") + .createOrReplaceTempView("t1") + Seq((1, "a", 40), (2, "b", 8)).toDF("id", "grp", "v").createOrReplaceTempView("t2") + + // The window requires ClusteredDistribution([id, grp]) with requireAllClusterKeys = false, + // so the DISTRIBUTE BY id child (HashPartitioning(id)) already satisfies it. The shuffle + // should be pushed down and its target relaxed to HashPartitioning(id) rather than the + // wider HashPartitioning(id, grp). + val sql = + """ + |SELECT id, grp, v, row_number() OVER (PARTITION BY id, grp ORDER BY v) AS rn + |FROM ( + | SELECT id, grp, v FROM (SELECT id, grp, v FROM t1 DISTRIBUTE BY id) + | UNION ALL + | SELECT id, grp, v FROM t2 + |) + |""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, expectedResult(sql)) + + val union = aqeInitialPlan(df).collect { case u: UnionExec => u }.head + assert(union.outputPartitioning.isInstanceOf[HashPartitioning], + s"the union should propagate a hash partitioning after the push-down:\n" + + df.queryExecution.executedPlan) + assert(union.outputPartitioning.asInstanceOf[HashPartitioning].expressions.length === 1, + s"the target should be relaxed to the single-column partitioning:\n" + + df.queryExecution.executedPlan) + } + } + } + + test("SPARK-60379: push down shuffle when children produce different partitionings") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("t1", "t2") { + Seq((1, "a", 10), (1, "a", 20), (2, "b", 5), (3, "c", 6)).toDF("id", "grp", "v") + .createOrReplaceTempView("t1") + Seq((1, "a", 40), (2, "b", 8)).toDF("id", "grp", "v").createOrReplaceTempView("t2") + + // The two children are hash-partitioned on different key sets (`id` vs `id, grp`), both + // of which satisfy the window's `ClusteredDistribution([id, grp])`. The rule should push + // the shuffle down rather than giving up on an empty intersection. + val sql = + """ + |SELECT id, grp, v, row_number() OVER (PARTITION BY id, grp ORDER BY v) AS rn + |FROM ( + | SELECT id, grp, v FROM (SELECT id, grp, v FROM t1 DISTRIBUTE BY id) + | UNION ALL + | SELECT id, grp, v FROM (SELECT id, grp, v FROM t2 DISTRIBUTE BY id, grp) + |) + |""".stripMargin + val df = spark.sql(sql) + checkAnswer(df, expectedResult(sql)) + assert(numShufflesOnUnion(aqeInitialPlan(df)) === 0, + s"shuffle should be pushed below the union:\n${df.queryExecution.executedPlan}") + + // Both partitionings tie on the number of children; the first one in child order wins. Review Comment: Thanks, verified in e3e43d9656c. The comment now names the byte ranking, and the new equal-byte test checks first-offered order in both child orders. <!-- SPARK_DEV_REVIEW_REPLY {"feedback_id":"inline:4361513972","thread_id":"inline:4361513972","verdict_sha256":"3e8febb23918757ca150c4722b1c9dc886490b9509809c986e163ebe64e247cd"} --> -- 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]
