[SPARK-58996][SQL][4.2] Fix SPJ partially clustered data correctness when EnsureRequirements re-runs

[ulysses-you]

What changes were proposed in this pull request?

Backport of SPARK-58996 (#58279, already merged to master, branch-4.x and branch-4.3) to branch-4.2.

EnsureRequirements is not idempotent for a storage-partitioned join under a partially clustered distribution: re-running it on a plan it already produced stacks a second GroupPartitionsExec over the first and duplicates rows. The fix makes the rule reuse the grouping an earlier pass inserted instead of re-deriving the alignment from an already-aligned layout:

• innermostGroupPartition descends only through a GroupPartitionsExec and the local SortExec this rule itself added; both the rewrite and the reads below go through it.
• rewriteGroupPartitions rewrites the innermost node and drops redundant groupings stacked above it, reproducing the plan a single pass would have produced.
• unwrapGroupPartitions peels to the pre-alignment plan so the statistics and original partition keys are read from that plan on every pass.
• applyGroupPartitions keeps the joinKeyPositions a reused node already holds (they were computed against the raw partition keys) instead of re-projecting a second time.
• withJoinKeyPositions, which every multi-child operator reaches and not just joins, reuses only a topmost GroupPartitionsExec.
• The ShuffleExchangeExec site strips every grouping this rule inserted before re-shuffling.

On branch-4.2 this additionally introduces a small PartitioningCollection companion object (with representativeOf and numKeyedPartitions) in partitioning.scala, mirroring master; branch-4.2's PartitioningCollection has no companion object yet.

Why are the changes needed?

The join returns duplicated rows when AQE re-runs EnsureRequirements. The re-run is not an AQE quirk: AdaptiveSparkPlanExec builds one rule instance, and ConvertSortMergeJoinToShuffledHashJoin and OptimizeSkewedJoin hand the whole tree back to it after rewriting some other join, all within one queryStagePreparationRules pass. A join child then arrives as SortExec(GroupPartitionsExec(...)); the partially clustered KeyedPartitioning reports isGrouped = false by design, so the distribution step adds a plain GroupPartitionsExec on top, and applyGroupPartitions rewrites that fresh outer node instead of the inherited one. A replicating grouping repeats every row, so the row set is emitted twice. Reproduction:

spark.sql.sources.v2.bucketing.enabled=true
spark.sql.sources.v2.bucketing.pushPartValues.enabled=true
spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled=true
spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold=100m

CREATE TABLE sp1 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO sp1 VALUES (1, 'aa'), (1, 'ab'), (2, 'bb');   -- two splits for id = 1

CREATE TABLE sp2 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO sp2 VALUES (1, 'p'), (2, 'q');

CREATE TABLE np1 (id BIGINT, data STRING); INSERT INTO np1 VALUES (7, 'x');
CREATE TABLE np2 (id BIGINT, data STRING); INSERT INTO np2 VALUES (7, 'y');

SELECT /*+ MERGE(a, b) */ a.id AS k FROM sp1 a JOIN sp2 b ON a.id = b.id
UNION ALL
SELECT c.id AS k FROM np1 c JOIN np2 d ON c.id = d.id;

returns (1, 1, 1, 1, 2, 7) instead of (1, 1, 2, 7).

Does this PR introduce any user-facing change?

Yes. It fixes wrong results (duplicated rows) for a storage-partitioned join under a partially clustered distribution when AQE re-runs EnsureRequirements. As on master, the partition-count fallback that picks the replicate side when there are no plan statistics now compares the pre-alignment split counts instead of the aligned report's distinct keys, which can change the replicated side on a first pass.

How was this patch tested?

• KeyGroupedPartitioningSuite: new tests for the row count, the replicate-side choice, and the subset-key join key positions under an EnsureRequirements re-run. The subset-key test sets spark.sql.requireAllClusterKeysForCoPartition=false, because on branch-4.2 a join on a subset of the partition keys only engages storage-partitioned join with that config off.
• EnsureRequirementsSuite: new unit tests for the local-sort reuse bound, topmost-only withJoinKeyPositions reuse, tag preservation, the shuffle site, the fallback counting, and the intentional single-child wrap.
• Both suites pass on branch-4.2.

[ulysses-you]

this is an additionally introducing for branch-4.2

[ulysses-you]

Merge Summary:

• merged into branch-4.2 bc8e972

Posted by merge_spark_pr.py

[ulysses-you]

thank you @peter-toth