[GLUTEN][Function] Support btrim function

apache · Mar 5, 2025 · 5b36b10 · 5b36b10
1 parent 86d0174
commit 5b36b10
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxStringFunctionsSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxStringFunctionsSuite.scala
@@ -316,6 +316,19 @@ class VeloxStringFunctionsSuite extends VeloxWholeStageTransformerSuite {
         s"from $LINEITEM_TABLE limit $LENGTH")(checkGlutenOperatorMatch[ProjectExecTransformer])
   }
 
+  test("btrim") {
+    runQueryAndCompare(
+      s"select l_orderkey, btrim(l_comment) " +
+        s"from $LINEITEM_TABLE limit $LENGTH")(checkGlutenOperatorMatch[ProjectExecTransformer])
+    runQueryAndCompare(
+      s"select l_orderkey, btrim('. abcdefg', l_comment) " +
+        s"from $LINEITEM_TABLE limit $LENGTH")(checkGlutenOperatorMatch[ProjectExecTransformer])
+    runQueryAndCompare(
+      s"select l_orderkey, btrim($NULL_STR_COL), " +
+        s"btrim($NULL_STR_COL, l_comment), btrim('. abcdefg', $NULL_STR_COL) " +
+        s"from $LINEITEM_TABLE limit $LENGTH")(checkGlutenOperatorMatch[ProjectExecTransformer])
+  }
+
   test("lpad") {
     runQueryAndCompare(
       s"select l_orderkey, lpad($NULL_STR_COL, 80) " +

diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
@@ -25,7 +25,7 @@ import org.apache.gluten.utils.DecimalArithmeticUtil
 import org.apache.spark.{SPARK_REVISION, SPARK_VERSION_SHORT}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.SQLConfHelper
-import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.{RuntimeReplaceable, StringTrimBoth, _}
 import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
 import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.execution.ScalarSubquery
@@ -359,6 +359,19 @@ object ExpressionConverter extends SQLConfHelper with Logging {
           children,
           s
         )
+      case s: RuntimeReplaceable =>
+        val (srcStr, trimStr) = s match {
+          case StringTrimBoth(srcStr, trimStr, child) => (srcStr, trimStr)
+        }
+        val children = trimStr
+          .map(replaceWithExpressionTransformer0(_, attributeSeq, expressionsMap))
+          .toSeq ++
+          Seq(replaceWithExpressionTransformer0(srcStr, attributeSeq, expressionsMap))
+        GenericExpressionTransformer(
+          substraitExprName,
+          children,
+          s
+        )
       case m: HashExpression[_] =>
         BackendsApiManager.getSparkPlanExecApiInstance.genHashExpressionTransformer(
           substraitExprName,

diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionMappings.scala b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionMappings.scala
@@ -21,7 +21,7 @@ import org.apache.gluten.config.GlutenConfig
 import org.apache.gluten.expression.ExpressionNames._
 import org.apache.gluten.sql.shims.SparkShimLoader
 
-import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.{StringTrimBoth, _}
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.execution.ScalarSubquery
@@ -80,6 +80,7 @@ object ExpressionMappings {
     Sig[StringTrimLeft](LTRIM),
     Sig[StringTrimRight](RTRIM),
     Sig[StringTrim](TRIM),
+    Sig[StringTrimBoth](BTRIM),
     Sig[StringLPad](LPAD),
     Sig[StringRPad](RPAD),
     Sig[StringReplace](REPLACE),

diff --git a/shims/common/src/main/scala/org/apache/gluten/expression/ExpressionNames.scala b/shims/common/src/main/scala/org/apache/gluten/expression/ExpressionNames.scala
@@ -107,6 +107,7 @@ object ExpressionNames {
   final val LTRIM = "ltrim"
   final val RTRIM = "rtrim"
   final val TRIM = "trim"
+  final val BTRIM = "btrim"
   final val LPAD = "lpad"
   final val RPAD = "rpad"
   final val REPLACE = "replace"