Update code

Author: beliefer

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala

@@ -50,14 +50,15 @@ case class AllDataTypesWithNonPrimitiveType(
     shortField: Short,
     byteField: Byte,
     booleanField: Boolean,
-    timestampField: Timestamp,
-    timestampNTZField: LocalDateTime,
     array: Seq[Int],
     arrayContainsNull: Seq[Option[Int]],
     map: Map[Int, Long],
     mapValueContainsNull: Map[Int, Option[Long]],
     data: (Seq[Int], (Int, String)))
 
+case class TimestampsWithNonPrimitiveType(
+    timestampField: Timestamp, timestampNTZField: LocalDateTime)
+
 case class BinaryData(binaryData: Array[Byte])
 
 case class Contact(name: String, phone: String)
@@ -69,8 +70,7 @@ abstract class OrcQueryTest extends OrcTest {
 
   test("Read/write All Types") {
     val data = (0 to 255).map { i =>
-      (s"$i", i, i.toLong, i.toFloat, i.toDouble, i.toShort, i.toByte, i % 2 == 0,
-        new Timestamp(i), LocalDateTime.of(2019, 3, 21, 0, 2, 3, 456000000 + i))
+      (s"$i", i, i.toLong, i.toFloat, i.toDouble, i.toShort, i.toByte, i % 2 == 0)
     }
 
     withOrcFile(data) { file =>
@@ -91,7 +91,6 @@ abstract class OrcQueryTest extends OrcTest {
     val data: Seq[AllDataTypesWithNonPrimitiveType] = (0 to 255).map { i =>
       AllDataTypesWithNonPrimitiveType(
         s"$i", i, i.toLong, i.toFloat, i.toDouble, i.toShort, i.toByte, i % 2 == 0,
-        new Timestamp(i), LocalDateTime.of(2019, 3, 21, 0, 2, 3, 456000000 + i),
         0 until i,
         (0 until i).map(Option(_).filter(_ % 3 == 0)),
         (0 until i).map(i => i -> i.toLong).toMap,
@@ -177,15 +176,13 @@ abstract class OrcQueryTest extends OrcTest {
       Option.empty[Long],
       Option.empty[Float],
       Option.empty[Double],
-      Option.empty[Boolean],
-      Option.empty[Timestamp],
-      Option.empty[LocalDateTime]
+      Option.empty[Boolean]
     ) :: Nil
 
     withOrcFile(data) { file =>
       checkAnswer(
         spark.read.orc(file),
-        Row(Seq.fill(7)(null): _*))
+        Row(Seq.fill(5)(null): _*))
     }
   }
 
@@ -720,6 +717,44 @@ abstract class OrcQuerySuite extends OrcQueryTest with SharedSparkSession {
       }
     }
   }
+
+  test("Read/write all timestamp types") {
+    val data = (0 to 255).map { i =>
+      (new Timestamp(i), LocalDateTime.of(2019, 3, 21, 0, 2, 3, 456000000 + i))
+    }
+
+    withOrcFile(data) { file =>
+      checkAnswer(
+        spark.read.orc(file),
+        data.toDF().collect())
+    }
+  }
+
+  test("Read/write all timestamp types with non-primitive type") {
+    val data: Seq[TimestampsWithNonPrimitiveType] = (0 to 255).map { i =>
+      TimestampsWithNonPrimitiveType(
+        new Timestamp(i), LocalDateTime.of(2019, 3, 21, 0, 2, 3, 456000000 + i))
+    }
+
+    withOrcFile(data) { file =>
+      checkAnswer(
+        spark.read.orc(file),
+        data.toDF().collect())
+    }
+  }
+
+  test("test for timestamp types: save and load case class RDD with `None`s as orc") {
+    val data = (
+      Option.empty[Timestamp],
+      Option.empty[LocalDateTime]
+    ) :: Nil
+
+    withOrcFile(data) { file =>
+      checkAnswer(
+        spark.read.orc(file),
+        Row(Seq.fill(2)(null): _*))
+    }
+  }
 }
 
 class OrcV1QuerySuite extends OrcQuerySuite {