Finish up Case.

Author: concretevitamin

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

@@ -202,3 +202,81 @@ case class If(predicate: Expression, trueValue: Expression, falseValue: Expressi
 
   override def toString = s"if ($predicate) $trueValue else $falseValue"
 }
+
+// TODO: is it a good idea to put this class in this file?
+// CASE WHEN a THEN b [WHEN c THEN d]* [ELSE e] END
+// When a = true, returns b; when c = true, return d; else return e
+case class Case(key: Option[Expression], branches: Seq[Expression]) extends Expression {
+  // Branches are considered in consecutive pairs (cond, val), and the last element
+  // is the val for the default catch-all case (w/o a companion condition, that is).
+
+  def children = key.toSeq ++ branches
+
+  override def nullable = branches
+    .sliding(2, 2)
+    .map {
+      case Seq(cond, value) => value.nullable
+      case Seq(elseValue) => elseValue.nullable
+    }
+    .reduce(_ || _)
+
+  def references = children.flatMap(_.references).toSet
+
+  override lazy val resolved = {
+    val allBranchesEqual = branches.sliding(2, 2).map {
+      case Seq(cond, value) => value.dataType
+      case Seq(elseValue) => elseValue.dataType
+    }.reduce(_ == _)
+    childrenResolved && allBranchesEqual
+  }
+
+  def dataType = {
+    if (!resolved) {
+      throw new UnresolvedException(this, "cannot resolve due to differing types in some branches")
+    }
+    branches(1).dataType
+  }
+
+  type EvaluatedType = Any
+
+  override def eval(input: Row): Any = {
+    def slidingCheck(expectedVal: Any): Any = {
+      branches.sliding(2, 2).foldLeft(None.asInstanceOf[Option[Any]]) {
+        case (Some(x), _) =>
+          Some(x)
+        case (None, Seq(cond, value)) =>
+          if (cond.eval(input) == true) Some(value.eval(input)) else None
+        case (None, Seq(elseValue)) =>
+          Some(elseValue.eval(input))
+      }.getOrElse(null)
+      // If all branches fail and an elseVal is not provided, the whole statement
+      // evaluates to null, according to Hive's semantics.
+    }
+    // Check if any branch's cond evaluates either to the key (if provided), or to true.
+    if (key.isDefined) {
+      slidingCheck(key.get.eval(input))
+    } else {
+      slidingCheck(true)
+    }
+  }
+
+  override def toString = {
+    var firstBranch = ""
+    var otherBranches = ""
+    if (key.isDefined) {
+      val keyString = key.get.toString
+      firstBranch = s"if ($keyString == ${branches(0)}) { ${branches(1)} }"
+      otherBranches = branches.sliding(2, 2).drop(1).map {
+        case Seq(cond, value) => s"\nelse if ($keyString == $cond) { $value }"
+        case Seq(elseValue) => s"\nelse { $elseValue }"
+      }.mkString
+    } else {
+      firstBranch = s"if (${branches(0)}) { ${branches(1)} }"
+      otherBranches = branches.sliding(2, 2).drop(1).map {
+        case Seq(cond, value) => s"\nelse if ($cond) { $value }"
+        case Seq(elseValue) => s"\nelse { $elseValue }"
+      }.mkString
+    }
+    firstBranch ++ otherBranches
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvaluationSuite.scala

@@ -35,7 +35,7 @@ class ExpressionEvaluationSuite extends FunSuite {
   /**
    * Checks for three-valued-logic.  Based on:
    * http://en.wikipedia.org/wiki/Null_(SQL)#Comparisons_with_NULL_and_the_three-valued_logic_.283VL.29
-   *
+   * I.e. in flat cpo "False -> Unknown -> True", OR is lowest upper bound, AND is greatest lower bound.
    * p       q       p OR q  p AND q  p = q
    * True    True    True    True     True
    * True    False   True    False    False

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

@@ -48,8 +48,9 @@ private[hive] case class AddJar(jarPath: String) extends Command
 
 private[hive] case class AddFile(filePath: String) extends Command
 
+// FIXME: add back private
 /** Provides a mapping from HiveQL statements to catalyst logical plans and expression trees. */
-private[hive] object HiveQl {
+object HiveQl {
   protected val nativeCommands = Seq(
     "TOK_DESCFUNCTION",
     "TOK_DESCTABLE",
@@ -798,6 +799,8 @@ private[hive] object HiveQl {
   val IN = "(?i)IN".r
   val DIV = "(?i)DIV".r
   val BETWEEN = "(?i)BETWEEN".r
+  val WHEN = "(?i)WHEN".r
+  val CASE = "(?i)CASE".r
 
   protected def nodeToExpr(node: Node): Expression = node match {
     /* Attribute References */
@@ -904,6 +907,12 @@ private[hive] object HiveQl {
     case Token(OR(), left :: right:: Nil) => Or(nodeToExpr(left), nodeToExpr(right))
     case Token(NOT(), child :: Nil) => Not(nodeToExpr(child))
 
+    /* Case statements */
+    case Token("TOK_FUNCTION", Token(WHEN(), Nil) :: branches) =>
+      Case(None, branches.map(nodeToExpr))
+    case Token("TOK_FUNCTION", Token(CASE(), Nil) :: branches) =>
+      Case(Some(nodeToExpr(branches(0))), branches.drop(1).map(nodeToExpr))
+
     /* Complex datatype manipulation */
     case Token("[", child :: ordinal :: Nil) =>
       GetItem(nodeToExpr(child), nodeToExpr(ordinal))