Support TRIM(BOTH/LEADING/TRAILING FROM str) format

apache · Jun 17, 2019 · a30fa00 · a30fa00
1 parent b7b4452
commit a30fa00
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 3 deletions.
diff --git a/sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4 b/sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4
@@ -691,7 +691,7 @@ primaryExpression
     | '(' query ')'                                                                            #subqueryExpression
     | qualifiedName '(' (setQuantifier? argument+=expression (',' argument+=expression)*)? ')'
        (OVER windowSpec)?                                                                      #functionCall
-    | qualifiedName '(' trimOption=(BOTH | LEADING | TRAILING) argument+=expression
+    | qualifiedName '(' trimOption=(BOTH | LEADING | TRAILING) (argument+=expression)?
       FROM argument+=expression ')'                                                            #functionCall
     | IDENTIFIER '->' expression                                                               #lambda
     | '(' IDENTIFIER (',' IDENTIFIER)+ ')' '->' expression                                     #lambda

diff --git a/...catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/...catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala
@@ -605,6 +605,12 @@ object StringTrim {
   usage = """
     _FUNC_(str) - Removes the leading and trailing space characters from `str`.
 
+    _FUNC_(BOTH FROM str) - Removes the leading and trailing space characters from `str`.
+
+    _FUNC_(LEADING FROM str) - Removes the leading space characters from `str`.
+
+    _FUNC_(TRAILING FROM str) - Removes the trailing space characters from `str`.
+
     _FUNC_(BOTH trimStr FROM str) - Remove the leading and trailing `trimStr` characters from `str`
 
     _FUNC_(LEADING trimStr FROM str) - Remove the leading `trimStr` characters from `str`
@@ -626,6 +632,12 @@ object StringTrim {
     Examples:
       > SELECT _FUNC_('    SparkSQL   ');
        SparkSQL
+      > SELECT _FUNC_(BOTH FROM '    SparkSQL   ');
+       SparkSQL
+      > SELECT _FUNC_(LEADING FROM '    SparkSQL   ');
+       SparkSQL
+      > SELECT _FUNC_(TRAILING FROM '    SparkSQL   ');
+           SparkSQL
       > SELECT _FUNC_('SL', 'SSparkSQLS');
        parkSQ
       > SELECT _FUNC_(BOTH 'SL' FROM 'SSparkSQLS');

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
@@ -1403,7 +1403,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
   }
 
   /**
-   * Create a (windowed) Function expression.
+   * Create a (windowed)/trim Function expression.
    */
   override def visitFunctionCall(ctx: FunctionCallContext): Expression = withOrigin(ctx) {
     def replaceFunctions(
@@ -1420,7 +1420,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
           case SqlBaseParser.LEADING => funcID.copy(funcName = "ltrim")
           case SqlBaseParser.TRAILING => funcID.copy(funcName = "rtrim")
           case _ => throw new ParseException("Function trim doesn't support with " +
-            s"type ${opt.getType}. Please use BOTH, LEADING or Trailing as trim type", ctx)
+            s"type ${opt.getType}. Please use BOTH, LEADING or TRAILING as trim type", ctx)
         }
       } else {
         funcID

diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala
@@ -717,6 +717,19 @@ class PlanParserSuite extends AnalysisTest {
       "SELECT TRIM(TRAILING 'c&^,.' FROM 'bc...,,,&&&ccc')",
       OneRowRelation().select('rtrim.function("c&^,.", "bc...,,,&&&ccc"))
     )
+
+    assertEqual(
+      "SELECT TRIM(BOTH FROM '  bunch o blanks  ')",
+      OneRowRelation().select('TRIM.function("  bunch o blanks  "))
+    )
+    assertEqual(
+      "SELECT TRIM(LEADING FROM '  bunch o blanks  ')",
+      OneRowRelation().select('ltrim.function("  bunch o blanks  "))
+    )
+    assertEqual(
+      "SELECT TRIM(TRAILING FROM '  bunch o blanks  ')",
+      OneRowRelation().select('rtrim.function("  bunch o blanks  "))
+    )
   }
 
   test("precedence of set operations") {