This is an automated email from the ASF dual-hosted git repository.
philo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new f8ad645c93 [TEST] Fix gluten test util getExecutedPlan (#8374)
f8ad645c93 is described below
commit f8ad645c93fb3205b0e689888f82480665ea667f
Author: Jinhai （陈锦海） <[email protected]>
AuthorDate: Mon Jan 6 14:48:11 2025 +0800
[TEST] Fix gluten test util getExecutedPlan (#8374)
---
.../org/apache/spark/sql/GlutenQueryTest.scala | 46 +++++++---------------
1 file changed, 15 insertions(+), 31 deletions(-)
diff --git
a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
index 8507233a57..b70ef26b3c 100644
--- a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
+++ b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
@@ -317,32 +317,23 @@ abstract class GlutenQueryTest extends PlanTest {
assert(executedPlan.exists(plan =>
plan.find(_.isInstanceOf[TransformSupport]).isDefined))
}
- /**
- * Get all the children plan of plans.
- * @param plans:
- * the input plans.
- * @return
- */
- def getChildrenPlan(plans: Seq[SparkPlan]): Seq[SparkPlan] = {
- if (plans.isEmpty) {
- return Seq()
- }
-
- val inputPlans: Seq[SparkPlan] = plans.map {
- case stage: ShuffleQueryStageExec => stage.plan
- case plan => plan
+ private def getExecutedPlan(plan: SparkPlan): Seq[SparkPlan] = {
+ val subTree = plan match {
+ case exec: AdaptiveSparkPlanExec =>
+ getExecutedPlan(exec.executedPlan)
+ case cmd: CommandResultExec =>
+ getExecutedPlan(cmd.commandPhysicalPlan)
+ case s: ShuffleQueryStageExec =>
+ getExecutedPlan(s.plan)
+ case plan =>
+ plan.children.flatMap(getExecutedPlan)
}
- var newChildren: Seq[SparkPlan] = Seq()
- inputPlans.foreach {
- plan =>
- newChildren = newChildren ++ getChildrenPlan(plan.children)
- // To avoid duplication of WholeStageCodegenXXX and its children.
- if (!plan.nodeName.startsWith("WholeStageCodegen")) {
- newChildren = newChildren :+ plan
- }
+ if (plan.nodeName.startsWith("WholeStageCodegen")) {
+ subTree
+ } else {
+ subTree :+ plan
}
- newChildren
}
/**
@@ -353,14 +344,7 @@ abstract class GlutenQueryTest extends PlanTest {
* A sequence of executed plans.
*/
def getExecutedPlan(df: DataFrame): Seq[SparkPlan] = {
- df.queryExecution.executedPlan match {
- case exec: AdaptiveSparkPlanExec =>
- getChildrenPlan(Seq(exec.executedPlan))
- case cmd: CommandResultExec =>
- getChildrenPlan(Seq(cmd.commandPhysicalPlan))
- case plan =>
- getChildrenPlan(Seq(plan))
- }
+ getExecutedPlan(df.queryExecution.executedPlan)
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
