[SPARK-4151][SQL] Add string operation function trim, ltrim, rtrim, length to support SparkSql (HiveQL)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala

@@ -103,6 +103,10 @@ class SqlParser extends AbstractSparkSQLParser {
   protected val UPPER = Keyword("UPPER")
   protected val WHEN = Keyword("WHEN")
   protected val WHERE = Keyword("WHERE")
+  protected val TRIM = Keyword("TRIM")
+  protected val LTRIM = Keyword("LTRIM")
+  protected val RTRIM = Keyword("RTRIM")
+  protected val LENGTH = Keyword("LENGTH")
 
   // Use reflection to find the reserved words defined in this class.
   protected val reservedWords =
@@ -283,6 +287,10 @@ class SqlParser extends AbstractSparkSQLParser {
     | MAX   ~ "(" ~> expression <~ ")" ^^ { case exp => Max(exp) }
     | UPPER ~ "(" ~> expression <~ ")" ^^ { case exp => Upper(exp) }
     | LOWER ~ "(" ~> expression <~ ")" ^^ { case exp => Lower(exp) }
+    | TRIM ~ "(" ~> expression <~ ")" ^^ { case exp => Trim(exp) }
+    | LTRIM ~ "(" ~> expression <~ ")" ^^ { case exp => Ltrim(exp) }
+    | RTRIM ~ "(" ~> expression <~ ")" ^^ { case exp => Rtrim(exp) }
+    | LENGTH ~ "(" ~> expression <~ ")" ^^ { case exp => Length(exp) }
     | IF ~ "(" ~> expression ~ ("," ~> expression) ~ ("," ~> expression) <~ ")" ^^
       { case c ~ t ~ f => If(c, t, f) }
     | CASE ~> expression.? ~ (WHEN ~> expression ~ (THEN ~> expression)).* ~

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala

@@ -23,7 +23,7 @@ import scala.collection.IndexedSeqOptimized
 
 
 import org.apache.spark.sql.catalyst.analysis.UnresolvedException
-import org.apache.spark.sql.catalyst.types.{BinaryType, BooleanType, DataType, StringType}
+import org.apache.spark.sql.catalyst.types._
 
 trait StringRegexExpression {
   self: BinaryExpression =>
@@ -71,7 +71,7 @@ trait StringRegexExpression {
   }
 }
 
-trait CaseConversionExpression {
+trait StringTransformationExpression {
   self: UnaryExpression =>
 
   type EvaluatedType = Any
@@ -92,6 +92,34 @@ trait CaseConversionExpression {
   }
 }
 
+/**
+ * This trait is use for string calculation that return Integer Type
+ * Functions that return Integer Type should use this trait
+ * eg: length(s), instr( string1, string2, start_position,nth_appearance )
+ */
+
+trait StringCalculationExpression {
+    self: UnaryExpression =>
+
+    type EvaluatedType = Any
+
+    def calc(v: String): Int
+
+    override def foldable: Boolean = child.foldable
+    def nullable: Boolean = child.nullable
+    def dataType: DataType = IntegerType
+
+    override def eval(input: Row): Any = {
+        val evaluated = child.eval(input)
+        if (evaluated == null) {
+            null
+        } else {
+            calc(evaluated.toString)
+        }
+    }
+}
+
+
 /**
  * Simple RegEx pattern matching function
  */
@@ -134,10 +162,51 @@ case class RLike(left: Expression, right: Expression)
   override def matches(regex: Pattern, str: String): Boolean = regex.matcher(str).find(0)
 }
 
+
+/**
+ * A function that strip whitespace (or other characters) from the beginning of a string
+ */
+case class Ltrim(child: Expression) extends UnaryExpression with StringTransformationExpression {
+
+  override def convert(v: String): String = v.dropWhile(_ == ' ')
+
+  override def toString() = s"Ltrim($child)"
+}
+
+/**
+ * A function that strip whitespace (or other characters) from the end of a string
+ */
+case class Rtrim(child: Expression) extends UnaryExpression with StringTransformationExpression {
+
+  override def convert(v: String): String = v.reverse.dropWhile(_ == ' ').reverse
+
+  override def toString() = s"Rtrim($child)"
+}
+
+/**
+ * A function that calculate the length of a string
+ */
+case class Length(child: Expression) extends UnaryExpression with StringCalculationExpression {
+
+  override def calc(v: String): Int = v.length()
+
+  override def toString() = s"Length($child)"
+}
+
+/**
+ * A function that trim the characters of a string
+ */
+case class Trim(child: Expression) extends UnaryExpression with StringTransformationExpression {
+
+  override def convert(v: String): String = v.trim()
+
+  override def toString() = s"Trim($child)"
+}
+
 /**
  * A function that converts the characters of a string to uppercase.
  */
-case class Upper(child: Expression) extends UnaryExpression with CaseConversionExpression {
+case class Upper(child: Expression) extends UnaryExpression with StringTransformationExpression {
 
   override def convert(v: String): String = v.toUpperCase()
 
@@ -147,7 +216,7 @@ case class Upper(child: Expression) extends UnaryExpression with CaseConversionE
 /**
  * A function that converts the characters of a string to lowercase.
  */
-case class Lower(child: Expression) extends UnaryExpression with CaseConversionExpression {
+case class Lower(child: Expression) extends UnaryExpression with StringTransformationExpression {
 
   override def convert(v: String): String = v.toLowerCase()
 

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -500,6 +500,56 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
         (2, "abc"),
         (3, null)))
   }
+  test("system function trim()") {
+    checkAnswer(
+      sql("SELECT N,TRIM(L) FROM untrimmedData"),
+      Seq(
+        (1, "Good"),
+        (2, "To"),
+        (3, "See"),
+        (4, "You !")))
+
+    checkAnswer(
+      sql("SELECT n, TRIM(s) FROM nullStrings"),
+      Seq(
+        (1, "abc"),
+        (2, "ABC"),
+        (3, null)))
+    }
+
+  test("system function ltrim()") {
+    checkAnswer(
+      sql("SELECT N,LTRIM(L) FROM untrimmedData"),
+      Seq(
+        (1, "Good "),
+        (2, "To "),
+        (3, "See"),
+        (4, "You ! ")))
+
+    checkAnswer(
+      sql("SELECT n, LTRIM(s) FROM nullStrings"),
+      Seq(
+        (1, "abc"),
+        (2, "ABC"),
+        (3, null)))
+    }
+
+  test("system function rtrim()") {
+    checkAnswer(
+      sql("SELECT N,RTRIM(L) FROM untrimmedData"),
+      Seq(
+        (1, " Good"),
+        (2, "To"),
+        (3, " See"),
+        (4, " You !")))
+
+    checkAnswer(
+      sql("SELECT n, RTRIM(s) FROM nullStrings"),
+      Seq(
+        (1, "abc"),
+        (2, "ABC"),
+        (3, null)))
+    }
 
   test("UNION") {
     checkAnswer(

sql/core/src/test/scala/org/apache/spark/sql/TestData.scala

@@ -72,6 +72,16 @@ object TestData {
 
   val emptyTableData = logical.LocalRelation('a.int, 'b.int)
 
+  case class UntrimmedData(N: Int, L: String)
+  val untrimmedData =
+    TestSQLContext.sparkContext.parallelize(
+      UntrimmedData(1, " Good ") ::
+      UntrimmedData(2, "To ") ::
+      UntrimmedData(3, " See") ::
+      UntrimmedData(4, " You ! ") :: Nil).toSchemaRDD
+  untrimmedData.registerTempTable("untrimmedData")
+
+
   case class UpperCaseData(N: Int, L: String)
   val upperCaseData =
     TestSQLContext.sparkContext.parallelize(

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala

@@ -869,6 +869,10 @@ private[hive] object HiveQl {
   val MIN = "(?i)MIN".r
   val UPPER = "(?i)UPPER".r
   val LOWER = "(?i)LOWER".r
+  val TRIM = "(?i)TRIM".r
+  val LTRIM = "(?i)LTRIM".r
+  val RTRIM = "(?i)RTRIM".r
+  val LENGTH = "(?i)LENGTH".r
   val RAND = "(?i)RAND".r
   val AND = "(?i)AND".r
   val OR = "(?i)OR".r
@@ -918,6 +922,10 @@ private[hive] object HiveQl {
     /* System functions about string operations */
     case Token("TOK_FUNCTION", Token(UPPER(), Nil) :: arg :: Nil) => Upper(nodeToExpr(arg))
     case Token("TOK_FUNCTION", Token(LOWER(), Nil) :: arg :: Nil) => Lower(nodeToExpr(arg))
+    case Token("TOK_FUNCTION", Token(TRIM(), Nil) :: arg :: Nil) => Trim(nodeToExpr(arg))
+    case Token("TOK_FUNCTION", Token(LTRIM(), Nil) :: arg :: Nil) => Ltrim(nodeToExpr(arg))
+    case Token("TOK_FUNCTION", Token(RTRIM(), Nil) :: arg :: Nil) => Rtrim(nodeToExpr(arg))
+    case Token("TOK_FUNCTION", Token(LENGTH(), Nil) :: arg :: Nil) => Length(nodeToExpr(arg))
 
     /* Casts */
     case Token("TOK_FUNCTION", Token("TOK_STRING", Nil) :: arg :: Nil) =>