[Iceberg] Make timestamptz return new Timestamp.MICROS logical type

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 2
+    "modification": 3
 }

CHANGES.md

@@ -74,6 +74,10 @@
 ## Breaking Changes
 
 * (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)).
+* [IcebergIO] Reading a `timestamptz` column will now return a `Timestamp.MICROS` Beam logical type, equivalent to a
+ `java.time.Instant` object, in order to preserve microseconds (the old Beam `Schema.FieldType#DATETIME` primitive type
+ truncates past milliseconds). Use pipeline option `--updateCompatibilityVersion=2.75.0` (or any older version) to
+ keep the old behavior ([#39344](https://github.com/apache/beam/issues/39344)).
 
 ## Deprecations
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java

@@ -932,6 +932,12 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso
           return java.time.Instant.parse(jsonBQString);
         }
       } else if (fieldType.isLogicalType(Timestamp.IDENTIFIER)) {
+        if (!jsonBQString.contains("UTC")) {
+          BigDecimal bd = new BigDecimal(jsonBQString);
+          long seconds = bd.longValue();
+          long nanos = bd.subtract(BigDecimal.valueOf(seconds)).movePointRight(9).longValue();
+          return java.time.Instant.ofEpochSecond(seconds, nanos);
+        }
         return VAR_PRECISION_FORMATTER.parse(jsonBQString, java.time.Instant::from);
       }
     }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilsTest.java

@@ -41,6 +41,7 @@
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
+import java.time.OffsetDateTime;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
@@ -1444,22 +1445,43 @@ public void testToBeamRow_timestampNanos_utcSuffix() {
 
   @Test
   @SuppressWarnings("JavaInstantGetSecondsGetNano")
-  public void testToBeamRow_timestampMicros_utcSuffix() {
+  public void testToBeamRow_timestampMicros() {
     Schema schema = Schema.builder().addLogicalTypeField("ts", Timestamp.MICROS).build();
 
     // BigQuery format with " UTC" suffix
     String timestamp = "2024-08-10 16:52:07.123456 UTC";
+    String parsableTimestamp = "2024-08-10T16:52:07.123456Z";
+    String negativeTimestamp = "1960-08-10T16:52:07.000123Z";
 
-    Row beamRow = BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", timestamp));
+    java.time.Instant instant = OffsetDateTime.parse(parsableTimestamp).toInstant();
+    String value = instant.getEpochSecond() + "." + instant.getNano() / 1000;
+    java.time.Instant negInstant = OffsetDateTime.parse(negativeTimestamp).toInstant();
+    String negValue =
+        BigDecimal.valueOf(negInstant.getEpochSecond())
+            .add(BigDecimal.valueOf(negInstant.getNano(), 9))
+            .toPlainString();
 
-    java.time.Instant actual = (java.time.Instant) beamRow.getValue("ts");
-    assertEquals(2024, actual.atZone(java.time.ZoneOffset.UTC).getYear());
-    assertEquals(8, actual.atZone(java.time.ZoneOffset.UTC).getMonthValue());
-    assertEquals(10, actual.atZone(java.time.ZoneOffset.UTC).getDayOfMonth());
-    assertEquals(16, actual.atZone(java.time.ZoneOffset.UTC).getHour());
-    assertEquals(52, actual.atZone(java.time.ZoneOffset.UTC).getMinute());
-    assertEquals(7, actual.atZone(java.time.ZoneOffset.UTC).getSecond());
-    assertEquals(123456000, actual.getNano());
+    List<TableRow> testRows =
+        Arrays.asList(
+            new TableRow().set("ts", timestamp),
+            new TableRow().set("ts", value),
+            new TableRow().set("negative", true).set("ts", negValue));
+
+    for (TableRow row : testRows) {
+      Row beamRow = BigQueryUtils.toBeamRow(schema, row);
+
+      java.time.Instant actual = (java.time.Instant) beamRow.getValue("ts");
+
+      assertEquals(
+          row.get("negative") == null ? 2024 : 1960,
+          actual.atZone(java.time.ZoneOffset.UTC).getYear());
+      assertEquals(8, actual.atZone(java.time.ZoneOffset.UTC).getMonthValue());
+      assertEquals(10, actual.atZone(java.time.ZoneOffset.UTC).getDayOfMonth());
+      assertEquals(16, actual.atZone(java.time.ZoneOffset.UTC).getHour());
+      assertEquals(52, actual.atZone(java.time.ZoneOffset.UTC).getMinute());
+      assertEquals(7, actual.atZone(java.time.ZoneOffset.UTC).getSecond());
+      assertEquals(row.get("negative") == null ? 123456000 : 123000, actual.getNano());
+    }
   }
 
   @Test

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java

@@ -25,6 +25,7 @@
 import java.util.Map;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.io.Read;
+import org.apache.beam.sdk.options.StreamingOptions;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.values.PBegin;
@@ -668,12 +669,23 @@ public PCollection<Row> expand(PBegin input) {
 
       Table table = TableCache.get(getCatalogConfig(), tableId);
 
+      @Nullable
+      String updateCompatibilityVersion =
+          input
+              .getPipeline()
+              .getOptions()
+              .as(StreamingOptions.class)
+              .getUpdateCompatibilityVersion();
+
       IcebergScanConfig scanConfig =
           IcebergScanConfig.builder()
               .setCatalogConfig(getCatalogConfig())
               .setScanType(IcebergScanConfig.ScanType.TABLE)
               .setTableIdentifier(tableId)
-              .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema()))
+              .setSchema(
+                  IcebergUtils.icebergSchemaToBeamSchema(
+                      table.schema(), updateCompatibilityVersion))
+              .setUpdateCompatibilityVersion(updateCompatibilityVersion)
               .setFromSnapshotInclusive(getFromSnapshot())
               .setToSnapshot(getToSnapshot())
               .setFromTimestamp(getFromTimestamp())

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java

@@ -164,7 +164,8 @@ public org.apache.iceberg.Schema recordIdSchema() {
 
   public Schema rowIdBeamSchema() {
     if (cachedRowIdBeamSchema == null) {
-      cachedRowIdBeamSchema = icebergSchemaToBeamSchema(recordIdSchema());
+      cachedRowIdBeamSchema =
+          icebergSchemaToBeamSchema(recordIdSchema(), getUpdateCompatibilityVersion());
     }
     return cachedRowIdBeamSchema;
   }
@@ -237,6 +238,9 @@ public Expression getFilter() {
   @Pure
   public abstract boolean getUseCdc();
 
+  @Pure
+  public abstract @Nullable String getUpdateCompatibilityVersion();
+
   @Pure
   public abstract @Nullable Boolean getStreaming();
 
@@ -335,6 +339,9 @@ public Builder setTableIdentifier(String... names) {
 
     public abstract Builder setUseCdc(boolean useCdc);
 
+    public abstract Builder setUpdateCompatibilityVersion(
+        @Nullable String updateCompatibilityVersion);
+
     public abstract Builder setStreaming(@Nullable Boolean streaming);
 
     public abstract Builder setPollInterval(@Nullable Duration pollInterval);

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java

@@ -40,7 +40,9 @@
 import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
 import org.apache.beam.sdk.schemas.logicaltypes.PassThroughLogicalType;
 import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
+import org.apache.beam.sdk.schemas.logicaltypes.Timestamp;
 import org.apache.beam.sdk.util.Preconditions;
+import org.apache.beam.sdk.util.construction.TransformUpgrader;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
@@ -83,7 +85,8 @@ private IcebergUtils() {}
           .put(MicrosInstant.IDENTIFIER, Types.TimestampType.withZone())
           .build();
 
-  private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
+  private static Schema.FieldType icebergTypeToBeamFieldType(
+      final Type type, @Nullable String updateCompatibilityVersion) {
     switch (type.typeId()) {
       case BOOLEAN:
         return Schema.FieldType.BOOLEAN;
@@ -102,7 +105,14 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
       case TIMESTAMP:
         Types.TimestampType ts = (Types.TimestampType) type.asPrimitiveType();
         if (ts.shouldAdjustToUTC()) {
-          return Schema.FieldType.DATETIME;
+          // timestamptz. The micros-precision Timestamp logical type preserves microseconds, while
+          // the legacy DATETIME (joda) mapping truncates to millis. Gated for update compatibility.
+          if (updateCompatibilityVersion != null
+              && !updateCompatibilityVersion.isEmpty()
+              && TransformUpgrader.compareVersions(updateCompatibilityVersion, "2.76.0") < 0) {
+            return Schema.FieldType.DATETIME;
+          }
+          return Schema.FieldType.logicalType(Timestamp.MICROS);
         }
         return Schema.FieldType.logicalType(SqlTypes.DATETIME);
       case STRING:
@@ -114,36 +124,51 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
       case DECIMAL:
         return Schema.FieldType.DECIMAL;
       case STRUCT:
-        return Schema.FieldType.row(icebergStructTypeToBeamSchema(type.asStructType()));
+        return Schema.FieldType.row(
+            icebergStructTypeToBeamSchema(type.asStructType(), updateCompatibilityVersion));
       case LIST:
-        return Schema.FieldType.array(icebergTypeToBeamFieldType(type.asListType().elementType()));
+        return Schema.FieldType.array(
+            icebergTypeToBeamFieldType(
+                type.asListType().elementType(), updateCompatibilityVersion));
       case MAP:
         return Schema.FieldType.map(
-            icebergTypeToBeamFieldType(type.asMapType().keyType()),
-            icebergTypeToBeamFieldType(type.asMapType().valueType()));
+            icebergTypeToBeamFieldType(type.asMapType().keyType(), updateCompatibilityVersion),
+            icebergTypeToBeamFieldType(type.asMapType().valueType(), updateCompatibilityVersion));
       default:
         throw new RuntimeException("Unrecognized Iceberg Type: " + type.typeId());
     }
   }
 
-  private static Schema.Field icebergFieldToBeamField(final Types.NestedField field) {
-    return Schema.Field.of(field.name(), icebergTypeToBeamFieldType(field.type()))
+  private static Schema.Field icebergFieldToBeamField(
+      final Types.NestedField field, @Nullable String updateCompatibilityVersion) {
+    return Schema.Field.of(
+            field.name(), icebergTypeToBeamFieldType(field.type(), updateCompatibilityVersion))
         .withNullable(field.isOptional());
   }
 
   /** Converts an Iceberg {@link org.apache.iceberg.Schema} to a Beam {@link Schema}. */
   public static Schema icebergSchemaToBeamSchema(final org.apache.iceberg.Schema schema) {
+    return icebergSchemaToBeamSchema(schema, null);
+  }
+
+  /**
+   * Converts an Iceberg {@link org.apache.iceberg.Schema} to a Beam {@link Schema}, accounting for
+   * update compatibility.
+   */
+  public static Schema icebergSchemaToBeamSchema(
+      final org.apache.iceberg.Schema schema, @Nullable String updateCompatibilityVersion) {
     Schema.Builder builder = Schema.builder();
     for (Types.NestedField f : schema.columns()) {
-      builder.addField(icebergFieldToBeamField(f));
+      builder.addField(icebergFieldToBeamField(f, updateCompatibilityVersion));
     }
     return builder.build();
   }
 
-  private static Schema icebergStructTypeToBeamSchema(final Types.StructType struct) {
+  private static Schema icebergStructTypeToBeamSchema(
+      final Types.StructType struct, @Nullable String updateCompatibilityVersion) {
     Schema.Builder builder = Schema.builder();
     for (Types.NestedField f : struct.fields()) {
-      builder.addField(icebergFieldToBeamField(f));
+      builder.addField(icebergFieldToBeamField(f, updateCompatibilityVersion));
     }
     return builder.build();
   }
@@ -198,7 +223,17 @@ static TypeAndMaxId beamFieldTypeToIcebergFieldType(
       String logicalTypeIdentifier = logicalType.getIdentifier();
       @Nullable Type type = BEAM_LOGICAL_TYPES_TO_ICEBERG_TYPES.get(logicalTypeIdentifier);
       if (type == null) {
-        throw new RuntimeException("Unsupported Beam logical type " + logicalTypeIdentifier);
+        if (beamType.isLogicalType(Timestamp.IDENTIFIER)) {
+          int precision = checkStateNotNull(logicalType.getArgument());
+          if (precision == Timestamp.MICROS.getArgument()) {
+            type = Types.TimestampType.withZone();
+          } else {
+            throw new UnsupportedOperationException(
+                "Unsupported Timestamp precision: " + precision);
+          }
+        } else {
+          throw new RuntimeException("Unsupported Beam logical type " + logicalTypeIdentifier);
+        }
       }
       return new TypeAndMaxId(--nestedFieldId, type);
     } else if (beamType.getTypeName().isCollectionType()) { // ARRAY or ITERABLE
@@ -613,21 +648,28 @@ private static Object getLogicalTypeValue(Object icebergValue, Schema.FieldType
         return LocalTime.parse(strValue);
       } else if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
         return LocalDateTime.parse(strValue);
+      } else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
+        return OffsetDateTime.parse(strValue).toInstant();
       }
     } else if (icebergValue instanceof Long) {
       if (type.isLogicalType(SqlTypes.TIME.getIdentifier())) {
         return DateTimeUtil.timeFromMicros((Long) icebergValue);
       } else if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
         return DateTimeUtil.timestampFromMicros((Long) icebergValue);
+      } else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
+        // timestamptz stored as micros since epoch -> java.time.Instant (micros preserved).
+        return DateTimeUtil.timestamptzFromMicros((Long) icebergValue).toInstant();
       }
     } else if (icebergValue instanceof Integer
         && type.isLogicalType(SqlTypes.DATE.getIdentifier())) {
       return DateTimeUtil.dateFromDays((Integer) icebergValue);
-    } else if (icebergValue instanceof OffsetDateTime
-        && type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
-      return ((OffsetDateTime) icebergValue)
-          .withOffsetSameInstant(ZoneOffset.UTC)
-          .toLocalDateTime();
+    } else if (icebergValue instanceof OffsetDateTime) {
+      OffsetDateTime odt = (OffsetDateTime) icebergValue;
+      if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
+        return odt.withOffsetSameInstant(ZoneOffset.UTC).toLocalDateTime();
+      } else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
+        return odt.toInstant();
+      }
     }
     // LocalDateTime, LocalDate, LocalTime
     return icebergValue;

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IncrementalScanSource.java

@@ -68,7 +68,9 @@ public PCollection<Row> expand(PBegin input) {
         .setCoder(KvCoder.of(ReadTaskDescriptor.getCoder(), ReadTask.getCoder()))
         .apply(Redistribute.arbitrarily())
         .apply("Read Rows From Tasks", ParDo.of(new ReadFromTasks(scanConfig)))
-        .setRowSchema(IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()));
+        .setRowSchema(
+            IcebergUtils.icebergSchemaToBeamSchema(
+                scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion()));
   }
 
   /** Continuously watches for new snapshots. */

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java

@@ -69,7 +69,9 @@ public void process(
         return;
       }
       FileScanTask task = fileScanTasks.get((int) l);
-      Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema());
+      Schema beamSchema =
+          IcebergUtils.icebergSchemaToBeamSchema(
+              scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion());
       try (CloseableIterable<Record> reader = ReadUtils.createReader(task, table, scanConfig)) {
 
         for (Record record : reader) {

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanSource.java

@@ -116,7 +116,9 @@ public void populateDisplayData(DisplayData.Builder builder) {
 
   @Override
   public Coder<Row> getOutputCoder() {
-    return RowCoder.of(IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()));
+    return RowCoder.of(
+        IcebergUtils.icebergSchemaToBeamSchema(
+            scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion()));
   }
 
   @Override

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanTaskReader.java

@@ -66,7 +66,10 @@ class ScanTaskReader extends BoundedSource.BoundedReader<Row> {
 
   public ScanTaskReader(ScanTaskSource source) {
     this.source = source;
-    this.beamSchema = icebergSchemaToBeamSchema(source.getScanConfig().getProjectedSchema());
+    this.beamSchema =
+        icebergSchemaToBeamSchema(
+            source.getScanConfig().getProjectedSchema(),
+            source.getScanConfig().getUpdateCompatibilityVersion());
   }
 
   @Override