dongjoon-hyun commented on PR #58659: URL: https://github.com/apache/spark/pull/58659#issuecomment-5670952320
Thank you for the thorough write-up, @peter-toth. I went through the whole diff. The direction (decide once in the planner, let the node carry the answer) looks right, and I could confirm that `pickCoPartitionTarget` and `keysCanSatisfy` preserve the previous decisions except for the widening you describe. I have one concern that I think needs to be addressed before merging, plus a few smaller ones. The first item is based on reading the code paths; I did not reproduce it by running a query. ### 1. `withNewChildInternal` keeps a stale `grouping` / `outputPartitioning` when AQE replaces the child `GroupPartitionsExec.withNewChildInternal` is `copy(child = newChild)`, so both derived fields survive a child rewrite. The factory doc says this holds "because nothing in the tree hands this node a child that reports a different partitioning", but two AQE rules do exactly that: - `OptimizeShuffleWithLocalRead.createProbeSideLocalRead` runs `transformDown` over the whole stage plan and replaces the probe-side `ShuffleQueryStageExec` of any `BroadcastHashJoinExec` with an `AQEShuffleReadExec`. For a one-mapper-per-task local read its `outputPartitioning` is the pre-shuffle child's partitioning, otherwise `UnknownPartitioning`. - `CoalesceShufflePartitions.collectCoalesceGroups` coalesces the shuffle stages under a node whose children need no compatible partitioning (e.g. a BHJ) independently, and for a `KeyedPartitioning` shuffle the coalesced `AQEShuffleReadExec` reports `UnknownPartitioning`. Scenario: a first AQE round plans an SPJ with a one-side keyed shuffle of `X` onto keyed table `Y`. On re-optimization `Y` turns out to be small and the join becomes a BHJ, with `X`'s materialized keyed shuffle stage as its probe side; the BHJ's output partitioning is that keyed layout. A later operator then puts a `GroupPartitionsExec` over the BHJ output: the push branch of a second SPJ join (an identity grouping keeps the keyed claim even over a marked layout), or a single-child operator on a subset of the keys under `allowJoinKeysSubsetOfPartitionKeys`. When `optimizeQueryStage` then applies one of the two rules above: - Before this PR, `copy` minted a fresh instance, the lazy vals were re-derived from the new child, the node reported `UnknownPartitioning` (`case o => o`), the parent's requirement was no longer satisfied, and `ValidateRequirements.validate(applied, ...)` reverted the rule. - With this PR, the node keeps reporting the old keyed partitioning, so validation passes, and at execution `grouping.partitions` indexes into a child RDD whose partition count and layout have changed (number of mappers, or the coalesced count). That is an exception at best and silently wrong join output at worst. The old safety net depended on re-derivation, which is what this PR removes. Suggestions: - Make `withNewChildInternal` re-derive when `newChild.outputPartitioning != child.outputPartitioning` (the common wrapper insertions keep the partitioning, so the perf win stays). That needs the derivation inputs back, i.e. keep `expectedPartitionKeys` (transient) rather than only `expectedKeyCount`, or recover them from `grouping.partitions`. The re-derivation also has to tolerate a child with no keyed partitioning without throwing, as the old `outputPartitioning` did, so that `ValidateRequirements` can reject the plan instead. - Add an AQE test where a `GroupPartitionsExec` sits over a keyed shuffle stage and a local read / coalesce is attempted. `KeyGroupedPartitioningSuite` has essentially one test with AQE enabled today. ### 2. `satisfies` can now project the partition keys on every call `keysSatisfy`'s second branch calls `numPartitionsProjectedOn` whenever `allowJoinKeysSubsetOfPartitionKeys` is on and `isFunctionOfClusterKeys` is false. That allocates a row per partition and hashes it with the uncached `hashCode`, and `satisfies` is asked by `EnsureRequirements`, `ValidateRequirements` and every AQE rule application. Previously that projection happened once per position set through the memo in `resolveKeyedPartitioning`, which this PR also removes, so the same set can now be projected up to three times there (`satisfies`, `eligible`, `maxBy`). Since `KeyedPartitioning` is immutable, a small transient memo on the position set inside the class would restore the old cost. ### 3. The assert in `withJoinKeyPositions` `withJoinKeyPositions` asserts `expectedKeyCount.isEmpty && reducers.isEmpty && !distributePartitions`. On a re-run, an aligned node from an earlier pass can pass through `resolveChild` unchanged and reach `shuffleToCoPartition` with a projected paired spec, and the planner then dies on the assert. The pushed positions index the layout `g` reports, so wrapping (`GroupPartitionsExec(g, Some(positions))`) is a sound fallback for that shape instead of asserting. ### 4. Nit: `outputPartitioning` is the only non-`@transient` derived field It used to be `@transient lazy val`, and every other derived parameter of the node is `@transient`. This matches `ShuffleExchangeExec`, so it may be intentional, but a short note would help. For what it is worth, things I checked and found fine: all keyed members of a `PartitioningCollection` share one `KeyLayout`, so the relaxed filter in `PartitioningCollection.createShuffleSpec` does not widen what `ValidateRequirements` accepts; `QueryPlan.doCanonicalize` normalizes the stored `KeyedPartitioning` through `mapExpressions`, so exchange reuse keeps working; and the DSv2 scan keeps pruned partitions as `None`, so runtime filters do not shift the indices the grouping holds. -- 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]
