fixup

apache · Mar 25, 2024 · e66f0d1 · e66f0d1
1 parent 3255ca3
commit e66f0d1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/gluten-core/src/main/scala/io/glutenproject/extension/columnar/EnumeratedTransform.scala b/gluten-core/src/main/scala/io/glutenproject/extension/columnar/EnumeratedTransform.scala
@@ -32,7 +32,7 @@ case class EnumeratedTransform(session: SparkSession, outputsColumnar: Boolean)
   with LogLevelUtil {
   import EnumeratedTransform._
 
-  private val cboRules = Seq(SingleNodeTransformCboRule())
+  private val cboRules = Seq(SingleNodeTransformCboRule(), ExchangeTransformCboRule())
   private val optimization = GlutenOptimization(cboRules)
 
   private val reqConvention = GlutenProperties.Conventions.ANY
@@ -60,4 +60,16 @@ object EnumeratedTransform {
 
     override def shape(): Shape[SparkPlan] = Shapes.fixedHeight(1)
   }
+
+  private case class ExchangeTransformCboRule() extends CboRule[SparkPlan] {
+    private val delegate = TransformPreOverrides.ExchangeTransformRule()
+
+
+    override def shift(node: SparkPlan): Iterable[SparkPlan] = {
+      val out = List(delegate.apply(node))
+      out
+    }
+
+    override def shape(): Shape[SparkPlan] = Shapes.fixedHeight(1)
+  }
 }
diff --git a/gluten-core/src/main/scala/io/glutenproject/extension/columnar/MiscColumnarRules.scala b/gluten-core/src/main/scala/io/glutenproject/extension/columnar/MiscColumnarRules.scala
@@ -109,7 +109,7 @@ object MiscColumnarRules {
     }
 
     // Exchange transformation.
-    private case class ExchangeTransformRule() extends Rule[SparkPlan] with LogLevelUtil {
+    case class ExchangeTransformRule() extends Rule[SparkPlan] with LogLevelUtil {
       override def apply(plan: SparkPlan): SparkPlan = plan match {
         case plan if TransformHints.isNotTransformable(plan) =>
           plan