tgravescs commented on a change in pull request #29224:
URL: https://github.com/apache/spark/pull/29224#discussion_r460120367
##########
File path:
sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
##########
@@ -731,6 +755,30 @@ class MyExtensions extends (SparkSessionExtensions =>
Unit) {
}
}
+object QueryPrepRuleHelper {
+ val myPrepTag: TreeNodeTag[String] = TreeNodeTag[String]("myPrepTag")
+ val myPrepTagValue: TreeNodeTag[String] =
TreeNodeTag[String]("myPrepTagValue")
+}
+
+// this rule will run during AQE query preparation and will write custom tags
to each node
+case class MyQueryStagePrepRule() extends Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = plan.transformDown {
+ case plan =>
+ plan.setTagValue(QueryPrepRuleHelper.myPrepTag,
QueryPrepRuleHelper.myPrepTagValue)
+ plan
+ }
+}
+
+// this rule will run during AQE query stage optimization and will verify
custom tags were
+// already written during query preparation phase
+case class MyNewQueryStageRule() extends Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = plan.transformDown {
+ case plan if !plan.isInstanceOf[AdaptiveSparkPlanExec] =>
+ assert(plan.getTagValue(QueryPrepRuleHelper.myPrepTag).get ==
QueryPrepRuleHelper.myPrepTagValue)
Review comment:
need to wrap this line
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
