fix test case queries

DonnyZone · DonnyZone · commit 5a04ddc77195 · 2017-08-17T15:50:26.000+08:00
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/QueryPartitionSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/QueryPartitionSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.hive
 
 import java.io.File
+import java.sql.Timestamp
 
 import com.google.common.io.Files
 import org.apache.hadoop.fs.FileSystem
@@ -71,22 +72,27 @@ class QueryPartitionSuite extends QueryTest with SQLTestUtils with TestHiveSingl
 
   test("SPARK-21739: Cast expression should initialize timezoneId " +
     "when it is called statically to convert something into TimestampType") {
-    // create table for test
-    sql("CREATE TABLE table_with_timestamp_partition(value int) PARTITIONED by (ts timestamp)")
-    sql("INSERT OVERWRITE TABLE table_with_timestamp_partition " +
-      "partition (ts = '2010-01-01 00:00:00.000') VALUES (1)")
-    sql("INSERT OVERWRITE TABLE table_with_timestamp_partition " +
-      "partition (ts = '2010-01-02 00:00:00.000') VALUES (2)")
-
-    // test for Cast expression in TableReader
-    checkAnswer(sql("select value from table_with_timestamp_partition"),
-      Seq(Row(1), Row(2)))
-
-    // test for Cast expression in HiveTableScanExec
-    checkAnswer(sql("select value from table_with_timestamp_partition " +
-      "where ts = '2010-01-02 00:00:00.000'"), Row(2))
-
-    sql("DROP TABLE IF EXISTS table_with_timestamp_partition")
+    withTable("table_with_timestamp_partition") {
+      // create table for test
+      sql("CREATE TABLE table_with_timestamp_partition(value int) PARTITIONED by (ts timestamp)")
+      sql("INSERT OVERWRITE TABLE table_with_timestamp_partition " +
+        "partition (ts = '2010-01-01 00:00:00.000') VALUES (1)")
+      sql("INSERT OVERWRITE TABLE table_with_timestamp_partition " +
+        "partition (ts = '2010-01-02 00:00:00.000') VALUES (2)")
+      sql("INSERT OVERWRITE TABLE table_with_timestamp_partition " +
+        "partition (ts = '2010-01-03 00:00:00.000') VALUES (3)")
+
+      // test for Cast expression in TableReader
+      checkAnswer(sql("select * from table_with_timestamp_partition"),
+        Seq(
+          Row(1, Timestamp.valueOf("2010-01-01 00:00:00.000")),
+          Row(2, Timestamp.valueOf("2010-01-02 00:00:00.000")),
+          Row(3, Timestamp.valueOf("2010-01-03 00:00:00.000"))))
+
+      // test for Cast expression in HiveTableScanExec
+      checkAnswer(sql("select value from table_with_timestamp_partition " +
+        "where ts = '2010-01-02 00:00:00.000'"), Row(2))
+    }
   }
 
 }
