diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index 187163cfc0b64..ba64a1da5577b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -432,6 +432,20 @@ case class CoalescedNullAwareHashPartitioning( copy(from = from.copy(expressions = newChildren)) } +// Describes why a keyed layout reports repeated partition keys on purpose, rather than leaving them +// for the grouping node that would settle them. +sealed trait UngroupingOrigin + +// One side of a partially clustered alignment: it keeps its splits and spreads them across the +// slots its key is expected in. +case object SPLIT_FOR_JOIN extends UngroupingOrigin + +// The other side of that alignment: it repeats each key's whole group in every slot of it. +case object REPLICATED_FOR_JOIN extends UngroupingOrigin + +// A spread an operator's ordering requirement is read through, which pairs with nothing. +case object SPLIT_FOR_ORDERING extends UngroupingOrigin + /** * The physical layout of the partitions a [[KeyedPartitioning]] describes, which is everything * about them except the expressions naming them. @@ -486,13 +500,23 @@ case class CoalescedNullAwareHashPartitioning( * one, clears the markers first. That is precision, not the * guarantee: without it the merge would spread the spurious * marker onto the accurate side. + * @param ungroupingOrigin Why the keys repeat on purpose, when a producer built them that way + * rather than leaving them for the grouping node that would settle them. + * A clustering refuses an ungrouped layout whatever this says: on its + * own such a layout never holds a key's rows together. The two + * alignment roles are read only where a pair is judged + * (`KeyedPartitioning.pairsUngrouped`). An ordering pairs nothing, + * which `None` and [[SPLIT_FOR_ORDERING]] answer. + * See [[UngroupingOrigin]]. */ + case class KeyLayout( @transient partitionKeys: Seq[InternalRowComparableWrapper], dataTypes: Seq[DataType], isGrouped: Boolean, isCollapsed: Boolean, - mayContainUnknownPartitionKeys: Boolean = false) { + mayContainUnknownPartitionKeys: Boolean = false, + ungroupingOrigin: Option[UngroupingOrigin] = None) { // The rows carry the types they are compared at, so the pair is checked against itself rather // than argued about. One row is enough: a layout's rows come from one wrapper factory, and the @@ -564,13 +588,18 @@ case class KeyLayout( * * == Distribution Satisfaction and Grouping == * Besides the default `satisfies()`, `KeyedPartitioning` answers a family of questions. They differ - * in what they let happen to the data before the distribution counts as met. Only * `keysMaySatisfy()` is asked from outside the class. The rest build it and `satisfies()` up. + * A layout a producer left ungrouped says so on the layout itself, and `satisfies()` reads it + * there (`UngroupingOrigin`). * * - `keysSatisfy()`: do the keys as they stand co-locate every cluster key, with nothing left for - * a node to project away? This is the strict question, and `satisfies()` is it plus `isGrouped`. - * Strict for a `ClusteredDistribution`; its `OrderedDistribution` arm is a gate on what may claim - * a global ordering, and the order itself is still the caller's to check. + * a node to project away? This is the strict question, and `satisfies()` is it plus `isGrouped` + * under a `ClusteredDistribution`, which an ungrouped layout serves only through the pairing + * its producer's stamp vouches for (`pairsUngrouped()`), never on its own. Its + * `OrderedDistribution` arm is a gate on what may claim a global ordering: an ordering reads + * the keys without pairing anything, so settling is not asked, but the keys do have to ascend + * (`keysSortedFor()`), and the two pairing roles are refused, built as they are for a + * clustering. * - `keysCanSatisfy()`: `keysSatisfy()`, or the keys co-locate them after a node has projected * away the expressions that carry none. Only * `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys` admits the second half. @@ -677,6 +706,7 @@ case class KeyedPartitioning( def isGrouped: Boolean = layout.isGrouped def isCollapsed: Boolean = layout.isCollapsed def mayContainUnknownPartitionKeys: Boolean = layout.mayContainUnknownPartitionKeys + def ungroupingOrigin: Option[UngroupingOrigin] = layout.ungroupingOrigin /** This partitioning over a changed layout, e.g. `withLayout(_.copy(isGrouped = false))`. */ def withLayout(f: KeyLayout => KeyLayout): KeyedPartitioning = copy(layout = f(layout)) @@ -695,7 +725,8 @@ case class KeyedPartitioning( * string is the only place a reader sees why a marked partitioning still shuffles. */ override protected def stringArgs: Iterator[Any] = - Iterator(expressions, partitionKeys, isGrouped, isCollapsed, mayContainUnknownPartitionKeys) + Iterator(expressions, partitionKeys, isGrouped, isCollapsed, + mayContainUnknownPartitionKeys) ++ ungroupingOrigin override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): KeyedPartitioning = @@ -794,7 +825,9 @@ case class KeyedPartitioning( def toGrouped: KeyedPartitioning = { // Unique keys need no dedup, only the sort. val uniqueKeys = if (isGrouped) partitionKeys else partitionKeys.distinct - withLayout(_.copy(partitionKeys = uniqueKeys.sorted(keyOrdering), isGrouped = true)) + // The keys are unique afterwards, so a reason to be ungrouped has been spent. + withLayout(_.copy(partitionKeys = uniqueKeys.sorted(keyOrdering), isGrouped = true, + ungroupingOrigin = None)) } /** @@ -854,10 +887,61 @@ case class KeyedPartitioning( reducers: Seq[Option[KeyReducer]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = KeyedPartitioning.reduceKeys(partitionKeys, keyDataTypes, reducers) - override def satisfies0(required: Distribution): Boolean = { - super.satisfies0(required) || (isGrouped && keysSatisfy(required)) + /** + * The ordering `o` reads these keys at. The single-reference-per-expression invariant of + * `supportsExpressions` gives one attribute per partition expression. Shared with + * `EnsureRequirements.resolveChild`, which sorts the expected keys with it. + */ + private[sql] def keyOrderingFor( + o: OrderedDistribution): Ordering[InternalRowComparableWrapper] = + RowOrdering.create(o.ordering, expressions.flatMap(_.references)) + .on((t: InternalRowComparableWrapper) => t.row) + + /** + * Whether the keys ascend in the sequence `o` reads them: the half of an ordering claim the + * key space does not carry. A grouped layout's keys are sorted by construction (`toGrouped` + * and the sources both sort with `groupedKeyRowOrdering`), but an ungrouped one has no + * producer that sorted it, so this check is what keeps `satisfies` from answering for a + * layout that cannot back the answer. `EnsureRequirements.resolveChild` asks the same + * question to decide keep-versus-spread, on this same construction. An empty zip is + * vacuously sorted, the answer for a single key. + */ + private[sql] def keysSortedFor(o: OrderedDistribution): Boolean = { + val keyOrdering = keyOrderingFor(o) + partitionKeys.zip(partitionKeys.drop(1)).forall { case (k1, k2) => keyOrdering.lteq(k1, k2) } } + override def satisfies0(required: Distribution): Boolean = required match { + // An ordering reads the keys without pairing partitions, so settling them is not what it + // asks: an ungrouped layout answers it on the keys alone, whether a producer stamped it for + // the ordering or nothing stamped it at all. The keys, though, have to ascend, and no + // producer vouches for that on an unstamped layout, so the check is here. The two pairing + // roles answer no: they were built for a clustering, not for an ordering. + case o: OrderedDistribution => + super.satisfies0(o) || + (ungroupingOrigin.forall(_ == SPLIT_FOR_ORDERING) && + keysSatisfy(o) && keysSortedFor(o)) + // A clustering asks that rows sharing a cluster key share a partition, which an ungrouped + // layout never gives on its own. Not even an alignment side: its claim is about the pair it + // was built for, read where a pair is judged (`pairsUngrouped`) and not here. + case _ => + super.satisfies0(required) || (isGrouped && keysSatisfy(required)) + } + + /** + * Whether this layout may pair as it stands despite being ungrouped: a producer stamped it as + * one side of an alignment, and the pairing being judged is the one the stamp vouches for. The + * caller holds the other side, and `KeyedShuffleSpec.isCompatibleWith` is where the two roles + * are asked to complement. `satisfies` refuses this layout, and rightly so: on its own it does + * not hold a key's rows together, so every consumer but that pairing needs the node that + * settles it. The count is asked here as it is asked there, since this is an admission of the + * member as it stands. + */ + private[sql] def pairsUngrouped(required: ClusteredDistribution): Boolean = + ungroupingOrigin.exists(r => r == SPLIT_FOR_JOIN || r == REPLICATED_FOR_JOIN) && + required.requiredNumPartitions.forall(_ == numPartitions) && + keysSatisfy(required) + /** * The positions of the partition expressions that cover a cluster key of `required`, and so * have to survive a projection. An expression covers one in two ways: @@ -1094,7 +1178,9 @@ object KeyedPartitioning { */ def concat(kps: Seq[KeyedPartitioning]): KeyedPartitioning = { val concatenatedKeys = kps.flatMap(_.partitionKeys) + val ungroupingOrigin = PartitioningCollection.concatUngrouping(kps.map(_.ungroupingOrigin)) kps.head.withLayout(_.copy( + ungroupingOrigin = ungroupingOrigin, partitionKeys = concatenatedKeys, // A child that has duplicates of its own puts them in the concatenation too, which answers // this without walking the keys. @@ -1371,8 +1457,9 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) // operation's still can, through the projection a `GroupPartitionsExec` performs. That is the // admission set this filter had before `satisfies` became strict, up to one shape it now also // keeps, a partition expression that *is* a cluster key, which `areKeysCompatible` turns away - // anyway. The set matters because `ValidateRequirements` builds a spec from a finished plan - // through here. + // anyway. The set matters because the planner's `shuffleToCoPartition` reads it through here to + // pick the layout the other children are laid out on, and a member this filter drops is one the + // collection cannot offer as that layout. // // Every admitted member stays, because `isCompatibleWith` answers for any of them and the // collection cannot know which one the other side matched. The cost is that @@ -1414,16 +1501,31 @@ object PartitioningCollection { private[sql] def numKeyedPartitions(partitioning: Partitioning): Option[Int] = representativeOf(partitioning).map(_.numPartitions) + /** + * The origin every one of `origins` carries, for a carrier merging them into one layout: the + * claim survives a merge only where all members make the same one. A disagreement, or no member + * at all, agrees on nothing. + */ + private[sql] def concatUngrouping( + origins: Seq[Option[UngroupingOrigin]]): Option[UngroupingOrigin] = + origins.distinct match { + case Seq(role) => role + case _ => None + } + /** * Whether `p` can serve `required` once a `GroupPartitionsExec` has projected a * [[KeyedPartitioning]]'s keys down to the cluster keys. `satisfies` asks whether it serves as it * stands, and only a keyed partitioning answers the two differently. * * A partitioning that is not grouped is not admitted, even though a node would also group it. - * This is the admission set `satisfies` gave the one caller before it became strict, and the - * caller feeds `ValidateRequirements` as well as the planner, so it does not widen what a - * finished plan is checked against. `EnsureRequirements.createKeyedShuffleSpecs` is where the - * planner asks the wider question, for a child it is about to group itself. + * This is the admission set `satisfies` gave the caller before it became strict, and it is what + * `PartitioningCollection.createShuffleSpec` answers with, which the planner's + * `shuffleToCoPartition` reads to lay the other children out on what this one reports. + * `ValidateRequirements` does not read it: `specsForPairing` is its admission, and it admits an + * ungrouped member only where partially clustered distribution reports one. + * `EnsureRequirements.createKeyedShuffleSpecs` is where the planner asks the wider question, for + * a child it is about to group itself. * * The partition count clause is here for consistency with the strict question. A node changes the * count, so the pre-grouping one is no prediction of it. @@ -1437,6 +1539,61 @@ object PartitioningCollection { case other => other.satisfies(required) } + /** + * The specs `p` offers for `distribution`, one per member that may serve it, each of them the + * layout that member reports (`reportedSpecOf`): the partitions it has, in the order it reports + * them, with the partition expressions that carry no cluster key left out. Leaving those out is + * what lets two members pair when the operation clusters on part of what they are partitioned + * on, and it invents nothing: a kept key is that partition's own, the count is the member's own, + * and the order is the member's own. + * + * A keyed member is admitted as it stands: on `satisfies`, the strict question, or, since a + * clustering is what a pair owes together, on the alignment stamp that vouches for exactly + * this pairing (`KeyedPartitioning.pairsUngrouped`; `isCompatibleWith` then asks the two + * sides' roles to complement). A member that is not keyed is asked on `satisfies` alone and + * has no projection to make. Compare the planner's admission of a member + * (`EnsureRequirements.createKeyedShuffleSpecs`): it asks the wider `keysMaySatisfy`, because + * the node it inserts settles the member afterwards, and it additionally requires the coverage + * of every operation key (`spark.sql.requireAllClusterKeysForCoPartition`), which this drops: + * a skew heuristic, since a member whose partitioning keys cover only a subset of the + * operation's keys is a sound pairing. + */ + private[sql] def specsForPairing( + p: Partitioning, + distribution: ClusteredDistribution): Seq[ShuffleSpec] = + flatten(p).flatMap { + case k: KeyedPartitioning => + Option.when(k.satisfies(distribution) || k.pairsUngrouped(distribution))( + reportedSpecOf(k, distribution)) + case other => + Option.when(other.satisfies(distribution))(other.createShuffleSpec(distribution)) + } + + /** + * The spec for a member a finished plan reports: its own layout with the partition expressions + * that carry no cluster key left out. No key is deduped and no key is re-sorted, so a side offers + * a view of the partitions the plan has rather than the layout + * `KeyedPartitioning.createShuffleSpec` builds for a node a planner would insert: that one dedups + * and sorts, and its count is the grouped one, which the plan does not hold. + * + * A projection here is a relabelling of the member's partitions, and this leans on the caller for + * that: only `keysSatisfy`'s projecting half admits such a member, and it holds the projection to + * one that merges no partition. The projecting half excludes a marked member, so a marked one is + * admitted only with its keys covering the clustering structurally, and the position count below + * comes out full for it: it leaves here unprojected, which its claim needs, since it routes the + * rows of an undeclared key by a hash over these keys in this order and dropping one breaks it. + */ + private def reportedSpecOf( + k: KeyedPartitioning, + distribution: ClusteredDistribution): KeyedShuffleSpec = { + val positions = k.positionsCoveringClusterKeys(distribution).toSeq + if (positions.size == k.expressions.length) { + KeyedShuffleSpec(k, distribution) + } else { + KeyedShuffleSpec(k.project(positions), distribution, Some(positions)) + } + } + /** * The uniform `mayContainUnknownPartitionKeys` marker of `p`'s keyed members, read from its * first keyed member, the same one `checkKeyedPartitioningInvariant` compares against. @@ -1471,11 +1628,17 @@ object PartitioningCollection { // already holds the canonical layout is returned as it is. That is what keeps repeated // `outputPartitioning` computations over deeply nested collections (e.g. chains of same-key // joins) O(1) per level. + // Interning below retypes the members that disagree with the agreed origin, and a retyped + // claim would be one nothing built. + val concatUngrouping = + PartitioningCollection.concatUngrouping(representatives.map(_.layout.ungroupingOrigin)) val canonicalLayout = representatives.map(_.layout) .find(l => l.isCollapsed == anyCollapsed && - l.mayContainUnknownPartitionKeys == anyUnknownKeys) + l.mayContainUnknownPartitionKeys == anyUnknownKeys && + l.ungroupingOrigin == concatUngrouping) .orElse(representatives.headOption.map(_.layout.copy( - isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys = anyUnknownKeys))) + isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys = anyUnknownKeys, + ungroupingOrigin = concatUngrouping))) .orNull def intern(p: Partitioning): Partitioning = representativeOf(p) match { @@ -1639,12 +1802,14 @@ case object SinglePartitionShuffleSpec extends LeafShuffleSpec { // disagree when the subset config projects them onto different key sets. // // `EnsureRequirements` never reaches this: a spec whose `canCreatePartitioning` is false is - // never the best one. The one production caller that can put a collection on the `other` side - // is `ValidateRequirements`' `specs.tail.forall(_.isCompatibleWith(specs.head))`, and there the - // stricter answer is the safer one. It does leave this direction stricter than the collection's - // own `exists`, against the symmetry this trait's doc assumes, but nothing observable follows: - // only `KeyedShuffleSpec` can make members disagree on `numPartitions`, and it has no - // `SinglePartitionShuffleSpec` case, so that direction is already false. + // never the best one. Nor does any other production caller put a collection on the `other` + // side: `pickCoPartitionTarget` flattens before it pairs a member, and `ValidateRequirements` + // offers leaf specs and pairs those, where it used to compare every child against the first + // one's whole spec. So what is left is the answer `ShuffleSpecSuite` pins. It stays `forall` + // rather than the `exists` a collection answers with, against the symmetry the trait doc + // assumes, but nothing observable follows: only `KeyedShuffleSpec` can make members disagree on + // `numPartitions`, and it has no `SinglePartitionShuffleSpec` case, so that direction is + // already false. case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith) } @@ -1922,6 +2087,13 @@ case class IdentityReducer(transform: TransformExpression) extends Reducer[Any, * does not: one `keyPositions` entry is empty there, which * `canCreatePartitioning` and `areKeysCompatible` both turn away. See * `KeyedPartitioning.createShuffleSpec`. + * The field's other producer is `PartitioningCollection.reportedSpecOf`, + * and there `Some` records a projection of the *report*: a view of the + * partitions the child already holds, which no node is owed. Such a spec + * is only ever read by `isCompatibleWith`, which ignores the field, so + * the "insert or compose a grouping node" reading the + * `EnsureRequirements` consumers give `Some` stays about planner-produced + * specs alone. */ case class KeyedShuffleSpec( partitioning: KeyedPartitioning, @@ -1962,18 +2134,29 @@ case class KeyedShuffleSpec( // Here we check: // 1. both distributions have the same number of clustering keys // 2. both partitioning have the same number of partitions - // 3. partition expressions from both sides are compatible, which means: - // 3.1 both sides have the same number of partition expressions - // 3.2 for each pair of partition expressions at the same index, the corresponding + // 3. where either side is ungrouped on purpose, the pair is one of each: the side that + // spreads a key's splits (`SPLIT_FOR_JOIN`) and the side that repeats the group + // (`REPLICATED_FOR_JOIN`). An unstamped ungrouped layout serves no clustering, so only + // `compatibleAsIs` can offer one here, and its builder settles the side with a grouping + // node. + // 4. partition expressions from both sides are compatible, which means: + // 4.1 both sides have the same number of partition expressions + // 4.2 for each pair of partition expressions at the same index, the corresponding // partition keys must share overlapping positions in their respective clustering keys. - // 3.3 each pair of partition expressions at the same index must describe one key space: two + // 4.3 each pair of partition expressions at the same index must describe one key space: two // bare references, the same transform function, or the two sides of one reduce. A pair // the join would first reduce onto one key space does not count, see // `areKeysCompatible`'s `allowReduce`. - // 4. the partition values from both sides are following the same order. + // 5. the partition values from both sides are following the same order. case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _) => distribution.clustering.length == otherDistribution.clustering.length && numPartitions == otherSpec.numPartitions && + ((partitioning.ungroupingOrigin, otherPartitioning.ungroupingOrigin) match { + case (None, None) => true + case (Some(l), Some(r)) => + Set(l, r) == Set(SPLIT_FOR_JOIN, REPLICATED_FOR_JOIN) + case _ => false + }) && areKeysCompatible(otherSpec, allowReduce = false) && // The reason the types are asked as well as the rows is on `describesSameKeys`, so the // next site comparing keys cannot forget the type clause. diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala index 86eec4be38707..fa77dc40a3153 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException} import org.apache.spark.sql.catalyst.dsl.expressions._ -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression} +import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, DirectShufflePartitionID, Expression, SortOrder, TransformExpression} import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.connector.catalog.functions.{FlipLowBitFunction, Reducer, ReducibleFunction, ScalarFunction} @@ -93,9 +93,8 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { test("SPARK-59289: createShuffleSpec drops a keyed member that is not grouped") { val a = AttributeReference("a", IntegerType)() val clustered = ClusteredDistribution(Seq(a)) - // A node would group this one too, but the admission set here is what a finished plan is - // checked against by `ValidateRequirements`, so it stays what the strict question admitted - // before a projection was allowed to answer it. + // A node would group this one too, but a collection is a layout the plan holds, so it stays + // what the strict question admitted before a projection was allowed to answer it. val ungrouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2))) val grouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2), InternalRow(3))) assert(!ungrouped.isGrouped && grouped.isGrouped, "test setup") @@ -114,6 +113,143 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { assert(specs.head.isInstanceOf[HashShuffleSpec], s"and it is the hash one, got $specs") } + test("SPARK-59671: the specs a side offers are the layouts it reports") { + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val cd = ClusteredDistribution(Seq(a)) + // An ungrouped member is offered where its layout says where the ungrouping comes from: one + // side of an alignment keeps its splits and spreads them. A finished plan is judged on the + // layouts it holds, so the reason is read off them rather than off a configuration. + val spread = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2))) + .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN))) + val specs = PartitioningCollection.specsForPairing(spread, cd) + assert(specs.size == 1, s"one member, one spec, got $specs") + val spec = specs.head.asInstanceOf[KeyedShuffleSpec] + assert((spec.partitioning eq spread) && spec.joinKeyPositions.isEmpty, + s"its own layout, no projection to make, got $specs") + + // A member that does not say why it is ungrouped is waiting for the node that would settle it, + // and a finished plan has none left to insert. + val unsettled = spread.withLayout(_.copy(ungroupingOrigin = None)) + assert(PartitioningCollection.specsForPairing(unsettled, cd).isEmpty, + "an ungrouped member is a plan only where a producer built one") + + // An origin is read for the distribution it is good for. A spread an ordering requirement is + // read through is no side of a clustering, and a clustering is what asks here. + val forOrdering = spread.withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_ORDERING))) + assert(PartitioningCollection.specsForPairing(forOrdering, cd).isEmpty, + "an origin a clustering does not read is not offered to it") + + // The side that repeats a key's whole group is the other half of that alignment, and it is + // offered the same way. + val repeat = spread.withLayout(_.copy(ungroupingOrigin = Some(REPLICATED_FOR_JOIN))) + assert(PartitioningCollection.specsForPairing(repeat, cd).size == 1) + + // A count the operation pinned is asked as it stands, which `satisfies` asks for every member, + // the ones whose layout says why they are ungrouped included. + val four = KeyedPartitioning(Seq(a), (1 to 4).map(InternalRow(_))) + .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN))) + assert(PartitioningCollection + .specsForPairing(four, ClusteredDistribution(Seq(a), requiredNumPartitions = Some(4))) + .size == 1, "the size the operation asks for is the one the member reports") + assert(PartitioningCollection + .specsForPairing(four, ClusteredDistribution(Seq(a), requiredNumPartitions = Some(3))) + .isEmpty, "a member of another size does not serve the distribution") + + // The subset permission applies where the operation's keys are a subset of the member's + // partitioning keys, so a member may carry an expression the operation does not cluster on. It + // is offered as its own partitions under the key the operation clusters on: that expression is + // left out, and no key is deduped or re-sorted for it. + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val source = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2, 2))) + val offered = PartitioningCollection.specsForPairing(source, cd) + .head.asInstanceOf[KeyedShuffleSpec] + assert(offered.partitioning.expressions === Seq(a) && + offered.partitioning.numPartitions === source.numPartitions && + offered.joinKeyPositions === Some(Seq(0)), + s"the member's own partitions under the operation's key, got $offered") + + // A marked layout is never offered through a projection: the projecting half of `keysSatisfy` + // excludes a marked member, so the keys an offered marked one carries are the keys its claim + // is over. + val marked = source.withLayout(_.copy(mayContainUnknownPartitionKeys = true)) + assert(PartitioningCollection.specsForPairing(marked, cd).isEmpty, + "a marked member is not offered through a projection") + } + + // A member whose keys do not cover the clustering is not offered, and a member that is not + // keyed is asked for its own spec. + val elsewhere = KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2))) + assert(PartitioningCollection.specsForPairing(elsewhere, cd).isEmpty) + assert(PartitioningCollection + .specsForPairing(HashPartitioning(Seq(a), 2), cd).head.isInstanceOf[HashShuffleSpec]) + } + + test("SPARK-59671: carriers keep one agreed origin and grouping spends it") { + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val keys = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def stamped(attr: AttributeReference, role: UngroupingOrigin): KeyedPartitioning = + KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin = Some(role))) + val spread = stamped(a, SPLIT_FOR_JOIN) + val repeat = stamped(a, REPLICATED_FOR_JOIN) + + // Grouping spends the claim: the keys are unique afterwards, so there are no splits left to + // spread and no group left to repeat, and a kept stamp would turn a sound pair of grouped + // sides away in `isCompatibleWith`, which reads a stamp as one side of an ungrouped pair. + assert(spread.toGrouped.isGrouped && spread.toGrouped.ungroupingOrigin.isEmpty) + + // A concatenation stands for the claim all of its members make: one role survives it, and + // two roles stand for nothing. + assert(KeyedPartitioning.concat(Seq(spread, spread)).ungroupingOrigin + .contains(SPLIT_FOR_JOIN)) + assert(KeyedPartitioning.concat(Seq(spread, repeat)).ungroupingOrigin.isEmpty) + + // A collection holds one canonical layout its members share, and it agrees the same way. An + // aligned pair's joined output is the disagreeing case: the zip of a spread side and a + // repeating side is a layout neither claim describes, so it comes out inert. + val agreed = PartitioningCollection.fromPartitionings( + Seq(spread, stamped(b, SPLIT_FOR_JOIN))) + assert(agreed.partitionings.forall(p => p.asInstanceOf[KeyedPartitioning] + .ungroupingOrigin.contains(SPLIT_FOR_JOIN)), "one role, one canonical layout") + val joined = PartitioningCollection.fromPartitionings(Seq(spread, repeat)) + assert(joined.partitionings.forall(_.asInstanceOf[KeyedPartitioning] + .ungroupingOrigin.isEmpty), "two roles agree on nothing") + } + + test("SPARK-59671: an ordering reads sorted keys, a clustering reads the pairing") { + val a = AttributeReference("a", IntegerType)() + def layout(keys: Seq[Int], role: Option[UngroupingOrigin]): KeyedPartitioning = + KeyedPartitioning(Seq(a), keys.map(InternalRow(_))) + .withLayout(_.copy(ungroupingOrigin = role)) + val sorted = Seq(1, 1, 2) + val ordering = OrderedDistribution(Seq(SortOrder(a, Ascending))) + val clustered = ClusteredDistribution(Seq(a)) + + withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") { + // An ordering pairs no partition with another, so it asks no settling, but it does ask + // the keys to ascend, which a grouped layout has by construction and an ungrouped one + // only where something sorted it. + assert(layout(sorted, None).satisfies(ordering)) + assert(layout(sorted, Some(SPLIT_FOR_ORDERING)).satisfies(ordering)) + assert(!layout(Seq(1, 2, 1, 2), None).satisfies(ordering), + "an ungrouped layout nothing sorted does not claim an ordering") + assert(!layout(Seq(2, 1), None).satisfies(ordering), + "a grouped layout whose keys do not ascend does not claim an ordering") + // The two pairing roles were built for a clustering, not for an ordering. + assert(!layout(sorted, Some(SPLIT_FOR_JOIN)).satisfies(ordering)) + assert(!layout(sorted, Some(REPLICATED_FOR_JOIN)).satisfies(ordering)) + // A clustering asks that a key's rows share a partition, which no ungrouped layout gives + // on its own: an alignment side's stamp is read where the pair is judged + // (`specsForPairing`, pinned above), not here. + Seq(None, Some(SPLIT_FOR_ORDERING), Some(SPLIT_FOR_JOIN), Some(REPLICATED_FOR_JOIN)).foreach { + role => assert(!layout(sorted, role).satisfies(clustered)) + } + assert(layout(Seq(1, 2), None).satisfies(clustered), + "a grouped layout answers a clustering as it stands") + } + } + protected def checkCompatible( left: ShuffleSpec, right: ShuffleSpec, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index ea6adcd814dc6..88893c7b79b49 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering import org.apache.spark.sql.catalyst.plans.QueryPlan -import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyLayout, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyLayout, KeyReducer, Partitioning, PartitioningCollection, REPLICATED_FOR_JOIN, UngroupingOrigin, UnknownPartitioning} import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper} import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan, SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -69,9 +69,13 @@ import org.apache.spark.sql.vectorized.ColumnarBatch * reducers rather than as their display names. A name embeds the exprId of the * attribute it transforms, so it would break canonical equality the same way the * reducers do. See `doCanonicalize`. - * @param distributePartitions When true, splits for a key are distributed across the expected - * partitions (padding with empty partitions). When false, all splits - * are replicated to every expected partition for that key. + * @param ungroupingOrigin Why this node's output is left ungrouped on purpose, stamped into + * the layout it reports. It also decides the routing: + * `SPLIT_FOR_JOIN` and `SPLIT_FOR_ORDERING` distribute each key's + * splits across the expected slots (padding with empty partitions), + * [[REPLICATED_FOR_JOIN]] groups the splits first and repeats each group + * into its expected slots, and `None` groups them, one partition per key. + * See [[UngroupingOrigin]]. * @param enableSortedMerge When true, uses [[SortedMergeCoalescedRDD]] to perform a k-way merge * of the coalesced partitions, preserving the child's output ordering * end-to-end. Set by [[EnsureRequirements]] when a parent operator @@ -89,10 +93,16 @@ case class GroupPartitionsExec( @transient joinKeyPositions: Option[Seq[Int]], @transient expectedKeyCount: Option[Int], @transient reducers: Option[Seq[Option[KeyReducer]]], - @transient distributePartitions: Boolean, + @transient ungroupingOrigin: Option[UngroupingOrigin], @transient enableSortedMerge: Boolean ) extends UnaryExecNode { + /** + * Whether each key's splits are distributed across its expected slots, or grouped into one + * partition per key first. Derived from the origin; see `@param ungroupingOrigin`. + */ + def distributePartitions: Boolean = ungroupingOrigin.exists(_ != REPLICATED_FOR_JOIN) + /** * The layout this node was planned to produce, but only while its child still reports the one it * was planned over. @@ -451,11 +461,15 @@ case class GroupPartitionsExec( return None } val positions = joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild)) + // Only a grouping node reaches here today: the one caller (`CombineAdjacentAggregation`) + // regroups the node between two aggregates, which is built unstamped, so forwarding the + // origin preserves rather than claims. A stamped node would also need its expected slots + // re-derived, which this path does not do. val regrouped = GroupPartitionsExec( child = newChild, joinKeyPositions = Option.when(positions != newChildKp.expressions.indices)(positions), reducers = reducers, - distributePartitions = distributePartitions, + ungroupingOrigin = ungroupingOrigin, enableSortedMerge = enableSortedMerge) regrouped.copyTagsFrom(this) Some(regrouped) @@ -523,7 +537,7 @@ private[sql] object GroupPartitionsExec { * * **Both are derived, and neither `copy` nor the generated `apply` re-derives them**, so a change * to `child`, `joinKeyPositions`, `expectedPartitionKeys` (stored as `expectedKeyCount`), - * `reducers` or `distributePartitions` has to come back through here. `enableSortedMerge` is not + * `reducers` or `ungroupingOrigin` has to come back through here. `enableSortedMerge` is not * an input to either, which is why `tryEnableSortedMerge` may `copy` it. * * Two other `copy` calls in this file are deliberate. `withNewChildInternal` carries both fields @@ -537,7 +551,7 @@ private[sql] object GroupPartitionsExec { joinKeyPositions: Option[Seq[Int]] = None, expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]] = None, reducers: Option[Seq[Option[KeyReducer]]] = None, - distributePartitions: Boolean = false, + ungroupingOrigin: Option[UngroupingOrigin] = None, enableSortedMerge: Boolean = false): GroupPartitionsExec = { // There must be a `KeyedPartitioning` in the child's output partitioning, as a // `GroupPartitionsExec` node is added to a plan only in that case. @@ -548,11 +562,11 @@ private[sql] object GroupPartitionsExec { // `representativeOf` found one, so the partitioning is an expression tree. val childExpr = childPartitioning.asInstanceOf[Partitioning with Expression] val grouping = computeGrouping( - childKp, joinKeyPositions, expectedPartitionKeys, reducers, distributePartitions) + childKp, joinKeyPositions, expectedPartitionKeys, reducers, ungroupingOrigin) GroupPartitionsExec(child, grouping, - computeOutputPartitioning(childExpr, grouping, joinKeyPositions, reducers), childPartitioning, - joinKeyPositions, expectedPartitionKeys.map(_.size), reducers, distributePartitions, - enableSortedMerge) + computeOutputPartitioning(childExpr, grouping, joinKeyPositions, reducers), + childPartitioning, joinKeyPositions, expectedPartitionKeys.map(_.size), reducers, + ungroupingOrigin, enableSortedMerge) } /** @@ -585,13 +599,14 @@ private[sql] object GroupPartitionsExec { // commits to a pairing and falls back to a shuffle when it happens. // // The marker read off `p` is the same bit the `else` branch publishes through - // `grouping.layout`, since `computeGrouping` copies the child's layout without touching it. A - // change to that copy has to move this guard with it. + // `grouping.layout`, since `computeGrouping` copies the child's layout and carries the + // marker over. A change to that copy has to move this guard with it. if (PartitioningCollection.keyedMarkerOf(p).contains(true) && !grouping.isIdentity) { UnknownPartitioning(grouping.partitions.size) } else { // One instance for every member, so they share it by reference. It already carries the - // child's marker, and the guard above is what lets that carry over. + // child's marker, and the guard above is what lets that carry over, and the origin this + // node's producer stamped, which `computeGrouping` wrote into it. val layout = grouping.layout p.transform { case k: KeyedPartitioning => @@ -680,7 +695,9 @@ private[sql] object GroupPartitionsExec { joinKeyPositions: Option[Seq[Int]], expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]], reducers: Option[Seq[Option[KeyReducer]]], - distributePartitions: Boolean): PartitionGrouping = { + ungroupingOrigin: Option[UngroupingOrigin]): PartitionGrouping = { + // The routing follows the origin, the same derivation the node itself reports. + val distributePartitions = ungroupingOrigin.exists(_ != REPLICATED_FOR_JOIN) // Project partition keys if join key positions are specified val (projectedDataTypes, projectedKeys) = joinKeyPositions.fold( @@ -740,16 +757,21 @@ private[sql] object GroupPartitionsExec { case ((_, Seq(single)), outputIndex) => single == outputIndex case _ => false } - // A `copy` of the child's layout rather than a fresh one, so the marker and anything the layout - // grows later are carried without this having to name them. Built here, where the child is in - // scope, so `computeOutputPartitioning` has nothing left to re-apply. + // A `copy` of the child's layout rather than a fresh one, so the marker and anything the + // layout grows later are carried without this having to name them. The origin is the one + // field not carried: the output claims what this node's own producer stamped, never what + // the child's claimed, and only while the shape holds the claim -- a grouping that settles + // every key leaves nothing repeating, so it spends the claim the way `toGrouped` does. + // Built here, where the child is in scope, so `computeOutputPartitioning` has nothing left + // to re-apply. PartitionGrouping( partitions, childKp.layout.copy( partitionKeys = partitions.map(_._1), dataTypes = reducedDataTypes, isGrouped = isGrouped, - isCollapsed = isCollapsed), + isCollapsed = isCollapsed, + ungroupingOrigin = if (isGrouped) None else ungroupingOrigin), isIdentity, numPrunedPartitions, numReplicatedPartitionReads) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index 1eea8a006ed22..0674047dd3b32 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -136,26 +136,21 @@ case class EnsureRequirements( case Some(resolution) => (distribution, resolution) match { case (o: OrderedDistribution, _) => - // OrderedDistribution requires grouped KeyedPartitioning with sorted keys - // according to the distribution's ordering. + // OrderedDistribution requires a KeyedPartitioning with keys sorted according to + // the distribution's ordering. val satisfyingKeyedPartitioning = resolution.fold(identity, _._1) - // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees - // one attribute per partition expression. - val attrs = satisfyingKeyedPartitioning.expressions.flatMap(_.references) - val keyRowOrdering = RowOrdering.create(o.ordering, attrs) - val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row) - val keys = satisfyingKeyedPartitioning.partitionKeys - // An empty zip is vacuously sorted, which is the answer for a single key. - if (keys.zip(keys.drop(1)).forall { case (k1, k2) => keyOrdering.lteq(k1, k2) }) { + if (satisfyingKeyedPartitioning.keysSortedFor(o)) { child } else { - // Use distributePartitions to spread splits across expected partitions - val sortedGroupedKeys = keys + // Spread the splits across the expected partitions, in the ordering's sequence + val sortedGroupedKeys = satisfyingKeyedPartitioning.partitionKeys .groupBy(identity).view.mapValues(_.size) - .toSeq.sortBy(_._1)(keyOrdering) + .toSeq.sortBy(_._1)(satisfyingKeyedPartitioning.keyOrderingFor(o)) GroupPartitionsExec(child, expectedPartitionKeys = Some(sortedGroupedKeys), - distributePartitions = true + // The keys stay ungrouped so that the ordering the operator reads is the one + // derived from them, which is no side of a pairing. + ungroupingOrigin = Some(SPLIT_FOR_ORDERING) ) } @@ -592,6 +587,9 @@ case class EnsureRequirements( right: SparkPlan, rightRequired: ClusteredDistribution): Option[Seq[SparkPlan]] = { parent match { + // A keyed alignment is planned for these two operators only. A `SortMergeAsOfJoinExec` is + // a `ShuffledJoin` and gets none: its matches are read in the order within a partition, + // which neither a spread side nor a repeating one preserves. case smj: SortMergeJoinExec => checkKeyGroupCompatible(left, leftRequired, right, rightRequired, smj.joinType) case sj: ShuffledHashJoinExec => @@ -904,13 +902,25 @@ case class EnsureRequirements( } // Now we need to push-down the common partition information to the `GroupPartitionsExec`s. + // + // Where `applyPartialClustering` holds, exactly one side repeats: `replicateRightSide` is + // the negation of `replicateLeftSide`, and the branch above is taken only when the side it + // picked may replicate for the join type. That split is the whole of the pairing's + // soundness: for a key, a partition holding part of it on the spread side holds all of it + // on the repeating side, so pairing the two index by index loses no match, and + // `ValidateRequirements` cannot tell such a pair from two sides that split the key between + // them. Each side stamps why it is left ungrouped (`UngroupingOrigin`), so a reader of the + // finished plan judges the pair on the layouts it holds; the stamp also decides the node's + // routing (`GroupPartitionsExec.distributePartitions`). ( GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions, Some(mergedPartitionKeys), leftReducers, - distributePartitions = applyPartialClustering && !replicateLeftSide), + ungroupingOrigin = Option.when(applyPartialClustering)( + if (replicateLeftSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)), GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions, Some(mergedPartitionKeys), rightReducers, - distributePartitions = applyPartialClustering && !replicateRightSide)) + ungroupingOrigin = Option.when(applyPartialClustering)( + if (replicateRightSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN))) } // The pairing is only worth committing to if both children still declare the same aligned key diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala index 1ac6b809fd250..4336b0847bc7e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala @@ -45,22 +45,33 @@ object ValidateRequirements extends Logging { assert(requiredChildDistributions.length == children.length) assert(requiredChildOrderings.length == children.length) + // A `ClusteredDistribution` is the one distribution an operator can owe its children + // together rather than one by one, so an operator whose children all owe one is judged on + // their mutual layout: the pairing below decides it, and the per-child distribution clause + // is skipped for it. Nothing is lost by the skip: `specsForPairing` admits a member only + // on its own answer, which is the clause's plus the alignment stamp a producer vouched for, + // and a side offering nothing fails the pairing by itself. Every other child, an operator + // with a single clustered child included, answers the clause for itself: a stamped layout + // is one side of a pair, and a consumer that is not the pairing's other side owes it the + // node that settles it. That is every such operator, not only a join: a cogroup zips + // corresponding partitions too. + val clusteredMultiChild = children.length > 1 && + requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution]) + val satisfied = children.zip(requiredChildDistributions.zip(requiredChildOrderings)).forall { case (child, (distribution, ordering)) - if !child.outputPartitioning.satisfies(distribution) + if (!clusteredMultiChild && !child.outputPartitioning.satisfies(distribution)) || !SortOrder.orderingSatisfies(child.outputOrdering, ordering) => logDebug(s"ValidateRequirements failed: $distribution, $ordering\n$plan") false case _ => true } - if (satisfied && children.length > 1 && - requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])) { - // Check the co-partitioning requirement. - val specs = children.map(_.outputPartitioning).zip(requiredChildDistributions).map { - case (p, d) => p.createShuffleSpec(d.asInstanceOf[ClusteredDistribution]) - } - if (specs.tail.forall(_.isCompatibleWith(specs.head))) { + if (satisfied && clusteredMultiChild) { + // Check the co-partitioning requirement. A pair aligned without grouping is one each + // side's layout answers for (`UngroupingOrigin`), so a pair no producer built is one + // whose sides offer nothing to pair and fails here. + if (satisfiesForPairing(children, requiredChildDistributions)) { true } else { logDebug(s"ValidateRequirements failed: children not co-partitioned in\n$plan") @@ -70,4 +81,31 @@ object ValidateRequirements extends Logging { satisfied } } + + /** + * Whether the sides of a multi-child clustered operator line up: every side offers the layouts + * it reports ([[PartitioningCollection.specsForPairing]]), and one member of the first side + * pairs with every other side. A plan holds what its members report, so no key is deduped and + * none is re-sorted to make a pair: a side is judged on the partitions it has, under the key + * the operation clusters on. + * + * `EnsureRequirements` asks the same question of a pair it takes as it stands by a different + * predicate: its `compatibleAsIs` reads two unprojected specs, while a member here may be + * relabelled onto the operation's cluster key (`reportedSpecOf`), and its `agreeingPairs` / + * `committed` path commits on the sides a reduce rebuilt rather than on the pair it picked. A + * finished plan is asked the strict question alone, which the reduced pair answers + * (`hasSameReducedKeys`). The coverage of every operation key + * (`spark.sql.requireAllClusterKeysForCoPartition`) is a skew heuristic and part of neither. + */ + private def satisfiesForPairing( + children: Seq[SparkPlan], + distributions: Seq[Distribution]): Boolean = { + val specs = children.zip(distributions).map { case (child, distribution) => + PartitioningCollection.specsForPairing( + child.outputPartitioning, distribution.asInstanceOf[ClusteredDistribution]) + } + specs.headOption.exists { firstSide => + firstSide.exists(head => specs.tail.forall(side => side.exists(_.isCompatibleWith(head)))) + } + } } 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 6cb17713e1c8a..951cf9b9f0b00 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 @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, import org.apache.spark.sql.catalyst.expressions.aggregate.Complete import org.apache.spark.sql.catalyst.plans.{Cross, ExistenceJoin, Inner, JoinType, LeftAnti, LeftSemi, LeftSingle} import org.apache.spark.sql.catalyst.plans.physical -import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN} import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryCatalystRuntimeFilterCatalog, InMemoryTableCatalog} import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.connector.distributions.Distributions @@ -45,10 +45,11 @@ import org.apache.spark.sql.execution.{ SparkPlan, UnionExec, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, ResultQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec} import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils import org.apache.spark.sql.execution.ui.SparkPlanGraphNode import org.apache.spark.sql.execution.window.{Final, Partial, WindowGroupLimitExec, WindowGroupLimitMode} @@ -7653,8 +7654,8 @@ class KeyGroupedPartitioningSuite test("SPARK-59120: reduced partition keys are read at the types they were built with") { // The join reduces the identity side onto the year key space, so its keys become `IntegerType` // years while the partitioning still reports `identity(ts)`, declaring `TimestampType`. With - // the subset opt-in on, AQE re-runs `createShuffleSpec` on the already reduced children through - // `ValidateRequirements`, which projects and sorts those keys. The mechanism is in + // the subset opt-in on, AQE re-runs `EnsureRequirements` over the already reduced children, + // and its `createShuffleSpec` projects and sorts those keys. The mechanism is in // `ShuffleSpecSuite`'s "createShuffleSpec sorts the projected keys at their built-with types". withTable("t_identity", "t_years") { createTsTable("t_identity", Array(identity("ts"))) @@ -9228,6 +9229,286 @@ class KeyGroupedPartitioningSuite } } + test("SPARK-59671: a partially clustered join leaves AQE's shuffle coalescing alone") { + // AQE validates a stage's whole candidate plan before accepting a shuffle-read change: on the + // base, a storage-partitioned join whose sides are aligned but not grouped kept every shuffle + // in its stage uncoalesced, unrelated ones included. Partially clustered distribution plans + // such a pair: the side that keeps its splits spreads them, and the other replicates its + // group across them, so both report repeated keys on purpose. The assertions below pin the + // read that the join's presence must leave alone. + val idCols = Array(Column.create("id", IntegerType), Column.create("data", StringType)) + createTable("pc1", idCols, Array(identity("id"))) + createTable("pc2", idCols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.pc1 VALUES (1, 'a1'), (2, 'a2'), (3, 'a3')") + // Key 1 twice: the side holding it keeps and spreads its two splits, which is what makes the + // pair ungrouped while its keys still line up. + sql("INSERT INTO testcat.ns.pc2 VALUES (1, 'b1'), (1, 'b1b'), (2, 'b2'), (4, 'b4')") + + withSQLConf( + SQLConf.V2_BUCKETING_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + "spark.sql.autoBroadcastJoinThreshold" -> "-1") { + val df = sql( + s""" + |SELECT /*+ MERGE(a) */ a.id, b.data + |FROM testcat.ns.pc1 a JOIN testcat.ns.pc2 b ON a.id = b.id + |UNION ALL + |SELECT count(*), cast(id % 2 AS STRING) + |FROM testcat.ns.pc1 GROUP BY id % 2 + |""".stripMargin) + checkAnswer(df, Seq(Row(1, "b1"), Row(1, "b1b"), Row(2, "b2"), Row(2, "1"), Row(1, "0"))) + + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val joins = collect(plan) { case j: ShuffledJoin => j } + assert(joins.size == 1, s"test setup: one storage-partitioned join:\n$plan") + assert(collectGroupPartitions(plan).exists { g => + PartitioningCollection.representativeOf(g.outputPartitioning).exists(!_.isGrouped) + }, s"test setup: the pair is spread, so a side repeats its keys:\n$plan") + // The join's side shuffles nothing (this suite's `collectShuffles` counts the exchanges a + // join reads through), and the chain shuffles once, for the aggregate: the stage holding + // both is the one whose read the join's presence could have kept uncoalesced. + assert(collectShuffles(plan).isEmpty, s"the join shuffles nothing:\n$plan") + assert(collectAllShuffles(plan).size === 1, + s"and the chain shuffles once, for the aggregate:\n$plan") + val aqeReads = collect(df.queryExecution.executedPlan) { case r: AQEShuffleReadExec => r } + assert(aqeReads.size === 1 && aqeReads.head.hasCoalescedPartition, + s"the aggregate's shuffle read must coalesce:\n${df.queryExecution.executedPlan}") + + // And the two share a stage, which is what makes the coalesce a decision the join can + // block: a read in a stage of its own would coalesce whatever the join did. + val finalStage = collect(df.queryExecution.executedPlan) { + case s: ResultQueryStageExec => s + } + assert(finalStage.size === 1 && finalStage.head.plan.exists { + case j: ShuffledJoin => true + case _ => false + }, s"test setup: the join is in the stage the read belongs to:\n" + + s"${df.queryExecution.executedPlan}") + } + } + + test("SPARK-59671: a three-table chain keeps its AQE coalescing") { + val idCols = Array(Column.create("id", IntegerType), Column.create("data", StringType)) + createTable("p3a", idCols, Array(identity("id"))) + createTable("p3b", idCols, Array(identity("id"))) + createTable("p3c", idCols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.p3a VALUES (1, 'a1'), (2, 'a2')") + sql("INSERT INTO testcat.ns.p3b VALUES (1, 'b1'), (1, 'b1b'), (2, 'b2')") + sql("INSERT INTO testcat.ns.p3c VALUES (1, 'c1'), (2, 'c2')") + withSQLConf( + SQLConf.V2_BUCKETING_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + "spark.sql.autoBroadcastJoinThreshold" -> "-1") { + val df = sql( + s""" + |SELECT /*+ MERGE(a, b), MERGE(a, c) */ a.id AS aid, b.id AS bid, c.data + |FROM testcat.ns.p3a a JOIN testcat.ns.p3b b ON a.id = b.id + |JOIN testcat.ns.p3c c ON a.id = c.id + |UNION ALL + |SELECT count(*), 0, 'x' FROM testcat.ns.p3a GROUP BY id % 2 + |""".stripMargin) + // The union coerces the join's `id` and the aggregate's `count(*)` to one type, bigint. + checkAnswer(df, Seq(Row(1L, 1, "c1"), Row(1L, 1, "c1"), Row(2L, 2, "c2"), + Row(1L, 0, "x"), Row(1L, 0, "x"))) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert(collectShuffles(plan).isEmpty, s"the chain must not shuffle:\n$plan") + assert(plan.exists(p => keyedPartitioningsOf(Seq(p)).size >= 2), + s"test setup: a side reports one keyed member per join key column:\n$plan") + assert(ValidateRequirements.validate(plan), s"the chain's pairing holds up:\n$plan") + // Each join zips a spread side against a repeating side: the collection it reports agrees + // on no origin, so nothing above may read the pair's claim off it, while the groupings + // below it hold the two complementary stamps. + val joins = collect(plan) { case j: ShuffledJoin => j } + assert(joins.nonEmpty, s"test setup: a storage-partitioned chain:\n$plan") + assert(joins.forall(j => keyedPartitioningsOf(Seq(j)).forall(_.ungroupingOrigin.isEmpty)), + s"a zipped pair reports no claim of its own:\n$plan") + val origins = collectAllGroupPartitions(plan).flatMap(_.ungroupingOrigin).distinct + assert(origins.contains(SPLIT_FOR_JOIN) && origins.contains(REPLICATED_FOR_JOIN), + s"one side spreads and the other repeats, got $origins:\n$plan") + assert(collect(df.queryExecution.executedPlan) { case r: AQEShuffleReadExec => r } + .exists(_.hasCoalescedPartition), + s"the aggregate's shuffle read must coalesce:\n${df.queryExecution.executedPlan}") + } + } + + test("SPARK-59671: an AQE local read over a split scan does not lose matches") { + // A lower join with a non-bucketed side cannot align, so it hash-shuffles the bucketed scan + // as its source reports it: ungrouped, with one key's rows across splits. AQE broadcasts the + // narrow side on its runtime size, and the wide side's shuffle, now a broadcast join's probe + // side, becomes a local read. A local read reports the partitioning from *before* the shuffle, + // which is that ungrouped keyed layout. The upper join then reads two layouts that agree key + // by key while neither holds a key's whole group: consuming them as an aligned pair would zip + // split i against split i and drop every match that crosses mappers. Validation has to refuse + // the pair so that AQE keeps the full shuffle read instead. + val idCols = Array(Column.create("id", IntegerType), Column.create("data", StringType)) + createTable("lrt", idCols, Array(identity("id"))) + createTable("lrs1", idCols, Array.empty[Transform]) + createTable("lrs2", idCols, Array.empty[Transform]) + // Key 1 in two splits, so that no single mapper holds all of its rows. The padding keeps the + // relation over the runtime broadcast threshold, which is what keeps the upper join a sort + // merge join reading both lower joins. + val pad = "x" * 4096 + sql(s"INSERT INTO testcat.ns.lrt VALUES (1, 't1a$pad')") + sql(s"INSERT INTO testcat.ns.lrt VALUES (1, 't1b$pad')") + sql("INSERT INTO testcat.ns.lrt VALUES (2, 't2a')") + sql("INSERT INTO testcat.ns.lrs1 VALUES (1, 's1a'), (2, 's1b')") + sql("INSERT INTO testcat.ns.lrs2 VALUES (1, 's2a'), (2, 's2b')") + + withSQLConf( + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + // Nothing broadcasts at plan time, so both lower joins plan as sort merge joins and + // shuffle their sides; AQE broadcasts the narrow sides on their runtime size alone. + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "1k", + // Coalescing would merge the wide side's shuffle into one slot first, and a local read + // over a single slot holds every row and loses nothing. Off, the read stays one slot per + // mapper, which is the shape that reports the scan's own ungrouped layout. + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + // More reducer slots than the wide side has mappers, which is when a local read keeps one + // partition per mapper. + SQLConf.SHUFFLE_PARTITIONS.key -> "4") { + val df = sql( + """ + |SELECT substring(a.data, 1, 3), substring(b.data, 1, 3) + |FROM (SELECT t1.id AS id, t1.data AS data FROM testcat.ns.lrt t1 + | JOIN testcat.ns.lrs1 s1 ON t1.id = s1.id) a + |JOIN (SELECT t2.id AS id, t2.data AS data FROM testcat.ns.lrt t2 + | JOIN testcat.ns.lrs2 s2 ON t2.id = s2.id) b + |ON a.id = b.id + |""".stripMargin) + // Every pair of key 1's rows across the two relations, and key 2's one pair. A read that + // kept one partition per mapper and paired them index by index would find only the two + // same-split pairs of key 1 and key 2's one. + checkAnswer(df, Seq(Row("t1a", "t1a"), Row("t1a", "t1b"), Row("t1b", "t1a"), + Row("t1b", "t1b"), Row("t2a", "t2a"))) + + // The shape the answer is about, so that it is not read as an accident: the wide side's + // shuffle carries a scan that reports one key across splits, and both narrow sides ended up + // broadcast, which is what turned the wide side's shuffle into a local-read candidate. + val plan = df.queryExecution.executedPlan + val shuffles = collectAllShuffles(plan) + // The two wide sides are one scan, so they share one exchange by reuse; each narrow side + // keeps its own, read whole as its broadcast's build input. + assert(shuffles.size == 3, + s"one shared wide-side shuffle and one per narrow side:\n${plan.treeString}") + val splitScans = shuffles + .flatMap(s => keyedPartitioningsOf(Seq(s.child))).filter(!_.isGrouped) + assert(splitScans.nonEmpty, + s"test setup: a shuffled scan reports a key across splits:\n${plan.treeString}") + assert(collect(plan) { case j: BroadcastHashJoinExec => j }.size == 2, + s"test setup: both lower joins broadcast at runtime:\n${plan.treeString}") + + // The refusal this test exists for: no local read may hand out the wide side's ungrouped + // keyed layout. A local read building a broadcast reads the whole relation and reports the + // scan's own unknown partitioning, which pairs nothing and is harmless. + assert(collect(plan) { case r: AQEShuffleReadExec if r.isLocalRead => r } + .forall(r => keyedPartitioningsOf(Seq(r)).isEmpty), + s"no local read reports a keyed layout:\n${plan.treeString}") + // And nothing re-pairs a partial read: the plan holds no grouping node at all. + assert(collectAllGroupPartitions(plan).isEmpty, + s"nothing re-pairs a partial read:\n${plan.treeString}") + } + } + + test("SPARK-59671: an ungrouped sort branch leaves AQE's shuffle coalescing alone") { + // The sort branch keeps the scan's splits as they stand: the catalog hands the partition keys + // over sorted, so no node settles them and the layout stays ungrouped under the sort. An + // ordering reads the keys without pairing anything, so that layout answers for itself. + // Before this, the finished plan failed validation on it and AQE kept the aggregate + // branch's shuffle read uncoalesced: the stage-wide symptom this PR is motivated by. + val idCols = Array(Column.create("id", IntegerType), Column.create("data", StringType)) + createTable("sot", idCols, Array(identity("id"))) + // Key 1 across two splits, so the layout the sort reads is ungrouped. + sql("INSERT INTO testcat.ns.sot VALUES (2, 'o2')") + sql("INSERT INTO testcat.ns.sot VALUES (1, 'o1a')") + sql("INSERT INTO testcat.ns.sot VALUES (1, 'o1b')") + + withTable("sos") { + sql("CREATE TABLE sos (k INT) USING parquet") + sql("INSERT INTO sos VALUES (1), (1), (2), (2), (3)") + withSQLConf( + SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val df = sql( + """ + |(SELECT id FROM testcat.ns.sot ORDER BY id) + |UNION ALL + |(SELECT count(*) FROM sos GROUP BY k) + |""".stripMargin) + // The union coerces the branch outputs to bigint. + val collected = df.collect().map(_.getLong(0)).toSeq + // The ordered branch comes first and stays ordered, its key's splits apart. + assert(collected.take(3) == Seq(1L, 1L, 2L), + s"the ordered branch keeps its order:\n${df.queryExecution.executedPlan}") + assert(collected.drop(3).sorted == Seq(1L, 2L, 2L), "one row per key group") + + val plan = df.queryExecution.executedPlan + // The symptom: the aggregate's shuffle read coalesces even though the finished plan + // holds an ungrouped layout. + assert(collect(plan) { case r: AQEShuffleReadExec => r }.exists(_.hasCoalescedPartition), + s"the aggregate's shuffle read coalesces:\n$plan") + // The shape it coalesces on: the sort reads the scan's own ungrouped layout with no + // grouping node in between, which is what the old validation refused the plan for. + assert(collectScans(plan).exists(s => s.outputPartitioning match { + case k: KeyedPartitioning => !k.isGrouped + case _ => false + }), s"test setup: the sort branch keeps an ungrouped layout:\n$plan") + assert(collectAllGroupPartitions(plan).isEmpty, + s"test setup: nothing settles the sort branch's keys:\n$plan") + } + } + } + + test("SPARK-59671: an aggregate over a left outer join's spread side groups it") { + // A left outer join reports its left side verbatim, and the producer's duplication gate + // makes that side the spread one: its layout holds a key's rows across partitions and + // carries the alignment stamp. The stamp vouches for the pair the join consumed and for + // nothing else, so the aggregate above must not read the layout as clustered: it gets a + // grouping node first, or key 1's two splits come out as two groups. + val idCols = Array(Column.create("id", IntegerType), Column.create("data", StringType)) + createTable("lo1", idCols, Array(identity("id"))) + createTable("lo2", idCols, Array(identity("id"))) + // The left side holds three splits against two, so the heuristic replicates the right and + // the left spreads, the only side a LeftOuter may duplicate. + sql("INSERT INTO testcat.ns.lo1 VALUES (1, 'l1a')") + sql("INSERT INTO testcat.ns.lo1 VALUES (1, 'l1b')") + sql("INSERT INTO testcat.ns.lo1 VALUES (2, 'l2')") + sql("INSERT INTO testcat.ns.lo2 VALUES (1, 'r1'), (2, 'r2')") + + withSQLConf( + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + "spark.sql.autoBroadcastJoinThreshold" -> "-1") { + val df = sql( + s""" + |SELECT id, count(*) FROM ( + | ${selectWithMergeJoinHint("l", "r")} l.id AS id + | FROM testcat.ns.lo1 l LEFT OUTER JOIN testcat.ns.lo2 r ON l.id = r.id + |) GROUP BY id + |""".stripMargin) + checkAnswer(df, Seq(Row(1, 2L), Row(2, 1L))) + + // The shape: the join reports its spread side verbatim, stamp included, and a grouping + // node settles it before the aggregate reads it. + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val joins = collect(plan) { case j: ShuffledJoin => j } + assert(joins.size == 1, s"test setup: one storage-partitioned join:\n$plan") + assert(keyedPartitioningsOf(Seq(joins.head)).exists(k => + !k.isGrouped && k.ungroupingOrigin.contains(SPLIT_FOR_JOIN)), + s"test setup: the join reports its spread side verbatim:\n$plan") + assert(collectAllGroupPartitions(plan).exists(g => + g.child.collectFirst { case _: SortMergeJoinExec => () }.isDefined && + (g.outputPartitioning match { + case k: KeyedPartitioning => k.isGrouped && k.ungroupingOrigin.isEmpty + case _ => false + })), + s"the aggregate's side is grouped over the join:\n$plan") + } + } } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala index fb9afa9cf3f3f..7bab293edf364 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala @@ -21,7 +21,7 @@ import org.apache.spark.{SparkContext, SparkException} import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder, TransformExpression} -import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyReducer, OrderedDistribution, Partitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN, UnknownPartitioning} import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer, DaysFunctionWithToYearsReducerWithLongResult, DaysToYearsReducerWithLongResult, Reducer, YearsFunctionWithToYearsReducerWithLongResult} import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge} @@ -50,7 +50,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession { val childKp = KeyedPartitioning(Seq(exprA, exprB), keys) .withLayout(_.copy(isCollapsed = childCollapsed)) GroupPartitionsExec(DummySparkPlan(outputPartitioning = childKp), joinKeyPositions, - expected, distributePartitions = distribute) + expected, ungroupingOrigin = Option.when(distribute)(SPLIT_FOR_JOIN)) .outputPartitioning.asInstanceOf[KeyedPartitioning] } def keyOf(a: Int): InternalRowComparableWrapper = @@ -590,7 +590,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession { val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1), row(2), row(2)))) val gpe = GroupPartitionsExec(child, expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 3, keyOf(3) -> 1)), - distributePartitions = true) + ungroupingOrigin = Some(SPLIT_FOR_JOIN)) gpe.execute() assert(gpe.metrics("numInputPartitions").value === 3) @@ -601,6 +601,52 @@ class GroupPartitionsExecSuite extends SharedSparkSession { assert(!gpe.metrics.contains("numReplicatedPartitionReads"), "distribute never replicates") } + test("SPARK-59671: the output claim is the node's own, never inherited") { + // The child's origin describes the shape its producer built. A node that groups those + // partitions settles them: its output carries only the claim its own producer stamped it + // with, none here, so no stale stamp above turns away a sound pair of grouped sides. + def keyOf(a: Int): InternalRowComparableWrapper = + InternalRowComparableWrapper(row(a), Seq(exprA)) + val child = ExecutableKeyedLeaf( + KeyedPartitioning(Seq(exprA), Seq(row(1), row(1), row(2))) + .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN)))) + + val grouped = GroupPartitionsExec(child) + val groupedOut = grouped.outputPartitioning.asInstanceOf[KeyedPartitioning] + assert(groupedOut.isGrouped && groupedOut.ungroupingOrigin.isEmpty) + assert(groupedOut.partitionKeys.map(_.row) == Seq(row(1), row(2))) + + // A stamped node stamps its own output whatever the child claimed, and its routing follows + // the stamp: a repeating side groups first. + val restamped = GroupPartitionsExec(child, + expectedPartitionKeys = Some(Seq(keyOf(1) -> 2, keyOf(2) -> 1)), + ungroupingOrigin = Some(REPLICATED_FOR_JOIN)) + assert(!restamped.distributePartitions, "a repeating side groups first") + val restampedOut = restamped.outputPartitioning.asInstanceOf[KeyedPartitioning] + assert(!restampedOut.isGrouped) + assert(restampedOut.ungroupingOrigin.contains(REPLICATED_FOR_JOIN)) + } + + test("SPARK-59671: an alignment that settles every key spends its stamp") { + // A stamp says why the keys repeat on purpose; a grouping where every key got one slot + // leaves nothing repeating, so the layout reports no claim even though the node was stamped + // for one -- and it keeps the ordering claim a grouped layout has always had. + def keyOf(a: Int): InternalRowComparableWrapper = + InternalRowComparableWrapper(row(a), Seq(exprA)) + val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1), row(2)))) + val gpe = GroupPartitionsExec(child, + expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 1)), + ungroupingOrigin = Some(SPLIT_FOR_JOIN)) + assert(gpe.ungroupingOrigin.contains(SPLIT_FOR_JOIN), + "the node keeps the stamp it was handed, and the routing it derives") + val out = gpe.outputPartitioning.asInstanceOf[KeyedPartitioning] + assert(out.isGrouped && out.ungroupingOrigin.isEmpty) + withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") { + assert(out.satisfies(OrderedDistribution(Seq(SortOrder(exprA, Ascending)))), + "a grouped layout keeps its ordering claim") + } + } + test("SPARK-59310: alignment prunes unmatched keys, pads missing ones") { // The expected keys carry key 1 and a key-3 slot the child does not hold, as an inner // join's intersection combined with the other side's layout would. The 2 splits of key 2 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 6b1314fb88f1b..6d3a7708a000a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.execution.{BinaryExecNode, DummySparkPlan, LeafExecNode, SafeForKWayMerge, SortExec, UnaryExecNode} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, GroupPartitionsExec} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeAsOfJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.python.FlatMapCoGroupsInPandasExec import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf @@ -2036,6 +2036,34 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } + test("SPARK-59671: an ordering spread carries its own reason and serves only the ordering") { + val exprKey = AttributeReference("k", IntegerType)() + // The keys arrive out of order for the ascending ordering the operator reads, and key 1 holds + // two splits. The spread that settles the order keeps the splits, so the layout it reports + // stays ungrouped, on purpose and for the ordering only. + val child = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(exprKey), + Seq(InternalRow(2), InternalRow(1), InternalRow(1)))) + val ordering = OrderedDistribution(Seq(SortOrder(exprKey, Ascending))) + + withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") { + val newChild = EnsureRequirements.apply(parentRequiring(child, ordering)).children.head + val groupings = groupPartitionsNodes(newChild) + assert(groupings.size == 1, + s"the unsorted keys are spread into order, got ${newChild.getClass.getSimpleName} " + + s"with ${groupings.size} grouping nodes") + assert(groupings.head.distributePartitions) + assert(groupings.head.ungroupingOrigin == Some(SPLIT_FOR_ORDERING)) + val layout = newChild.outputPartitioning.asInstanceOf[KeyedPartitioning] + assert(!layout.isGrouped) + assert(layout.partitionKeys.map(_.row.getInt(0)) == Seq(1, 1, 2)) + // The spread serves the ordering it was built for... + assert(layout.satisfies(ordering)) + // ...and its stamp is no pairing role: the same layout must not serve a clustering. + assert(!layout.satisfies(ClusteredDistribution(Seq(exprKey)))) + } + } + test("SPARK-58968: a projection whose resulting count matches requiredNumPartitions still " + "groups") { val exprN = AttributeReference("n", IntegerType)() @@ -2597,6 +2625,71 @@ class EnsureRequirementsSuite extends SharedSparkSession { ((false, true))) } + test("SPARK-59671: the alignment stamps one repeating side, only where the join type may") { + // The replicate side is the one with fewer pre-alignment partitions (the dummy plans carry no + // `logicalLink`, which forces the fallback). A join type may duplicate one side, both, or + // neither, and the stamps follow: a side that may not repeat groups instead, and a join type + // that may repeat neither side builds no spread pair at all. + def stampsOf(smj: SparkPlan): Seq[(Option[UngroupingOrigin], Boolean)] = + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true") { + EnsureRequirements.apply(smj).children.map { child => + val gpe = child.collectFirst { case g: GroupPartitionsExec => g }.get + (gpe.ungroupingOrigin, + gpe.outputPartitioning.asInstanceOf[KeyedPartitioning].isGrouped) + } + } + def smj(joinType: JoinType, leftKeys: Seq[InternalRow], + rightKeys: Seq[InternalRow]): SortMergeJoinExec = + SortMergeJoinExec(Seq(exprA), Seq(exprB), joinType, None, + DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprA), leftKeys)), + DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprB), rightKeys))) + val threeSplits = Seq(InternalRow(1), InternalRow(1), InternalRow(1)) + val twoKeys = Seq(InternalRow(1), InternalRow(2)) + + // The right side holds fewer partitions and is picked to repeat. An inner join may duplicate + // either side and a left outer join the right one, so both get one spread and one repeating + // side. + assert(stampsOf(smj(Inner, threeSplits, twoKeys)) === + Seq((Some(SPLIT_FOR_JOIN), false), (Some(REPLICATED_FOR_JOIN), false))) + assert(stampsOf(smj(LeftOuter, threeSplits, twoKeys)) === + Seq((Some(SPLIT_FOR_JOIN), false), (Some(REPLICATED_FOR_JOIN), false))) + + // The mirrored left outer join picks the left side, which it may not duplicate: the spread + // is skipped and both sides group. + assert(stampsOf(smj(LeftOuter, Seq(InternalRow(1)), threeSplits)) === + Seq((None, true), (None, true))) + + // A full outer join may duplicate neither side: both group, whichever is picked. + assert(stampsOf(smj(FullOuter, threeSplits, twoKeys)) === + Seq((None, true), (None, true))) + } + + test("SPARK-59671: an as-of join gets no keyed alignment, so its sides shuffle") { + // `EnsureRequirements.checkKeyGroupCompatible` plans a keyed alignment for the sort-merge and + // shuffled-hash joins only, and an as-of join is a `ShuffledJoin` outside that list: its + // matches are read in the order within a partition, which neither a spread nor a repeating + // side preserves. The same shapes an inner join stamps and aligns must fall back to a + // shuffle here. + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true") { + val asOf = SortMergeAsOfJoinExec(Seq(exprA), Seq(exprB), Seq(exprA), Seq(exprB), + GreaterThan(exprA, exprB), exprA, Inner, None, + DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprA), + Seq(InternalRow(1), InternalRow(1), InternalRow(1)))), + DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprB), + Seq(InternalRow(1), InternalRow(2))))) + EnsureRequirements.apply(asOf).children.foreach { child => + assert(groupPartitionsNodes(child).isEmpty, + s"an as-of join builds no stamped side:\n$child") + assert(child.collectFirst { case s: ShuffleExchangeExec => s }.isDefined, + s"the side shuffles onto a hash layout:\n$child") + } + } + } + test("SPARK-58996: a single-child operator over a partially clustered layout still gets " + "grouped") { // A partially clustered `GroupPartitionsExec` reports a non-grouped partitioning by design, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala index 4e02a10eb41ec..b3002b4504067 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala @@ -17,12 +17,17 @@ package org.apache.spark.sql.execution.exchange -import org.apache.spark.sql.catalyst.expressions.{Ascending, SortOrder} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Literal, SortOrder} +import org.apache.spark.sql.catalyst.optimizer.BuildLeft import org.apache.spark.sql.catalyst.plans.Inner -import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, SinglePartition} -import org.apache.spark.sql.execution.SortExec -import org.apache.spark.sql.execution.joins.SortMergeJoinExec +import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HashPartitioning, KeyedPartitioning, Partitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SinglePartition, SPLIT_FOR_JOIN, UngroupingOrigin} +import org.apache.spark.sql.execution.{CoGroupExec, DummySparkPlan, SortExec, SparkPlan} +import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, ObjectType} class ValidateRequirementsSuite extends SharedSparkSession { @@ -158,4 +163,303 @@ class ValidateRequirementsSuite extends SharedSparkSession { testNestedJoin(Seq((2, 2), (1, 1)), Seq((2, 2)), Seq(5, 5, 5), false) testNestedJoin(Seq((2, 2), (1, 1)), Seq((2, 5)), Seq(5, 5, 5), false) } + + test("SPARK-59671: a co-partitioning operator judges keyed children by their pairing") { + // The sides of a storage-partitioned join aligned for skew repeat their spread keys on purpose: + // neither satisfies a clustering on its own, and each side's layout says where that ungrouping + // comes from, so the pairing is what the operator reads. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def side( + attr: AttributeReference, + role: UngroupingOrigin, + keys: Seq[InternalRow]): DummySparkPlan = + DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin = Some(role)))) + val left = side(a, SPLIT_FOR_JOIN, rows) + val right = side(b, REPLICATED_FOR_JOIN, rows) + val join = ShuffledHashJoinExec(Seq(a), Seq(b), Inner, BuildLeft, None, left, right) + assert(ValidateRequirements.validate(join), + s"a spread side and one that repeats the group pair, which is the requirement:\n$join") + + // The same key rows in another order do not pair: position by position is the contract. + val off = side(b, REPLICATED_FOR_JOIN, Seq(InternalRow(1), InternalRow(2), InternalRow(1))) + assert(!ValidateRequirements.validate(join.copy(right = off)), + "the same keys in a different order are not aligned") + + // Nor does a keyed side pair with one that never pairs: a hashed side matches neither the keys + // nor the layout of a keyed one. + val hashed = DummySparkPlan(outputPartitioning = HashPartitioning(Seq(b), 3)) + assert(!ValidateRequirements.validate(join.copy(right = hashed)), + "a keyed side does not pair with a hashed one") + + // A side that does not say why it is ungrouped is one no producer built: it is waiting for the + // node that would settle it, and a finished plan has none left to insert. + val unsettled = DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(b), rows)) + assert(!ValidateRequirements.validate(join.copy(right = unsettled)), + s"an ungrouped side is admitted only where a producer built one:\n$join") + + // Nor do two sides of the same role pair: an alignment is one of each, and two spreads leave a + // key's rows split between them with nothing repeating the group. + assert(!ValidateRequirements.validate(join.copy(right = side(b, SPLIT_FOR_JOIN, rows))), + "two sides that spread the same keys are not an alignment") + + // And a pair each side satisfies on its own is still refused when the sides do not line up, + // which is the operator's requirement: the pairing, not the per-side answer. + val leftKeys = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))) + val rightKeys = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(b), Seq(InternalRow(4), InternalRow(5), InternalRow(6)))) + assert(leftKeys.outputPartitioning.satisfies(ClusteredDistribution(Seq(a))) && + rightKeys.outputPartitioning.satisfies(ClusteredDistribution(Seq(b))), + "test setup: each side answers its own clustering") + assert(!ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(b), Inner, BuildLeft, None, leftKeys, rightKeys)), + "the sides do not line up, and nothing else the operator reads says otherwise") + } + + test("SPARK-59671: a keyed pair is judged on the layouts the plan holds") { + // Neither side's keys are deduped or re-sorted to make the pair line up: reading it through + // `createShuffleSpec` would do both, so two sides whose layouts disagree as they stand would + // answer as though they agreed. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def side(attr: AttributeReference, role: UngroupingOrigin, + keys: Seq[InternalRow]): DummySparkPlan = + DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin = Some(role)))) + val left = side(a, SPLIT_FOR_JOIN, rows) + // The same keys in another order, which a projection onto distinct sorted keys would + // normalize away. + val off = side(b, REPLICATED_FOR_JOIN, Seq(InternalRow(1), InternalRow(2), InternalRow(1))) + assert(!ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(b), Inner, BuildLeft, None, left, off)), + "the sides report different layouts, and nothing normalizes them") + + // And a count: the other side holds two partitions, and neither side has the node that would + // make the two agree. + val twoKeys = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2)))) + assert(!ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(b), Inner, BuildLeft, None, left, twoKeys)), + "a child holding three partitions does not pair with one holding two") + } + + test("SPARK-59671: a pair is judged on the key the operation clusters on") { + // The subset permission applies where the operation's keys are a subset of the source's + // partitioning keys: a side partitioned on `[a, b]` can serve an operator on `[a]`. Such a side + // is judged on its own partitions under `[a]`, which is what its own key already gives it when + // dropping `b` merges no partition, whether or not a node stands over it. No key is deduped and + // none is re-sorted, so the pairs that line up are the ones whose keys line up as reported. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val c = AttributeReference("c", IntegerType)() + val d = AttributeReference("d", IntegerType)() + withSQLConf( + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false") { + def joinOnKeys(left: SparkPlan, right: SparkPlan): SparkPlan = + ShuffledHashJoinExec(Seq(a), Seq(c), Inner, BuildLeft, None, left, right) + + // The projection's result, which is what such a plan reports from the side it grouped. + val projected = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2)))) + val projectedRight = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(c), Seq(InternalRow(1), InternalRow(2)))) + assert(ValidateRequirements.validate(joinOnKeys(projected, projectedRight)), + "the layout the grouping node leaves is what the pair is judged on") + + // The source's keys, one step earlier: the second expression is the operation's to drop, and + // what is left is the same pair, so it stands. + val source = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2, 2)))) + val sourceRight = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(c, d), Seq(InternalRow(1, 1), InternalRow(2, 2)))) + assert(ValidateRequirements.validate(joinOnKeys(source, sourceRight)), + s"the key the operation clusters on is what the side offers:\n$source") + + // And it is offered in the order it is reported, so a side whose keys run the other way is + // not aligned: a spec that sorted them would call these two the same pair. + val reversed = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a, b), Seq(InternalRow(2, 2), InternalRow(1, 1)))) + assert(!ValidateRequirements.validate(joinOnKeys(reversed, sourceRight)), + s"the keys are compared as reported, not as sorted:\n$reversed") + } + } + + test("SPARK-59671: a single clustered child still owes its own grouping") { + // The pairing stands in for the per-child check only where children pair with each other. + // An operator with a single clustered child, an aggregate over a join output say, is judged + // per side: the ungrouped keyed layout does not satisfy it until a grouping stands under. + val a = AttributeReference("a", IntegerType)() + val child = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1)))) + val parent = DummySparkPlan( + children = Seq(child), + requiredChildDistribution = Seq(ClusteredDistribution(Seq(a))), + requiredChildOrdering = Seq(Nil)) + assert(!ValidateRequirements.validate(parent), "an ungrouped child fails a lone clustered slot") + } + + test("SPARK-59671: a partially clustered pair planned by the rule passes validation") { + // Partially clustered distribution spreads a side, so its replicate side reports a + // non-grouped layout whose keys repeat on purpose, and the validator's per-side check + // refused such a pair. AQE validates a stage's whole candidate plan before accepting a + // shuffle-read change, so the plan's stage took no coalescing either. The pairing takes + // it: the spread is deliberate, and the two sides agree index by index. + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true") { + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val left = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))) + val right = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2)))) + val plan = new EnsureRequirements().apply( + SortMergeJoinExec(Seq(a), Seq(b), Inner, None, left, right)) + val join = plan.collectFirst { case j: SortMergeJoinExec => j } + .getOrElse(fail(s"expected the join back:\n${plan.treeString}")) + def groupingAt(plan: SparkPlan): Option[GroupPartitionsExec] = plan match { + case g: GroupPartitionsExec => Some(g) + case s: SortExec if !s.global => groupingAt(s.child) + case _ => None + } + val sideGroupings = join.children.flatMap(groupingAt) + assert(sideGroupings.size == 2, + s"test setup: both sides are aligned:\n${plan.treeString}") + assert(sideGroupings.exists { g => + PartitioningCollection.representativeOf(g.outputPartitioning).exists(!_.isGrouped) + }, s"test setup: a side keeps its splits, so its keys repeat:\n${plan.treeString}") + assert(ValidateRequirements.validate(plan), + s"a spread pair that agrees on its keys is accepted:\n${plan.treeString}") + } + } + + test("SPARK-59671: a collection of keyed members is judged by its pairing") { + // A side that reports several keyed alternatives (a projection over a join keeps one per join + // key column) is judged on the members the admission keeps: a side whose members all carry the + // same origin for their ungrouped keys is a pair the producer aligns, and it holds up when the + // member keyed on the join key is the one that lines up with the other side. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def keyed(attr: AttributeReference, role: UngroupingOrigin): KeyedPartitioning = + KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin = Some(role))) + def side(partitioning: Partitioning): DummySparkPlan = + DummySparkPlan(outputPartitioning = partitioning) + val bothAlternatives = side(PartitioningCollection.fromPartitionings( + Seq(keyed(a, REPLICATED_FOR_JOIN), keyed(b, REPLICATED_FOR_JOIN)))) + assert(ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(a), Inner, BuildLeft, None, bothAlternatives, side(keyed(a, SPLIT_FOR_JOIN)))), + "the member keyed on the join key pairs with the other side") + + // The pairing still has to be there: a side whose members are keyed on something else offers + // nothing to pair with. + val wrongKeys = side(PartitioningCollection.fromPartitionings( + Seq(keyed(b, REPLICATED_FOR_JOIN), keyed(b, REPLICATED_FOR_JOIN)))) + assert(!ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(a), Inner, BuildLeft, None, wrongKeys, side(keyed(a, SPLIT_FOR_JOIN)))), + "a side offering no member keyed on the join keys does not pair") + } + + test("SPARK-59671: a multi-child clustered operator is judged on its children's pairing") { + // Every operator whose children all owe a `ClusteredDistribution` reads one layout they have to + // hold together, whichever operator it is: a cogroup zips corresponding partitions, so two + // sides that each satisfy the distribution on their own are not enough. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def cogroupOf(left: SparkPlan, right: SparkPlan): SparkPlan = CoGroupExec( + (key: Any, l: Iterator[Any], r: Iterator[Any]) => Nil, + Literal(1), Literal(1), Literal(1), Seq(a), Seq(b), Seq(a), Seq(b), Nil, Nil, + AttributeReference("obj", ObjectType(classOf[AnyRef]))(), left, right) + def grouped(attr: AttributeReference, keys: Seq[Int]): DummySparkPlan = DummySparkPlan( + outputOrdering = Seq(SortOrder(attr, Ascending)), + outputPartitioning = KeyedPartitioning(Seq(attr), keys.map(InternalRow(_)))) + + // A side that does not say why it is ungrouped is one no producer built, which is the whole of + // what the validator can read: the operator kind is the producer's to check, and it is checked + // where the alignment is built. + val unsettled = DummySparkPlan( + outputOrdering = Seq(SortOrder(a, Ascending)), + outputPartitioning = KeyedPartitioning(Seq(a), rows)) + assert(!ValidateRequirements.validate(cogroupOf(unsettled, unsettled)), + "an ungrouped side no producer built is refused, whoever reads it") + + // What a cogroup shares with a join is the mutual check on the layouts its sides report. + def cogroup(aKeys: Seq[Int], bKeys: Seq[Int]): SparkPlan = + cogroupOf(grouped(a, aKeys), grouped(b, bKeys)) + assert(ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(1, 2))), + "two sides holding the same grouped layout are read as they stand") + assert(!ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(1, 2, 3))), + "a side holding three partitions does not pair with one holding two") + assert(!ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(2, 1))), + "the same keys in another order are not aligned") + } + + test("SPARK-59671: a side is judged on the member that pairs, whichever one it is") { + // A side that reports several keyed alternatives offers all of them: the one keyed on the + // join key is the one that lines up, and it need not be the first the side reports. Reading a + // single member would refuse this pair, which the planner builds whenever a projection keeps + // two key columns. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def keyed(attr: AttributeReference, role: UngroupingOrigin): KeyedPartitioning = + KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin = Some(role))) + val left = DummySparkPlan(outputPartitioning = + PartitioningCollection.fromPartitionings( + Seq(keyed(b, REPLICATED_FOR_JOIN), keyed(a, REPLICATED_FOR_JOIN)))) + val right = DummySparkPlan(outputPartitioning = keyed(a, SPLIT_FOR_JOIN)) + assert(ValidateRequirements.validate(ShuffledHashJoinExec( + Seq(a), Seq(a), Inner, BuildLeft, None, left, right)), + "the second member is the one keyed on the join key, and it pairs") + } + + test("SPARK-59671: the ordering requirement is asked on top of the pairing") { + // A pair can line up and still owe its operator an ordering: the pairing covers the + // distribution clause alone, so a sort-merge join over the aligned sides, with nothing + // ordering them, is refused. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2)) + def side(attr: AttributeReference, role: UngroupingOrigin): DummySparkPlan = + DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin = Some(role)))) + val left = side(a, SPLIT_FOR_JOIN) + val right = side(b, REPLICATED_FOR_JOIN) + assert(left.outputOrdering.isEmpty && right.outputOrdering.isEmpty, + "test setup: nothing orders the sides") + assert(!ValidateRequirements.validate( + SortMergeJoinExec(Seq(a), Seq(b), Inner, None, left, right)), + "the sides pair, and still owe the join its ordering") + } + + test("SPARK-59671: a collapsed pair is not admitted on its pairing alone") { + // A layout whose keys were collapsed (a partition standing for several of the source's) serves + // a clustering through a grouping node only where that grouping is permitted, and the + // permission is a config. Without it nothing admits such a member, and a side offering nothing + // has nothing to pair with. + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + def collapsed(attr: AttributeReference): DummySparkPlan = DummySparkPlan( + outputPartitioning = KeyedPartitioning(Seq(attr), + Seq(InternalRow(1), InternalRow(1), InternalRow(2))) + .withLayout(_.copy(isCollapsed = true))) + def pair: SparkPlan = ShuffledHashJoinExec(Seq(a), Seq(b), Inner, BuildLeft, None, + collapsed(a), collapsed(b)) + + // Whether a collapsed layout may be grouped is the producer's permission to apply, and it + // applies it where the stamp goes: the validator reads no configuration, so the same pair + // answers the same way either way. + Seq("false", "true").foreach { subset => + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> subset) { + assert(!ValidateRequirements.validate(pair), + s"a pair no producer built is refused whatever the permission says:\n$pair") + } + } + } }