[SPARK-58968][SQL] Fix SPJ allowKeysSubsetOfPartitionKeys correctness for non-join operators

[peter-toth]

What changes were proposed in this pull request?

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 KeyedPartitionings by isGrouped and then had two
branches that each had to insert a GroupPartitionsExec. This PR changes the classification to
be about what still has to happen to the data:

• splitKeyedPartitionings now takes the required distribution and returns
  (other, satisfying, needsGrouping). needsGrouping pairs each partitioning with the
  partition expression positions the node has to project to.
• A new projectionPositions helper computes those positions, from the joinKeyPositions that
  KeyedPartitioning.createShuffleSpec already derives.
• The four-way match on the required distribution collapses to three cases, with a single place
  that inserts GroupPartitionsExec.

projectionPositions reports "no projection needed" in two cases. When projecting would not
merge any partition that coalescing duplicate partition keys alone leaves apart, every operation
key already lives on a single partition, so the partitioning satisfies the distribution as it is;
keeping it is also better than projecting, because KeyedPartitioning([id, name]) and
KeyedPartitioning([id]) then describe the same number of partitions and only the first lets a
downstream operator co-partition on name too. And when no partition expression's reference is
a cluster key at all there is no position to project to - requireAllClusterKeys matches at the
expression level, so a transform or nested-field partition expression can satisfy the distribution
while KeyedShuffleSpec.keyPositions yields nothing, and projecting to an empty position list
would collapse every partition into one.

The classification asks satisfies for an already grouped partitioning, not groupedSatisfies,
because only satisfies also enforces Distribution.requiredNumPartitions - and satisfies is how
this rule used to reach a grouped partitioning. A non-grouped one is left to groupedSatisfies
alone, also as before. Why that requirement is enforced on this path and not on the others is worth
a separate look.

For an operator that co-partitions more than one child no projection is done here: the multi-child
block below owns it, through checkKeyGroupCompatible for a storage-partitioned join and through
withJoinKeyPositions otherwise. Projecting inline as well would apply the projection twice, the
second time with positions that index into the unprojected partition expressions. Note that
withJoinKeyPositions is handed the best spec's positions for every child, so that delegation is
only correct when both sides' positions agree - a pre-existing limitation, unchanged here.

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 left as a follow-up.

The KeyedPartitioning scaladoc is updated as well: it described nonGroupedSatisfies /
groupedSatisfies as methods "called on non-grouped KPs", and taught isGrouped as the axis this
PR replaces.

Why are the changes needed?

With v2BucketingAllowKeysSubsetOfPartitionKeys enabled, KeyedPartitioning.groupedSatisfies
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; 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 - it 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')/(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.

Collapsing the two branches is what makes the second case impossible to forget again, but it also
opens two ways for a single insertion point to be wrong, which is what the two "no projection
needed" cases above guard: inserting a node that merges nothing, and projecting to no position at
all. Neither can happen on master, where a grouped partitioning got no node.

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 -
a window PARTITION BY, and the single-pass aggregate shapes (FlatMapGroupsInBatchExec,
ArrowAggregatePythonExec, MapGroupsExec). A two-phase SQL aggregate was already correct: 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.

The config is disabled by default, so nothing changes unless it is turned on.

How was this patch tested?

Added regression tests in KeyGroupedPartitioningSuite:

• window top-k over PARTITION BY a subset of the partition keys
• window top-k over a duplicated PARTITION BY key
• 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

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 under a ClusteredDistribution carrying requiredNumPartitions,
  asserting the count is still honoured with a shuffle
• a years(k)-partitioned KeyedPartitioning under a requireAllClusterKeys distribution,
  asserting no node is inserted where there is no position to project to

The three window top-k tests and the non-grouped plain-window test fail without the
EnsureRequirements change. The other four pass on master already and guard this PR's single
insertion point: each of them fails when its own condition is removed and passes under the other
ablations.

KeyGroupedPartitioningSuite, EnsureRequirementsSuite, PlannerSuite,
ProjectedOrderingAndPartitioningSuite, AdaptiveQueryExecSuite, the execution.exchange and
RemoveRedundant* suites and the connector distribution suites pass.

The three window tests come from #58245.

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

Generated-by: Claude Code

[ulysses-you]

I'm fine with this fix, the main change is same that adding an extra GroupPartitionExec with projected key position for non-join operators. This fix makes change inlines the method splitKeyedPartitionings and return the needsGrouping which including the key positions.