andygrove opened a new pull request, #7231: URL: https://github.com/apache/datafusion-comet/pull/7231
## Which issue does this PR close? Closes #7208. ## Rationale for this change With AQE on, `REBALANCE(k)` followed by `GROUP BY k` returns duplicate groups when the data is skewed, and a `PARTITION BY k` window restarts `row_number()` inside a key. Spark's `OptimizeSkewInRebalancePartitions` splits the skewed shuffle partition. In plain Spark that is safe, because `ValidateRequirements` rejects the split when an operator above needs `ClusteredDistribution(k)`. Comet's aggregate, window and join operators declared no `requiredChildDistribution`, so the check always passed and the native operators ran over the pieces of a key separately. ## What changes are included in this PR? `CometBaseAggregateExec`, `CometWindowExec`, `CometHashJoinExec` and `CometSortMergeJoinExec` now return `originalPlan.requiredChildDistribution`, so each declares the same requirement as the Spark operator it replaces. `requiredChildOrdering` and the broadcast hash join are unchanged. ## How are these changes tested? New test in `CometExecSuite`, `AQE does not split a skewed rebalance partition under operators that need clustering`. It uses the data from the issue and checks the `GROUP BY` query, the `row_number()` window query, and a self-join with sort-merge join on and off. Each query is compared with Spark, and the plan is asserted to have no skewed `AQEShuffleReadExec`. The test fails without the operator changes and passes with them. The full `CometExecSuite` passes on the default profile (Spark 4.1): 177 tests, 0 failures. -- 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]
