sunchao opened a new pull request, #7357: URL: https://github.com/apache/datafusion-comet/pull/7357
## Which issue does this PR close? Addresses the inherited logical-link case in #6034. The broader audit remains open. ## Rationale for this change Spark’s adaptive query execution (AQE) can switch a join to broadcast when an aggregation makes the build side small enough. With a mixed `COUNT(*)` and `COUNT(DISTINCT ...)`, the first shuffle can be above the broadcast threshold and the second well below it. Comet can lose the link to the second stage and keep a sort-merge join. ## What changes are included in this PR? Keep the current AQE stage on a reused Comet operator when the link is direct, or when the inherited stage wraps the operator’s original logical node. Compare nodes by identity, including through nested stage wrappers. When restoring the original link, copy its direct and inherited tags on the operator itself. Calling Spark’s `setLogicalLink` here promotes an inherited tag to a direct one and propagates it into children that can already belong to a newer stage. ## How are these changes tested? The SQL regression joins two synthetic ranges. Static broadcasting is disabled and AQE gets a 16 KiB broadcast threshold. On stock Spark 3.5.9, changing only the planner rule gave these results. | Regression check | Apache main | This PR | | --- | --- | --- | | Initial native broadcast joins | 0 | 0 | | Final join | Sort-merge | Native broadcast hash | | Rows matching the expected aggregates | 16 of 16 | 16 of 16 | The fixed plan’s first build shuffle reported 196,096 bytes and the second reported 1,280 bytes. The comparison measures plan selection and result parity. Rule tests also cover nested stage identity and restoring stale direct and inherited links. -- 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]
