Changed parser to support case when function.

Author: ravipesala

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

@@ -128,6 +128,11 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
   protected val UNION = Keyword("UNION")
   protected val UPPER = Keyword("UPPER")
   protected val WHERE = Keyword("WHERE")
+  protected val CASE = Keyword("CASE")
+  protected val WHEN = Keyword("WHEN")
+  protected val THEN = Keyword("THEN")
+  protected val ELSE = Keyword("ELSE")
+  protected val END = Keyword("END")
 
   // Use reflection to find the reserved words defined in this class.
   protected val reservedWords =
@@ -333,6 +338,24 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
     IF ~> "(" ~> expression ~ "," ~ expression ~ "," ~ expression <~ ")" ^^ {
       case c ~ "," ~ t ~ "," ~ f => If(c,t,f)
     } |
+    CASE ~> opt(expression) ~ (WHEN ~ expression ~ THEN ~ expression).* ~
+      opt(ELSE ~> expression) <~ END ^^ {
+       case c ~ l ~ el =>
+         var caseWhenExpr = l.map{x =>
+          x match {
+           case w ~ we ~ t ~ te =>
+             c match {
+               case Some(e) => Seq(EqualTo(e, we), te)
+               case None => Seq(we, te)
+             }
+          }
+        }.toSeq.reduce(_ ++ _)
+        caseWhenExpr = el match {
+          case Some(e) => caseWhenExpr ++ Seq(e)
+          case None => caseWhenExpr
+        }
+        CaseWhen(caseWhenExpr)
+    } |
     (SUBSTR | SUBSTRING) ~> "(" ~> expression ~ "," ~ expression <~ ")" ^^ {
       case s ~ "," ~ p => Substring(s,p,Literal(Integer.MAX_VALUE))
     } |

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

@@ -680,9 +680,20 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
       sql("SELECT CAST(TRUE AS STRING), CAST(FALSE AS STRING) FROM testData LIMIT 1"),
       ("true", "false") :: Nil)
   }
-  
+
   test("SPARK-3371 Renaming a function expression with group by gives error") {
     registerFunction("len", (s: String) => s.length)
     checkAnswer(
-      sql("SELECT len(value) as temp FROM testData WHERE key = 1 group by len(value)"), 1)}    
+      sql("SELECT len(value) as temp FROM testData WHERE key = 1 group by len(value)"), 1)
+  }
+
+  test("SPARK-3813 CASE a WHEN b THEN c [WHEN d THEN e]* [ELSE f] END") {
+    checkAnswer(
+      sql("SELECT CASE key WHEN 1 THEN 1 ELSE 0 END FROM testData WHERE key = 1 group by key"), 1)
+  }
+
+  test("SPARK-3813 CASE WHEN a THEN b [WHEN c THEN d]* [ELSE e] END") {
+    checkAnswer(
+      sql("SELECT CASE WHEN key=1 THEN 1 ELSE 2 END FROM testData WHERE key = 1 group by key"), 1)
+  }
 }