[SPARK-59050][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins

[ulysses-you]

What changes were proposed in this pull request?

The one-side shuffle optimization (spark.sql.sources.v2.bucketing.shuffle.enabled) re-shuffles a join side onto the keyed side's declared KeyedPartitioning. KeyGroupedPartitioner silently routes rows whose key is not among the declared partition keys to arbitrary partitions, so the re-shuffled side's KeyedPartitioning may not actually cover all its rows.

• Add KeyedPartitioning.mayContainUnknownPartitionKeys, set on every partitioning produced by KeyedShuffleSpec.createPartitioning, and require KeyedShuffleSpec.areKeysCompatible to only co-partition such a partitioning with a side whose keys are a subset of the declared keys (in the same order when both sides are flagged). Since the subset test compares key values, it is only applied between partitionings using the same partition function per position: an identity-vs-transform pair holds raw values on one side and transform outputs on the other, so their key sets are not comparable.
• Propagate the marker through PartitioningPreservingUnaryExecNode, GroupPartitionsExec and UnionExec. Because a boolean marker cannot express which keys are trustworthy, these sites drop the keyed claim entirely whenever the declared key set changes: a projection that drops a key position (coarsening the declared keys), a union whose merged keys are a superset of a flagged leg's keys, and GroupPartitionsExec with reducers (whose declared keys are the reduced keys an out-of-set key can reduce into).

Why are the changes needed?

A storage-partitioned join where the preserved side is one-side-shuffled (e.g. a RIGHT OUTER JOIN t with a keyed a and a non-keyed t holding keys a does not) produces an output whose declared partitioning omits those keys. A following storage-partitioned join trusts the declaration and silently loses the matches -- wrong results with no error.

Reproduction

Only spark.sql.sources.v2.bucketing.shuffle.enabled is needed (the other SPJ toggles are at their defaults). The keyed tables need a catalog that reports partitioning (here testcat is the in-memory test catalog).

-- 1) config
SET spark.sql.sources.v2.bucketing.shuffle.enabled = true;

-- 2) tables + data
-- a: keyed on id, keys {1, 2}
CREATE TABLE testcat.ns.a (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO testcat.ns.a VALUES (1, 'a1'), (2, 'a2');

-- t: v1 parquet, keys {1, 2, 3}  (contains a key `a` does not have)
CREATE TABLE t (id BIGINT, data STRING) USING parquet;
INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3');

-- u: keyed on id, keys {1, 2, 3}
CREATE TABLE testcat.ns.u (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO testcat.ns.u VALUES (1, 'u1'), (2, 'u2'), (3, 'u3');

-- 3) query: the RIGHT OUTER join preserves t's id=3 row; the following
--    storage-partitioned join must still match it against u.
SELECT r.id, u.data
FROM (
  SELECT t.id AS id
  FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id
) r
JOIN testcat.ns.u u ON r.id = u.id;

Result

	id
Expected	1, 2, 3
Actual (before this PR)	1, 2 <- id=3's match silently lost

Turning spark.sql.sources.v2.bucketing.shuffle.enabled off (or disabling SPJ entirely) returns all three rows, confirming this is an SPJ-path bug rather than a query semantics issue.

Mechanism

1. The first join one-side-shuffles t onto a's declared keys {1, 2}. t's id=3 is not among the declared keys, so KeyGroupedPartitioner silently routes it to an arbitrary partition, while the shuffle still declares the {1, 2} layout.
2. a RIGHT OUTER JOIN t preserves that misplaced id=3 row, so the join output's declared partitioning ({1, 2}) no longer matches its data.
3. The second join trusts the declared layout when planning its storage-partitioned join against u ({1, 2, 3}) and never co-locates the id=3 rows, silently dropping the match.

The fix marks such one-side-shuffled partitionings with mayContainUnknownPartitionKeys and restricts areKeysCompatible to only co-partition them with a side whose keys are a subset of the declared keys (in the same order when both sides are flagged), so the second join falls back to a shuffle and returns correct results. Downstream operators that change the declared key set (a key-dropping projection, a union with several legs, GroupPartitionsExec with reducers) drop the keyed claim entirely, since the marker alone cannot express which keys the claim still covers.

Two further wrong-results shapes the review measured are fixed the same way and covered by repro tests: a Project that drops a key position coarsens the declared key set (4 expected, 2 returned before the fix), and a UNION ALL leg that declares exactly the out-of-set key (4 expected, 3 returned). All three repro tests fail before the fix and pass after.

Does this PR introduce any user-facing change?

Yes: it fixes a wrong-results bug in spark.sql.sources.v2.bucketing.shuffle.enabled. Storage-partitioned joins whose preserved side is one-side-shuffled against a smaller keyed side, followed by another storage-partitioned join on a larger key set, now fall back to shuffles and return correct results. Previously they could silently lose matches.

The buggy config ships since 4.0.0, so this is a candidate for backporting to the 4.x lines.

How was this patch tested?

• KeyGroupedPartitioningSuite: identity and bucket one-side-shuffle repros (assert correct answers, shuffle count and mayContainUnknownPartitionKeys flags), outer-join variants, the key-dropping-project and union repros (union covered with both a disjoint and an overlapping second leg), subset-keyed partner compatibility, and a keyed-preserved-side regression (the one-side shuffle is still used when sound).
• ShuffleSpecSuite: unit tests for the areKeysCompatible subset/order rules and the same-partition-function requirement on the unknown-keyed path.
• Confirmed the bug-repro tests fail without the fix (wrong results) and pass with it.
• Regression suites all pass: KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, ShuffleSpecSuite, DistributionSuite, ValidateRequirementsSuite, ProjectedOrderingAndPartitioningSuite.

Was this patch authored or co-authored using generative AI tooling?

Yes. Generated-by: Claude Code.

[ulysses-you]

cc @peter-toth @cloud-fan if you have time to take a look, thank you

[peter-toth]

I can take a look at this tomorrow.

[peter-toth]

Thanks for the PR, @ulysses-you!

I reproduced the bug and confirmed all six new tests fail on base, so the mechanism and the fix's direction both hold: a KeyedPartitioning from KeyedShuffleSpec.createPartitioning describes only its declared keys, and restricting areKeysCompatible to a subset partner is the right guarantee. What a boolean cannot express is which keys are trustworthy, so the guarantee stops holding the moment a node changes the declared key set. I measured two queries that still return wrong results with this PR applied: a Project that drops a key position, and a UNION ALL whose other leg contributes exactly the out-of-set key. Both are this same bug, reached through two of the propagation sites the PR adds. Note that tightening areKeysCompatible would not catch either one -- in both the partner's key set is exactly equal to the declared set, so the gate has to be at the producer. The common fix is to drop the keyed partitioning entirely, not just the marker, wherever the declared key set changes; findings 1, 2 and 4 each spell out the site.

Blocking

• 1. Projection coarsens the declared key set: PartitioningPreservingUnaryExecNode carries the marker across a key projection, so an out-of-set full key can land inside the projected declared set. Measured: 4 expected rows, 2 returned, with only shuffle.enabled on.
• 2. Union widens the declared key set: mergedKeys is a superset of the flagged leg's keys, so another leg can declare the very key the flagged leg holds out-of-set. Measured: 4 expected rows, 3 returned.
• 3. Description overstates the fix: the description says the marker is propagated so "a downstream storage-partitioned operator never trusts a layout that may not cover the data". Findings 1 and 2 are counterexamples. Either close them here, or name the shapes that stay broken in the description and in the mayContainUnknownPartitionKeys scaladoc.

Non-blocking

• 4. Reducers coarsen the keys the same way: GroupPartitionsExec declares reduced keys, and an out-of-set key can reduce into the declared set. Derived from reading, not measured.
• 5. Subset test compares keys from different domains: isExpressionCompatible above it admits identity-vs-transform pairs, whose partitionKeys are then raw values on one side and transform outputs on the other, so declared.contains compares unrelated numbers.
• 6. Affects-version and backport: spark.sql.sources.v2.bucketing.shuffle.enabled ships since 4.0.0, but SPARK-59050 lists 5.0.0 only and the PR says nothing about backporting a silent wrong-results fix. Does the bug reach the 4.x lines, and is a backport planned?

Minor

• 7. Single-element loop: Seq("RIGHT OUTER").foreach iterates once; LEFT OUTER and FULL OUTER are spelled out separately below it.
• 8. Test names lack the ticket id: 89 of this suite's 110 tests carry a SPARK-xxxxx: prefix; the six new ones don't.
• 9. FULL OUTER block asserts less than its siblings: it omits the collectGroupPartitions(...).isEmpty check that would actually pin the comment's claim.

[peter-toth]

Finding 1. The marker travels through the projection but the declared key set does not. When projectablePositions.length < numPositions, sharedKeys is keySource.projectKeys(...), so two different full keys collapse onto one projected key -- and a key that was outside the declared set before the projection can be inside it after. areKeysCompatible's subset test then reads as a guarantee that no longer holds.

Measured on 83af168, only spark.sql.sources.v2.bucketing.shuffle.enabled=true and AQE off (as in the new tests):

-- a: keyed on (id, k), keys {(1,x),(2,x),(3,x),(4,x)};  u: keyed on id, keys {1,2,3,4}
-- t: v1 parquet with (1,z),(2,z),(3,z),(4,z) -- every (id,k) is out-of-set, no id is
SELECT r.id, u.k
FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t
      ON a.id = t.id AND a.k = t.k) r
JOIN testcat.ns.u u ON r.id = u.id

Expected 4 rows, got 2 ([1,u1], [3,u3]); with spark.sql.sources.v2.bucketing.enabled=false all 4 come back. The Project [id] above the RIGHT OUTER join drops key position 1, so the flagged partitioning declares {1,2,3,4}, u's keys are a subset, the second SortMergeJoin storage-partitions with no exchange, and the t rows whose (id,k) hashed to the wrong partition never meet their u match.

The marker is only sound while the key set survives verbatim, so drop the keyed claim when it doesn't -- right after projectablePositions is computed:

    if (projectablePositions.length < numPositions &&
        kps.exists(_.mayContainUnknownPartitionKeys)) {
      return LazyList.empty
    }

[ulysses-you]

Fixed in 355d388: projectKeyedPartitionings now returns an empty list when the projection drops a key position (projectablePositions.length < numPositions) and any input KP is flagged, so the keyed claim is dropped entirely instead of being carried across the coarsened declared set. Added a wrong-results repro test (SPARK-59050: SPJ: project dropping a key position drops the unknown-keyed claim) that fails on the previous commit with 4 expected / 2 returned, matching your measurement, and passes after the fix.

[peter-toth]

Finding 2. mergedKeys is the concatenation of every leg's keys, so the merged declared set is a superset of the flagged leg's. A key that the flagged leg holds out-of-set can be contributed as a declared key by another leg, and areKeysCompatible then accepts a partner that holds it.

Measured on 83af168, same single config. s is keyed on id with key {3} -- exactly the key t has and a does not:

-- a keyed {1,2};  t v1 parquet {1,2,3};  s keyed {3};  u keyed {1,2,3}
SELECT r.id, u.data
FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id
      UNION ALL
      SELECT id FROM testcat.ns.s) r
JOIN testcat.ns.u u ON r.id = u.id

Expected 4 rows (id=3 matches u twice, once via t and once via s), got 3 -- t's id=3 is lost. The union declares [1,2,3], which equals u's keys, so the second join storage-partitions with no GroupPartitionsExec and no exchange.

This is the shape "SPJ: union preserves the unknown-partition-keys marker" covers, but with a deliberately disjoint leg (s = {4,5}), which is why it passes.

The merged claim can't be trusted once a leg is flagged, so fall back before building it:

      if (compatible) {
        if (kps.exists(_.mayContainUnknownPartitionKeys)) {
          return super.outputPartitioning
        }

A laxer rule would keep the merged partitioning when every flagged leg's key set already equals mergedKeys, but with more than one leg that is only the all-legs-identical case. The existing union test would then assert the opposite -- that the keyed partitioning is dropped.

[ulysses-you]

Fixed in 355d388: UnionExec now returns super.outputPartitioning whenever any leg is flagged, before building mergedKeys. Since the merged set is the concatenation of every leg's keys, with multiple legs it is always larger than any single leg's declared set, so there is no sound way to keep the merged claim. Updated the union test accordingly: it now asserts the merged keyed partitioning is dropped (no GroupPartitionsExec, the union side re-shuffles), and covers both a disjoint second leg and your wider shape where the second leg declares exactly the out-of-set key (fails 4 expected / 3 returned before the fix, passes after).

[peter-toth]

Finding 4. Same shape as findings 1 and 2, one hop earlier. With reducers defined, groupedPartitions is keyed by the reduced keys, so this partitioning declares a coarsened set -- and an out-of-set key can reduce into it. A flagged side with identity keys {0,1,2,3} that holds an out-of-set id=4, reduced by bucket(4, .), declares {0,1,2,3} again, and bucket(4,4)=0 is one of them, while the id=4 rows sit in nonNegativeMod(hash(4), 4).

I did not build this one -- it needs spark.sql.sources.v2.bucketing.allowCompatibleTransforms=true on top of the shuffle config, and the reducer has to be picked for the flagged side. The experiment that settles it: a keyed identity(id) with ids 0..3, t v1 holding id 4, u keyed bucket(4, id) holding id 4, then a RIGHT OUTER JOIN t joined to u on id. Same fix shape as findings 1 and 2: don't propagate the keyed claim when reducers.isDefined.

[ulysses-you]

Guarded in 355d388: GroupPartitionsExec.outputPartitioning returns UnknownPartitioning(0) when reducers.isDefined and any input KP is flagged. Note this is deliberately not super.outputPartitioning -- for a UnaryExecNode that is the child's flagged keyed partitioning, which is equally untrustworthy here since the regrouped data layout no longer matches the child's declared keys. I did not build the end-to-end experiment (it needs allowCompatibleTransforms plus the reducer to be picked for the flagged side), so this is a conservative guard rather than a repro-tested fix -- happy to add the experiment as a follow-up if you'd like.

[peter-toth]

Finding 5. This compares the two sides' partitionKeys directly, but isExpressionCompatible just above admits an AttributeReference against a TransformExpression (and two different-but-compatible transforms) when allowCompatibleTransforms is on. In those cases the two key sequences live in different domains -- raw id values on one side, bucket ids on the other -- so declared.contains is comparing unrelated numbers.

It can reject a sound pairing, and it can accept an unsound one: flagged identity keys {0,1,2,3} "contain" bucket keys {0,1}, so the pairing passes while the question that matters (are the flagged side's out-of-set ids in the partition their bucket label points at?) is never asked. Also from reading, not measured.

EnsureRequirements computes leftReducedKeys / rightReducedKeys a few lines below the areKeysCompatible call, so comparing in the reduced key space is available; refusing the marker path outright unless both sides are the same function would also close it.

[ulysses-you]

Fixed in 355d388 (the 'refuse the marker path unless both sides are the same function' option): when either side is flagged, areKeysCompatible now requires the partition expressions to match per position -- two leaf attributes, or TransformExpression.isSameFunction -- and returns false otherwise, so the subset comparison always happens within one key domain. Unflagged pairs keep the existing behavior (the identity-vs-transform pair stays admissible and EnsureRequirements reconciles the domains via reducers). Unit-tested in ShuffleSpecSuite.

[peter-toth]

Finding 7. This iterates a single element -- and LEFT OUTER and FULL OUTER each get their own block below rather than joining the loop. Either inline the body and drop the loop, or fold the two later blocks in (they differ in the assertions, so inlining is probably simpler).

[ulysses-you]

Inlined in 355d388 -- the RIGHT OUTER block now stands on its own alongside the FULL OUTER and LEFT OUTER blocks.

[peter-toth]

Finding 8. 89 of this suite's 110 tests are prefixed with their ticket id; these six aren't. Worth SPARK-59050: SPJ: ... on each so they are greppable from the ticket.

[ulysses-you]

Done in 355d388 -- all six new tests are now prefixed with SPARK-59050:.

[peter-toth]

Finding 9. The other three blocks in this test also assert collectGroupPartitions(...).isEmpty. The comment here says FULL OUTER is safe because it exposes UnknownPartitioning, and that is exactly what the missing assertion would pin -- as written, nothing fails if a future change lets the downstream join storage-partition on this layout.

[ulysses-you]

Added in 355d388 -- the FULL OUTER block now also asserts collectGroupPartitions(...).isEmpty, pinning the comment's claim.

[ulysses-you]

Thanks for the careful review, @peter-toth -- and for measuring the two extra wrong-results shapes. I've pushed 355d388 addressing the findings; inline replies are on the individual threads. On the two body findings that had no thread:

• 3 (description overstated the fix): rewrote the description. It now says explicitly that a boolean marker cannot express which keys are trustworthy, so the sites that change the declared key set drop the keyed claim entirely: a key-dropping projection, a union whose merged keys are a superset of a flagged leg's keys, and GroupPartitionsExec with reducers. Findings 1, 2 and 4 are all closed in this PR (each with a repro or a guard), so no shapes are knowingly left broken.
• 6 (affects-version and backport): you're right -- spark.sql.sources.v2.bucketing.shuffle.enabled ships since 4.0.0 and a silent wrong-results fix is worth backporting. I've added a note to the description and will update SPARK-59050's affected versions and open backport PRs to the 4.x lines once this merges.

One honest gap: finding 4 (reducers + flagged) is implemented as a conservative guard but without an end-to-end repro test, since the experiment needs allowCompatibleTransforms plus the reducer to be picked for the flagged side -- flagged in my inline reply if you'd like it added.

[peter-toth]

Re-checked through 948df55. Findings 1, 2, 5, 7, 8 and 9 are resolved — I re-measured the project and union repros against the current base (0-of-4 and 3-of-4 rows there, correct here) — and 6 is answered. Finding 4's guard landed, and I built the end-to-end experiment it was missing: it throws at planning time on a query that returns correct results on master, so it comes back as finding 11. I also found a fourth site that coarsens the declared key set and keeps the marker, which I missed in round 1.

One coordination note: my own #58351 renames isNarrowed to isCollapsed and redefines it as actual key collapse, and it edits the same constructor and the same four propagation sites, so whichever of the two lands second needs a rebase. It changes what the fix for finding 10 should look like and what finding 12 should not do, so both are written with it in mind.

Blocking

• 11. Reducers guard throws at planning (new): GroupPartitionsExec returns UnknownPartitioning(0), and the 0 breaks PartitioningCollection's uniform-numPartitions requirement, so the join above it throws. Measured on a query that returns four correct rows on master. The guard also only ever fires on a spurious marker, because finding 5's same-function rule already keeps reducers away from a flagged spec. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:84]
• 10. Projected join keys coarsen the declared set (late catch): createShuffleSpec projects the declared keys under allowKeysSubsetOfPartitionKeys and carries the marker onto the coarsened set — the same hazard finding 1 fixed for Project. Measured: 4 expected rows, 0 returned. GroupPartitionsExec's joinKeyPositions projection does the same one hop later. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:664]
• 3. Description overstates the fix (regressed): the rewritten description says these sites "drop the keyed claim entirely whenever the declared key set changes". Finding 10 is a counterexample at two more sites — close them, or name them in the description and in the mayContainUnknownPartitionKeys scaladoc.

Non-blocking

• 12. A spurious marker costs a sound SPJ (new): kps.exists(...) reads the marker off an inner join's mixed-flag PartitioningCollection, where the join has already dropped every out-of-set row. Measured: a key-dropping projection over such a join plans 1 shuffle on master and 2 here. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala:140]

Minor

• 13. Layout invariant still stated absolutely (new): the new @param qualifies what the == Partition Keys == narrative asserts flatly. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:549]

[peter-toth]

Finding 11. Two problems on this line.

UnknownPartitioning(0) claims zero partitions while this node produces groupedPartitions.size of them. Every other operator that gives up on describing its partitioning reports the real count — PartitioningPreservingUnaryExecNode uses UnknownPartitioning(child.outputPartitioning.numPartitions), ShuffledJoin's FullOuter arm the same. The 0 breaks PartitioningCollection's uniform-numPartitions requirement as soon as an inner join above this node builds one from both sides.

Measured on 948df55 with spark.sql.sources.v2.bucketing.shuffle.enabled=true and spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled=true:

-- a keyed bucket(4, id), ids 0..3;  t v1 parquet, ids 0..3;  u keyed bucket(2, id), ids 0..3
SELECT * FROM testcat.ns.a a JOIN t ON a.id = t.id JOIN testcat.ns.u u ON a.id = u.id
ORDER BY a.id

master returns the four rows; this commit throws:

java.lang.IllegalArgumentException: requirement failed: PartitioningCollection requires all of its partitionings have the same numPartitions.
  at org.apache.spark.sql.catalyst.plans.physical.PartitioningCollection.<init>(partitioning.scala:859)
  at org.apache.spark.sql.catalyst.plans.physical.PartitioningCollection$.fromPartitionings(partitioning.scala:978)
  at org.apache.spark.sql.execution.joins.ShuffledJoin.outputPartitioning(ShuffledJoin.scala:73)
  at ...EnsureRequirements.ensureDistributionAndOrdering(EnsureRequirements.scala:69)

The left side is GroupPartitions ... Reducers: [BucketReducer(2)] over the first join, reporting 0; the right side is a GroupPartitionsExec over u reporting 2. The ORDER BY is only there to make something above the top join ask for its outputPartitioning; a third join or an aggregate does it too.

Suggested change

	return UnknownPartitioning(0)
	return UnknownPartitioning(groupedPartitions.size)

The second problem is why the guard fires at all. It needs reducers defined and a flagged KP in the child's partitioning, and finding 5's fix makes those nearly exclusive: when the spec is flagged, areKeysCompatible now requires isSameFunction per position, and a same-function pair has no reducer (BucketFunction.reducer returns null when gcd == thisNumBuckets, DaysFunction.reducer(DaysFunction) returns null). The only path I found is the one above — the first join is InnerLike, so its outputPartitioning is a PartitioningCollection holding a's unflagged KP and the re-shuffled side's flagged one, createKeyedShuffleSpec's collectFirst picks the unflagged one so areKeysCompatible never sees a marker and the reducer is computed, while p.exists here still finds the flagged sibling. In that shape the marker is spurious: the inner join has already dropped every row whose key is outside the declared set, so the reduced claim is sound. Worth scoping the predicate to the KP the reduction is actually about, or dropping the branch and recording in the scaladoc why a flagged spec cannot carry reducers (a third-party ReducibleFunction returning a self-reducer is the remaining hole).

[peter-toth]

Finding 10. This projects the declared keys and carries the marker onto the projected set — the same coarsening finding 1 fixed in projectKeyedPartitionings. Two full keys collapse onto one projected key, so a key that was outside the declared set can be inside the projected one, and areKeysCompatible's subset test reads as a guarantee that no longer holds. I missed this site in round 1; my list of producers named the projection, the union and the reducers, and not this one.

Measured on 948df55 with spark.sql.sources.v2.bucketing.shuffle.enabled=true and spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled=true, AQE off as in the new tests:

-- a keyed (id, k) = {(1,x),(2,x),(3,x),(4,x)};  u keyed id = {1,2,3,4}
-- t v1 parquet = (1,z),(2,z),(3,z),(4,z) -- every (id,k) is out-of-set, no id is
SELECT r.id, r.k, u.data
FROM (SELECT t.id AS id, t.k AS k FROM testcat.ns.a a RIGHT OUTER JOIN t
      ON a.id = t.id AND a.k = t.k) r
JOIN testcat.ns.u u ON r.id = u.id

r.k stays in the outer select list, so the projection keeps both positions and your new guard in projectKeyedPartitionings does not fire. Expected 4 rows, 0 returned — the same as base, so this shape is not fixed. The second join storage-partitions on the coarsened {1,2,3,4}: GroupPartitions JoinKeyPositions: [0] ExpectedPartitionKeys: 4 on both sides, and no exchange between them.

GroupPartitionsExec.outputPartitioning coarsens the same way one hop later (sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:88-90): with joinKeyPositions defined it declares the projected keys and keeps the marker, and only the reducers case drops the claim.

On the fix: please do not route it through isNarrowed. The hazard is about keys that are not in the declared list at all — an out-of-set full key projecting into the projected declared set — while the flag only describes how distinct the declared keys are among themselves. Today's isNarrowed would happen to be true here because it records provenance, but #58351 redefines it as actual key collapse and renames it isCollapsed, and in the repro above the four declared keys stay four distinct projected keys, so it reads false. A direct test is both clearer and stable across that change: refuse the marker path whenever joinKeyPositions.length < expressions.length. EnsureRequirements.createKeyedShuffleSpec.tryCreate is the natural gate — it already returns an Option and still holds the unprojected arity, so returning None sends the child down the ordinary shuffle path. Doing it here by returning the unprojected result instead would be shorter, but it risks createPartitioning's clustering(positionSet.head) on a position that maps to no clustering key, which is why the projection exists.

[peter-toth]

Finding 12. kps can hold a flagged and an unflagged KeyedPartitioning at the same time: PartitioningCollection's invariant only requires a shared partitionKeys reference and equal arity, and ShuffledJoin.outputPartitioning builds exactly that for InnerLike — the keyed side's own KP plus the re-shuffled side's flagged one. In that collection the marker is spurious: the inner join has already dropped every row whose key is outside the declared set, so the surviving rows do obey the declared layout. exists drops a keyed claim that is sound.

Measured on 948df55 with only spark.sql.sources.v2.bucketing.shuffle.enabled=true:

-- a keyed (id, k) = {(1,x)..(4,x)};  t v1 parquet = (1,x)..(4,x);  u keyed id = {1..4}
SELECT r.id, u.data
FROM (SELECT t.id AS id FROM testcat.ns.a a JOIN t ON a.id = t.id AND a.k = t.k) r
JOIN testcat.ns.u u ON r.id = u.id

Four correct rows on master and here, but master plans one shuffle (the first join's one-side shuffle) and this commit plans two — the second join loses its storage-partitioned join.

Filtering kps to the unflagged ones before this guard, and keeping the drop when none remain, would preserve it. The all-flagged case genuinely has to drop: with both join sides re-shuffled onto equal declared keys, an out-of-set key hashes to the same index on both sides, so the join can emit matched rows whose key is not declared. The p.exists in GroupPartitionsExec has the same scoping question (finding 11).

One thing to watch when this meets #58351: that PR normalizes isCollapsed across a collection by OR and adds a require that the members agree. Please do not extend either to this marker. isCollapsed describes the shared physical layout, so every member naming that layout is equally coarse; this marker describes which rows the side that produced it may hold, and an inner join filters those rows away. Requiring agreement would make the mixed collection above impossible to represent and would lock in the lost SPJ.

[peter-toth]

Finding 13. This @param qualifies an invariant that the == Partition Keys == section states flatly a hundred lines above (sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:463): "partition i holds key partitionKeys(i). A consumer must therefore treat the given order as authoritative rather than re-derive one." That paragraph is where a consumer looks for the layout contract, so one sentence there — a partition may also hold rows whose key is not declared at all when mayContainUnknownPartitionKeys is set — keeps the two in sync. #58351 restructures this class doc around key collapse, so land the sentence wherever the layout paragraph ends up after the rebase.