Update gluten-substrait/src/main/scala/org/apache/gluten/execution/Wh…

…oleStageTransformer.scala Co-authored-by: Hongze Zhang <[email protected]>
apache · Sep 3, 2024 · 44bf8e9 · 44bf8e9
1 parent ffe058a
commit 44bf8e9
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
@@ -361,14 +361,11 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
         plan: SparkPlan,
         buffer: ArrayBuffer[TransformSupport]): Unit = {
       plan match {
-        case node: TransformSupport =>
-          if (plan.children.forall(!_.isInstanceOf[TransformSupport])) {
+        case node: TransformSupport if node.children.forall(!_.isInstanceOf[TransformSupport]) =>
             buffer.append(node)
-          } else {
-            plan.children
-              .filter(_.isInstanceOf[TransformSupport])
+        case node: TransformSupport =>
+            node.children
               .foreach(collectTransformSupportLeaves(_, buffer))
-          }
         case _ =>
       }
     }