Add timestamp tests

Author: MaxGekk

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SpecialDatetimeValuesSuite.scala

@@ -17,14 +17,16 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import java.time.LocalDate
+import java.time.{Instant, LocalDate, LocalDateTime, LocalTime, ZoneId}
 
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Literal}
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
-import org.apache.spark.sql.types.DateType
+import org.apache.spark.sql.catalyst.util.DateTimeConstants.MICROS_PER_MINUTE
+import org.apache.spark.sql.catalyst.util.DateTimeUtils.{instantToMicros, localDateTimeToMicros}
+import org.apache.spark.sql.types.{AtomicType, DateType, TimestampNTZType, TimestampType}
 
 class SpecialDatetimeValuesSuite extends PlanTest {
   object Optimize extends RuleExecutor[LogicalPlan] {
@@ -55,4 +57,45 @@ class SpecialDatetimeValuesSuite extends PlanTest {
       assert(expected === lits.toSet)
     }
   }
+
+  private def testSpecialTs(tsType: AtomicType, expected: Set[Long], zoneId: ZoneId): Unit = {
+    val in = Project(Seq(
+      Alias(Cast(Literal("epoch"), tsType, Some(zoneId.getId)), "epoch")(),
+      Alias(Cast(Literal("now"), tsType, Some(zoneId.getId)), "now")(),
+      Alias(Cast(Literal("tomorrow"), tsType, Some(zoneId.getId)), "tomorrow")(),
+      Alias(Cast(Literal("yesterday"), tsType, Some(zoneId.getId)), "yesterday")()),
+      LocalRelation())
+
+    val plan = Optimize.execute(in.analyze).asInstanceOf[Project]
+    val lits = new scala.collection.mutable.ArrayBuffer[Long]
+    plan.transformAllExpressions { case e: Literal if e.dataType == tsType =>
+      lits += e.value.asInstanceOf[Long]
+      e
+    }
+    assert(lits.forall(ts => expected.exists(ets => Math.abs(ets -ts) <= MICROS_PER_MINUTE)))
+  }
+
+  test("special timestamp_ltz values") {
+    testSpecialDatetimeValues { zoneId =>
+      val expected = Set(
+        Instant.ofEpochSecond(0),
+        Instant.now(),
+        Instant.now().atZone(zoneId).`with`(LocalTime.MIDNIGHT).plusDays(1).toInstant,
+        Instant.now().atZone(zoneId).`with`(LocalTime.MIDNIGHT).minusDays(1).toInstant
+      ).map(instantToMicros)
+      testSpecialTs(TimestampType, expected, zoneId)
+    }
+  }
+
+  test("special timestamp_ntz values") {
+    testSpecialDatetimeValues { zoneId =>
+      val expected = Set(
+        LocalDateTime.of(1970, 1, 1, 0, 0),
+        LocalDateTime.now(),
+        LocalDateTime.now().`with`(LocalTime.MIDNIGHT).plusDays(1),
+        LocalDateTime.now().`with`(LocalTime.MIDNIGHT).minusDays(1)
+      ).map(localDateTimeToMicros)
+      testSpecialTs(TimestampNTZType, expected, zoneId)
+    }
+  }
 }