diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index a45ea80c2ee5..6c9c9fce24f0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -920,26 +920,153 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup Seq(firstPartitioning) ++ convertedOtherPartitionings } - // Compares two leaf partitionings for union pass-through equivalence. Callers pass leaf - // partitionings only; a `PartitioningCollection` is flattened to its members by - // `outputPartitioning` before reaching here. + // Compares two leaf partitionings for the co-located pass-through in `outputPartitioning`. + // Callers pass leaf partitionings only; a `PartitioningCollection` is flattened to its members + // there. `KeyedPartitioning` never reaches here: it is merged per key position by + // `mergeKeyedPartitionings` and filtered out of the pass-through candidates. private def comparePartitioning(left: Partitioning, right: Partitioning): Boolean = { (left, right) match { case (SinglePartition, SinglePartition) => true case (l: HashPartitioningLike, r: HashPartitioningLike) => l == r - // For `KeyedPartitioning`, only the partition expressions must match (the other child's - // expressions have already been remapped to the first child's attributes by - // `prepareOutputPartitioning`). The partition keys are intentionally not compared here: - // children typically carry different key sets, and `outputPartitioning` merges them. - case (l: KeyedPartitioning, r: KeyedPartitioning) => - l.expressions.length == r.expressions.length && - l.expressions.zip(r.expressions).forall { case (le, re) => le.semanticEquals(re) } // Note: two `RangePartitioning`s with even same ordering and number of partitions // are not equal, because they might have different partition bounds. case _ => false } } + /** + * Merges the children's `KeyedPartitioning`s into the alternatives this union can report, or an + * empty sequence if they cannot be merged. + * + * A `UnionExec` concatenates its children's partitions in order (one child's partitions after + * another's), so a merged `KeyedPartitioning` describes the concatenation of the children's + * partition keys, one key per physical output partition. The merge works per key position, the + * way [[PartitioningPreservingUnaryExecNode]] narrows a projected `KeyedPartitioning` rather than + * dropping it: a position survives when every child partitions by the same expression for it, the + * first child's candidates are narrowed to the surviving positions, and each child's keys are + * projected to its own positions for those expressions before being concatenated. Children that + * agree on fewer positions than they partition by therefore still merge, at a coarser + * granularity, and children that list the same expressions in a different order merge by + * permutation. + * + * @param keyedCandidates each child's `KeyedPartitioning`s, with expressions already remapped to + * the first child's attributes; every child must offer at least one + * @param toUnionOutput maps the first child's attributes to this union's own output attributes + */ + private def mergeKeyedPartitionings( + keyedCandidates: Seq[Seq[KeyedPartitioning]], + toUnionOutput: Partitioning => Partitioning): Seq[Partitioning] = { + val headCandidates = keyedCandidates.head + + // Where each partition expression sits within a child. All of a child's candidates are namings + // of the same positions, so an expression's position does not depend on which candidate named + // it; a child that repeats an expression keeps its first position. Keyed on the canonicalized + // expression so that a `TransformExpression` matches as a whole -- `bucket(4, id)` must not + // match `years(id)` or `bucket(8, id)` merely because they reference the same column. + val otherPositions = keyedCandidates.tail.map { candidates => + val positions = mutable.Map.empty[Expression, Int] + candidates.foreach(_.expressions.zipWithIndex.foreach { case (e, i) => + positions.getOrElseUpdate(e.canonicalized, i) + }) + positions + } + + // The first child's key positions that every other child also partitions by, each with the + // namings admitted at that position and the positions those namings occupy in the other + // children. A position must be dropped unless *every* other child has it: a child that does + // not partition by an expression has no key column to contribute for it. + val retained = headCandidates.head.expressions.indices.flatMap { headPos => + val located = ExpressionSet(headCandidates.map(_.expressions(headPos))).toSeq.flatMap { e => + val positions = otherPositions.map(_.get(e.canonicalized)) + Option.when(positions.forall(_.isDefined))(e.canonicalized -> positions.map(_.get)) + } + located.headOption.map { case (_, otherPos) => + // Namings that would project the other children at different positions imply different + // merged keys, and all the alternatives returned below must share one key list, so keep + // only those agreeing with the first survivor. Which naming that is follows the first + // child's candidate order and is otherwise arbitrary: when a position admits several + // namings landing on different positions of another child, a different pick could keep + // that child's positions distinct where this one collapses them. Collapsing stays correct + // and is reported as narrowing below, so this only ever costs granularity. + (headPos, located.collect { case (e, p) if p == otherPos => e }.toSet, otherPos) + } + } + if (retained.isEmpty) { + return Nil + } + val retainedHeadPositions = retained.map(_._1) + + // Filter and narrow the first child's candidates instead of rebuilding them from the admitted + // namings: its candidate list was already bounded by `EXPRESSION_PROJECTION_CANDIDATE_LIMIT` + // when `PartitioningPreservingUnaryExecNode` produced it, and re-deriving the combinations + // could resurrect ones that limit had dropped. So no separate bound is needed here -- the + // result cannot be larger than the list it filters. A candidate is dropped when it names an + // expression the other children lack at a surviving position; the narrowing then collapses + // candidates that differed only at a dropped position, hence the deduplication. + val seen = mutable.Set.empty[Seq[Expression]] + val narrowedExpressions = headCandidates.flatMap { kp => + val admitted = retained.forall { + case (headPos, namings, _) => namings.contains(kp.expressions(headPos).canonicalized) + } + if (admitted) { + val narrowed = retainedHeadPositions.map(kp.expressions) + Option.when(seen.add(narrowed.map(_.canonicalized)))(narrowed) + } else { + None + } + }.toList + if (narrowedExpressions.isEmpty) { + // Every candidate names something the other children lack at some surviving position. This + // can only happen when the surviving namings are split across candidates (one names the + // survivor at position 0, another at position 1), which needs the children to alias their + // partition columns in incompatible ways; recombining them is what the paragraph above + // deliberately avoids. + return Nil + } + + // Each child's own positions for the surviving expressions, in the first child's order. Two + // surviving positions can land on the *same* position of another child, because that child may + // name one key column with two different expressions across its candidates -- the two sides of + // an equi-join, or two aliases of one column. The merged keys stay truthful in that case (the + // namings are equal-valued), but that child contributes fewer distinct key columns than it + // partitions by, which is a narrowing and has to be reported as one below. + val childPositions = keyedCandidates.indices.map { childIndex => + if (childIndex == 0) retainedHeadPositions else retained.map(_._3(childIndex - 1)) + } + + // Each child's keys projected to its own positions, then concatenated in child order to match + // the layout `doExecute` produces. `projectKeys` takes an arbitrary position sequence, so a + // child that lists the expressions in a different order is simply permuted, and a child whose + // positions are already aligned skips the projection entirely. + val mergedKeys = keyedCandidates.zip(childPositions).flatMap { case (candidates, positions) => + // Any candidate serves as the key source: a child's candidates share one `partitionKeys` + // reference and name the same positions. + val keySource = candidates.head + if (positions == keySource.expressions.indices) { + keySource.partitionKeys + } else { + keySource.projectKeys(positions)._2 + } + } + val isGrouped = mergedKeys.distinct.size == mergedKeys.size + // Narrowed if any child contributes fewer distinct key columns than it partitions by -- its + // partitions that held distinct keys now share one -- or if any child's partitioning was + // already narrowed, which keeps the flag sticky as `PartitioningPreservingUnaryExecNode` does. + // Counting *distinct* positions matters: comparing against the number of surviving positions + // would miss a child two of them collapse onto. Read across every candidate of every child, and + // uniform across the alternatives returned: like the keys and `isGrouped`, it describes the + // merged layout rather than one naming of it, and reading it as `false` would silently drop + // `GroupPartitionsExec`'s skew protection. + val isNarrowed = keyedCandidates.zip(childPositions).exists { case (candidates, positions) => + candidates.exists(_.isNarrowed) || + candidates.head.expressions.length > positions.distinct.length + } + + narrowedExpressions.map { expressions => + toUnionOutput(KeyedPartitioning(expressions, mergedKeys, isGrouped, isNarrowed)) + } + } + override def outputPartitioning: Partitioning = { if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) { return super.outputPartitioning @@ -957,30 +1084,38 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup case _ => p } - // Case A: every child is a single `KeyedPartitioning`. A `UnionExec` concatenates its - // children's partitions in order (one child's partitions after another's), so the merged - // `KeyedPartitioning` carries the concatenation of the children's partition keys, one key - // per physical output partition. Children usually hold different key sets, so the merged - // keys often contain duplicates and `isGrouped` is false; a downstream `GroupPartitionsExec` - // regroups partitions that share a key. This concatenation (numPartitions = sum) is a - // distinct physical strategy from the co-located pass-through below (numPartitions = N), so - // it is kept as a separate case and never folded into a `PartitioningCollection`. - if (partitionings.forall(_.isInstanceOf[KeyedPartitioning])) { - val kps = partitionings.map(_.asInstanceOf[KeyedPartitioning]) - val headKp = kps.head - // The `KeyedPartitioning`s must agree on the partition expressions to merge. - val compatible = kps.forall(comparePartitioning(_, headKp)) - if (compatible) { - val mergedKeys = kps.flatMap(_.partitionKeys) - val mergedExpressions = headKp.expressions.map(_.transform { - case a: Attribute if attributeMap.contains(a) => attributeMap(a) - }) - val isGrouped = mergedKeys.distinct.size == mergedKeys.size - val isNarrowed = kps.exists(_.isNarrowed) - return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed) - } else { - return super.outputPartitioning + // Case A: every child offers at least one `KeyedPartitioning`, either directly or as a member + // of a `PartitioningCollection` (a `ProjectExec` that aliases a partition key to several output + // names reports one `KeyedPartitioning` per alias, and an inner join reports one per side). The + // merge concatenates the children's partition keys, so `numPartitions` is their sum; children + // usually hold different key sets, so the merged keys often contain duplicates and `isGrouped` + // is false, and a downstream `GroupPartitionsExec` regroups partitions that share a key. That + // concatenation is a distinct physical strategy from the co-located pass-through below + // (numPartitions = N), so Case A and Case B never contribute to the same + // `PartitioningCollection`. See `mergeKeyedPartitionings` for how the merge itself works. + // + // Case A is tried first, so a child offering a `KeyedPartitioning` never reaches Case B. That + // only shadows a co-locatable candidate if some child's collection holds both a + // `KeyedPartitioning` and a `HashPartitioningLike`/`SinglePartition`, which no producer builds + // today: `EnsureRequirements` cannot leave a join with one keyed and one hash side + // (`HashShuffleSpec.isCompatibleWith` rejects a `KeyedShuffleSpec`, and + // `KeyedShuffleSpec.createPartitioning` yields another `KeyedPartitioning`), so a join's + // collection is homogeneous and projection preserves that. If a mixed collection ever becomes + // reachable, revisit which case should win rather than leaving it to this ordering. + val keyedCandidates = partitionings.map { p => + PartitioningCollection.flatten(p).collect { case kp: KeyedPartitioning => kp } + } + if (keyedCandidates.forall(_.nonEmpty)) { + val merged = mergeKeyedPartitionings(keyedCandidates, toUnionOutput) + if (merged.length == 1) { + return merged.head + } else if (merged.length > 1) { + return PartitioningCollection.fromPartitionings(merged) } + // The children agree on no key position. Falling through to Case B rather than returning + // `super.outputPartitioning` outright is only for uniformity: Case B needs every child to + // offer a co-locatable candidate too, which per the paragraph above cannot happen while + // every child also offers a `KeyedPartitioning`, so it yields `super.outputPartitioning`. } // Case B: treat each child's partitioning as a set of candidate partitionings (a @@ -1222,22 +1357,32 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup override def usedInputs: AttributeSet = AttributeSet.empty protected override def doExecute(): RDD[InternalRow] = { - outputPartitioning match { - case _: UnknownPartitioning | _: KeyedPartitioning => - // An `UnknownPartitioning` union simply concatenates its children. A - // `KeyedPartitioning` union does the same: its merged partition keys describe the - // concatenated layout (one key per physical partition), and a downstream - // `GroupPartitionsExec` regroups partitions that share a key. This differs from an - // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware - // union RDD interleaves same-index partitions across children. - sparkContext.union(children.map(_.execute())) - case _ => - // This union has a known, index-co-locatable partitioning, i.e., its children have the - // same partitioning in semantics so this union can choose not to change the partitioning - // by using a custom partitioning aware union RDD. - val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty) - new SQLPartitioningAwareUnionRDD( - sparkContext, nonEmptyRdds, outputPartitioning.numPartitions) + val partitioning = outputPartitioning + // An `UnknownPartitioning` union simply concatenates its children. A `KeyedPartitioning` + // union does the same, whether it reports one `KeyedPartitioning` or a + // `PartitioningCollection` of alternatives: the merged partition keys describe the + // concatenated layout (one key per physical partition), and a downstream + // `GroupPartitionsExec` regroups partitions that share a key. This differs from an + // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware + // union RDD interleaves same-index partitions across children. + // + // The two never mix in one `PartitioningCollection`: Case A of `outputPartitioning` + // contributes only `KeyedPartitioning`s and Case B filters them out. The `require` is a + // backstop -- a mixed collection has no correct arm here, because its members would disagree + // on whether `numPartitions` is the sum over the children or the count of one of them, so it + // must fail loudly instead of silently picking one layout. + val members = PartitioningCollection.flatten(partitioning) + val keyed = members.count(_.isInstanceOf[KeyedPartitioning]) + require(keyed == 0 || keyed == members.length, + s"UnionExec cannot mix KeyedPartitioning with co-locatable partitionings: $partitioning") + if (partitioning.isInstanceOf[UnknownPartitioning] || keyed > 0) { + sparkContext.union(children.map(_.execute())) + } else { + // This union has a known, index-co-locatable partitioning, i.e., its children have the + // same partitioning in semantics so this union can choose not to change the partitioning + // by using a custom partitioning aware union RDD. + val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty) + new SQLPartitioningAwareUnionRDD(sparkContext, nonEmptyRdds, partitioning.numPartitions) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index faefbab6c8ca..1d5e91063fd8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -1659,6 +1659,60 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP } } + test("SPARK-58594: union partitioning - PartitioningCollection of keyed partitionings") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[InMemoryCatalog].getName) { + withTable("testcat.ns.k1", "testcat.ns.k2") { + sql("CREATE TABLE testcat.ns.k1 (id bigint, data string) PARTITIONED BY (id)") + sql("INSERT INTO testcat.ns.k1 VALUES (1, 'a1'), (2, 'a2')") + sql("CREATE TABLE testcat.ns.k2 (id bigint, data string) PARTITIONED BY (id)") + sql("INSERT INTO testcat.ns.k2 VALUES (2, 'b2'), (3, 'b3')") + + // Each child selects `id` twice, so its `ProjectExec` reports one `KeyedPartitioning` per + // alias inside a `PartitioningCollection` instead of a single `KeyedPartitioning`. + def unionDF: DataFrame = sql( + """SELECT id, id AS k, data FROM testcat.ns.k1 + |UNION ALL + |SELECT id, id AS k, data FROM testcat.ns.k2 + |""".stripMargin) + + val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + unionDF.collect() + } + + Seq(true, false).foreach { enabled => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) { + val union = unionDF + val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + assert( + unionExec.head.children.forall(_.outputPartitioning.isInstanceOf[ + PartitioningCollection]), + "each child should report a PartitioningCollection of KeyedPartitionings") + + val partitioning = unionExec.head.outputPartitioning + if (enabled) { + // Both children offer `identity(k)` and `identity(id)`, so the union merges the two + // collections member-wise and keeps both alternatives. + val collection = partitioning.asInstanceOf[PartitioningCollection] + assert(collection.partitionings.forall(_.isInstanceOf[KeyedPartitioning]), + s"expected a collection of KeyedPartitionings but got $partitioning") + assert(collection.partitionings.size == 2, + "one merged KeyedPartitioning per alias the children agree on") + // The merged keys are the concatenation of the children's keys: [1, 2] ++ [2, 3]. + assert(collection.numPartitions == 4) + } else { + assert(partitioning.isInstanceOf[UnknownPartitioning]) + } + + checkAnswer(union, correctResult) + } + } + } + } + } + test("SPARK-58317: union partitioning - PartitioningCollection child intersects to single") { withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index a732294786b2..a4982390d17a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -4505,4 +4505,499 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with assert(shuffles.isEmpty, "should not contain any shuffle") checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5))) } + + test("SPARK-58594: union partitioning - two surviving positions collapsing onto one position " + + "of another leg counts as narrowing") { + // The second leg is an SPJ that selects both sides of one join key (`t3.k1 AS a`, `t4.k1 AS b`) + // and only one side of the other (`t3.k2 AS c`), so it reports `PC(KP([a, c]), KP([b, c]))`: + // `a` and `b` are two namings of its *first* key position. The first leg partitions by (a, b), + // so both of its positions survive -- but both project the second leg at position 0, dropping + // that leg's `k2`. The merged keys stay truthful, since the join condition forces + // `t4.k1 = t3.k1`, yet the second leg contributes one distinct key column where it partitions + // by two, so the merge must report `isNarrowed`. Counting surviving positions instead of each + // leg's *distinct* projected positions reports `false` here, which lets `groupedSatisfies` + // accept an ungrouped, narrowed partitioning without `allowKeysSubsetOfPartitionKeys` and drops + // the skew protection that flag exists to demand. + // + // `max(c)` is what keeps `c` alive. Without it `ColumnPruning` drops `c`, the second leg's own + // projection can no longer express its second key position, and that leg arrives already + // `isNarrowed = true` -- so the assertion below would hold for an entirely different reason. + val cols = Array( + Column.create("k1", LongType), + Column.create("k2", LongType), + Column.create("k3", LongType)) + withTable("t1", "t3", "t4") { + createTable("t1", cols, Array(identity("k1"), identity("k2"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 10, 100), (2, 20, 200)") + createTable("t3", cols, Array(identity("k1"), identity("k2"))) + sql("INSERT INTO testcat.ns.t3 VALUES (5, 50, 500), (5, 60, 600)") + createTable("t4", cols, Array(identity("k1"), identity("k2"))) + sql("INSERT INTO testcat.ns.t4 VALUES (5, 50, 501), (5, 60, 601)") + + val unionSql = + """SELECT a, b, max(c) AS m FROM ( + | SELECT k1 AS a, k2 AS b, k3 AS c FROM testcat.ns.t1 + | UNION ALL + | SELECT /*+ MERGE(t3, t4) */ t3.k1 AS a, t4.k1 AS b, t3.k2 AS c + | FROM testcat.ns.t3 JOIN testcat.ns.t4 + | ON t3.k1 = t4.k1 AND t3.k2 = t4.k2 + |) u + |GROUP BY a, b + |""".stripMargin + + Seq(true, false).foreach { allowSubset => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> allowSubset.toString) { + val df = sql(unionSql) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(kp.numPartitions == 4, "both legs keep two partitions each") + assert(!kp.isGrouped, + "the second leg's two partitions both merge to the key (5, 5)") + assert(kp.isNarrowed, + "the second leg contributes one distinct key column but partitions by two") + + // `collectShuffles`/`collectGroupPartitions` only look inside a `SortMergeJoinExec` + // subtree, and the nodes of interest here sit above the union, so collect directly. + val shuffles = collect(plan) { case s: ShuffleExchangeExec => s } + val groupPartitions = collect(plan) { case g: GroupPartitionsExec => g } + if (allowSubset) { + assert(shuffles.isEmpty, "the config opts in to the skew risk") + assert(groupPartitions.nonEmpty, + "GroupPartitionsExec must coalesce the two partitions merging to (5, 5)") + } else { + assert(shuffles.nonEmpty, + "a narrowed, ungrouped partitioning must not satisfy the aggregate " + + "without the config") + } + checkAnswer(df, Seq(Row(1, 10, 100), Row(2, 20, 200), Row(5, 5, 60))) + } + } + } + } + + test("SPARK-58594: storage-partitioned join over a union of three legs, one of them permuted") { + // Three legs pin the per-leg position bookkeeping: the middle leg lists the key expressions in + // the opposite order, so each leg needs its own projection positions. Were the leg-to-positions + // indexing off by one, the legs' keys would be projected with each other's positions and the + // merged key list would mix (id, dept) rows with (dept, id) ones. + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + withTable("t1", "t2", "t3", "t4") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1')") + createTable("t2", cols, Array(identity("dept"), identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (2, 'y', 'b2')") + createTable("t3", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t3 VALUES (3, 'z', 'c3')") + createTable("t4", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t4 VALUES " + + "(1, 'x', 'd1'), (2, 'y', 'd2'), (3, 'z', 'd3')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t4) */ u.id, u.dept, u.data, t4.data AS t4data + |FROM ( + | SELECT id, dept, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, dept, data FROM testcat.ns.t2 + | UNION ALL + | SELECT id, dept, data FROM testcat.ns.t3 + |) u + |JOIN testcat.ns.t4 ON u.id = t4.id AND u.dept = t4.dept + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + assert(union.children.length == 3, + "the three legs must not be collapsed into nested unions") + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(kp.expressions.length == 2, "all three legs partition by both columns") + assert(kp.numPartitions == 3, "one partition per leg") + assert(!kp.isNarrowed, "a permutation drops no position") + assert(kp.isGrouped) + + assert(collectShuffles(plan).isEmpty, + "no shuffle: every leg's keys are projected into the first leg's order") + checkAnswer(df, Seq( + Row(1, "x", "a1", "d1"), Row(2, "y", "b2", "d2"), Row(3, "z", "c3", "d3"))) + } + } + } + + test("SPARK-58594: storage-partitioned join over a union of storage-partitioned joins") { + // An inner `ShuffledJoin` reports `PartitioningCollection(left, right)`, so an SPJ inside a + // union leg makes that leg report a collection of `KeyedPartitioning`s -- one per join side. + // Without merging collections the union reports `UnknownPartitioning` and both sides of the + // outer join shuffle. Selecting both join keys (`t1.id` and `t2.id AS k`) is what keeps the + // collection alive: projecting only one key would narrow each leg back to a single + // `KeyedPartitioning` that the existing merge already handles. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + val partitions = Array(identity("id")) + withTable("t1", "t2", "t3", "t4", "t5") { + createTable("t1", cols, partitions) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')") + createTable("t2", cols, partitions) + sql("INSERT INTO testcat.ns.t2 VALUES (1, 'b1'), (2, 'b2')") + createTable("t3", cols, partitions) + sql("INSERT INTO testcat.ns.t3 VALUES (3, 'c3'), (4, 'c4')") + createTable("t4", cols, partitions) + sql("INSERT INTO testcat.ns.t4 VALUES (3, 'd3'), (4, 'd4')") + createTable("t5", cols, partitions) + sql("INSERT INTO testcat.ns.t5 VALUES (1, 'e1'), (2, 'e2'), (3, 'e3'), (4, 'e4')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t5) */ u.id, u.k, t5.data + |FROM ( + | SELECT /*+ MERGE(t1, t2) */ t1.id, t2.id AS k + | FROM testcat.ns.t1 JOIN testcat.ns.t2 ON t1.id = t2.id + | UNION ALL + | SELECT /*+ MERGE(t3, t4) */ t3.id, t4.id AS k + | FROM testcat.ns.t3 JOIN testcat.ns.t4 ON t3.id = t4.id + |) u + |JOIN testcat.ns.t5 ON u.id = t5.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + assert(union.children.forall(_.outputPartitioning.isInstanceOf[ + physical.PartitioningCollection]), + "each union leg must report a PartitioningCollection for this test to be meaningful") + + val collection = union.outputPartitioning match { + case pc: physical.PartitioningCollection => pc + case other => fail(s"expected a PartitioningCollection of KeyedPartitionings, got $other") + } + val kps = collection.partitionings.map(_.asInstanceOf[physical.KeyedPartitioning]) + assert(kps.forall(_.numPartitions == 4), + "each merged descriptor carries the concatenation of both legs' keys") + assert(kps.forall(_.isGrouped), "disjoint leg keys merge without duplicates") + + assert(collectShuffles(plan).isEmpty, + "no shuffle: the merged descriptor for `id` matches t5") + assert(collectGroupPartitions(plan).isEmpty, + "no GroupPartitionsExec: merged keys are already grouped and aligned") + checkAnswer(df, Seq( + Row(1, 1, "e1"), Row(2, 2, "e2"), Row(3, 3, "e3"), Row(4, 4, "e4"))) + } + } + } + + test("SPARK-58594: storage-partitioned join over union: children report a " + + "PartitioningCollection of KeyedPartitionings") { + // Each leg selects `id` twice (`id` and `id AS k`), so its `ProjectExec` reports one + // `KeyedPartitioning` per alias inside a `PartitioningCollection` rather than a single + // `KeyedPartitioning`. The union merges the collections member-wise and keeps both + // alternatives, so the SMJ can pick the one it needs. The join is on `u.id`, which is the + // second alternative (aliases come first), so keeping only one member would not do: `u.k` has + // to stay in the outer projection, otherwise `ColumnPruning` drops `id AS k` and each leg + // collapses back to a single `KeyedPartitioning` that Case A already handled. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + val partitions = Array(identity("id")) + withTable("t1", "t2", "t3") { + createTable("t1", cols, partitions) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')") + createTable("t2", cols, partitions) + sql("INSERT INTO testcat.ns.t2 VALUES (3, 'b3'), (4, 'b4')") + createTable("t3", cols, partitions) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'), (4, 'c4')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.k, u.data, t3.data AS t3data + |FROM ( + | SELECT id, id AS k, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, id AS k, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + assert(union.children.forall(_.outputPartitioning.isInstanceOf[ + physical.PartitioningCollection]), + "each union leg must report a PartitioningCollection for this test to be meaningful") + + val collection = union.outputPartitioning match { + case pc: physical.PartitioningCollection => pc + case other => fail(s"expected a PartitioningCollection of KeyedPartitionings, got $other") + } + assert(collection.partitionings.size == 2, + "one merged KeyedPartitioning per alias both legs agree on") + val kps = collection.partitionings.map(_.asInstanceOf[physical.KeyedPartitioning]) + assert(kps.forall(_.numPartitions == 4), + "each merged descriptor carries the concatenation of both legs' keys") + assert(kps.forall(_.isGrouped), "disjoint leg keys merge without duplicates") + assert(kps.map(_.isNarrowed).distinct.size == 1, + "the alternatives name one layout, so they must agree on isNarrowed") + + assert(collectShuffles(plan).isEmpty, + "no shuffle: the merged descriptor for `id` matches t3") + assert(collectGroupPartitions(plan).isEmpty, + "no GroupPartitionsExec: merged keys are already grouped and aligned") + checkAnswer(df, Seq( + Row(1, 1, "a1", "c1"), Row(2, 2, "a2", "c2"), + Row(3, 3, "b3", "c3"), Row(4, 4, "b4", "c4"))) + } + } + } + + test("SPARK-58594: storage-partitioned join over union: merged PartitioningCollection needs " + + "grouping") { + // Same collection-per-leg shape as above, but the legs share the key 2, so every merged + // descriptor has a duplicate key and `isGrouped` is false. That routes the plan through + // `GroupPartitionsExec`, which resolves the child's partitioning with + // `collectFirst { case k: KeyedPartitioning => k }` -- i.e. it groups on one alternative and + // has to stay correct for the others, which is sound only because all members of the merged + // collection share one `partitionKeys` reference. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + val partitions = Array(identity("id")) + withTable("t1", "t2", "t3") { + createTable("t1", cols, partitions) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')") + createTable("t2", cols, partitions) + sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3')") + createTable("t3", cols, partitions) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.k, u.data, t3.data AS t3data + |FROM ( + | SELECT id, id AS k, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, id AS k, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val collection = union.outputPartitioning match { + case pc: physical.PartitioningCollection => pc + case other => fail(s"expected a PartitioningCollection of KeyedPartitionings, got $other") + } + val kps = collection.partitionings.map(_.asInstanceOf[physical.KeyedPartitioning]) + assert(kps.forall(_.numPartitions == 4), "merged keys are [1, 2] ++ [2, 3]") + assert(kps.forall(!_.isGrouped), + "the key 2 appears in both legs, so the merge is ungrouped") + + assert(collectShuffles(plan).isEmpty, "no shuffle: grouping the merged keys matches t3") + assert(collectGroupPartitions(plan).nonEmpty, + "GroupPartitionsExec must coalesce the two partitions holding key 2") + checkAnswer(df, Seq( + Row(1, 1, "a1", "c1"), Row(2, 2, "a2", "c2"), + Row(2, 2, "b2", "c2"), Row(3, 3, "b3", "c3"))) + } + } + } + + test("SPARK-58594: storage-partitioned join over union: PartitioningCollection under AQE") { + // The SPARK-58317 pass-through has an AQE variant; cover the merged keyed collection too, since + // AQE re-plans the union's parent from stage output statistics. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + val partitions = Array(identity("id")) + withTable("t1", "t2", "t3") { + createTable("t1", cols, partitions) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')") + createTable("t2", cols, partitions) + sql("INSERT INTO testcat.ns.t2 VALUES (3, 'b3'), (4, 'b4')") + createTable("t3", cols, partitions) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'), (4, 'c4')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.k, u.data, t3.data AS t3data + |FROM ( + | SELECT id, id AS k, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, id AS k, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id + |""".stripMargin) + checkAnswer(df, Seq( + Row(1, 1, "a1", "c1"), Row(2, 2, "a2", "c2"), + Row(3, 3, "b3", "c3"), Row(4, 4, "b4", "c4"))) + assert(collectShuffles(df.queryExecution.executedPlan).isEmpty, + "the merged keyed collection must survive AQE planning") + } + } + } + + test("SPARK-58594: storage-partitioned join over union: legs agreeing on fewer key positions, " + + "narrowed keys stay distinct") { + // t1 is partitioned by (id, dept) and t2 by (id) alone, so the legs agree only on `id`. The + // merge narrows to that position and projects t1's two-column keys down to it. t1's keys stay + // distinct once projected ([1, 2]), so the merged descriptor is grouped and no config is + // needed: `groupedSatisfies` only gates a narrowed partitioning when it is also ungrouped. + // `u.dept` has to appear in the *outer* projection, not just the legs: `ColumnPruning` pushes + // through the union, and pruning a partition column makes `V2ScanPartitioningAndOrdering` drop + // the KeyedPartitioning at the scan, leaving the union nothing to merge (the follow-up + // SPARK-46367 called out). + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + val joinCols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2", "t3") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1'), (2, 'y', 'a2')") + createTable("t2", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (3, 'p', 'b3'), (4, 'q', 'b4')") + createTable("t3", joinCols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'), (4, 'c4')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.dept, u.data, t3.data AS t3data + |FROM ( + | SELECT id, dept, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, dept, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(kp.expressions.length == 1, "narrowed to the single position both legs partition by") + assert(kp.numPartitions == 4, "keys are t1's projected to `id` ([1, 2]) plus t2's ([3, 4])") + assert(kp.isNarrowed, "t1 lost its `dept` position") + assert(kp.isGrouped, "projecting t1's keys to `id` leaves them distinct") + + assert(collectShuffles(plan).isEmpty, "no shuffle: the narrowed descriptor matches t3") + checkAnswer(df, Seq( + Row(1, "x", "a1", "c1"), Row(2, "y", "a2", "c2"), + Row(3, "p", "b3", "c3"), Row(4, "q", "b4", "c4"))) + } + } + } + + test("SPARK-58594: storage-partitioned join over union: legs agreeing on fewer key positions, " + + "narrowed keys collide") { + // As above, but t1 holds two partitions differing only in `dept`, so projecting its keys to + // `id` duplicates the key 1. The merge is then narrowed *and* ungrouped, which carries the same + // skew risk as using a subset of the partition keys for a join, so `groupedSatisfies` refuses + // it unless `allowKeysSubsetOfPartitionKeys` is enabled -- matching how SPARK-46367 gates a + // narrowing projection. + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + val joinCols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2", "t3") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1'), (1, 'y', 'a2')") + createTable("t2", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (2, 'p', 'b2')") + createTable("t3", joinCols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2')") + + Seq(true, false).foreach { allowSubset => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> allowSubset.toString) { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.dept, u.data, t3.data AS t3data + |FROM ( + | SELECT id, dept, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, dept, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(kp.numPartitions == 3, "t1's two partitions plus t2's one") + assert(kp.isNarrowed && !kp.isGrouped, + "projecting t1's keys to `id` collides on 1, so the merge is narrowed and ungrouped") + + if (allowSubset) { + assert(collectShuffles(plan).isEmpty, + "the config opts in to the skew risk, so the narrowed descriptor is used") + assert(collectGroupPartitions(plan).nonEmpty, + "GroupPartitionsExec must coalesce t1's two partitions onto the key 1") + } else { + assert(collectShuffles(plan).nonEmpty, + "without the config a narrowed, ungrouped partitioning must not satisfy the join") + } + checkAnswer(df, Seq( + Row(1, "x", "a1", "c1"), Row(1, "y", "a2", "c1"), Row(2, "p", "b2", "c2"))) + } + } + } + } + + test("SPARK-58594: storage-partitioned join over union: legs list the same key expressions in " + + "a different order") { + // t1 is partitioned by (id, dept) and t2 by (dept, id). Both positions survive, but they sit at + // opposite positions in the two legs, so t2's keys are permuted rather than narrowed -- + // `projectKeys` takes an arbitrary position sequence, so that costs nothing extra. Nothing is + // dropped, so the merge is not narrowed. + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + withTable("t1", "t2", "t3") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1'), (2, 'y', 'a2')") + createTable("t2", cols, Array(identity("dept"), identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (3, 'z', 'b3'), (4, 'w', 'b4')") + createTable("t3", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t3 VALUES " + + "(1, 'x', 'c1'), (2, 'y', 'c2'), (3, 'z', 'c3'), (4, 'w', 'c4')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = sql( + """SELECT /*+ MERGE(u, t3) */ u.id, u.dept, u.data, t3.data AS t3data + |FROM ( + | SELECT id, dept, data FROM testcat.ns.t1 + | UNION ALL + | SELECT id, dept, data FROM testcat.ns.t2 + |) u + |JOIN testcat.ns.t3 ON u.id = t3.id AND u.dept = t3.dept + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(kp.expressions.length == 2, "both positions survive; only their order differs") + assert(kp.numPartitions == 4) + assert(!kp.isNarrowed, "a permutation drops no position") + assert(kp.isGrouped) + + assert(collectShuffles(plan).isEmpty, + "no shuffle: t2's keys are permuted into t1's order, so the merge matches t3") + checkAnswer(df, Seq( + Row(1, "x", "a1", "c1"), Row(2, "y", "a2", "c2"), + Row(3, "z", "b3", "c3"), Row(4, "w", "b4", "c4"))) + } + } + } }