[GLUTEN-7581][VL] Code cleanup for GlutenColumnarRule

apache · Oct 18, 2024 · 023b568 · 023b568
1 parent 00b71e8
commit 023b568
Showing 1 changed file with 21 additions and 38 deletions.
diff --git a/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenColumnarRule.scala b/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenColumnarRule.scala
@@ -58,38 +58,6 @@ object GlutenColumnarRule {
     override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
       copy(child = newChild)
   }
-
-  object OutputsColumnarTester {
-    def wrap(plan: SparkPlan): SparkPlan = {
-      if (plan.supportsColumnar) {
-        DummyColumnarOutputExec(plan)
-      } else {
-        DummyRowOutputExec(plan)
-      }
-    }
-
-    def inferOutputsColumnar(plan: SparkPlan): Boolean = plan match {
-      case DummyRowOutputExec(_) => false
-      case RowToColumnarExec(DummyRowOutputExec(_)) => true
-      case DummyColumnarOutputExec(_) => true
-      case ColumnarToRowExec(DummyColumnarOutputExec(_)) => false
-      case _ =>
-        throw new IllegalStateException(
-          "This should not happen. Please leave a issue at" +
-            " https://github.com/apache/incubator-gluten.")
-    }
-
-    def unwrap(plan: SparkPlan): SparkPlan = plan match {
-      case DummyRowOutputExec(child) => child
-      case RowToColumnarExec(DummyRowOutputExec(child)) => child
-      case DummyColumnarOutputExec(child) => child
-      case ColumnarToRowExec(DummyColumnarOutputExec(child)) => child
-      case _ =>
-        throw new IllegalStateException(
-          "This should not happen. Please leave a issue at" +
-            " https://github.com/apache/incubator-gluten.")
-    }
-  }
 }
 
 case class GlutenColumnarRule(
@@ -109,16 +77,31 @@ case class GlutenColumnarRule(
    */
   final override def preColumnarTransitions: Rule[SparkPlan] = plan => {
     // To infer caller's property: ApplyColumnarRulesAndInsertTransitions#outputsColumnar.
-    OutputsColumnarTester.wrap(plan)
+    if (plan.supportsColumnar) {
+      DummyColumnarOutputExec(plan)
+    } else {
+      DummyRowOutputExec(plan)
+    }
   }
 
   override def postColumnarTransitions: Rule[SparkPlan] = plan => {
-    val outputsColumnar = OutputsColumnarTester.inferOutputsColumnar(plan)
-    val unwrapped = OutputsColumnarTester.unwrap(plan)
-    val vanillaPlan = Transitions.insertTransitions(unwrapped, outputsColumnar)
+    val (childPlan, outputsColumnar) = plan match {
+      case DummyRowOutputExec(child) =>
+        (child, false)
+      case RowToColumnarExec(DummyRowOutputExec(child)) =>
+        (child, true)
+      case DummyColumnarOutputExec(child) =>
+        (child, true)
+      case ColumnarToRowExec(DummyColumnarOutputExec(child)) =>
+        (child, false)
+      case _ =>
+        throw new IllegalStateException(
+          "This should not happen. Please leave a issue at" +
+            " https://github.com/apache/incubator-gluten.")
+    }
+    val vanillaPlan = Transitions.insertTransitions(childPlan, outputsColumnar)
     val applier = applierBuilder.apply(session)
-    val out = applier.apply(vanillaPlan, outputsColumnar)
-    out
+    applier.apply(vanillaPlan, outputsColumnar)
   }
 
 }