Skip to content

Commit ac30205

Browse files
ravipesalamarmbrus
authored andcommitted
[SPARK-3813][SQL] Support "case when" conditional functions in Spark SQL.
"case when" conditional function is already supported in Spark SQL but there is no support in SqlParser. So added parser support to it. Author : ravipesala ravindra.pesalahuawei.com Author: ravipesala <[email protected]> Closes apache#2678 from ravipesala/SPARK-3813 and squashes the following commits: 70c75a7 [ravipesala] Fixed styles 713ea84 [ravipesala] Updated as per admin comments 709684f [ravipesala] Changed parser to support case when function.
1 parent bc3b6cb commit ac30205

File tree

2 files changed

+27
-2
lines changed

2 files changed

+27
-2
lines changed

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,13 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
7777
protected val BETWEEN = Keyword("BETWEEN")
7878
protected val BY = Keyword("BY")
7979
protected val CACHE = Keyword("CACHE")
80+
protected val CASE = Keyword("CASE")
8081
protected val CAST = Keyword("CAST")
8182
protected val COUNT = Keyword("COUNT")
8283
protected val DESC = Keyword("DESC")
8384
protected val DISTINCT = Keyword("DISTINCT")
85+
protected val ELSE = Keyword("ELSE")
86+
protected val END = Keyword("END")
8487
protected val EXCEPT = Keyword("EXCEPT")
8588
protected val FALSE = Keyword("FALSE")
8689
protected val FIRST = Keyword("FIRST")
@@ -122,11 +125,13 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
122125
protected val SUBSTRING = Keyword("SUBSTRING")
123126
protected val SUM = Keyword("SUM")
124127
protected val TABLE = Keyword("TABLE")
128+
protected val THEN = Keyword("THEN")
125129
protected val TIMESTAMP = Keyword("TIMESTAMP")
126130
protected val TRUE = Keyword("TRUE")
127131
protected val UNCACHE = Keyword("UNCACHE")
128132
protected val UNION = Keyword("UNION")
129133
protected val UPPER = Keyword("UPPER")
134+
protected val WHEN = Keyword("WHEN")
130135
protected val WHERE = Keyword("WHERE")
131136

132137
// Use reflection to find the reserved words defined in this class.
@@ -333,6 +338,15 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
333338
IF ~> "(" ~> expression ~ "," ~ expression ~ "," ~ expression <~ ")" ^^ {
334339
case c ~ "," ~ t ~ "," ~ f => If(c,t,f)
335340
} |
341+
CASE ~> expression.? ~ (WHEN ~> expression ~ (THEN ~> expression)).* ~
342+
(ELSE ~> expression).? <~ END ^^ {
343+
case casePart ~ altPart ~ elsePart =>
344+
val altExprs = altPart.flatMap {
345+
case we ~ te =>
346+
Seq(casePart.fold(we)(EqualTo(_, we)), te)
347+
}
348+
CaseWhen(altExprs ++ elsePart.toList)
349+
} |
336350
(SUBSTR | SUBSTRING) ~> "(" ~> expression ~ "," ~ expression <~ ")" ^^ {
337351
case s ~ "," ~ p => Substring(s,p,Literal(Integer.MAX_VALUE))
338352
} |

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -680,9 +680,20 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
680680
sql("SELECT CAST(TRUE AS STRING), CAST(FALSE AS STRING) FROM testData LIMIT 1"),
681681
("true", "false") :: Nil)
682682
}
683-
683+
684684
test("SPARK-3371 Renaming a function expression with group by gives error") {
685685
registerFunction("len", (s: String) => s.length)
686686
checkAnswer(
687-
sql("SELECT len(value) as temp FROM testData WHERE key = 1 group by len(value)"), 1)}
687+
sql("SELECT len(value) as temp FROM testData WHERE key = 1 group by len(value)"), 1)
688+
}
689+
690+
test("SPARK-3813 CASE a WHEN b THEN c [WHEN d THEN e]* [ELSE f] END") {
691+
checkAnswer(
692+
sql("SELECT CASE key WHEN 1 THEN 1 ELSE 0 END FROM testData WHERE key = 1 group by key"), 1)
693+
}
694+
695+
test("SPARK-3813 CASE WHEN a THEN b [WHEN c THEN d]* [ELSE e] END") {
696+
checkAnswer(
697+
sql("SELECT CASE WHEN key=1 THEN 1 ELSE 2 END FROM testData WHERE key = 1 group by key"), 1)
698+
}
688699
}

0 commit comments

Comments
 (0)