peter-toth opened a new pull request, #58469: URL: https://github.com/apache/spark/pull/58469
### What changes were proposed in this pull request? This builds on two changes that are now on master. #58351 (SPARK-59057) renamed `KeyedPartitioning.isNarrowed` to `isCollapsed` and split `groupedSatisfies` into `keysSatisfy` and `mayGroupToSatisfy`. Two things here follow from that. The classification asks whether a member can satisfy the distribution once a `GroupPartitionsExec` is allowed, which is `keysSatisfy` for a grouped member and `mayGroupToSatisfy` for a non-grouped one, since only the second may coalesce duplicate keys and only it needs the collapse permission. That composition is a new `KeyedPartitioning.keysMaySatisfy`, which keeps `keysSatisfy` private. And the tests that build a `KeyedPartitioning` state `isCollapsed` explicitly, because that parameter lost its default. #58420 (SPARK-59120) made every reader of a `KeyedPartitioning`'s partition keys take its types from `keyDataTypes`, the schema the keys were written under, rather than from the partition expressions. This PR projects those keys on more paths, so it needs that fix underneath it. Without it, projecting the keys of a partitioning whose keys a reducer rewrote throws `ClassCastException` at planning. This is an alternative to #58245, which fixes the same JIRA by adding the projection to one of the two branches below. `EnsureRequirements` split a child's `KeyedPartitioning`s by `isGrouped` and then had two branches that each had to insert a `GroupPartitionsExec`. This PR classifies by what still has to happen to the data instead. - `splitKeyedPartitionings` now takes the required distribution and answers two questions, in this order. Whether a non-`KeyedPartitioning` member already satisfies it, and if not, how a `KeyedPartitioning` member can. As it is, or after a `GroupPartitionsExec` projecting to the partition expression positions returned with it, with `None` positions when the node only has to coalesce duplicate partition keys. At most one partitioning comes back, because the caller acts on a single one of them. The order is deliberate, since the keyed half is the one that projects partition keys, so it only runs when no plain member is enough. - A new `clusterKeyPositions` helper derives those positions from the required clustering, and a new `KeyedPartitioning.numPartitionsProjectedOn` answers how many partitions a projection onto them would leave. - The four-way match collapses to three cases, because the two arms that each had to insert a `GroupPartitionsExec` become one. It now matches on the distribution and the resolution together, since two of the three arms are decided by the resolution rather than by the distribution. The `OrderedDistribution` arm keeps its own insertion, as on `master`. `clusterKeyPositions` keeps a partition expression when it is one of the operation keys. `keysSatisfy` recognises that at the *reference* level, where a `bucket(4, a)` transform covers the cluster key `a`, and also at the *expression* level, where the cluster key is the partition expression itself. Both are honoured here, which is why the positions are derived from the clustering rather than taken from `KeyedShuffleSpec.keyPositions`. That answers only the first, which is the right question for a storage-partitioned join but would drop a partition expression that is itself an operation key. The expression-level test is what keeps a position whose expression is clustered on while its references are not. No production path builds such a clustering today. `V2ExpressionUtils.toCatalystTransformOpt` maps `IdentityTransform` to the resolved attribute itself, so the reference-level test matches the same position anyway, and `DistributionAndOrderingUtils.prepareQuery` maps `resolveTransfo rmExpression` over a write's clustering, so a `TransformExpression` does not survive into one. The test builds the shape by hand, with a table partitioned by `(id, years(ts))` and clustered on those same two expressions. The check is here because `keysSatisfy` already accepts that shape, so deriving the positions any other way would make the two disagree. A member whose expressions cover no operation key at all is skipped rather than projected to no position, which would put every partition into one. Only a partitioning whose expressions have no references can get there. A DSv2 scan cannot report one, because `supportsExpressions` refuses it, but nothing rejects it at `KeyedPartitioning` construction. A partitioning that satisfies the distribution is kept as it is when nothing is left for a node to do, which holds in two ways. Either the projection drops no position, so `satisfies` is not the over-claim this fix is about. Or a position is dropped but the projection merges nothing, so every operation key already lives on a single partition. Keeping the partitioning is then better than projecting, because `KeyedPartitioning([id, name])` and `KeyedPartitioning([id])` describe the same number of partitions and only the first lets a downstream operator co-partition on `name` too. This is where the collapse gives the single insertion point a new way to be wrong, namely inserting a node that merges nothing, which cannot happen on `master`, where a grouped partitioning got no node at all. That question is asked of every member of the child's partitioning, not just of the one a node would be built from, because a child can report a whole `PartitioningCollection` and its members can disagree about which positions are operation keys. An inner join is where they do. Its `outputPartitioning` is the two sides' partitionings, and unlike `AliasAwareOutputExpression` it does not enumerate the mixed combinations, so a window keyed on one side's first key column and the other side's remaining ones sees one member covering position 0 and one covering the rest. The first needs no node, the second would coalesce for nothing. Among the members that do need a node, the one whose projection leaves the most partitions is used, and that answer is exact rather than a heuristic. Only the position sets not contained in another one have to be projected. Projecting to fewer positions can merge partitions but never split them, so a contained set can never leave more partitions than the set containing it, and can only tie. Dropping it therefore costs no parallelism, and on a tie it settles the choice toward the wider set, which still names the keys the narrower one would have dropped. In the ordinary case one set contains all the others and a single projection settles it, and only sets that genuinely disagree, neither containing the other, each cost one. That matters because the projection is the expensive step here, since `KeyedPartitioning.projectKeys` allocates a row per input partition and `InternalRowComparableWrapper.hashCode` is uncached. So the count is memoized per position set as well, a set covering ev ery position needs no projection at all, being the identity on the key values, and a single surviving candidate with no required count needs none either, which is the shape the default config produces. Keeping one member per position set costs nothing either, because `PartitioningCollection` guarantees its members share the `partitionKeys` reference and their arity, so position `i` addresses the same key column in all of them. `Distribution.requiredNumPartitions` needs care, because a `GroupPartitionsExec` derives its number of partitions from the partition keys it is handed rather than from the operator. The requirement is therefore checked against the count such a node *would* produce, not against the count the partitioning happens to have now, and it filters the candidates rather than vetoing the winner. One that would land on the required number must not lose the ranking to one that cannot honour it and send the whole child to a shuffle instead. As-is satisfaction goes through `satisfies`, which enforces the count on its own. For an operator that co-partitions more than one child no projection is done here, because that belongs to the multi-child block below, which a storage-partitioned join reaches through `checkKeyGroupCompatible` and anything else through `withJoinKeyPositions`. Projecting inline as well would leave that block deriving positions from an already projected partitioning and applying them to the unprojected keys. `KeyedPartitioning.satisfies` is not touched, so nothing outside `EnsureRequirements` changes behaviour. It does still answer `true` for a partitioning that needs a projection first, which means `ValidateRequirements` cannot catch a missing `GroupPartitionsExec`. Giving that check the strict test directly, without changing what `satisfies` answers, is a follow-up we are working on. The `OrderedDistribution` arm also loses a `MatchError`. It tested that the partition keys are sorted with `partitionKeys.sliding(2)`, which yields one short window for a single-key partitioning, and `case Seq(k1, k2)` cannot match it, so planning threw. A v2 table whose rows all share one partition value reports such a partitioning, since `DataSourceV2ScanExecBase` has no single-partition short-circuit, and with `v2BucketingAllowSorting` on, a global sort on the partition key reaches this branch. The test is now a zip of the keys with their successors, which is vacuously true for one key. Pre-existing, but the branch is rewritten here anyway. Two scaladocs are corrected as well. `KeyedPartitioning` taught `isGrouped` as the axis this PR replaces. `GroupPartitionsExec.joinKeyPositions` described its projection as being "for join compatibility", and it now carries the projection for a single-child operator too. #### Backport to `branch-4.2` Everything above is #58262's description, unchanged. This is what differs on this branch. **`keysMaySatisfy` is not added, and `EnsureRequirements` asks `groupedSatisfies` directly.** #58351 split `groupedSatisfies` into `keysSatisfy`, the key matching, and `mayGroupToSatisfy`, that matching plus permission to coalesce a collapsed partitioning, and `keysMaySatisfy` composes the two. Neither the split nor the permission is on this branch: `isCollapsed` and its gate arrive in 4.3.0 with SPARK-46367, and this branch's `groupedSatisfies` is exactly what `master` calls `keysSatisfy`. So `keysMaySatisfy`'s two arms coincide here. The class doc describes the two questions this branch has rather than `master`'s four. **Four smaller deviations.** - `nonGroupedSatisfies` becomes private, as on `master`: the rewrite removes its last caller outside the class. - `splitKeyedPartitionings` keeps a local recursion instead of `PartitioningCollection.flatten`, which is not on this branch. - The shared `exprA` .. `exprD` fixtures in `EnsureRequirementsSuite` become attributes, where this branch has `Literal`s, because the new tests match partition expressions against cluster keys. SPARK-57038 made the same change on `master`, and its absence here is also why this branch's planner reads the partition expressions through `collectLeaves()` where `master` reads `references`. - The config is spelled `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS` here, the `JOIN_` was dropped from the name later. Read `v2BucketingAllowKeysSubsetOfPartitionKeys` above as that. **Two of the 24 tests are not carried over**, because the shapes they cover do not exist on this branch: - `SPARK-58968: keysMaySatisfy asks the collapse gate of a non-grouped partitioning only`, in `ProjectedOrderingAndPartitioningSuite`. It contrasts `keysMaySatisfy` with `mayGroupToSatisfy` under the collapse gate, and none of the three is here, so that suite is untouched. - `SPARK-58968: window top-k over union output partitioning coalesces partitions`. `UnionExec` does not merge `KeyedPartitioning`s on this branch, so the union reports no keyed partitioning for the window to see and no `GroupPartitionsExec` is inserted. The two `SPARK-46367` test tidy-ups do not apply either, since those tests are not on this branch. **The measurements below are `master`'s.** On this branch, 22 tests and **14 of them fail on the base** at `773f49a0456`: the three window tests, the transform position, the `MatchError`, four of the six `requiredNumPartitions` tests, and all five multi-member ones. `SPARK-58968: a partitioning that covers no operation key is shuffled, not projected` passes on this branch's base, where it failed on `master`; it stays as a guard on the skip path. Ran `DistributionSuite` and `ShuffleSpecSuite` with 20 tests, and `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `ValidateRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite` and `PlannerSuite` with 263. `dev/lint-scala` is clean. **Nothing goes below this branch.** `GroupPartitionsExec.scala` does not exist on `branch-4.1`, so 4.1 and lower are unaffected. ### Why are the changes needed? With `v2BucketingAllowKeysSubsetOfPartitionKeys` enabled, `KeyedPartitioning.keysSatisfy` only requires that some operation key overlaps the partition attributes. A partitioning grouped on `(id, name)` therefore reports that it satisfies `ClusteredDistribution([id])` while rows sharing an `id` still sit on separate partitions. A storage-partitioned join projects the keys down to the operation keys in `checkKeyGroupCompatible`, and for a non-join operator nothing did. `isGrouped` is the wrong thing to classify on, because it only says the *full* partition keys are unique and says nothing about whether the *projected* keys are. A `GroupPartitionsExec` is needed either to coalesce duplicate partition keys, or to project down to the operation keys, or both, and splitting by provenance put those two reasons in different branches. Both branches returned wrong results. 1. A grouped partitioning over a superset of the operation keys got no node at all, so a top-k window over `PARTITION BY id` on an `(id, name)`-partitioned table surfaced `id=1` twice, once per `(1,'aa')` and `(1,'bb')` partition. 2. A non-grouped partitioning over a superset got a node without a projection, so it coalesced by the full partition keys and left the operation key split. Any source reporting more than one split per partition value hits this. `SUM(price) OVER (PARTITION BY id)` over the same table with two splits for `(1,'aa')` returned 25.0 and 20.0 instead of 45.0. Both reach back to 4.2.0, where `GroupPartitionsExec` and this classification were introduced. A reduced storage-partitioned join reaches the same wrong result through a third shape. Joining an `identity(ts)`-partitioned table to a `years(ts)`-partitioned one under `v2BucketingAllowCompatibleTransforms` leaves both sides grouped on `(year, bucket)`, so two rows sharing a `ts` in different buckets sit on separate partitions. `SUM(v) OVER (PARTITION BY ts)` then returns 10 and 20 on `master` where the answer is 30 and 30, and the projection this PR inserts is what merges the two partitions. Before #58420 that same query threw `ClassCastException` at planning once the keys were read, so the wrong answer only became observable when that fix landed. ### Does this PR introduce _any_ user-facing change? Yes, it fixes a data correctness issue. With `v2BucketingAllowKeysSubsetOfPartitionKeys` enabled, a single-child operator whose keys are a subset of the partition keys now produces correct results, namely a window `PARTITION BY` and the single-pass aggregate shapes (`FlatMapGroupsInBatchExec`, `ArrowAggregatePythonExec`, `MapGroupsExec`). A two-phase SQL aggregate was already correct, because its partial `HashAggregate` is a `PartitioningPreservingUnaryExecNode`, so it narrows `KP([id, name])` to `KP([id])` before the final aggregate sees it. A cogroup is unaffected either way, because the projection there is left to the multi-child block, exactly as before. Two further changes are not gated on that config. The `requiredNumPartitions` rule applies whether `allowKeysSubsetOfPartitionKeys` is on or off, since a scan still needs `spark.sql.sources.v2.bucketing.enabled` to report a `KeyedPartitioning` at all, and `master` has no count check on the grouping path. A non-grouped `KeyedPartitioning` with 3 partitions and 2 distinct keys under `ClusteredDistribution([k], requiredNumPartitions = Some(3))` got a `GroupPartitionsExec` with 2 partitions on `master` and now gets a shuffle with 3. I could not find a query where such a distribution meets a `KeyedPartitioning` today, so this is robustness rather than a reachable wrong result. Only `StatefulOperatorPartitioning` and `AQEUtils` ever set the requirement. The `AQEUtils` one fires only over a `HashPartitioning` child. `StatefulOperatorPartitioning` sets it through a plain `ClusteredDistribution` when `spark.sql.streaming.statefulOperator.useStrictDistribution` is off, and a streaming scan never reports a `KeyedPartitioning`, because `Micr oBatchScanExec`, `ContinuousScanExec` and `RealTimeStreamScanExec` all leave `keyGroupedPartitioning` at `None`. The initial-state child of `flatMapGroupsWithState` and `transformWithState` is a batch plan though, so it can be a scan that does. That operator co-partitions its two children, so the shape is the multi-child block's to decide either way, and what changes here is only that a candidate which cannot honour the count no longer wins the resolution. The requirement is part of the `ClusteredDistribution` contract though, the two could meet more widely as this area grows, and honouring it costs nothing where nobody asks for a count, so it seemed the wrong thing to keep ignoring. The `MatchError` fix is gated on `v2BucketingAllowSorting` instead, which is also off by default. `explain` gains one label where the new projection happens. A `GroupPartitionsExec` inserted for a single-child operator now carries `joinKeyPositions`, so the node prints `JoinKeyPositions: [...]` where it printed nothing before. The name is historical, and the scaladoc now says the projection is not join-specific. This needs `allowKeysSubsetOfPartitionKeys` on, so no golden file moves. ### How was this patch tested? Added regression tests in `KeyGroupedPartitioningSuite`: - window top-k over `PARTITION BY` a subset of the partition keys, for both `PARTITION BY id` and the duplicated `PARTITION BY id, id` - window top-k over union output partitioning - a plain window over a subset of the partition keys on a non-grouped `KeyedPartitioning`, asserting the inserted node projects to the operation key rather than only coalescing - no `GroupPartitionsExec` and no shuffle when projecting to the operation keys merges nothing - a window over an inner join's two-member `PartitioningCollection`, keyed on the left side's first key column and the right side's remaining two, asserting no node is inserted because the left member needs none, even though the right one covers more operation keys - a window over a join that reduced one side's keys onto the other side's key space, asserting the two rows sharing a `ts` end up on one partition and in `EnsureRequirementsSuite`: - a `FlatMapCoGroupsInPandasExec` over `(n, i)`-partitioned children grouped on `i`, asserting both sides are grouped on `i` and not on `n` - a grouped `KeyedPartitioning` whose count differs from `requiredNumPartitions`, asserting the count is still honoured with a shuffle - an `(n, i)`-partitioned `KeyedPartitioning` whose count matches `requiredNumPartitions` but which needs a projection, asserting it falls back to a shuffle rather than to a node that would break the count it just passed - a projecting `KeyedPartitioning` whose *post-projection* count matches `requiredNumPartitions`, asserting it still groups on the operation key with no shuffle - a non-grouped `KeyedPartitioning` whose *post-grouping* count matches `requiredNumPartitions`, asserting it still groups without a shuffle - the same count rule with `allowKeysSubsetOfPartitionKeys` left at its default, asserting the shuffle - an `(id, years(ts))`-partitioned `KeyedPartitioning` clustered on those same two expressions, with and without `requireAllClusterKeys`, asserting no node is inserted when a cluster key is the partition expression itself - a `(bucket(4, a), b)`-partitioned `KeyedPartitioning` clustered on `a` alone, asserting the node projects to the bucket position. This is the one test where the reference-level match decides a kept position, and it is the shape the single-reference soundness argument is about. `master` inserts no node at all here, so the partitions sharing a bucket stay apart. - a single-partition `KeyedPartitioning` under `OrderedDistribution`, asserting planning no longer throws a `MatchError` - a `KeyedPartitioning` with no partition expressions, asserting it is still kept as it is, since without that guard it reaches the shuffle branch, where `UnspecifiedDistribution.createPartitioning` throws - a `KeyedPartitioning` whose expressions reference no column, asserting the child is shuffled rather than projected to no position at all - a `PartitioningCollection` whose two members cover a different number of operation keys, asserting the wider one supplies the projection - two members whose position sets are nested and whose projections leave the same number of partitions, asserting the containing set still wins, so the projection keeps naming the key the other would have dropped - two members whose position sets are nested, where only the narrower one's projection lands on `requiredNumPartitions`, asserting it is used rather than pruned and lost to the wider one - two members whose position sets are not nested and where the narrower one leaves more partitions, asserting the narrower projection wins over the wider coverage - two members covering one position each whose projections leave different numbers of partitions, asserting the one leaving the most supplies the projection - two members covering one position each whose projections leave the same number of partitions, asserting the one the child reports first wins, in both collection orders and in `ProjectedOrderingAndPartitioningSuite` a grouped and collapsed `KeyedPartitioning`, asserting `keysMaySatisfy` accepts it where `mayGroupToSatisfy` refuses it, because a grouped partitioning has no duplicate keys left for a node to coalesce. Seventeen of the twenty-four fail without the production change in this commit, measured on `master`. Those are the four window tests, the reference-free expressions, the transform position, the `MatchError`, four of the six `requiredNumPartitions` tests, all five multi-member ones and the `keysMaySatisfy` one, which does not even compile there. The other seven pass already and guard this PR's own insertion decision. Removing the condition each one covers makes it fail, with one exception worth naming. `the candidate covering the most operation keys wins` survives the removal of either the containment prune or the ranking, because either mechanism alone picks the same winner, so it earns its place by documenting the shape rather than by pinning a single condition. The join test is one of the seven. The base leaves that plan alone too, and it is here because an earlier revision of this PR inserted a node there. Removing the `isCoPartitioned` guard also fails two pre-existing SPJ te sts. `DistributionSuite` and `ShuffleSpecSuite` pass with 29 tests, and `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite` and `PlannerSuite` pass with 295. The window `PARTITION BY` tests come from #58245. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code Co-authored-by: Xiduo You <[email protected]> -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
