Use date_trunc for DATE field in dynamic partition overwrite (#1310)

GoogleCloudDataproc · Oct 23, 2024 · 2875a56 · 2875a56
1 parent 14fa9b8
commit 2875a56
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/...nnector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java b/...nnector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java
@@ -735,12 +735,24 @@ static String getQueryForTimePartitionedTable(
       TimePartitioning timePartitioning) {
     TimePartitioning.Type partitionType = timePartitioning.getType();
     String partitionField = timePartitioning.getField();
+    FieldList allFields = destinationDefinition.getSchema().getFields();
+    Optional<LegacySQLTypeName> partitionFieldType =
+        allFields.stream()
+            .filter(field -> partitionField.equals(field.getName()))
+            .map(field -> field.getType())
+            .findFirst();
+    // Using timestamp_trunc on a DATE field results in $cast_to_DATETIME which prevents pruning
+    // when required_partition_filter is set, as it is not considered a monotonic function
+    String truncFuntion =
+        (partitionFieldType.isPresent() && partitionFieldType.get().equals(LegacySQLTypeName.DATE))
+            ? "date_trunc"
+            : "timestamp_trunc";
     String extractedPartitionedSource =
-        String.format("timestamp_trunc(`%s`, %s)", partitionField, partitionType.toString());
+        String.format("%s(`%s`, %s)", truncFuntion, partitionField, partitionType.toString());
     String extractedPartitionedTarget =
         String.format(
-            "timestamp_trunc(`%s`.`%s`, %s)", "target", partitionField, partitionType.toString());
-    FieldList allFields = destinationDefinition.getSchema().getFields();
+            "%s(`target`.`%s`, %s)", truncFuntion, partitionField, partitionType.toString());
+
     String commaSeparatedFields =
         allFields.stream().map(Field::getName).collect(Collectors.joining("`,`", "`", "`"));