ulysses-you opened a new pull request, #58561:
URL: https://github.com/apache/spark/pull/58561
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR:
https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g.,
'[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a
faster review.
7. If you want to add a new configuration, please read the guideline first
for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the
guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
Backport of SPARK-58996 (apache/spark#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
```
```sql
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`.
--
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]
