Added support for DateType/TimestampType.

Updated ExpressionEvalHelper to avoid conversion.
apache · Jul 27, 2015 · 24a3e46 · 24a3e46
1 parent 9989064
commit 24a3e46
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
@@ -255,6 +255,10 @@ public Object get(int ordinal, DataType dataType) {
       return getDouble(ordinal);
     } else if (dataType instanceof DecimalType) {
       return getDecimal(ordinal);
+    } else if (dataType instanceof DateType) {
+      return getInt(ordinal);
+    } else if (dataType instanceof TimestampType) {
+      return getLong(ordinal);
     } else if (dataType instanceof StringType) {
       return getUTF8String(ordinal);
     } else if (dataType instanceof StructType) {

diff --git a/...alyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala b/...alyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala
@@ -160,17 +160,20 @@ trait ExpressionEvalHelper {
       expected: Any,
       inputRow: InternalRow = EmptyRow): Unit = {
 
-    val plan = generateProject(
+    val project = generateProject(
       GenerateUnsafeProjection.generate(Alias(expression, s"Optimized($expression)")() :: Nil),
       expression)
-
-    val unsafeRow = plan(inputRow)
-    // UnsafeRow cannot be compared with GenericInternalRow directly
-    val actual = FromUnsafeProjection(expression.dataType :: Nil)(unsafeRow)
-    val expectedRow = InternalRow(expected)
-    if (actual != expectedRow) {
-      val input = if (inputRow == EmptyRow) "" else s", input: $inputRow"
-      fail(s"Incorrect Evaluation: $expression, actual: $actual, expected: $expectedRow$input")
+    val out = project(inputRow)
+    val input = if (inputRow == EmptyRow) "" else s", input: $inputRow"
+
+    if (expected == null) {
+      if (!out.isNullAt(0)) {
+        val actual = out.get(0, expression.dataType)
+        fail(s"Incorrect Evaluation: $expression, actual: $actual, expected: $expected$input")
+      }
+    } else if (out.get(0, expression.dataType) != expected) {
+      val actual = out.get(0, expression.dataType)
+      fail(s"Incorrect Evaluation: $expression, actual: $actual, expected: $expected$input")
     }
   }
 
@@ -200,8 +203,7 @@ trait ExpressionEvalHelper {
     plan = generateProject(
       GenerateUnsafeProjection.generate(Alias(expression, s"Optimized($expression)")() :: Nil),
       expression)
-    actual = FromUnsafeProjection(expression.dataType :: Nil)(
-      plan(inputRow)).get(0, expression.dataType)
+    actual = plan(inputRow).get(0, expression.dataType)
     assert(checkResult(actual, expected))
   }
 }