[Managed Iceberg] Add support for TIMESTAMP, TIME, and DATE types

.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

@@ -64,6 +64,7 @@
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 * [Managed Iceberg] Support creating tables if needed ([#32686](https://github.com/apache/beam/pull/32686))
+* [Managed Iceberg] Add support for TIMESTAMP, TIME, and DATE types ([#32688](https://github.com/apache/beam/pull/32688))
 
 ## New Features / Improvements
 

sdks/java/io/iceberg/hive/src/test/java/org/apache/beam/sdk/io/iceberg/hive/IcebergHiveCatalogIT.java

@@ -32,6 +32,7 @@
 import org.apache.beam.sdk.io.iceberg.hive.testutils.HiveMetastoreExtension;
 import org.apache.beam.sdk.managed.Managed;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
@@ -64,7 +65,10 @@
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.OutputFile;
 import org.apache.iceberg.parquet.Parquet;
+import org.apache.iceberg.util.DateTimeUtil;
 import org.apache.thrift.TException;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Rule;
@@ -100,6 +104,10 @@ public class IcebergHiveCatalogIT {
           .addArrayField("arr_long", Schema.FieldType.INT64)
           .addRowField("row", NESTED_ROW_SCHEMA)
           .addNullableRowField("nullable_row", NESTED_ROW_SCHEMA)
+          .addDateTimeField("datetime_tz")
+          .addLogicalTypeField("datetime", SqlTypes.DATETIME)
+          .addLogicalTypeField("date", SqlTypes.DATE)
+          .addLogicalTypeField("time", SqlTypes.TIME)
           .build();
 
   private static final SimpleFunction<Long, Row> ROW_FUNC =
@@ -127,6 +135,10 @@ public Row apply(Long num) {
               .addValue(LongStream.range(1, num % 10).boxed().collect(Collectors.toList()))
               .addValue(nestedRow)
               .addValue(num % 2 == 0 ? null : nestedRow)
+              .addValue(new DateTime(num).withZone(DateTimeZone.forOffsetHoursMinutes(3, 25)))
+              .addValue(DateTimeUtil.timestampFromMicros(num))
+              .addValue(DateTimeUtil.dateFromDays(Integer.parseInt(strNum)))
+              .addValue(DateTimeUtil.timeFromMicros(num))
               .build();
         }
       };

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

@@ -150,7 +150,16 @@
  *     <td> DOUBLE </td> <td> DOUBLE </td>
  *   </tr>
  *   <tr>
- *     <td> DATETIME </td> <td> STRING </td>
+ *     <td> SqlTypes.DATETIME </td> <td> TIMESTAMP </td>
+ *   </tr>
+ *   <tr>
+ *     <td> DATETIME </td> <td> TIMESTAMPTZ </td>
+ *   </tr>
+ *   <tr>
+ *     <td> SqlTypes.DATE </td> <td> DATE </td>
+ *   </tr>
+ *   <tr>
+ *     <td> SqlTypes.TIME </td> <td> TIME </td>
  *   </tr>
  *   <tr>
  *     <td> ITERABLE </td> <td> LIST </td>
@@ -166,6 +175,29 @@
  *   </tr>
  * </table>
  *
+ * <p><b>Note:</b> {@code SqlTypes} are Beam logical types.
+ *
+ * <h3>Note on timestamps</h3>
+ *
+ * <p>For an existing table, the following Beam types are supported for both {@code timestamp} and
+ * {@code timestamptz}:
+ *
+ * <ul>
+ *   <li>{@code SqlTypes.DATETIME} --> Using a {@link java.time.LocalDateTime} object
+ *   <li>{@code DATETIME} --> Using a {@link org.joda.time.DateTime} object
+ *   <li>{@code INT64} --> Using a {@link Long} representing micros since EPOCH
+ *   <li>{@code STRING} --> Using a timestamp {@link String} representation (e.g. {@code
+ *       "2024-10-08T13:18:20.053+03:27"})
+ * </ul>
+ *
+ * <p><b>Note</b>: If you expect Beam to create the Iceberg table at runtime, please provide {@code
+ * SqlTypes.DATETIME} for a {@code timestamp} column and {@code DATETIME} for a {@code timestamptz}
+ * column. If the table does not exist, Beam will treat {@code STRING} and {@code INT64} at
+ * face-value and create equivalent column types.
+ *
+ * <p>For Iceberg reads, the connector will produce Beam {@code SqlTypes.DATETIME} types for
+ * Iceberg's {@code timestamp} and {@code DATETIME} types for {@code timestamptz}.
+ *
  * <h3>Dynamic Destinations</h3>
  *
  * <p>Managed Iceberg supports writing to dynamic destinations. To do so, please provide an

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

@@ -20,12 +20,18 @@
 import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
 
 import java.nio.ByteBuffer;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
 import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
@@ -34,15 +40,13 @@
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.DateTimeUtil;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.DateTime;
-import org.joda.time.DateTimeZone;
+import org.joda.time.Instant;
 
-/** Utilities for converting between Beam and Iceberg types. */
+/** Utilities for converting between Beam and Iceberg types, made public for user's convenience. */
 public class IcebergUtils {
-  // This is made public for users convenience, as many may have more experience working with
-  // Iceberg types.
-
   private IcebergUtils() {}
 
   private static final Map<Schema.TypeName, Type> BEAM_TYPES_TO_ICEBERG_TYPES =
@@ -54,6 +58,14 @@ private IcebergUtils() {}
           .put(Schema.TypeName.DOUBLE, Types.DoubleType.get())
           .put(Schema.TypeName.STRING, Types.StringType.get())
           .put(Schema.TypeName.BYTES, Types.BinaryType.get())
+          .put(Schema.TypeName.DATETIME, Types.TimestampType.withZone())
+          .build();
+
+  private static final Map<String, Type> BEAM_LOGICAL_TYPES_TO_ICEBERG_TYPES =
+      ImmutableMap.<String, Type>builder()
+          .put(SqlTypes.DATE.getIdentifier(), Types.DateType.get())
+          .put(SqlTypes.TIME.getIdentifier(), Types.TimeType.get())
+          .put(SqlTypes.DATETIME.getIdentifier(), Types.TimestampType.withoutZone())
           .build();
 
   private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
@@ -69,9 +81,15 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
       case DOUBLE:
         return Schema.FieldType.DOUBLE;
       case DATE:
+        return Schema.FieldType.logicalType(SqlTypes.DATE);
       case TIME:
-      case TIMESTAMP: // TODO: Logical types?
-        return Schema.FieldType.DATETIME;
+        return Schema.FieldType.logicalType(SqlTypes.TIME);
+      case TIMESTAMP:
+        Types.TimestampType ts = (Types.TimestampType) type.asPrimitiveType();
+        if (ts.shouldAdjustToUTC()) {
+          return Schema.FieldType.DATETIME;
+        }
+        return Schema.FieldType.logicalType(SqlTypes.DATETIME);
       case STRING:
         return Schema.FieldType.STRING;
       case UUID:
@@ -151,6 +169,14 @@ static TypeAndMaxId beamFieldTypeToIcebergFieldType(
       // other types.
       return new TypeAndMaxId(
           --nestedFieldId, BEAM_TYPES_TO_ICEBERG_TYPES.get(beamType.getTypeName()));
+    } else if (beamType.getTypeName().isLogicalType()) {
+      String logicalTypeIdentifier =
+          checkArgumentNotNull(beamType.getLogicalType()).getIdentifier();
+      @Nullable Type type = BEAM_LOGICAL_TYPES_TO_ICEBERG_TYPES.get(logicalTypeIdentifier);
+      if (type == null) {
+        throw new RuntimeException("Unsupported Beam logical type " + logicalTypeIdentifier);
+      }
+      return new TypeAndMaxId(--nestedFieldId, type);
     } else if (beamType.getTypeName().isCollectionType()) { // ARRAY or ITERABLE
       Schema.FieldType beamCollectionType =
           Preconditions.checkArgumentNotNull(beamType.getCollectionElementType());
@@ -227,8 +253,6 @@ static TypeAndMaxId beamFieldTypeToIcebergFieldType(
    *
    * <p>The following unsupported Beam types will be defaulted to {@link Types.StringType}:
    * <li>{@link Schema.TypeName.DECIMAL}
-   * <li>{@link Schema.TypeName.DATETIME}
-   * <li>{@link Schema.TypeName.LOGICAL_TYPE}
    */
   public static org.apache.iceberg.Schema beamSchemaToIcebergSchema(final Schema schema) {
     List<Types.NestedField> fields = new ArrayList<>(schema.getFieldCount());
@@ -282,12 +306,20 @@ private static void copyFieldIntoRecord(Record rec, Types.NestedField field, Row
         Optional.ofNullable(value.getDouble(name)).ifPresent(v -> rec.setField(name, v));
         break;
       case DATE:
-        throw new UnsupportedOperationException("Date fields not yet supported");
+        Optional.ofNullable(value.getLogicalTypeValue(name, LocalDate.class))
+            .ifPresent(v -> rec.setField(name, v));
+        break;
       case TIME:
-        throw new UnsupportedOperationException("Time fields not yet supported");
+        Optional.ofNullable(value.getLogicalTypeValue(name, LocalTime.class))
+            .ifPresent(v -> rec.setField(name, v));
+        break;
       case TIMESTAMP:
-        Optional.ofNullable(value.getDateTime(name))
-            .ifPresent(v -> rec.setField(name, v.getMillis()));
+        Object val = value.getValue(name);
+        if (val == null) {
+          break;
+        }
+        Types.TimestampType ts = (Types.TimestampType) field.type().asPrimitiveType();
+        rec.setField(name, getIcebergTimestampValue(val, ts.shouldAdjustToUTC()));
         break;
       case STRING:
         Optional.ofNullable(value.getString(name)).ifPresent(v -> rec.setField(name, v));
@@ -322,6 +354,55 @@ private static void copyFieldIntoRecord(Record rec, Types.NestedField field, Row
     }
   }
 
+  /**
+   * Returns the appropriate value for an Iceberg timestamp field
+   *
+   * <p>If `timestamp`, we resolve incoming values to a {@link LocalDateTime}.
+   *
+   * <p>If `timestamptz`, we resolve to a UTC {@link OffsetDateTime}. Iceberg already resolves all
+   * incoming timestamps to UTC, so there is no harm in doing it from our side.
+   *
+   * <p>Valid types are:
+   *
+   * <ul>
+   *   <li>{@link SqlTypes.DATETIME} --> {@link LocalDateTime}
+   *   <li>{@link Schema.FieldType.DATETIME} --> {@link Instant}
+   *   <li>{@link Schema.FieldType.INT64} --> {@link Long}
+   *   <li>{@link Schema.FieldType.STRING} --> {@link String}
+   * </ul>
+   */
+  private static Object getIcebergTimestampValue(Object beamValue, boolean shouldAdjustToUtc) {
+    // timestamptz
+    if (shouldAdjustToUtc) {
+      if (beamValue instanceof LocalDateTime) { // SqlTypes.DATETIME
+        return OffsetDateTime.of((LocalDateTime) beamValue, ZoneOffset.UTC);
+      } else if (beamValue instanceof Instant) { // FieldType.DATETIME
+        return DateTimeUtil.timestamptzFromMicros(((Instant) beamValue).getMillis() * 1000L);
+      } else if (beamValue instanceof Long) { // FieldType.INT64
+        return DateTimeUtil.timestamptzFromMicros((Long) beamValue);
+      } else if (beamValue instanceof String) { // FieldType.STRING
+        return OffsetDateTime.parse((String) beamValue).withOffsetSameInstant(ZoneOffset.UTC);
+      } else {
+        throw new UnsupportedOperationException(
+            "Unsupported Beam type for Iceberg timestamp with timezone: " + beamValue.getClass());
+      }
+    }
+
+    // timestamp
+    if (beamValue instanceof LocalDateTime) { // SqlType.DATETIME
+      return beamValue;
+    } else if (beamValue instanceof Instant) { // FieldType.DATETIME
+      return DateTimeUtil.timestampFromMicros(((Instant) beamValue).getMillis() * 1000L);
+    } else if (beamValue instanceof Long) { // FieldType.INT64
+      return DateTimeUtil.timestampFromMicros((Long) beamValue);
+    } else if (beamValue instanceof String) { // FieldType.STRING
+      return LocalDateTime.parse((String) beamValue);
+    } else {
+      throw new UnsupportedOperationException(
+          "Unsupported Beam type for Iceberg timestamp with timezone: " + beamValue.getClass());
+    }
+  }
+
   /** Converts an Iceberg {@link Record} to a Beam {@link Row}. */
   public static Row icebergRecordToBeamRow(Schema schema, Record record) {
     Row.Builder rowBuilder = Row.withSchema(schema);
@@ -345,16 +426,17 @@ public static Row icebergRecordToBeamRow(Schema schema, Record record) {
         case FLOAT: // Iceberg and Beam both use float
         case DOUBLE: // Iceberg and Beam both use double
         case STRING: // Iceberg and Beam both use String
-        case BOOLEAN: // Iceberg and Beam both use String
+        case BOOLEAN: // Iceberg and Beam both use boolean
         case ARRAY:
         case ITERABLE:
         case MAP:
           rowBuilder.addValue(icebergValue);
           break;
         case DATETIME:
-          // Iceberg uses a long for millis; Beam uses joda time DateTime
-          long millis = (long) icebergValue;
-          rowBuilder.addValue(new DateTime(millis, DateTimeZone.UTC));
+          // Iceberg uses a long for micros.
+          // Beam DATETIME uses joda's DateTime, which only supports millis,
+          // so we do lose some precision here
+          rowBuilder.addValue(getBeamDateTimeValue(icebergValue));
           break;
         case BYTES:
           // Iceberg uses ByteBuffer; Beam uses byte[]
@@ -369,13 +451,59 @@ public static Row icebergRecordToBeamRow(Schema schema, Record record) {
           rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, nestedRecord));
           break;
         case LOGICAL_TYPE:
-          throw new UnsupportedOperationException(
-              "Cannot convert iceberg field to Beam logical type");
+          rowBuilder.addValue(getLogicalTypeValue(icebergValue, field.getType()));
+          break;
         default:
           throw new UnsupportedOperationException(
               "Unsupported Beam type: " + field.getType().getTypeName());
       }
     }
     return rowBuilder.build();
   }
+
+  private static DateTime getBeamDateTimeValue(Object icebergValue) {
+    long micros;
+    if (icebergValue instanceof OffsetDateTime) {
+      micros = DateTimeUtil.microsFromTimestamptz((OffsetDateTime) icebergValue);
+    } else if (icebergValue instanceof LocalDateTime) {
+      micros = DateTimeUtil.microsFromTimestamp((LocalDateTime) icebergValue);
+    } else if (icebergValue instanceof Long) {
+      micros = (long) icebergValue;
+    } else if (icebergValue instanceof String) {
+      return DateTime.parse((String) icebergValue);
+    } else {
+      throw new UnsupportedOperationException(
+          "Unsupported Iceberg type for Beam type DATETIME: " + icebergValue.getClass());
+    }
+    return new DateTime(micros / 1000L);
+  }
+
+  private static Object getLogicalTypeValue(Object icebergValue, Schema.FieldType type) {
+    if (icebergValue instanceof String) {
+      String strValue = (String) icebergValue;
+      if (type.isLogicalType(SqlTypes.DATE.getIdentifier())) {
+        return LocalDate.parse(strValue);
+      } else if (type.isLogicalType(SqlTypes.TIME.getIdentifier())) {
+        return LocalTime.parse(strValue);
+      } else if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
+        return LocalDateTime.parse(strValue);
+      }
+    } 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 (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();
+    }
+    // LocalDateTime, LocalDate, LocalTime
+    return icebergValue;
+  }
 }