Handle Java 8 types for dates and timestamps when compiling filters (G…

…oogleCloudDataproc#1025) Co-authored-by: Thomas Powell <[email protected]>
suryasoma · Jul 21, 2023 · 741bdb3 · 741bdb3
1 parent 0e69e08
commit 741bdb3
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/...uery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkFilterUtils.java b/...uery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkFilterUtils.java
@@ -21,6 +21,8 @@
 import com.google.common.collect.ImmutableList;
 import java.sql.Date;
 import java.sql.Timestamp;
+import java.time.Instant;
+import java.time.LocalDate;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.Optional;
@@ -315,10 +317,10 @@ static String compileValue(Object value, char arrayStart, char arrayEnd) {
     if (value instanceof String) {
       return "'" + escape((String) value) + "'";
     }
-    if (value instanceof Date) {
+    if (value instanceof Date || value instanceof LocalDate) {
       return "DATE '" + value + "'";
     }
-    if (value instanceof Timestamp) {
+    if (value instanceof Timestamp || value instanceof Instant) {
       return "TIMESTAMP '" + value + "'";
     }
     if (value instanceof Object[]) {

diff --git a/...-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkFilterUtilsTest.java b/...-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkFilterUtilsTest.java
@@ -21,6 +21,10 @@
 import java.sql.Date;
 import java.sql.Timestamp;
 import java.text.ParseException;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Optional;
@@ -201,6 +205,16 @@ public void testDateFilters() throws ParseException {
         .isEqualTo("`datefield` IN (DATE '2020-09-01', DATE '2020-11-03')");
   }
 
+  @Test
+  public void testDateFilters_java8Time() {
+    assertThat(
+            SparkFilterUtils.compileFilter(
+                In.apply(
+                    "datefield",
+                    new Object[] {LocalDate.of(2020, 9, 1), LocalDate.of(2020, 11, 3)})))
+        .isEqualTo("`datefield` IN (DATE '2020-09-01', DATE '2020-11-03')");
+  }
+
   @Test
   public void testTimestampFilters() throws ParseException {
     Timestamp ts1 = Timestamp.valueOf("2008-12-25 15:30:00");
@@ -210,6 +224,15 @@ public void testTimestampFilters() throws ParseException {
             "`tsfield` IN (TIMESTAMP '2008-12-25 15:30:00.0', TIMESTAMP '2020-01-25 02:10:10.0')");
   }
 
+  @Test
+  public void testTimestampFilters_java8Time() {
+    Instant ts1 = LocalDateTime.of(2008, 12, 25, 15, 30, 0).toInstant(ZoneOffset.UTC);
+    Instant ts2 = LocalDateTime.of(2020, 1, 25, 2, 10, 10).toInstant(ZoneOffset.UTC);
+    assertThat(SparkFilterUtils.compileFilter(In.apply("tsfield", new Object[] {ts1, ts2})))
+        .isEqualTo(
+            "`tsfield` IN (TIMESTAMP '2008-12-25T15:30:00Z', TIMESTAMP '2020-01-25T02:10:10Z')");
+  }
+
   @Test
   public void testFiltersWithNestedOrAnd_1() {
     if (dataFormat == ARROW && !pushAllFilters) {