From 826d2aedba42bc87f6d1caad91ede0b7027964ec Mon Sep 17 00:00:00 2001
From: Ivan Kukrkic <ivan.kukrkic@databricks.com>
Date: Thu, 26 Sep 2024 19:17:59 +0200
Subject: [PATCH] Adding more test cases.

---
 .../jdbc/v2/PostgresIntegrationSuite.scala    | 15 ++++++
 .../apache/spark/sql/jdbc/JDBCV2Suite.scala   | 51 +++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala
index d47dd522544ba..3e32649c25fd7 100644
--- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala
+++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala
@@ -145,6 +145,21 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCT
       assert(df.collect().toSet === expectedResult)
     }
 
+    testFilterPushdown("millenium", "2001-01-01 00:00:00.0",
+      Set(
+        Row(1, Timestamp.valueOf("2024-01-01 01:01:01.0")),
+        Row(2, Timestamp.valueOf("2024-02-02 02:02:02.0"))
+      ))
+    testFilterPushdown("century", "2001-01-01 00:00:00.0",
+      Set(
+        Row(1, Timestamp.valueOf("2024-01-01 01:01:01.0")),
+        Row(2, Timestamp.valueOf("2024-02-02 02:02:02.0"))
+      ))
+    testFilterPushdown("decade", "2020-01-01 00:00:00.0",
+      Set(
+        Row(1, Timestamp.valueOf("2024-01-01 01:01:01.0")),
+        Row(2, Timestamp.valueOf("2024-02-02 02:02:02.0"))
+      ))
     testFilterPushdown("YEAR", "2024-01-01 00:00:00.0",
       Set(
         Row(1, Timestamp.valueOf("2024-01-01 01:01:01.0")),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
index 12b641e580a45..78df155fcbb10 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
@@ -1675,6 +1675,57 @@ class JDBCV2Suite extends QueryTest with SharedSparkSession with ExplainSuiteHel
       "PushedFilters: [(DATE_TRUNC('MicroseconD', TIME1)) = 1725560625000000]"
     checkPushedInfo(df9, expectedPlanFragment9)
     checkAnswer(df9, Seq(Row("adam")))
+
+    val df10 = sql("SELECT name FROM h2.test.datetime WHERE " +
+      "DATE_TRUNC('MicroSecondS', time1) = timestamp'2024-09-05 11:23:45.000000'")
+    checkFiltersRemoved(df10)
+    val expectedPlanFragment10 =
+      "PushedFilters: [(DATE_TRUNC('MicroSecondS', TIME1)) = 1725560625000000]"
+    checkPushedInfo(df10, expectedPlanFragment10)
+    checkAnswer(df10, Seq(Row("adam")))
+
+    val df11 = sql("SELECT name FROM h2.test.datetime WHERE " +
+      "DATE_TRUNC('MILLENNIUM', time1) = timestamp'2001-01-01 00:00:00.000000'")
+    checkFiltersRemoved(df11)
+    val expectedPlanFragment11 =
+      "PushedFilters: [(DATE_TRUNC('MILLENNIUM', TIME1)) = 978336000000000]"
+    checkPushedInfo(df11, expectedPlanFragment11)
+    checkAnswer(df11, Seq(Row("adam"), Row("alex"), Row("amy")))
+
+    val df12 = sql("SELECT name FROM h2.test.datetime WHERE " +
+      "DATE_TRUNC('Century', time1) = timestamp'2001-01-01 00:00:00.000000'")
+    checkFiltersRemoved(df12)
+    val expectedPlanFragment12 =
+      "PushedFilters: [(DATE_TRUNC('Century', TIME1)) = 978336000000000]"
+    checkPushedInfo(df12, expectedPlanFragment12)
+    checkAnswer(df12, Seq(Row("adam"), Row("alex"), Row("amy")))
+
+    val df13 = sql("SELECT name FROM h2.test.datetime WHERE " +
+      "DATE_TRUNC('DECADE', time1) = timestamp'2020-01-01 00:00:00.000000'")
+    checkFiltersRemoved(df13)
+    val expectedPlanFragment13 =
+      "PushedFilters: [(DATE_TRUNC('DECADE', TIME1)) = 1577865600000000]"
+    checkPushedInfo(df13, expectedPlanFragment13)
+    checkAnswer(df13, Seq(Row("adam"), Row("alex"), Row("amy")))
+
+    val df14 = sql("SELECT name FROM h2.test.datetime WHERE " +
+      "DATE_TRUNC('week', time1) = timestamp'2024-09-01 00:00:00.0'")
+    checkFiltersRemoved(df14)
+    val expectedPlanFragment14 =
+      "PushedFilters: [(DATE_TRUNC('week', TIME1)) = 1725174000000000]"
+    checkPushedInfo(df14, expectedPlanFragment14)
+    checkAnswer(df14, Seq(Row("adam")))
+
+    val e = intercept[SparkException] {
+      checkAnswer(
+        sql("SELECT name FROM h2.test.datetime WHERE " +
+        "DATE_TRUNC('invalid time unit', time1) = timestamp'2024-09-01 00:00:00.0'"),
+        Seq.empty
+      )
+    }
+    assert(e.getMessage.contains(
+      "Invalid value \"invalid time unit\" for parameter \"date-time field\"")
+    )
   }
 
   test("scan with filter push-down with misc functions") {