From 122e126684c1dff40d263936ffb00c80417b400e Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sat, 29 Aug 2026 11:03:23 +0200 Subject: [PATCH] [SPARK-59057][SQL] Make KeyedPartitioning.isNarrowed mean actual key collapse, and rename it to isCollapsed ### What changes were proposed in this pull request? `KeyedPartitioning` carries a flag that gates whether `GroupPartitionsExec` may coalesce its duplicate partition keys without `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`. Today the flag records *provenance*, meaning "a projection dropped key positions, or my input already had some dropped". The gate's own comment describes something else, a *collapse*: "partitions that held distinct keys in the original finer-grained partitioning". The two are not the same, and the gate reads the wrong one. This changes the flag to mean what the gate needs, and renames it from `isNarrowed` to `isCollapsed`. A collapse is a projection or a reduction mapping keys that were distinct in the input onto the same key, so one partition here stands for several of the original ones. Dropping key positions no longer sets the flag on its own. The projected keys have to actually lose distinctness. It stays sticky, since neither grouping nor a further projection can make a partitioning finer again. The gate still rests on the same two conditions, but each now sits where it belongs. `isCollapsed` says the collapse happened, and it gates `mayGroupToSatisfy`, the question `EnsureRequirements` asks before it inserts a `GroupPartitionsExec`. The second condition is that there is still something left to merge, and that one is the caller. `EnsureRequirements` asks the question only of partitionings that are not grouped. Once the keys are unique, grouping merges nothing and there is nothing left to gate, however coarse the partitioning already is, so `satisfies` needs no gate at all. That splits the old `groupedSatisfies` in two, because its two callers were asking different questions through one method: * `keysSatisfy(required)`: whether the partition keys match what the distribution asks for, ignoring duplicate keys. `satisfies` is this plus `isGrouped`. * `mayGroupToSatisfy(required)`: `keysSatisfy` plus a check that coalescing the duplicates is allowed. `EnsureRequirements` asks this of non-grouped partitionings. That check is a conjunct rather than a case inside the key matching. SPARK-58974 established it there by hoisting it above the `requireAllClusterKeys` branch, since the risk it guards does not depend on which key sets count as matching. `OrderedDistribution` stays exempt, because `GroupPartitionsExec` pads that path out to the expected split counts instead of coalescing. Producers: * `PartitioningPreservingUnaryExecNode` goes through a new `KeyedPartitioning.project`, which computes `isGrouped` and `isCollapsed` together from the projected keys. The flag is inherited, or set when the projection leaves fewer distinct keys than the input had. A projection that keeps every position returns the partitioning unchanged, with no pass over the keys. * `UnionExec`'s keyed merge moves to `KeyedPartitioning.concat`, which joins `apply`, `project` and `toGrouped` as the fourth rule for how the flags travel. The rule itself is unchanged: it ORs the children's flags, so a child that really did collapse its keys marks the union. What changes for a union is upstream of it, since the children's flags are now precise. * `GroupPartitionsExec`, `KeyedPartitioning.toGrouped`, `KeyedPartitioning.createShuffleSpec` and `KeyedShuffleSpec.createPartitioning` all propagate the flag. Each of them used to build a partitioning through the 3-argument constructor, which defaulted the flag to `false`. That default is gone from the parameter, so every producer now has to state it. * `createShuffleSpec` projects onto the operation keys, so it computes its own collapse with `project`. `GroupPartitionsExec` decides both answers from the key groups it keeps, by asking whether a group spans more than one of its own child's partition keys. It does not compare against the aligned key list, because that list is the one both join sides agreed on. Such a list can be missing keys this side had (partition filtering) or repeat them (padding), and neither is a collapse. `GroupPartitionsExec` computes the flag once, from one member of the child's partitioning. `PartitioningCollection` normalizes the flag across its members by OR, alongside the `partitionKeys` interning it already did, and its invariant check enforces the result. The flag describes the shared physical layout rather than one member's naming of it, and consumers read it 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. Uniformity also means a consumer can read one representative instead of every member. The coverage there is structural. Once the flag is uniform, the mixed collection that could slip through cannot be built, so the test asserts the invariant rather than a plan shape. This supersedes https://github.com/apache/spark/pull/58316 (SPARK-59026), which restores the flag in two of the producers above. Propagating it is that PR's finding, and its unit test is taken here with credit. Its end-to-end test is not taken. Its table has unique ids, so dropping the second key column keeps every key distinct, and that is exactly the scenario this change reclassifies as no collapse. The two cannot both land as written. The class doc is reorganized around the new meaning. It leads with what a partition key is and what `EnsureRequirements` uses the keys for, and it gains a "Key Collapse" section that tells a collapse apart from a grouping. That section also carries why a collapsed partitioning is still reported rather than dropped to `UnknownPartitioning`, which was nowhere in the code. Both the method split and `project` point the same way as item 5 of https://github.com/apache/spark/pull/58262#issuecomment-5438648022, which is that the SPJ decisions in `EnsureRequirements` are worth untangling on their own. `project` settles one half of that here. Applying a projection to a partitioning now happens in one place instead of once per producer. The other half is left to the follow-up. Which positions to keep is still asked in two places, `KeyedShuffleSpec.keyPositions` (where `createShuffleSpec` takes its `joinKeyPositions` from) and the overlap test in `keysSatisfy`. A comment there now records that the two are the same test, and that the answer only becomes true once the projection is applied. ### Why are the changes needed? Provenance leaves the gate open in three ways, which is why this is filed as a bug: * **The flag is dropped wherever a partitioning is rebuilt.** That happens in `toGrouped`, in `createShuffleSpec`'s projected partitioning, in `KeyedShuffleSpec.createPartitioning` and in `GroupPartitionsExec`, so the gate cannot see the risk even when it exists. This is the defect https://github.com/apache/spark/pull/58316 (SPARK-59026) reports, and it is fixed here. * **A reduction never sets the flag at all.** Under `allowCompatibleTransforms`, a `bucket(4, id)` side reduced onto a `bucket(2, id)` join maps four distinct keys onto two, which is a real collapse. It drops no key position, so provenance misses the whole class of them and the opt-in is bypassed. * **A `PartitioningCollection` whose members disagreed could be entered through the plain one.** `EnsureRequirements` looks for any single member that a `GroupPartitionsExec` would make satisfy the distribution, so a collection holding one collapsed and one plain member reached one regardless. Normalizing the flag across members closes that. Provenance also over-refuses, and it does so on one of the commonest shapes. `!isGrouped` has causes that have nothing to do with dropping key positions: * A data source reports one partition key per input split, so a table with several splits for the same partition value already has duplicate keys before any projection. * `UnionExec` computes `isGrouped` over the concatenation of its children's keys, so two children that are individually key-distinct make it false by overlapping with each other. In both cases grouping merges only partitions that already shared a key. That is what `GroupPartitionsExec` does for any partitioning that never went through a projection, and it needs no opt-in. Today the mere presence of a narrowing projection anywhere upstream turns that into a refusal, so a plan gets a shuffle that protects nothing. Measured on the multi-split shape. A table partitioned by `(id, dept)` has two splits for the same `(1, 'x')` value, is projected down to `id`, and is joined on it with the opt-in off. Before: no `GroupPartitionsExec` and 2 shuffles, and the projected partitioning reports the flag set. After: 1 `GroupPartitionsExec`, 0 shuffles, flag clear. ### Does this PR introduce _any_ user-facing change? It should reach `branch-4.3` and `branch-4.x` as well as master, since 4.3.0 is where the flag first ships. Yes, a plan-level change. Storage-partitioned operations now proceed without `allowKeysSubsetOfPartitionKeys.enabled` in the cases above, where they previously fell back to a shuffle. Query results are unchanged. No migration guide entry is needed, because the flag and its gate arrived in 4.3.0 (SPARK-46367), which is unreleased, so no released version behaves the old way. Propagating the flag through `createShuffleSpec` also makes the spec's projected partitioning compare equal to the child's in one more case. That case is the opt-in on, every position selected, and the source already collapsed and grouped. `EnsureRequirements`' "child partitionings not modified" fast path then fires where it previously did not. The partitions are the same, since that path is only reached when both sides' keys already match, but the existing `GroupPartitionsExec` keeps `expectedPartitionKeys = None`, so `explain` loses its `ExpectedPartitionKeys` line. Reducers are a second case. With `allowCompatibleTransforms`, a `bucket(4, id)` side reduced onto a `bucket(2, id)` join really does collapse its keys, so its `GroupPartitionsExec` output now carries the flag where it did not before, and a coalescing further up the plan needs the opt-in. That direction is a narrowing, not a widening, and it is what the flag is supposed to say. Planning cost was measured on the worst case for the collapse test. The partitioning is projected ten times over, each projection dropping one key position, which is what a plan with stacked `Project`s asks for. Every projection drops a position, so the identity fast path never fires, and no projection collapses a key, so the walk never stops early. Over a 50k-split, 25k-distinct-key partitioning with 12-position keys, 20 evaluations of that chain took 345 ms for the provenance formula and 1291 ms for this one, i.e. about 5 ms per projection that drops a position. The obvious two-pass form of the same test, one `distinct` over the projected keys and another over the source keys, took 2137 ms in the same harness, which is why `project` walks the two key lists in lockstep instead. A projection that keeps every position pays nothing, since it cannot collapse anything. The trade is deliberate: a few milliseconds of planning buys the removal of a runtime shuffle in the cases above, and it also closes the three holes that let the gate through. ### How was this patch tested? * New unit test in `ProjectedOrderingAndPartitioningSuite` for the multi-split shape. A source with duplicate keys, projected down, is ungrouped but not collapsed, and `mayGroupToSatisfy` accepts it with the opt-in off. * New end-to-end test in `KeyGroupedPartitioningSuite` for the same shape through a real plan, asserting the grouping happens and the shuffles disappear. * New unit test that a collapsed member anywhere in a `PartitioningCollection` marks every member of it, including through a nested collection, so the outcome does not depend on which join side the collapse came from. * New operator-level test in `GroupPartitionsExecSuite` that the output flag reports what the node's grouping merges. It covers the replicating and the distributing mode, and the case where the join's agreed keys drop the merged one. * New end-to-end test for the shuffle-template chain from #58316, with the expectation these semantics call for. Its `items` has unique ids, so dropping `name` collapses nothing, and the final aggregate may group the union's overlapping keys with the opt-in off. That keeps the chain covered while pinning the reclassification. * New end-to-end test that reducing keys onto a coarser transform collapses them. With `allowCompatibleTransforms`, an `identity(item_id)` side reduced onto `bucket(4, id)` maps several ids onto one bucket, and the flag now says so. * New end-to-end test that filtering partition keys out is not a collapse. With `partitionFilter` on, an inner join plans both sides on the intersection of their keys, and the side that lost a key must not report a collapse. Otherwise the sticky flag costs a shuffle above a later union. * One existing expectation flipped, which is the contract change. In `SPARK-46367: narrowing projection with duplicate keys ...`, the scenario whose projected keys stay distinct now asserts the flag is clear. * Every new expectation is guarded by an ablation, verified one at a time. The old provenance formula fails four of them: the multi-split unit test, its end-to-end counterpart, the flipped `SPARK-46367` scenario and the shuffle-template chain. Deciding the flag from a key count taken after the alignment, instead of from the key groups kept, fails the partition-filter test. Dropping the collection normalization fails its unit test. Dropping the `GroupPartitionsExec` term fails the reducer test. * `DistributionSuite` covers the propagation through `toGrouped` and `KeyedShuffleSpec.createPartitioning`. That test comes from #58316. * Also ran `DistributionSuite`, `ShuffleSpecSuite`, `KeyGroupedPartitioningSuite`, `ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `PlannerSuite`, `WriteDistributionAndOrderingSuite`, `DataFrameSetOperationsSuite`, `AdaptiveQueryExecSuite`, `CoalesceShufflePartitionsSuite` and the TPC-DS plan stability suites. No golden file changed. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../plans/physical/partitioning.scala | 398 +++++++++++------- .../apache/spark/sql/internal/SQLConf.scala | 8 +- .../sql/catalyst/DistributionSuite.scala | 19 +- .../AliasAwareOutputExpression.scala | 19 +- .../execution/basicPhysicalOperators.scala | 6 +- .../datasources/v2/GroupPartitionsExec.scala | 67 ++- .../exchange/EnsureRequirements.scala | 14 +- ...taSourceV2CatalystRuntimeFilterSuite.scala | 2 +- .../DistributionAndOrderingSuiteBase.scala | 4 +- .../KeyGroupedPartitioningSuite.scala | 288 ++++++++++++- ...rojectedOrderingAndPartitioningSuite.scala | 120 ++++-- .../v2/GroupPartitionsExecSuite.scala | 33 ++ 12 files changed, 726 insertions(+), 252 deletions(-) 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 688968bda16bf..4973cf04db8a4 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. Every partitioning that never went through a projection is 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 @@ -561,10 +596,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) } /** @@ -583,42 +656,31 @@ case class KeyedPartitioning( KeyedPartitioning.reduceKeys(partitionKeys, expressionDataTypes, 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 { @@ -634,6 +696,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) { @@ -641,14 +717,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 @@ -669,7 +743,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 = { @@ -832,11 +926,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 { @@ -850,16 +956,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 @@ -870,19 +971,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") } } } @@ -920,6 +1018,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 @@ -929,38 +1037,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)) } @@ -1445,7 +1548,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 9398d14eb94e8..d15fc0fd1af65 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 6f671bfcb80be..7b1b3d7d8a92b 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.expressionDataTypes, keyedPartitioning.partitionKeys) - )(keyedPartitioning.projectKeys) + (childKp.expressionDataTypes, 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 c2051ac91b062..b2e863e330820 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 c0a1faa032635..506f781ed471a 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 @@ -499,7 +499,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 948de99959544..4a58d77b4af74 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. @@ -4946,6 +4953,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))