fix

apache · Jul 10, 2024 · 8cf51dc · 8cf51dc
1 parent 014a401
commit 8cf51dc
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/...ore/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala b/...ore/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala
@@ -17,7 +17,7 @@
 package org.apache.gluten.extension.columnar
 
 import org.apache.gluten.GlutenConfig
-import org.apache.gluten.extension.columnar.OffloadOthers.ReplaceSingleNode
+import org.apache.gluten.extension.columnar.MiscColumnarRules.TransformPreOverrides
 import org.apache.gluten.extension.columnar.rewrite.RewriteSparkPlanRulesManager
 
 import org.apache.spark.sql.catalyst.expressions.SortOrder
@@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.{SortExec, SparkPlan}
  * SortAggregate with the same key. So, this rule adds local sort back if necessary.
  */
 object EnsureLocalSortRequirements extends Rule[SparkPlan] {
-  private lazy val offload = new ReplaceSingleNode()
+  private lazy val offload = TransformPreOverrides.apply()
 
   private def addLocalSort(
       originalChild: SparkPlan,
@@ -48,7 +48,7 @@ object EnsureLocalSortRequirements extends Rule[SparkPlan] {
         // The sort can not be offloaded
         rewrittenPlan
       } else {
-        rewrittenPlan.transform { case p => offload.doReplace(p) }
+        offload.apply(rewrittenPlan)
       }
     }
   }