From cf3fa3a1bb2fd398e8d823007aee790b775d68a0 Mon Sep 17 00:00:00 2001
From: Grigory Pomadchin <gr.pomadchin@gmail.com>
Date: Sat, 5 Mar 2022 16:39:09 -0500
Subject: [PATCH] Adjust tests and explanations to limitations of the current
 Instant encoding

---
 .../src/test/scala/frameless/ColumnTests.scala    | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/dataset/src/test/scala/frameless/ColumnTests.scala b/dataset/src/test/scala/frameless/ColumnTests.scala
index d651f8261..aa1449690 100644
--- a/dataset/src/test/scala/frameless/ColumnTests.scala
+++ b/dataset/src/test/scala/frameless/ColumnTests.scala
@@ -1,6 +1,6 @@
 package frameless
 
-import java.time.Instant
+import java.time.{Instant, Period, Duration}
 import org.scalacheck.Prop._
 import org.scalacheck.{Arbitrary, Gen, Prop}, Arbitrary.arbitrary
 import org.scalatest.matchers.should.Matchers
@@ -13,7 +13,14 @@ final class ColumnTests extends TypedDatasetSuite with Matchers {
   private implicit object OrderingImplicits {
     implicit val sqlDateOrdering: Ordering[SQLDate] = Ordering.by(_.days)
     implicit val sqlTimestmapOrdering: Ordering[SQLTimestamp] = Ordering.by(_.us)
-    implicit val arbInstant: Arbitrary[Instant] = Arbitrary(Gen.choose[Instant](Instant.EPOCH, Instant.now))
+    implicit val periodOrdering: Ordering[Period] = Ordering.by(p => (p.getYears, p.getMonths, p.getDays))
+    /**
+      * DateTimeUtils.instantToMicros supports dates starting 1970-01-01T00:00:00Z, which is Instant.EPOCH.
+      * This function also overflows on Instant.MAX, to be sure it never overflows we use Instant.MAX / 4.
+      * For implementation details check the org.apache.spark.sql.catalyst.util.DateTimeUtils.instantToMicros function details.
+      */
+    implicit val arbInstant: Arbitrary[Instant] = Arbitrary(Gen.choose[Instant](Instant.EPOCH, Instant.ofEpochMilli(Instant.MAX.getEpochSecond / 4)))
+    implicit val arbPeriod: Arbitrary[Period] = Arbitrary(Gen.chooseNum(0, Int.MaxValue).map(l => Period.of(l, l, l)))
   }
 
   test("select('a < 'b, 'a <= 'b, 'a > 'b, 'a >= 'b)") {
@@ -44,6 +51,8 @@ final class ColumnTests extends TypedDatasetSuite with Matchers {
     check(forAll(prop[SQLTimestamp] _))
     check(forAll(prop[String] _))
     check(forAll(prop[Instant] _))
+    check(forAll(prop[Duration] _))
+    check(forAll(prop[Period] _))
   }
 
   test("between") {
@@ -71,6 +80,8 @@ final class ColumnTests extends TypedDatasetSuite with Matchers {
     check(forAll(prop[SQLTimestamp] _))
     check(forAll(prop[String] _))
     check(forAll(prop[Instant] _))
+    check(forAll(prop[Duration] _))
+    check(forAll(prop[Period] _))
   }
 
   test("toString") {