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 6c585e0434323..2e7e011725469 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 @@ -438,33 +438,22 @@ case class CoalescedNullAwareHashPartitioning( * Represents a partitioning where rows are split across partitions based on transforms defined by * `expressions`. * - * == Usage Forms == - * `KeyedPartitioning` is used in two distinct forms: - * - * 1. '''As outputPartitioning''': When used as a node's output partitioning (e.g., in - * `BatchScanExec` or `GroupPartitionsExec`), the `partitionKeys` are typically in sorted order - * because data sources produce them that way and `GroupPartitionsExec` sorts while grouping. - * Sorted order is not a hard requirement, but it is a useful property: when both sides of a - * storage-partitioned join report sorted keys, `EnsureRequirements` can often match them - * without inserting an additional `GroupPartitionsExec`. The keys may no longer be sorted after - * a narrowing projection through `PartitioningPreservingUnaryExecNode` or after `UnionExec` - * concatenates its children's keys; `EnsureRequirements` reconciles both sides either via - * `GroupPartitionsExec` with `expectedPartitionKeys`, or -- when the chosen shuffle spec is a - * `KeyedShuffleSpec` and the other child's spec is not compatible with it -- by shuffling that - * other child onto these keys in the order given here. - * - * 2. '''In KeyedShuffleSpec''': When used within `KeyedShuffleSpec`, the `partitionKeys` may not - * be in sorted order, and consumers must not assume otherwise. - * * == Partition Keys == - * - `partitionKeys`: The partition keys, one per partition. May contain duplicates initially - * (ungrouped state), but becomes unique after `GroupPartitionsExec` applies grouping. + * A partition key is a property of the partition it belongs to. `partitionKeys(i)` is the constant + * value that `expressions` takes for every row in partition `i`, so there is one key per partition. + * Keys may repeat while the partitioning is ungrouped. They become unique once + * `GroupPartitionsExec` has grouped it. + * + * `EnsureRequirements` uses the keys for three things: * - * `partitionKeys` is a physical layout indexed by partition id, not a set: partition `i` holds key - * `partitionKeys(i)`. A consumer must therefore treat the given order as authoritative rather than - * re-derive one. In particular `ShuffleExchangeExec` builds its `KeyGroupedPartitioner` from this - * order, so that a side shuffled onto a `KeyedPartitioning` lands in the same partitions as the - * side that declared it. + * - Deciding whether the partitioning meets a required distribution, for instance whether a + * group-by can run on these partitions as they are. + * - Pairing up the partitions of two children that hold the same key, which is what lets a join + * read both sides without shuffling either. + * - Laying another child out on these keys, when it cannot pair the partitions up. + * `ShuffleExchangeExec` builds a `KeyGroupedPartitioner` from this order, so the shuffled side + * lands in the same partitions as the side that declared the keys. This is also why a consumer + * must keep the order given here. * * == Grouping State == * A KeyedPartitioning can be in two states: @@ -474,31 +463,81 @@ case class CoalescedNullAwareHashPartitioning( * splits for the same partition value. * * - '''Grouped''' (when `isGrouped == true`): `partitionKeys` contains only unique values, with - * each partition having a distinct key. This occurs when: (1) a data source natively produces - * unique partition keys, or (2) `GroupPartitionsExec` coalesces partitions with duplicate keys. + * each partition having a distinct key. A data source can report unique partition keys natively, + * or a `GroupPartitionsExec` can coalesce the partitions that share a key. * * == Distribution Satisfaction and Grouping == - * Besides the default `satisfies()`, `KeyedPartitioning` exposes two additional methods: + * Besides the default `satisfies()`, `KeyedPartitioning` answers three questions. They differ in + * what they let happen to the data before the distribution counts as met. * - * - `nonGroupedSatisfies()`: as-is satisfaction (without inserting `GroupPartitionsExec`). It is - * the default `Partitioning` implementation, so for a `ClusteredDistribution` it is always false. - * - `groupedSatisfies()`: whether the distribution would be satisfied once the duplicate partition - * keys are coalesced. It has two callers, and they ask different questions: - * - `EnsureRequirements` calls it on non-grouped KPs to ask whether inserting a - * `GroupPartitionsExec` would help. When it returns true, the distribution is NOT yet - * satisfied -- `EnsureRequirements` will insert one to coalesce the duplicate keys. - * - `satisfies0()` calls it on grouped KPs. Since `nonGroupedSatisfies()` is false for a - * `ClusteredDistribution`, this is the only route by which a grouped KP satisfies one, and no - * grouping is involved: the keys are already unique. - * - * That second caller is why the narrowing guard in `groupedSatisfies()` is a conjunction with - * `!isGrouped`. A narrowed KP whose projected keys stayed distinct is grouped, and dropping the - * `!isGrouped` term would stop it from satisfying a `ClusteredDistribution` and cost it a shuffle, - * even though grouping it would merge nothing. + * - `nonGroupedSatisfies()`: is the distribution met by the partitioning as it stands, with no node + * inserted? It is the default `Partitioning` implementation, so for a `ClusteredDistribution` it + * is always false. + * - `keysSatisfy()`: do the partition keys match what the distribution asks for? Duplicate keys are + * ignored, so this asks about the key expressions alone. `satisfies()` is this plus `isGrouped`, + * and nothing more. + * - `mayGroupToSatisfy()`: may `EnsureRequirements` insert a `GroupPartitionsExec` to meet the + * distribution? It is asked of non-grouped partitionings only, where such a node coalesces the + * duplicate keys. So the answer is `keysSatisfy()`, plus whether that coalescing is allowed. + * "Key Collapse" below says when it is not. * * For `OrderedDistribution`, `GroupPartitionsExec` must also sort the partition keys to meet the * ordering requirement. * + * == Key Collapse == + * Two things happen to partition keys, and only the first is a loss of granularity: + * + * - '''Key collapse''': a projection or a reduction maps two keys that were distinct onto the same + * new key. `[(1, 'a'), (1, 'b'), (2, 'c')]` projected onto the first position gives `[1, 1, 2]`, + * so three distinct keys became two. `isCollapsed` records this. + * - '''Grouping''': `GroupPartitionsExec` physically combines the partitions that share a key. + * `[1, 1, 2]` becomes `[1, 2]`. `isGrouped` says the keys are unique, not how they became unique. + * + * Nothing downstream makes a partitioning finer, so once the flag is set, every partitioning + * derived from that one inherits it. + * + * Whether a grouping needs `allowKeysSubsetOfPartitionKeys` follows from that difference: + * + * - '''Grouping without a collapse''' merges only partitions that already shared a key. A source + * reporting several splits per key produces those, and so does a union of children whose keys + * overlap. A partitioning that went through neither a projection nor a reduction is always in + * this case, so no opt-in is needed. + * - '''Grouping after a collapse''' merges partitions that the finer-grained partitioning held + * apart. The two `1` partitions above came from the different keys `(1, 'a')` and `(1, 'b')`, so + * the merged partition holds more data than any partition the source declared. This is what the + * opt-in exists to gate. + * - '''Grouping for an `OrderedDistribution`''' is not gated at all. `GroupPartitionsExec` pads + * that path out to the expected split counts instead of coalescing, so it merges nothing. + * + * A collapsed partitioning is therefore in one of two states, and only the first is gated: + * + * - '''Collapsed and ungrouped''': duplicate keys remain, so grouping them would merge partitions + * the source held apart. `mayGroupToSatisfy()` refuses a `ClusteredDistribution` unless the + * config is on. + * - '''Collapsed and grouped''': the keys are already unique, so grouping merges nothing and + * `satisfies()` accepts a `ClusteredDistribution` whatever the config says. A partitioning + * reaches this state by being grouped with the config on, or by having its keys reduced onto a + * coarser transform. + * + * A collapsed partitioning is still reported as it is. The gate above refuses one route only, and + * that route is meeting a `ClusteredDistribution` by grouping. It refuses nothing when the operator + * asks for `UnspecifiedDistribution`, nothing once the keys are unique, and nothing on the path + * that lays another child out on these keys, which never reads the flag. Which of those applies + * depends on the required distribution, and the producer of the partitioning does not know it. + * Reporting `UnknownPartitioning` would give up all of them, and make the plan shape depend on a + * config. + * + * == Key Order == + * `partitionKeys` is usually in ascending order, but nothing guarantees it, so a consumer must not + * assume it. A projection that drops key positions can leave the keys unsorted. So can a + * `UnionExec` that concatenates its children's keys. The keys inside a `KeyedShuffleSpec` are not + * sorted at all. + * + * Sorted keys are still worth having. Data sources produce them that way, and `GroupPartitionsExec` + * sorts while grouping. When both sides of a storage-partitioned join report them, + * `EnsureRequirements` can often match the two sides without inserting an additional + * `GroupPartitionsExec`. + * * == Example == * Consider a data source with partition transform `[years(ts_col)]` and 4 input splits: * @@ -508,8 +547,9 @@ case class CoalescedNullAwareHashPartitioning( * partitionKeys: [0, 1, 2, 2] // partitions 2 and 3 share a key * numPartitions: 4 * isGrouped: false - * satisfies(ClusteredDistribution(...)) == false // isGrouped guards groupedSatisfies - * groupedSatisfies(ClusteredDistribution(...)) == true // would satisfy after grouping + * keysSatisfy(ClusteredDistribution(...)) == true // the keys match + * satisfies(ClusteredDistribution(...)) == false // but they are not unique yet + * mayGroupToSatisfy(...) == true // grouping would settle it * }}} * * '''After GroupPartitionsExec''' (grouped): @@ -518,7 +558,8 @@ case class CoalescedNullAwareHashPartitioning( * partitionKeys: [0, 1, 2] // duplicates removed * numPartitions: 3 * isGrouped: true - * satisfies(ClusteredDistribution(...)) == true // satisfies now + * keysSatisfy(ClusteredDistribution(...)) == true // the keys still match + * satisfies(ClusteredDistribution(...)) == true // and now they are unique * }}} * * @param expressions Partition transform expressions (e.g., `years(col)`, `bucket(10, col)`). @@ -528,22 +569,16 @@ case class CoalescedNullAwareHashPartitioning( * guaranteed after projection. May contain duplicates when ungrouped. * @param isGrouped Whether partition keys are unique (no duplicates). Computed on first * creation, then preserved through copy operations to avoid recomputation. - * @param isNarrowed Whether this partitioning was derived from a finer-grained one by dropping key - * positions (e.g. via `PartitioningPreservingUnaryExecNode`). When true and the - * keys are no longer unique, `GroupPartitionsExec` may merge partitions that held - * distinct keys in the original partitioning, carrying the same skew risk as - * `allowKeysSubsetOfPartitionKeys`. "May", because the condition is a proxy: the - * duplicate keys can also come from a source that reports several splits per - * partition key, in which case grouping merges only same-key partitions. Such a - * partitioning can only satisfy `ClusteredDistribution` by being grouped, and - * `groupedSatisfies` refuses that unless the config is enabled, regardless of - * `requireAllClusterKeysForDistribution`. + * @param isCollapsed Whether a projection or a reduction mapped keys that were distinct in the + * partitioning this one was derived from onto the same key, so one key here can + * stand for several of the original ones. Sticky. See "Key Collapse" above for + * what it gates and how it travels. */ case class KeyedPartitioning( expressions: Seq[Expression], @transient partitionKeys: Seq[InternalRowComparableWrapper], isGrouped: Boolean, - isNarrowed: Boolean = false) extends Expression with Partitioning with Unevaluable { + isCollapsed: Boolean) extends Expression with Partitioning with Unevaluable { override val numPartitions = partitionKeys.length override def children: Seq[Expression] = expressions @@ -602,10 +637,48 @@ case class KeyedPartitioning( @transient lazy val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row) - def toGrouped: KeyedPartitioning = { - val groupedPartitionKeys = partitionKeys.distinct.sorted(keyOrdering) + /** + * Projects this partitioning onto the key positions in `positions`, whose order becomes the order + * of the projected expressions and key fields. The projected `isGrouped` and `isCollapsed` are + * computed together, because dropping a position can map keys that were distinct here onto the + * same key. Keeping every position changes no key, so both answers are inherited as they are, + * with no pass over the keys. `GroupPartitionsExec` decides the same two answers from the key + * groups it keeps, so it does not use this. + */ + def project(positions: Seq[Int]): KeyedPartitioning = { + if (positions == expressions.indices) { + this + } else { + // One pass answers both questions, walking the projected keys alongside the keys they came + // from. Two different source keys landing on one projected key is the collapse, and it is + // also what makes the projected keys non-unique, so the walk stops at the first one. The + // source keys are never hashed, only compared where a projected key repeats. + val projectedKeys = projectKeys(positions)._2 + val sourceOf = + mutable.HashMap.empty[InternalRowComparableWrapper, InternalRowComparableWrapper] + var collapses = false + val projectedIter = projectedKeys.iterator + val sourceIter = partitionKeys.iterator + while (projectedIter.hasNext && !collapses) { + val projected = projectedIter.next() + val source = sourceIter.next() + sourceOf.put(projected, source) match { + case Some(previous) if previous != source => collapses = true + case _ => + } + } + copy( + expressions = positions.map(expressions), + partitionKeys = projectedKeys, + isGrouped = !collapses && sourceOf.size == projectedKeys.length, + isCollapsed = isCollapsed || collapses) + } + } - new KeyedPartitioning(expressions, groupedPartitionKeys, isGrouped = true) + def toGrouped: KeyedPartitioning = { + // Unique keys need no dedup, only the sort. + val uniqueKeys = if (isGrouped) partitionKeys else partitionKeys.distinct + copy(partitionKeys = uniqueKeys.sorted(keyOrdering), isGrouped = true) } /** @@ -624,42 +697,31 @@ case class KeyedPartitioning( KeyedPartitioning.reduceKeys(partitionKeys, keyDataTypes, reducers) override def satisfies0(required: Distribution): Boolean = { - nonGroupedSatisfies(required) || (isGrouped && groupedSatisfies(required)) + nonGroupedSatisfies(required) || (isGrouped && keysSatisfy(required)) } def nonGroupedSatisfies(required: Distribution): Boolean = super.satisfies0(required) - def groupedSatisfies(required: Distribution): Boolean = { + /** The first of the three questions the class doc lists. */ + private def keysSatisfy(required: Distribution): Boolean = { required match { case c @ ClusteredDistribution(requiredClustering, requireAllClusterKeys, _, _) => - // Both branches below are gated by the same switch, so read it once. - val allowKeysSubsetOfPartitionKeys = - SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys - - if (isNarrowed && !isGrouped && !allowKeysSubsetOfPartitionKeys) { - // A narrowed, non-grouped partitioning carries the same skew risk as using a subset of - // partition keys for a join: GroupPartitionsExec may merge partitions that held distinct - // keys in the original finer-grained partitioning. Require the same config to opt in. - // - // Checked before the `requireAllClusterKeys` branch, because the risk is independent of - // which key sets count as matching (SPARK-58974). - // - // The `!isGrouped` term is required, not redundant: a narrowing projection whose - // projected keys stayed distinct is grouped, and `satisfies0` routes such a partitioning - // through here (`nonGroupedSatisfies` is false for a `ClusteredDistribution`). Grouping - // it would merge nothing, so refusing it would only cost a shuffle. - false - } else if (requireAllClusterKeys) { + if (requireAllClusterKeys) { // Checks whether this partitioning is partitioned on exactly same clustering keys of // `ClusteredDistribution`. c.areAllClusterKeysMatched(expressions) } else { // We'll need to find leaf attributes from the partition expressions first. - lazy val attributes = AttributeSet.fromAttributeSets(expressions.map(_.references)) - - if (allowKeysSubsetOfPartitionKeys) { - // check that operation keys (required clustering keys) - // overlap with partition keys (KeyedPartitioning attributes) + val attributes = AttributeSet.fromAttributeSets(expressions.map(_.references)) + + if (SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys) { + // The operation keys may be a subset of the partition keys, so one partition expression + // covering one of them is enough. Partitions can then still hold rows that share an + // operation key. What makes the distribution true is the projection onto the covering + // positions, plus the grouping that follows it. Given the single reference per + // expression that `supportsExpressions` enforces, the test below is exactly + // `KeyedShuffleSpec.keyPositions.exists(_.nonEmpty)`, where `createShuffleSpec` takes + // its `joinKeyPositions` from. Consolidating the two is left to a follow-up. requiredClustering.exists(x => attributes.exists(_.semanticEquals(x))) && expressions.forall(_.references.size == 1) } else { @@ -675,6 +737,20 @@ case class KeyedPartitioning( } } + /** + * The third of the three questions the class doc lists. Ask it only of a partitioning that is not + * grouped, since a grouped one has nothing to coalesce and `satisfies` is the question for it. + */ + def mayGroupToSatisfy(required: Distribution): Boolean = { + val mayCoalesce = required match { + case _: ClusteredDistribution => + !isCollapsed || SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys + case _ => true + } + // The permission is the cheap half, so it is asked first. + mayCoalesce && keysSatisfy(required) + } + override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = { val result = KeyedShuffleSpec(this, distribution) if (SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys) { @@ -682,14 +758,12 @@ case class KeyedPartitioning( // `KeyedPartitioning` grouped on the operation keys, and use that as // the returned shuffle spec. val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2) - val projectedExpressions = joinKeyPositions.map(expressions) - val projectedKeys = projectKeys(joinKeyPositions)._2 - // Sort the distinct projected keys the same way `GroupPartitionsExec` does (both sort with - // `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the keyed side is grouped - // and the other side is re-shuffled using this spec, the two `KeyedPartitioning`s carry the - // same keys in a different order and `PartitioningCollection.fromPartitionings` rejects them. - val projectedPartitioning = - new KeyedPartitioning(projectedExpressions, projectedKeys, isGrouped = false).toGrouped + // `toGrouped` sorts the keys the same way `GroupPartitionsExec` does (both sort with + // `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the keyed side is + // grouped and the other side is re-shuffled using this spec, the two `KeyedPartitioning`s + // carry the same keys in a different order and `PartitioningCollection.fromPartitionings` + // rejects them. + val projectedPartitioning = project(joinKeyPositions).toGrouped result.copy(partitioning = projectedPartitioning, joinKeyPositions = Some(joinKeyPositions)) } else { result @@ -710,7 +784,27 @@ object KeyedPartitioning { InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(dataTypes) val comparablePartitionKeys = partitionKeys.map(comparableKeyWrapperFactory) val isGrouped = comparablePartitionKeys.distinct.size == comparablePartitionKeys.size - new KeyedPartitioning(expressions, comparablePartitionKeys, isGrouped) + // Built from scratch, so it is the layout everything else is compared against. + new KeyedPartitioning(expressions, comparablePartitionKeys, isGrouped, isCollapsed = false) + } + + /** + * Concatenates partitionings that agree on their expressions, which is what a `UnionExec` does to + * its children's partitions. The result reports one key per output partition, so its keys are the + * children's keys in child order. + * + * Keys repeating across children is not a collapse. Only a child's own collapse carries over, + * since such a key still stands for several finer-grained ones in the concatenation. + */ + def concat(kps: Seq[KeyedPartitioning]): KeyedPartitioning = { + val concatenatedKeys = kps.flatMap(_.partitionKeys) + kps.head.copy( + partitionKeys = concatenatedKeys, + // A child that has duplicates of its own puts them in the concatenation too, which answers + // this without walking the keys. + isGrouped = kps.forall(_.isGrouped) && + concatenatedKeys.distinct.length == concatenatedKeys.length, + isCollapsed = kps.exists(_.isCollapsed)) } def supportsExpressions(expressions: Seq[Expression]): Boolean = { @@ -873,11 +967,23 @@ case class RangePartitioning(ordering: Seq[SortOrder], numPartitions: Int) * in this collection do not need to be equivalent, which is useful for * Outer Join operators. * - * [[KeyedPartitioning]]s within a `PartitioningCollection` describe the same physical partitioning - * and so must share the same `partitionKeys` reference, differing only in their `expressions` (with - * matching arity). Use [[PartitioningCollection.fromPartitionings]] to build a collection from - * independently-computed partitionings (e.g. join `outputPartitioning`); it interns `partitionKeys` - * references (including across nested collections) so the invariant holds. + * [[KeyedPartitioning]]s within a `PartitioningCollection` describe the same physical partitioning. + * The constructor therefore requires all of them to share the same `partitionKeys` reference and + * `isCollapsed` flag, and to have matching expression arity. Only their `expressions` differ. + * + * Use [[PartitioningCollection.fromPartitionings]] to build one from independently-computed + * partitionings, such as a join's `outputPartitioning`. Its inputs need not agree on `isCollapsed`. + * Each member carries the history of the child it came from, so one side can have collapsed its + * keys in a projection while the other reports what its source declared. `fromPartitionings` ORs + * the flags, including across nested collections, which is right because the members name one + * shared layout, and one coarse member makes that layout coarse. Uniformity matters because + * consumers read the flag off a single member. `satisfies0` and `EnsureRequirements` accept when + * any one member satisfies the distribution, so a member that under-reported the collapse would let + * the gate through. + * + * The key lists are required rather than reconciled. `fromPartitionings` interns the reference when + * they are structurally equal, since members whose keys differ describe different layouts, which + * one collection cannot stand for. */ case class PartitioningCollection(partitionings: Seq[Partitioning]) extends Expression with Partitioning with Unevaluable { @@ -891,16 +997,11 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) /** * First [[KeyedPartitioning]] reachable from this collection through direct members or nested * collections, if any. Since every collection validates the invariant on construction, this - * single representative stands for all [[KeyedPartitioning]]s in the subtree: they all share - * its `partitionKeys` reference and expression arity. The invariant check forces this lazy val - * during construction, so it is only recomputed after deserialization. + * single representative stands for all [[KeyedPartitioning]]s in the subtree. The invariant check + * forces this lazy val during construction, so it is only recomputed after deserialization. */ @transient private[physical] lazy val firstKeyedPartitioning: Option[KeyedPartitioning] = - partitionings.view.map { - case k: KeyedPartitioning => Some(k) - case pc: PartitioningCollection => pc.firstKeyedPartitioning - case _ => None - }.collectFirst { case Some(k) => k } + partitionings.view.flatMap(PartitioningCollection.representativeOf).headOption /** * Nested collections already enforced the invariant on their own construction, so comparing one @@ -911,19 +1012,16 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) */ private def checkKeyedPartitioningInvariant(): Unit = { firstKeyedPartitioning.foreach { first => - partitionings.foreach { p => - val representative = p match { - case k: KeyedPartitioning => k - case pc: PartitioningCollection => pc.firstKeyedPartitioning.orNull - case _ => null - } - if (representative != null && (representative ne first)) { - require(representative.expressions.length == first.expressions.length, + partitionings.iterator.flatMap(PartitioningCollection.representativeOf).foreach { rep => + if (rep ne first) { + require(rep.expressions.length == first.expressions.length, "All KeyedPartitionings in a PartitioningCollection must have matching expression " + "arity") - require(representative.partitionKeys eq first.partitionKeys, + require(rep.partitionKeys eq first.partitionKeys, "All KeyedPartitionings in a PartitioningCollection must share the same " + "partitionKeys reference") + require(rep.isCollapsed == first.isCollapsed, + "All KeyedPartitionings in a PartitioningCollection must agree on isCollapsed") } } } @@ -961,6 +1059,16 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) } object PartitioningCollection { + /** + * One [[KeyedPartitioning]] standing for every one in this partitioning, if there is any. By the + * invariant in the class doc, any of them describes the layout. + */ + private[physical] def representativeOf(p: Partitioning): Option[KeyedPartitioning] = p match { + case k: KeyedPartitioning => Some(k) + case pc: PartitioningCollection => pc.firstKeyedPartitioning + case _ => None + } + /** * Builds a [[PartitioningCollection]], unifying the `partitionKeys` reference across all * [[KeyedPartitioning]]s (including those in nested collections). Use this when combining @@ -970,38 +1078,33 @@ object PartitioningCollection { * Note: this can't be implemented with `TreeNode.transform`. */ def fromPartitionings(partitionings: Seq[Partitioning]): PartitioningCollection = { + // See the class doc for why the flag is normalized by OR rather than required to agree. One + // representative per member is enough, because every collection agrees on the flag internally + // by this same construction, and only a member that disagrees is rebuilt. + val anyCollapsed = partitionings.exists(representativeOf(_).exists(_.isCollapsed)) + var canonicalKeys: Seq[InternalRowComparableWrapper] = null - def intern(p: Partitioning): Partitioning = p match { - case k: KeyedPartitioning => - if (canonicalKeys == null) { - canonicalKeys = k.partitionKeys - k - } else if (k.partitionKeys ne canonicalKeys) { - require(k.partitionKeys == canonicalKeys, - "All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys") - k.copy(partitionKeys = canonicalKeys) + // A partitioning with no `KeyedPartitioning` in it has nothing to normalize, and one that + // already agrees on both the keys and the flag 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. + def intern(p: Partitioning): Partitioning = representativeOf(p) match { + case None => p + case Some(representative) => + if (canonicalKeys == null) canonicalKeys = representative.partitionKeys + if ((representative.partitionKeys eq canonicalKeys) && + representative.isCollapsed == anyCollapsed) { + p } else { - k - } - case pc: PartitioningCollection => - pc.firstKeyedPartitioning match { - // No KeyedPartitioning anywhere in this subtree: nothing to intern. Returning the - // collection as-is keeps repeated outputPartitioning computations over deeply nested - // collections (e.g. chains of same-key joins) O(1) per level. - case None => pc - case Some(representative) if canonicalKeys == null => - canonicalKeys = representative.partitionKeys - pc - // The collection's own invariant guarantees all its KeyedPartitionings share the - // representative's `partitionKeys` reference, so reference-equality of the - // representative's keys means the whole subtree is already interned. - case Some(representative) if representative.partitionKeys eq canonicalKeys => pc - case Some(representative) => - require(representative.partitionKeys == canonicalKeys, - "All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys") - new PartitioningCollection(pc.partitionings.map(intern)) + require(representative.partitionKeys == canonicalKeys, + "All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys") + p match { + case keyed: KeyedPartitioning => + keyed.copy(partitionKeys = canonicalKeys, isCollapsed = anyCollapsed) + case pc: PartitioningCollection => + new PartitioningCollection(pc.partitionings.map(intern)) + } } - case other => other } new PartitioningCollection(partitionings.map(intern)) } @@ -1492,7 +1595,10 @@ case class KeyedShuffleSpec( te.copy(children = te.children.map(_ => clustering(positionSet.head))) case (_, positionSet) => clustering(positionSet.head) } - KeyedPartitioning(newExpressions, partitioning.partitionKeys, partitioning.isGrouped) + // The shuffled side is laid out on this side's partition keys, so it inherits the flag. That + // is conservative rather than strictly true, and it can only ever add a shuffle: a later + // grouping of the shared key set carries the collapsed side's risk. + partitioning.copy(expressions = newExpressions) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 2e866508c8675..8bfa055aa4518 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2515,9 +2515,11 @@ object SQLConf { "tables. At planning time, Spark will group the partitions by only those keys that are " + "in the operation's keys. That is currently enabled only if " + s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key} is false. This config also gates " + - "grouping a partitioning that was narrowed to a subset of its keys and whose keys are no " + - "longer distinct, which carries the same risk of skew; that applies regardless of " + - s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}." + "grouping a partitioning whose keys collapsed, that is, where a projection or a " + + "reduction mapped keys that were distinct in the source onto the same key, so that " + + "grouping them would produce a partition larger than any the source declared. That " + + s"applies regardless of ${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}. It does not " + + "apply to duplicate keys the source itself reported, which are grouped without this config." ) .version("4.0.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala index 2e9e1270b3fc5..5c3afe3072137 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala @@ -401,11 +401,11 @@ class DistributionSuite extends SparkFunSuite { assert(!nonGroupedKP.isGrouped) // satisfies() must return false: the partitions are not yet grouped. checkSatisfied(nonGroupedKP, ClusteredDistribution(Seq(x)), false) - // groupedSatisfies() returns true: it CAN satisfy once GroupPartitionsExec groups them. - assert(nonGroupedKP.groupedSatisfies(ClusteredDistribution(Seq(x)))) + // mayGroupToSatisfy() returns true, because grouping them makes it satisfy. + assert(nonGroupedKP.mayGroupToSatisfy(ClusteredDistribution(Seq(x)))) // Grouped: all distinct keys, so isGrouped=true and satisfies() delegates to - // groupedSatisfies(). + // keysSatisfy(). val groupedKP = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2), InternalRow(3))) assert(groupedKP.isGrouped) checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true) @@ -466,4 +466,17 @@ class DistributionSuite extends SparkFunSuite { } assert(arityMismatch.getMessage.contains("matching expression arity")) } + + test("SPARK-59057: toGrouped and KeyedShuffleSpec.createPartitioning keep isCollapsed sticky") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + + val collapsedKP = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(1), InternalRow(2))) + .copy(isCollapsed = true) + assert(collapsedKP.toGrouped.isCollapsed, "toGrouped must keep isCollapsed sticky") + + val spec = KeyedShuffleSpec(collapsedKP, ClusteredDistribution(Seq(x))) + val created = spec.createPartitioning(Seq(y)).asInstanceOf[KeyedPartitioning] + assert(created.isCollapsed, "createPartitioning must keep isCollapsed sticky") + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala index e5d293f6e2193..6fca19fbf0a90 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala @@ -131,18 +131,10 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode if (projectablePositions.isEmpty) return LazyList.empty - // All input KPs share the same partitionKeys by invariant; use the first as the key source. - val keySource = kps.head - val sharedKeys = - if (projectablePositions.length == numPositions) keySource.partitionKeys - else keySource.projectKeys(projectablePositions)._2 - - val isGrouped = sharedKeys.distinct.size == sharedKeys.size - // A KP is narrowed if this node drops positions, or if the input KPs were already narrowed - // (i.e. came from a finer-grained partitioning). The flag must be sticky: a subsequent - // PartitioningPreservingUnaryExecNode that passes all positions through would otherwise - // recompute isNarrowed=false, silently dropping the protection. - val isNarrowed = projectablePositions.length < numPositions || keySource.isNarrowed + // All input KPs share the same partitionKeys and isCollapsed flag by invariant, so the first + // one projects the keys and both flags for every combination below. Only the expressions + // differ. + val projected = kps.head.project(projectablePositions) // Cross-product the per-position alternatives to produce all concrete KPs. // Note: generateCartesianProduct expects thunks () => Seq[T], but wrapping LazyLists in thunks @@ -151,8 +143,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode // so all cross-product combinations are distinct by construction. MultiTransform.generateCartesianProduct( projectablePositions.map(i => () => alternativesPerPosition(i))) - .map(projectedExprs => - new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isNarrowed)) + .map(projectedExprs => projected.copy(expressions = projectedExprs)) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index bea86501e6f3a..f0186af264d90 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -986,11 +986,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // The `KeyedPartitioning`s must agree on the partition expressions to merge. val compatible = kps.forall(comparePartitioning(_, headKp)) if (compatible) { - val mergedKeys = kps.flatMap(_.partitionKeys) - val mergedExpressions = headKp.expressions - val isGrouped = mergedKeys.distinct.size == mergedKeys.size - val isNarrowed = kps.exists(_.isNarrowed) - return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed) + return KeyedPartitioning.concat(kps) } else { return super.outputPartitioning } 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 22e7649886edd..df9a2f5a0a905 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 @@ -68,21 +68,20 @@ case class GroupPartitionsExec( child.outputPartitioning match { case p: Partitioning with Expression => // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they - // can only differ in `expressions`; their `partitionKeys` reference is shared (enforced by - // `PartitioningCollection`), so `groupedPartitions` is computed only once. - val partitionKeys = groupedPartitions.map(_._1) + // can only differ in `expressions`. Their `partitionKeys` reference and `isCollapsed` flag + // are shared (enforced by `PartitioningCollection`), so the grouping is computed once. + val partitionKeys = grouping.partitions.map(_._1) p.transform { case k: KeyedPartitioning => val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions)) - KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped = isGrouped) + KeyedPartitioning( + projectedExpressions, partitionKeys, grouping.isGrouped, grouping.isCollapsed) }.asInstanceOf[Partitioning] case o => o } } - /** - * Aligns partitions based on `expectedPartitionKeys` and clustering mode. - */ + /** Aligns partitions based on `expectedPartitionKeys` and clustering mode. */ private def alignToExpectedKeys(keyMap: Map[InternalRowComparableWrapper, Seq[Int]]) = { var isGrouped = true val alignedPartitions = expectedPartitionKeys.get.flatMap { case (key, numSplits) => @@ -119,15 +118,14 @@ case class GroupPartitionsExec( * 3. Grouping input partition indices by their (possibly projected/reduced) keys * 4. Sorting or distributing based on whether partial clustering is enabled * - * Returns a tuple of (partitions, isGrouped) where: - * - partitions: sequence of (partitionKey, inputPartitionIndices) pairs representing - * how input partitions should be grouped together - * - isGrouped: whether the output partitioning is grouped (no duplicates in partition keys) + * `isCollapsed` says whether the output stands for more than one of the child's partition keys. + * Two things set it: the child's own flag, and a merge this node performs. */ - @transient private lazy val groupedPartitionsTuple = { - // There must be a `KeyedPartitioning` in child's output partitioning as a - // `GroupPartitionsExec` node is added to a plan only in that case. - val keyedPartitioning = child.outputPartitioning + @transient private lazy val grouping: PartitionGrouping = { + // There must be a `KeyedPartitioning` in the child's output partitioning, as a + // `GroupPartitionsExec` node is added to a plan only in that case. Any member will do, see + // `outputPartitioning` above. + val childKp = child.outputPartitioning .asInstanceOf[Partitioning with Expression] .collectFirst { case k: KeyedPartitioning => k } .getOrElse( @@ -136,8 +134,8 @@ case class GroupPartitionsExec( // Project partition keys if join key positions are specified val (projectedDataTypes, projectedKeys) = joinKeyPositions.fold( - (keyedPartitioning.keyDataTypes, keyedPartitioning.partitionKeys) - )(keyedPartitioning.projectKeys) + (childKp.keyDataTypes, childKp.partitionKeys) + )(childKp.projectKeys) // Reduce keys if reducers are specified val (reducedDataTypes, reducedKeys) = reducers.fold((projectedDataTypes, projectedKeys))( @@ -145,17 +143,38 @@ case class GroupPartitionsExec( val keyToPartitionIndices = reducedKeys.zipWithIndex.groupMap(_._1)(_._2) - if (expectedPartitionKeys.isDefined) { + val (partitions, isGrouped) = if (expectedPartitionKeys.isDefined) { alignToExpectedKeys(keyToPartitionIndices) } else { (groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes), true) } + + // Both cheap terms come first, so the scan below runs only where a merge is possible. A + // grouping that left the keys as they are groups the child's own key values, and one of those + // groups can only ever cover the one key it was built from. + val keysChanged = + joinKeyPositions.exists(_.length < childKp.expressions.length) || reducers.isDefined + val isCollapsed = childKp.isCollapsed || keysChanged && { + // The groups this node keeps are the ones that can merge keys of the child, and asking the + // child's keys rather than its partitions is what tells such a merge from a source that + // reports several splits per key. + val keptGroups = expectedPartitionKeys match { + case Some(expected) => expected.view.flatMap { case (key, _) => + keyToPartitionIndices.get(key) + } + case None => keyToPartitionIndices.values.view + } + val childKeys = childKp.partitionKeys.toArray + keptGroups.exists { group => + val first = childKeys(group.head) + group.tail.exists(childKeys(_) != first) + } + } + PartitionGrouping(partitions, isGrouped, isCollapsed) } @transient lazy val groupedPartitions: Seq[(InternalRowComparableWrapper, Seq[Int])] = - groupedPartitionsTuple._1 - - @transient lazy val isGrouped: Boolean = groupedPartitionsTuple._2 + grouping.partitions @transient private lazy val hasCoalescing: Boolean = groupedPartitions.exists(_._2.size > 1) @@ -335,6 +354,12 @@ case class GroupPartitionsExec( } } +/** What a [[GroupPartitionsExec]] computes once and reports from several members. */ +private case class PartitionGrouping( + partitions: Seq[(InternalRowComparableWrapper, Seq[Int])], + isGrouped: Boolean, + isCollapsed: Boolean) + /** * A PartitionCoalescer that groups partitions according to a pre-computed grouping plan. * 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 36857539fb78c..e5cc7b285cdcf 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 @@ -73,20 +73,20 @@ case class EnsureRequirements( child } else { // Check KeyedPartitioning satisfaction conditions - val groupedSatisfies = grouped.find(_.satisfies(distribution)) + val satisfyingGrouped = grouped.find(_.satisfies(distribution)) val nonGroupedSatisfiesAsIs = nonGrouped.exists(_.nonGroupedSatisfies(distribution)) - val nonGroupedSatisfiesWhenGrouped = nonGrouped.find(_.groupedSatisfies(distribution)) + val groupableNonGrouped = nonGrouped.find(_.mayGroupToSatisfy(distribution)) // Check if any KeyedPartitioning satisfies the distribution - if (groupedSatisfies.isDefined || nonGroupedSatisfiesAsIs - || nonGroupedSatisfiesWhenGrouped.isDefined) { + if (satisfyingGrouped.isDefined || nonGroupedSatisfiesAsIs + || groupableNonGrouped.isDefined) { distribution match { case o: OrderedDistribution => // OrderedDistribution requires grouped KeyedPartitioning with sorted keys // according to the distribution's ordering. - // Find any KeyedPartitioning that satisfies via groupedSatisfies. + // Find any KeyedPartitioning that satisfies, grouped or groupable. val satisfyingKeyedPartitioning = - groupedSatisfies.orElse(nonGroupedSatisfiesWhenGrouped).get + satisfyingGrouped.orElse(groupableNonGrouped).get // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees // one attribute per partition expression. val attrs = satisfyingKeyedPartitioning.expressions.flatMap(_.references) @@ -107,7 +107,7 @@ case class EnsureRequirements( ) } - case _ if groupedSatisfies.isDefined => + case _ if satisfyingGrouped.isDefined => // Grouped KeyedPartitioning already satisfies child diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala index 11e7461bc7f14..bdd108cc9d52a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -594,7 +594,7 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { val partitioning = KeyedPartitioning( Seq(partAttr), Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))), - isGrouped = false) + isGrouped = false, isCollapsed = false) def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = { val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala index 98ed8bf0b3d0d..ac06262047863 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala @@ -53,8 +53,8 @@ abstract class DistributionAndOrderingSuiteBase plan: QueryPlan[T]): Partitioning = partitioning match { case HashPartitioning(exprs, numPartitions) => HashPartitioning(exprs.map(resolveAttrs(_, plan)), numPartitions) - case KeyedPartitioning(expressions, partitionKeys, isGrouped, _) => - KeyedPartitioning(expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped) + case kp: KeyedPartitioning => + kp.copy(expressions = kp.expressions.map(resolveAttrs(_, plan))) case PartitioningCollection(partitionings) => PartitioningCollection(partitionings.map(resolvePartitioning(_, plan))) case RangePartitioning(ordering, numPartitions) => 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 9c2aa8afab5e7..fc13372bba096 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 @@ -607,6 +607,14 @@ class KeyGroupedPartitioningSuite } } + /** Every `KeyedPartitioning` these nodes report, flattening partitioning collections. */ + protected def keyedPartitioningsOf( + nodes: Seq[SparkPlan]): Seq[physical.KeyedPartitioning] = { + nodes.map(_.outputPartitioning) + .flatMap(physical.PartitioningCollection.flatten) + .collect { case kp: physical.KeyedPartitioning => kp } + } + test("partitioned join: exact distribution (same number of buckets) from both sides") { val customers_partitions = Array(bucket(4, "customer_id")) val orders_partitions = Array(bucket(4, "customer_id")) @@ -4243,7 +4251,7 @@ class KeyGroupedPartitioningSuite test("SPARK-46367: narrowing projection requires allowKeysSubsetOfPartitionKeys") { // items is partitioned by (id, name). The subquery projects away 'name', narrowing - // KeyedPartitioning([id, name]) -> KeyedPartitioning([id]) with isNarrowed=true. + // KeyedPartitioning([id, name]) -> KeyedPartitioning([id]) with isCollapsed=true. // Because id=1 maps to two original partitions ("aa" and "bb"), isGrouped=false. // GroupPartitionsExec would merge them, carrying the same skew risk as subset partition // keys -- so SPJ requires allowKeysSubsetOfPartitionKeys to be enabled. @@ -4290,9 +4298,9 @@ class KeyGroupedPartitioningSuite test("SPARK-46367: narrowing projection with distinct projected keys does not require " + "allowKeysSubsetOfPartitionKeys") { // items is partitioned by (id, name) but each id value is unique, so projecting away 'name' - // produces KeyedPartitioning([id]) with isNarrowed=true but isGrouped=true. - // Because no two original partitions share the same projected key, GroupPartitionsExec does not - // merge any partitions -- no skew risk -- so SPJ works without config. + // produces KeyedPartitioning([id]) with isGrouped=true and isCollapsed=false. No two original + // partitions share the same projected key, so the projection lost no distinct key and + // GroupPartitionsExec would merge nothing. There is no skew risk, so SPJ works without config. val items_partitions = Array(identity("id"), identity("name")) createTable(items, itemsColumns, items_partitions) sql(s"INSERT INTO testcat.ns.$items VALUES " + @@ -4321,7 +4329,7 @@ class KeyGroupedPartitioningSuite val shuffles = collectShuffles(df.queryExecution.executedPlan) assert(shuffles.isEmpty, - "should not add shuffle: narrowed KP remains grouped so no skew risk") + "should not add shuffle: the projected KP stays grouped, so there is no skew risk") checkAnswer(df, Seq(Row(1, 42.0f), Row(2, 11.0f), Row(3, 19.5f))) } @@ -4331,7 +4339,7 @@ class KeyGroupedPartitioningSuite "with allowKeysSubsetOfPartitionKeys") { // Table partitioned by (id, name): id=1 maps to two distinct partition keys (1,'aa') and // (1,'bb'). The partial HashAggregate (a PartitioningPreservingUnaryExecNode) projects away - // 'name', narrowing the KP from [id,name] to KP([id], isNarrowed=true, isGrouped=false). + // 'name', collapsing KP([id,name]) to KP([id], isCollapsed=true, isGrouped=false). // By default a shuffle is required; with allowKeysSubsetOfPartitionKeys enabled, // EnsureRequirements inserts GroupPartitionsExec to coalesce both id=1 partitions so the final // aggregate sees all id=1 partial results in one task -- correct and shuffle-free. @@ -4367,7 +4375,7 @@ class KeyGroupedPartitioningSuite "with allowKeysSubsetOfPartitionKeys") { // Same narrowing mechanism as the aggregate test: the partial HashAggregate (a // PartitioningPreservingUnaryExecNode) for the inner GROUP BY id, price projects away 'name', - // narrowing KP([id,name]) to KP([id], isNarrowed=true, isGrouped=false). With + // collapsing KP([id,name]) to KP([id], isCollapsed=true, isGrouped=false). With // allowKeysSubsetOfPartitionKeys enabled, EnsureRequirements inserts GroupPartitionsExec to // coalesce both id=1 partitions for the final aggregate. The window PARTITION BY id then sees // KP([id], isGrouped=true) from the aggregate output and needs no further exchange. @@ -4869,14 +4877,14 @@ class KeyGroupedPartitioningSuite checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5))) } - test("SPARK-58974: the narrowing skew guard applies regardless of requireAllClusterKeys") { + test("SPARK-58974: the collapse skew guard applies regardless of requireAllClusterKeys") { // Same shape as "SPARK-46367: narrowing projection requires allowKeysSubsetOfPartitionKeys": // items is partitioned by (id, name) and id=1 maps to two of its partitions, so projecting // `name` away narrows [id, name] to [id] and collapses the keys to [1, 1, 2]. // - // The narrowing guard describes a skew risk that does not depend on which key sets count as + // The collapse skew guard describes a risk that does not depend on which key sets count as // matching, so it must apply for either value of `requireAllClusterKeys`. It used to sit inside - // the `requireAllClusterKeys = false` arm of `groupedSatisfies`, so with that setting enabled + // the `requireAllClusterKeys = false` arm of the key matching, so with that setting enabled // such a partitioning was grouped anyway -- taking exactly the exposure // `allowKeysSubsetOfPartitionKeys` exists to gate, with nobody opting in. createTable(items, itemsColumns, Array(identity("id"), identity("name"))) @@ -4915,11 +4923,10 @@ class KeyGroupedPartitioningSuite val df = sql(query) val plan = df.queryExecution.executedPlan - val narrowed = collect(plan) { case p: ProjectExec => p } - .map(_.outputPartitioning) - .collect { case kp: physical.KeyedPartitioning if kp.isNarrowed => kp } - assert(narrowed.exists(!_.isGrouped), - "this test needs a narrowed, ungrouped partitioning to be meaningful") + val collapsed = keyedPartitioningsOf(collect(plan) { case p: ProjectExec => p }) + .filter(_.isCollapsed) + assert(collapsed.exists(!_.isGrouped), + "this test needs a collapsed, ungrouped partitioning to be meaningful") // Count over the whole plan, so a shuffle or a grouping anywhere in it is caught, not // only the ones inside the join subtree. @@ -4932,9 +4939,9 @@ class KeyGroupedPartitioningSuite assert(shuffles.isEmpty, s"$settings: the opt-in must avoid the shuffles") } else { // `requireAllClusterKeys` says which key sets count as matching; it does not authorise - // coalescing a narrowed partitioning, so the guard behaves the same either way. + // coalescing a collapsed partitioning, so the guard behaves the same either way. assert(groupPartitions.isEmpty, - s"$settings must not coalesce a narrowed partitioning without " + + s"$settings must not coalesce a collapsed partitioning without " + "allowKeysSubsetOfPartitionKeys") // Both sides are shuffled, unless `v2BucketingShuffleEnabled` lets the refused side be // laid out on the other side's declared partition keys. @@ -5041,6 +5048,253 @@ class KeyGroupedPartitioningSuite } } } + test("SPARK-59057: several splits per partition key are grouped without " + + "allowKeysSubsetOfPartitionKeys") { + // The scan reports one partition key per split, so a table with two splits for the same + // (id, dept) value already has duplicate keys before any projection. Dropping dept maps + // (1, 'x'), (1, 'x') onto 1 and (2, 'y') onto 2: two distinct keys before, two after, so the + // projection collapsed nothing. Grouping merges only the two splits that already shared a key, + // which is what happens for any partitioning that never went through a projection, so it must + // not need the opt-in. + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + val t2cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1'), (1, 'x', 'a2'), (2, 'y', 'a3')") + createTable("t2", t2cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (1, 'b1'), (2, 'b2')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + // `dept RLIKE ...` has no V2 translation, so the scan must output dept and the Project + // above the Filter is what drops it. + val df = sql( + """SELECT /*+ MERGE(u, t2) */ u.id, t2.data + |FROM (SELECT id FROM testcat.ns.t1 WHERE dept RLIKE 'x|y') u + |JOIN testcat.ns.t2 ON u.id = t2.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val projected = keyedPartitioningsOf(collect(plan) { case p: ProjectExec => p }) + assert(projected.exists(kp => !kp.isGrouped && !kp.isCollapsed), + "this test needs an ungrouped partitioning that the projection did not collapse") + + assert(collectAllGroupPartitions(plan).nonEmpty, + "the duplicate splits must be grouped, no opt-in needed") + assert(collectAllShuffles(plan).isEmpty, "grouping must replace the shuffles") + checkAnswer(df, Seq(Row(1, "b1"), Row(1, "b1"), Row(2, "b2"))) + } + } + } + + test("SPARK-59057: filtering partition keys out is not a key collapse") { + // With partition filtering on, an inner join plans both sides on the intersection of their + // partition keys, so a side that had more keys ends up with fewer than it reported. That is + // pruning rather than merging, and counting it as a collapse would make the sticky flag refuse + // grouping further up the plan. Note that these nodes neither project nor reduce, so this pins + // the gate that keeps a pruning-only node from reporting a collapse at all. The node that does + // prune a collapsed key is the next test. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2", "t3") { + createTable("t1", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2'), (3, 'a3')") + createTable("t2", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (1, 'b1'), (2, 'b2')") + createTable("t3", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + val df = sql( + """SELECT id, count(*) AS cnt FROM ( + | SELECT /*+ MERGE(t1, t2) */ t1.id FROM testcat.ns.t1 JOIN testcat.ns.t2 + | ON t1.id = t2.id + | UNION ALL + | SELECT id FROM testcat.ns.t3 + |) GROUP BY id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val groupPartitions = collectAllGroupPartitions(plan) + val keyed = keyedPartitioningsOf(groupPartitions) + assert(keyed.nonEmpty, "the join must be planned as a storage-partitioned join") + assert(keyed.forall(!_.isCollapsed), + "key 3 was filtered out, not merged into another partition") + assert(collectAllShuffles(plan).isEmpty, + "nothing collapsed, so the aggregate above the union must not need a shuffle") + checkAnswer(df, Seq(Row(1, 2), Row(2, 2))) + } + } + } + + test("SPARK-59057: a shuffle built from a template that collapsed nothing needs no opt-in") { + // The chain apache#58316 built for the shuffle-template path, with the expectation the collapse + // semantics call for. items is partitioned by (id, name) with *unique* ids, so projecting + // `name` away drops a key position but loses no distinct key, so nothing collapsed. purchases + // reports no partitioning, so with v2BucketingShuffleEnabled its side is shuffled using the + // projected partitioning as the template. A RIGHT OUTER join then exposes only that shuffled + // side, and the union with t3 (overlapping keys) makes the merged partitioning ungrouped. The + // duplicate keys there come from the two children holding the same ids, not from the + // projection, so the final aggregate may group them with the opt-in off. + createTable(items, itemsColumns, Array(identity("id"), identity("name"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array.empty) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + createTable("t3", Array(Column.create("id", LongType)), Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1), (2)") + + val query = + s""" + |SELECT id, COUNT(*) AS cnt FROM ( + | ${selectWithMergeJoinHint("sub", "p")} + | p.item_id AS id + | FROM (SELECT id FROM testcat.ns.$items WHERE name >= 'aa') sub + | RIGHT OUTER JOIN testcat.ns.$purchases p + | ON sub.id = p.item_id + | UNION ALL + | SELECT id FROM testcat.ns.t3 + |) GROUP BY id + |""".stripMargin + + Seq(false, true).foreach { allowSubset => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> allowSubset.toString) { + val plan = sql(query).queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(!kp.isGrouped, "keys 1 and 2 repeat across the union children") + assert(!kp.isCollapsed, + "the ids were unique, so dropping `name` collapsed nothing to carry down this chain") + assert(collectAllGroupPartitions(plan).nonEmpty, + s"allowSubset=$allowSubset: grouping merges only partitions that already shared a key") + assert(collectAllShuffles(plan).size == 1, + s"allowSubset=$allowSubset: only the purchases side shuffles") + checkAnswer(sql(query), Seq(Row(1L, 2L), Row(2L, 2L), Row(3L, 1L))) + } + } + } + + test("SPARK-59057: a collapse is reported when the splits are distributed, not replicated") { + // Under partially clustered distribution `GroupPartitionsExec` spreads a key's splits over one + // partition each instead of replicating them, so the partitions it emits never hold more than + // one split. The collapse has to be read off the key groups it keeps, or this whole mode + // silently reports nothing collapsed. Here dropping `name` maps the distinct keys (1, 'aa') + // and (1, 'bb') onto key 1, which is exactly the state the gate exists for. + createTable(items, itemsColumns, Array(identity("id"), identity("name"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'bb', 41.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'cc', 10.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(identity("item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-02' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, i.name, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val distributing = collectAllGroupPartitions(plan).filter(_.distributePartitions) + assert(distributing.nonEmpty, "this test needs a distributing GroupPartitionsExec") + val keyed = keyedPartitioningsOf(distributing) + assert(keyed.nonEmpty && keyed.forall(_.isCollapsed), + "dropping `name` merged two distinct keys, however the splits are laid out afterwards") + checkAnswer(df, Seq(Row(1, "aa", 42.0), Row(1, "aa", 44.0), + Row(1, "bb", 42.0), Row(1, "bb", 44.0), Row(2, "cc", 11.0))) + } + } + + test("SPARK-59057: a collapse confined to keys the other side filters out is not one") { + // items is partitioned by `identity(id)` with ids 0, 4, 5 and purchases by `bucket(4, item_id)` + // holding only bucket 1, so items' keys are reduced onto buckets [0, 0, 1], three distinct ids + // onto two buckets. Bucket 0 is then dropped by partition filtering, since purchases has + // no rows for it, so the only key the grouping outputs is bucket 1, covering the single id 5. + // Counting distinct keys before the filtering reports a collapse (2 < 3) for a partitioning + // where nothing was merged, and the flag is sticky, so it would keep saying so downstream. + createTable(items, itemsColumns, Array(identity("id"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(0, 'aa', 39.0, cast('2020-01-01' as timestamp)), " + + s"(4, 'bb', 40.0, cast('2020-01-01' as timestamp)), " + + s"(5, 'cc', 41.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(bucket(4, "item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(5, 42.0, cast('2020-01-01' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val keyed = keyedPartitioningsOf(collectAllGroupPartitions(plan)) + assert(keyed.nonEmpty, "the reduced join must be planned as a storage-partitioned join") + assert(keyed.forall(!_.isCollapsed), + "the only surviving key covers one source id, so nothing was merged") + checkAnswer(df, Seq(Row(5, 42.0))) + } + } + + test("SPARK-59057: reducing keys onto a coarser transform collapses keys") { + // `identity(item_id)` reduced onto `bucket(4, id)` maps several ids onto one bucket, so keys + // that were distinct in the source land on the same key. One output partition then covers what + // were separate id partitions. That is a collapse, and unlike a projection it drops no key + // position, so the old provenance flag missed it. + createTable(items, itemsColumns, Array(bucket(4, "id"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(0, 'aa', 39.0, cast('2020-01-01' as timestamp)), " + + s"(4, 'bb', 40.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(identity("item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(0, 42.0, cast('2020-01-01' as timestamp)), " + + s"(4, 44.0, cast('2020-01-15' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val keyed = keyedPartitioningsOf(collectAllGroupPartitions(plan)) + assert(keyed.nonEmpty, "the reduced join must be planned as a storage-partitioned join") + assert(keyed.exists(_.isCollapsed), + "the side whose keys were reduced onto a coarser transform reports a collapse") + checkAnswer(df, Seq(Row(0, 42.0), Row(4, 44.0))) + } + } + } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala index 408f116540bba..a0d76f8d5ccc7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala @@ -275,7 +275,7 @@ class ProjectedOrderingAndPartitioningSuite project.outputPartitioning match { case kp: KeyedPartitioning => assert(kp.expressions === Seq(x), - "narrowed partitioning must keep the projected expression") + "the projected partitioning must keep the projected expression") assert(kp.numPartitions === 4, "partition count must be preserved") case other => @@ -299,7 +299,7 @@ class ProjectedOrderingAndPartitioningSuite case pc: PartitioningCollection => val kps = pc.partitionings.map(_.asInstanceOf[KeyedPartitioning]) assert(kps.forall(_.expressions.length == 1), - "all narrowed KPs must have 1 expression") + "all projected KPs must have 1 expression") assert(kps.map(_.expressions.head.asInstanceOf[Attribute].name).toSet === Set("x", "x_alias"), "both the original and aliased attribute must appear") @@ -329,11 +329,11 @@ class ProjectedOrderingAndPartitioningSuite case pc: PartitioningCollection => val kps = pc.partitionings.map(_.asInstanceOf[KeyedPartitioning]) assert(kps.forall(_.expressions.length == 2), - "narrowed KPs must have 2 expressions (z dropped, x and y kept)") + "projected KPs must have 2 expressions (z dropped, x and y kept)") assert(kps.map(_.expressions.map(_.asInstanceOf[Attribute].name)).toSet === Set(Seq("x", "y"), Seq("x_alias", "y"))) assert(kps.tail.forall(_.partitionKeys eq kps.head.partitionKeys), - "all narrowed KPs must share the same partitionKeys object") + "all projected KPs must share the same partitionKeys object") case other => fail(s"Expected PartitioningCollection, got $other") } @@ -359,7 +359,7 @@ class ProjectedOrderingAndPartitioningSuite case kp: KeyedPartitioning => assert(kp.expressions.map(_.asInstanceOf[Attribute].name) === Seq("y", "z"), "expressions must follow original KP position order [y, z], not output order [z, y]") - assert(kp.isNarrowed, "dropping x must mark the KP as narrowed") + assert(kp.isCollapsed, "dropping x maps (1,1,1) and (2,1,1) onto the same key (1,1)") assert(!kp.isGrouped, "projected keys have duplicate (1,1) entries") case other => fail(s"Expected KeyedPartitioning, got $other") @@ -387,13 +387,13 @@ class ProjectedOrderingAndPartitioningSuite case pc: PartitioningCollection => val kps = pc.partitionings.map(_.asInstanceOf[KeyedPartitioning]) assert(kps.forall(_.expressions.length == 2), - "narrowed KPs must have 2 expressions (x dropped)") + "projected KPs must have 2 expressions (x dropped)") assert(kps.map(_.expressions.map(_.asInstanceOf[Attribute].name)).toSet === Set(Seq("y", "z"), Seq("y", "z_alias")), "expressions must follow original KP position order [y, z/z_alias], not output order") assert(kps.tail.forall(_.partitionKeys eq kps.head.partitionKeys), - "all narrowed KPs must share the same partitionKeys object") - assert(kps.forall(_.isNarrowed), "all KPs must be marked as narrowed") + "all projected KPs must share the same partitionKeys object") + assert(kps.forall(_.isCollapsed), "all KPs must be marked as collapsed") assert(kps.forall(!_.isGrouped), "projected keys have duplicate (1,1) entries") case other => fail(s"Expected PartitioningCollection, got $other") @@ -445,7 +445,7 @@ class ProjectedOrderingAndPartitioningSuite // Scenario 1: projected keys have duplicates (x-values: 1, 1, 2) -> isGrouped=false. // GroupPartitionsExec would merge the two x=1 partitions, carrying the same skew risk as - // allowKeysSubsetOfPartitionKeys. EnsureRequirements calls groupedSatisfies() directly. + // allowKeysSubsetOfPartitionKeys. EnsureRequirements calls mayGroupToSatisfy() directly. val keys2d = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1)) val project = ProjectExec(Seq(x), DummyLeafExecWithPartitioning(output = Seq(x, y), @@ -454,10 +454,10 @@ class ProjectedOrderingAndPartitioningSuite withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { project.outputPartitioning match { case kp: KeyedPartitioning => - assert(!kp.isGrouped, "narrowed keys must have duplicates (1 appears twice)") - assert(kp.isNarrowed, "projection must mark the KP as narrowed") - assert(!kp.groupedSatisfies(ClusteredDistribution(Seq(x))), - "narrowed ungrouped KP must not satisfy via groupedSatisfies without config") + assert(!kp.isGrouped, "collapsed keys must have duplicates (1 appears twice)") + assert(kp.isCollapsed, "dropping y maps (1,1) and (1,2) onto the same key 1") + assert(!kp.mayGroupToSatisfy(ClusteredDistribution(Seq(x))), + "collapsed ungrouped KP must not satisfy via mayGroupToSatisfy without config") case other => fail(s"Expected KeyedPartitioning, got $other") } } @@ -465,15 +465,17 @@ class ProjectedOrderingAndPartitioningSuite withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { project.outputPartitioning match { case kp: KeyedPartitioning => - assert(kp.groupedSatisfies(ClusteredDistribution(Seq(x))), - "narrowed ungrouped KP must satisfy via groupedSatisfies when config is enabled") + assert(kp.mayGroupToSatisfy(ClusteredDistribution(Seq(x))), + "collapsed ungrouped KP must satisfy via mayGroupToSatisfy when config is enabled") case other => fail(s"Expected KeyedPartitioning, got $other") } } // Scenario 2: projected keys are distinct (x-values: 1, 2, 3) -> isGrouped=true. // Each projected key maps to exactly one original partition so GroupPartitionsExec does not - // merge any partitions. No skew risk: must satisfy ClusteredDistribution regardless of config. + // merge any partitions. Nothing collapsed either. The projection dropped a position but kept + // every key distinct, so the partitioning is not coarser than the layout it came from. There is + // no skew risk, so it must satisfy ClusteredDistribution regardless of config. val keys2dDistinct = Seq(InternalRow(1, 1), InternalRow(2, 2), InternalRow(3, 3)) val projectDistinct = ProjectExec(Seq(x), DummyLeafExecWithPartitioning(output = Seq(x, y), @@ -483,15 +485,67 @@ class ProjectedOrderingAndPartitioningSuite projectDistinct.outputPartitioning match { case kp: KeyedPartitioning => assert(kp.isGrouped, "distinct projected keys must be grouped") - assert(kp.isNarrowed, "projection must mark the KP as narrowed") + assert(!kp.isCollapsed, + "dropping y lost no distinct key, so this partitioning is not coarser than its source") assert(kp.satisfies(ClusteredDistribution(Seq(x))), - "grouped narrowed KP must satisfy ClusteredDistribution without config (no merging)") + "a grouped KP must satisfy ClusteredDistribution without config (no merging)") case other => fail(s"Expected KeyedPartitioning, got $other") } } } - test("SPARK-58974: the narrowing guard applies for either value of requireAllClusterKeys") { + test("SPARK-59057: PartitioningCollection normalizes isCollapsed across its members") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + val z = AttributeReference("z", IntegerType)() + val keys = Seq(InternalRow(1), InternalRow(2)) + + def allCollapsed(p: Partitioning): Boolean = + PartitioningCollection.flatten(p).collect { case kp: KeyedPartitioning => kp } + .forall(_.isCollapsed) + + // See the `PartitioningCollection` class doc for why a collapsed member marks the others. + val collection = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(x), keys), + KeyedPartitioning(Seq(y), keys).copy(isCollapsed = true))) + assert(allCollapsed(collection), "a collapsed member must mark the whole collection") + + // Nested collections are normalized too, so a collapsed sibling reaches into them. + val nested = PartitioningCollection.fromPartitionings(Seq( + PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(x), keys), KeyedPartitioning(Seq(y), keys))), + KeyedPartitioning(Seq(z), keys).copy(isCollapsed = true))) + assert(allCollapsed(nested), + "a collapsed sibling must mark the members of a nested collection") + } + + test("SPARK-59057: a projection that keeps every distinct key collapses nothing") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + + // The source reports two splits for the same partition value, so its keys already contain a + // duplicate before any projection. Dropping y maps (1,1),(1,1) onto 1 and (2,2) onto 2: two + // distinct keys before, two after. Grouping would merge only the two splits that already + // shared a key, which is what GroupPartitionsExec does for any partitioning and needs no + // opt-in, so `mayGroupToSatisfy` must accept it with the config off. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 1), InternalRow(2, 2)) + val project = ProjectExec(Seq(x), + DummyLeafExecWithPartitioning(output = Seq(x, y), + partitioning = KeyedPartitioning(Seq(x, y), keys))) + + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + project.outputPartitioning match { + case kp: KeyedPartitioning => + assert(!kp.isGrouped, "the duplicate key comes from the source's two splits") + assert(!kp.isCollapsed, "the projection lost no distinct key") + assert(kp.mayGroupToSatisfy(ClusteredDistribution(Seq(x))), + "grouping merges only splits that already shared a key, so no opt-in is needed") + case other => fail(s"Expected KeyedPartitioning, got $other") + } + } + } + + test("SPARK-58974: the collapse skew guard applies regardless of requireAllClusterKeys") { val x = AttributeReference("x", IntegerType)() val y = AttributeReference("y", IntegerType)() @@ -503,7 +557,7 @@ class ProjectedOrderingAndPartitioningSuite DummyLeafExecWithPartitioning(output = Seq(x, y), partitioning = KeyedPartitioning(Seq(x, y), keys))) val kp = project.outputPartitioning.asInstanceOf[KeyedPartitioning] - assert(kp.isNarrowed && !kp.isGrouped) + assert(kp.isCollapsed && !kp.isGrouped) Seq(true, false).foreach { requireAll => // Both values on purpose: the whole claim of the fix is that the guard answers the same @@ -512,18 +566,18 @@ class ProjectedOrderingAndPartitioningSuite // `requireAllClusterKeys` branch on its own would accept it. val required = ClusteredDistribution(Seq(x), requireAllClusterKeys = requireAll) withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { - assert(!kp.groupedSatisfies(required), - s"requireAllClusterKeys=$requireAll must not group a narrowed partitioning whose keys " + + assert(!kp.mayGroupToSatisfy(required), + s"requireAllClusterKeys=$requireAll must not group a collapsed partitioning whose keys " + "are no longer distinct") } withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { - assert(kp.groupedSatisfies(required), + assert(kp.mayGroupToSatisfy(required), s"requireAllClusterKeys=$requireAll: the opt-in must allow the grouping") } } } - test("SPARK-46367: isNarrowed is sticky across chained PartitioningPreservingUnaryExecNodes") { + test("SPARK-46367: isCollapsed is sticky across chained PartitioningPreservingUnaryExecNodes") { val x = AttributeReference("x", IntegerType)() val y = AttributeReference("y", IntegerType)() @@ -537,10 +591,10 @@ class ProjectedOrderingAndPartitioningSuite outerProject.outputPartitioning match { case kp: KeyedPartitioning => assert(!kp.isGrouped, "duplicate keys must survive the second hop") - assert(kp.isNarrowed, - "isNarrowed must be sticky: a second hop that keeps all positions must not reset it") - assert(!kp.groupedSatisfies(ClusteredDistribution(Seq(x))), - "narrowed ungrouped KP must still not satisfy ClusteredDistribution without config " + + assert(kp.isCollapsed, + "isCollapsed must be sticky: a second hop that keeps all positions must not reset it") + assert(!kp.mayGroupToSatisfy(ClusteredDistribution(Seq(x))), + "collapsed ungrouped KP must still not satisfy ClusteredDistribution without config " + "after a second PartitioningPreservingUnaryExecNode hop") case other => fail(s"Expected KeyedPartitioning, got $other") } @@ -572,7 +626,7 @@ class ProjectedOrderingAndPartitioningSuite } assert(kp.partitionKeys eq child.partitioning.asInstanceOf[KeyedPartitioning].partitionKeys, "partition keys must be unchanged") - assert(!kp.isNarrowed, "same number of positions: not narrowed") + assert(!kp.isCollapsed, "no position dropped: nothing collapsed") case other => fail(s"Expected KeyedPartitioning, got $other") } } @@ -580,7 +634,7 @@ class ProjectedOrderingAndPartitioningSuite test("SPARK-46367: narrowing projection drops transform when its column is absent") { // KP([bucket(32, id), years(ts)], keys2d) through Project(id) -- ts is dropped. // bucket(32, id) is projectable (id in output); years(ts) is not (ts absent). - // Result: KP([bucket(32, id)], keys1d, isNarrowed=true, isGrouped=false). + // Result: KP([bucket(32, id)], keys1d, isCollapsed=true, isGrouped=false). val id = AttributeReference("id", IntegerType)() val ts = AttributeReference("ts", IntegerType)() val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) @@ -601,7 +655,7 @@ class ProjectedOrderingAndPartitioningSuite assert(te.children.head.asInstanceOf[Attribute].name === "id") case other => fail(s"Expected TransformExpression, got $other") } - assert(kp.isNarrowed, "dropping years(ts) position must mark the KP as narrowed") + assert(kp.isCollapsed, "dropping years(ts) maps (0,2020) and (0,2021) onto bucket 0") assert(!kp.isGrouped, "projected bucket keys (0,1,0) have duplicates") case other => fail(s"Expected KeyedPartitioning, got $other") } @@ -610,7 +664,7 @@ class ProjectedOrderingAndPartitioningSuite test("SPARK-46367: alias substitution rewrites years transform while preserving bucket") { // KP([bucket(32, id), years(ts)], keys2d) through Project(id, ts as ts_alias). // bucket(32, id) keeps id (no alias for id); years(ts) is rewritten to years(ts_alias). - // Result: KP([bucket(32, id), years(ts_alias)], keys2d) -- not narrowed. + // Result: KP([bucket(32, id), years(ts_alias)], keys2d), nothing collapsed. val id = AttributeReference("id", IntegerType)() val ts = AttributeReference("ts", IntegerType)() val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) @@ -641,7 +695,7 @@ class ProjectedOrderingAndPartitioningSuite } assert(kp.partitionKeys eq child.partitioning.asInstanceOf[KeyedPartitioning].partitionKeys, "partition keys must be unchanged") - assert(!kp.isNarrowed, "both positions projected: not narrowed") + assert(!kp.isCollapsed, "both positions projected: nothing collapsed") case other => fail(s"Expected KeyedPartitioning, got $other") } } 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 be772bc3f28e3..46ee475d41f0a 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,6 +21,7 @@ import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyedShuffleSpec, Partitioning, PartitioningCollection, UnknownPartitioning} +import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -35,6 +36,38 @@ class GroupPartitionsExecSuite extends SharedSparkSession { private def row(a: Int): InternalRow = InternalRow.fromSeq(Seq(a)) private def row(a: Int, b: Int): InternalRow = InternalRow.fromSeq(Seq(a, b)) + test("SPARK-59057: the output flag reports what this node's grouping merges") { + // Keys [(1,1), (1,2), (2,1)] projected onto position 0 give [1, 1, 2]. The first two groups + // cover keys the child held apart, the third does not. + val keys = Seq(row(1, 1), row(1, 2), row(2, 1)) + def gpe(joinKeyPositions: Option[Seq[Int]], + expected: Option[Seq[(InternalRowComparableWrapper, Int)]] = None, + distribute: Boolean = false, + childCollapsed: Boolean = false): KeyedPartitioning = { + val childKp = KeyedPartitioning(Seq(exprA, exprB), keys).copy(isCollapsed = childCollapsed) + GroupPartitionsExec(DummySparkPlan(outputPartitioning = childKp), joinKeyPositions, + expected, distributePartitions = distribute) + .outputPartitioning.asInstanceOf[KeyedPartitioning] + } + def keyOf(a: Int): InternalRowComparableWrapper = + InternalRowComparableWrapper(row(a), Seq(exprA)) + + assert(!gpe(None).isCollapsed, + "no projection, so every group covers the one key it was built from") + assert(gpe(Some(Seq(0))).isCollapsed, "keys (1,1) and (1,2) are merged into key 1") + assert(gpe(None, childCollapsed = true).isCollapsed, "the child's flag is sticky") + + // The keys the join agreed on decide it. Keeping key 1 keeps the merge, keeping only key 2 + // does not, and that holds however the splits of a kept key are laid out afterwards. Key 1 has + // two child splits, which is the split count `EnsureRequirements` derives for it. + Seq(false, true).foreach { distribute => + assert(gpe(Some(Seq(0)), Some(Seq(keyOf(1) -> 2, keyOf(2) -> 1)), distribute).isCollapsed, + s"distributePartitions=$distribute: the merged key 1 survives") + assert(!gpe(Some(Seq(0)), Some(Seq(keyOf(2) -> 1)), distribute).isCollapsed, + s"distributePartitions=$distribute: only key 2 survives, and it merges nothing") + } + } + test("SPARK-56241: non-coalescing passes through child ordering unchanged") { // Each partition has a distinct key — no coalescing happens. val partitionKeys = Seq(row(1), row(2), row(3))