erenavsarogullari opened a new pull request, #54242: URL: https://github.com/apache/spark/pull/54242
### What changes were proposed in this pull request? Spark Adaptive Query Execution(AQE) framework coalesces the eligible shuffle partitions (e.g: empty/small-sized shuffle partitions) during the query execution. This optimization requires to be created coalesce groups for related Shuffle Stages (e.g: SortMergeJoin can have 1 ShuffleQueryStage per join-leg) to guarantee that both SMJ legs having the same num of partitions. To create coalesce groups for related ShuffleStages, Spark Plan Tree needs to be traversed by finding ShuffleQueryStages. https://issues.apache.org/jira/browse/SPARK-46590 has fixed incorrect coalesce grouping problem by adding BinaryExecNode Support for SparkPlan Tree traversal. This PR aims to introduce following complementary improvements on the top of SPARK-46590: **1-** Adding warning log message to ShufflePartitionsUtil.coalescePartitionsWithoutSkew() when numOfPartitions of ShuffleStages in the same coalesce group are not equal. This is required for the consistency because ShufflePartitionsUtil.coalescePartitionsWithSkew() logs warning message for the same case, **2-** Adding problematic shuffleStageIds to warning messages when numOfPartitions of ShuffleStages in the same coalesce group are not equal. This info can help for troubleshooting. **3-** Aligning the warning logs for specially for both ShufflePartitionsUtil.coalescePartitionsWithoutSkew() and coalescePartitionsWithSkew() cases **4-** 2 new UT cases are being added: Current UT Cases cover following use cases: ``` skewed SMJ under Union under BNLJ, skewed SMJ under Union under CartesianProduct ``` This PR also adds following new UT cases: **4.1-** skewed SMJ under Union under BHJ, **4.2-** non-skewed SMJ under Union under BHJ ### Why are the changes needed? - Warning message is useful for troubleshooting when numOfPartitions of ShuffleStages in the same coalesce group are not equal, - Problematic `ShuffleStageIds` can also help for the trouble shooting, - Additional UT coverage is also useful to verify coalesce grouping logic to avoid SPARK-46590 kinds of issues. ### Does this PR introduce _any_ user-facing change? Yes, adding new warning message when numOfPartitions of ShuffleStages in the same coalesce group are not equal ### How was this patch tested? Added 2 new UT cases for existing use-cases to test coalesce grouping logic such as: - skewed SMJ under Union under BHJ, - non-skewed SMJ under Union under BHJ ### Was this patch authored or co-authored using generative AI tooling? No -- 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]
