ulysses-you opened a new pull request, #58279: URL: https://github.com/apache/spark/pull/58279
### What changes were proposed in this pull request? `EnsureRequirements` is not idempotent for a storage-partitioned join that uses a partially clustered distribution: re-running it on a plan it already produced stacks a second `GroupPartitionsExec` on top of the first, and the result duplicates rows. AQE hands the whole plan back to `EnsureRequirements` after rewriting some other join -- `ConvertSortMergeJoinToShuffledHashJoin` and `OptimizeSkewedJoin` both do this. On that second pass a join child is `SortExec(GroupPartitionsExec(...))` rather than a bare scan. A partially clustered `KeyedPartitioning` reports `isGrouped = false` by design, so the distribution step treats it as satisfied "only after grouping" and adds a plain `GroupPartitionsExec` on top; `applyGroupPartitions` then writes the join's `expectedPartitionKeys` and `distributePartitions` into that fresh outer node. The alignment is therefore re-derived from an already-aligned layout. The inner node replicates an input partition across the expected partitions, and the outer node concatenates those replicas back into a single partition before replicating again, so every row of the replicated side is emitted twice: ``` outer GroupPartitions groupedPartitions = [([1],[0,1]), ([1],[0,1]), ([2],[2])] inner GroupPartitions groupedPartitions = [([1],[0]), ([1],[0]), ([2],[1])] ``` This adds `rewriteGroupPartitions`, which descends to the innermost `GroupPartitionsExec` (through a local `SortExec` and through a redundant grouping), rewrites it, and drops what sits above it, reproducing the plan a single pass would have produced. Two things bound the change: * Only a *local* `SortExec` is traversed. A global one requires `OrderedDistribution`, which a `KeyedPartitioning` can satisfy (behind `spark.sql.sources.v2.bucketing.sorting.enabled`) through a `GroupPartitionsExec` built to emit the partition keys in sorted order; reusing that node would overwrite its `expectedPartitionKeys` and clear `distributePartitions`. * Dropping a grouping is safe only because this is reached from `checkKeyGroupCompatible`, which runs for joins alone. An operator with a single child -- an aggregate or a window over a partially clustered join -- genuinely needs its non-grouped input grouped and never gets here (see the partially-clustered aggregate and window tests in `KeyGroupedPartitioningSuite`). `withJoinKeyPositions` is routed through the same helper for consistency. It has no reproduction of its own: instrumentation shows it only ever receives a bare `GroupPartitionsExec` on the paths the suites exercise, where old and new code behave identically. ### Why are the changes needed? The join returns duplicated rows. Reproduction, against a DSv2 catalog that reports `KeyGroupedPartitioning`: ``` 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'), (2, 'bb'); ALTER TABLE sp1 ADD COLUMN extra STRING; -- opens a second split for id = 1 INSERT INTO sp1 VALUES (1, 'ab', 'x'); 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)`; the correct answer is `(1, 1, 2, 7)`. Two details are load-bearing. The `np1`/`np2` branch exists only to create a materialized shuffle stage, which is what makes `ConvertSortMergeJoinToShuffledHashJoin` fire and re-run `EnsureRequirements` over the whole plan -- the storage-partitioned join has no shuffle of its own, so it cannot trigger the re-run by itself. And `id = 1` must map to two input splits: in the in-memory test catalog the `ALTER TABLE` between the two inserts changes the write schema and opens the second split, while on a real connector this corresponds to a partition value with more than one data file whose files are not combined into a single task. A single pass is self-consistent: within one `ensureDistributionAndOrdering` call the distribution step creates the `GroupPartitionsExec` and `applyGroupPartitions` rewrites that same node. The stacking only appears once the rule is applied to its own output. ### 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`. For the query above, the result changes from `(1, 1, 1, 1, 2, 7)` to the correct `(1, 1, 2, 7)`. ### How was this patch tested? New tests: * `KeyGroupedPartitioningSuite`, "partially clustered join keeps its row count when EnsureRequirements re-runs" -- the end-to-end query above, whose row count was wrong before the fix. * `EnsureRequirementsSuite`, "only a local sort is looked through when reusing GroupPartitionsExec" -- covers the bare, local-sort and global-sort cases, so a `GroupPartitionsExec` serving a global sort's `OrderedDistribution` is never reused. * `EnsureRequirementsSuite`, "a single-child operator over a partially clustered layout still gets grouped" -- guards the case where stacking a second grouping is legitimate. Existing suites run locally: `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`, `EnsureRequirementsSuite`, `PlannerSuite`, `AdaptiveQueryExecSuite`, `DataFrameJoinSuite`, `DisableUnnecessaryBucketedScanWithoutHiveSupportSuite(AE)`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code -- 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]
