sunchao commented on code in PR #7357:
URL: https://github.com/apache/datafusion-comet/pull/7357#discussion_r4515456090
##########
spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala:
##########
@@ -1429,7 +1429,21 @@ class CometExecRuleSuite extends CometTestBase {
}
}
- test("CometExecRule repairs ordinary and inherited logical links from the
original plan") {
+ test("CometExecRule preserves an inherited AQE stage for its original
logical subtree") {
+ withAdaptiveAggregateConf {
+ val aggregate = createAdaptiveAggregate()
+ val original = aggregate.originalPlan.logicalLink.get
+ val stage = LogicalQueryStage(LogicalQueryStage(original, aggregate),
aggregate)
+ aggregate.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
+ aggregate.setTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG, stage)
+
+ val repaired = applyCometExecRule(aggregate)
+ assert(repaired.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isEmpty)
+ assert(repaired.logicalLink.exists(_ eq stage))
+ }
+ }
Review Comment:
Updated the existing nested-stage test to exercise both direct and inherited
links on the original operator. In both cases it checks that the Comet operator
keeps the exact outer stage and its direct tag stays empty. The neighboring
stale-link test covers the original with no tag.
I also tried a direct-only lookup on the original. The direct case passed
and the inherited case failed at the stage identity assertion on Spark 3.5.9
and 4.1.3. With the PR code, the full planner-rule suite passed on both
versions.
[86f96624a8](https://github.com/apache/datafusion-comet/pull/7357/commits/86f96624a8bcfc5ceb7b265e8a0922f7c14167ce)
--
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]
