(IcebergIO) Support PartitionSpec/SortOrder on dynamic table creation via IcebergIO

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

@@ -18,10 +18,12 @@
 package org.apache.beam.sdk.io.iceberg;
 
 import java.io.Serializable;
+import java.util.List;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.ValueInSingleWindow;
 import org.apache.iceberg.catalog.TableIdentifier;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public interface DynamicDestinations extends Serializable {
 
@@ -34,6 +36,14 @@ public interface DynamicDestinations extends Serializable {
   String getTableStringIdentifier(ValueInSingleWindow<Row> element);
 
   static DynamicDestinations singleTable(TableIdentifier tableId, Schema inputSchema) {
-    return new OneTableDynamicDestinations(tableId, inputSchema);
+    return new OneTableDynamicDestinations(tableId, inputSchema, null, null);
+  }
+
+  static DynamicDestinations singleTable(
+      TableIdentifier tableId,
+      Schema inputSchema,
+      @Nullable List<String> partitionFields,
+      @Nullable List<String> sortFields) {
+    return new OneTableDynamicDestinations(tableId, inputSchema, partitionFields, sortFields);
   }
 }

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

@@ -409,6 +409,10 @@ public abstract static class WriteRows extends PTransform<PCollection<Row>, Iceb
 
     abstract @Nullable Map<String, String> getWriteProperties();
 
+    abstract @Nullable List<String> getPartitionFields();
+
+    abstract @Nullable List<String> getSortFields();
+
     abstract Builder toBuilder();
 
     @AutoValue.Builder
@@ -429,6 +433,10 @@ abstract static class Builder {
 
       abstract Builder setWriteProperties(Map<String, String> writeProperties);
 
+      abstract Builder setPartitionFields(List<String> partitionFields);
+
+      abstract Builder setSortFields(List<String> sortFields);
+
       abstract WriteRows build();
     }
 
@@ -483,6 +491,26 @@ public WriteRows withWriteProperties(Map<String, String> writeProperties) {
       return toBuilder().setWriteProperties(writeProperties).build();
     }
 
+    /**
+     * Defines the desired Partition Spec to be applied when the Iceberg table must be dynamically
+     * created, e.g. `bucket(id_field, 32)` or `day(timestamp_field)`
+     *
+     * <p>See: https://iceberg.apache.org/spec/#partitioning
+     */
+    public WriteRows withPartitionFields(List<String> partitionFields) {
+      return toBuilder().setPartitionFields(partitionFields).build();
+    }
+
+    /**
+     * Defines the desired Sort Order to be applied when the Iceberg table must be dynamically
+     * created, e.g. `int_field desc` or `bucket(modulo_5, 4) asc nulls last`
+     *
+     * <p>See: https://iceberg.apache.org/spec/#sorting
+     */
+    public WriteRows withSortOrder(List<String> sortFields) {
+      return toBuilder().setSortFields(sortFields).build();
+    }
+
     @Override
     public IcebergWriteResult expand(PCollection<Row> input) {
       List<?> allToArgs = Arrays.asList(getTableIdentifier(), getDynamicDestinations());
@@ -494,7 +522,10 @@ public IcebergWriteResult expand(PCollection<Row> input) {
       if (destinations == null) {
         destinations =
             DynamicDestinations.singleTable(
-                Preconditions.checkNotNull(getTableIdentifier()), input.getSchema());
+                Preconditions.checkNotNull(getTableIdentifier()),
+                input.getSchema(),
+                getPartitionFields(),
+                getSortFields());
       }
 
       // Assign destinations before re-windowing to global in WriteToDestinations because

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

@@ -23,20 +23,24 @@
 import java.io.IOException;
 import java.io.ObjectInput;
 import java.io.ObjectOutput;
+import java.util.List;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.ValueInSingleWindow;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 class OneTableDynamicDestinations implements DynamicDestinations, Externalizable {
   // TableId represented as String for serializability
   private transient @MonotonicNonNull String tableIdString;
 
   private transient @MonotonicNonNull TableIdentifier tableId;
   private transient @MonotonicNonNull Schema rowSchema;
+  private transient @Nullable List<String> partitionFields;
+  private transient @Nullable List<String> sortFields;
 
   @VisibleForTesting
   TableIdentifier getTableIdentifier() {
@@ -46,9 +50,15 @@ TableIdentifier getTableIdentifier() {
     return tableId;
   }
 
-  OneTableDynamicDestinations(TableIdentifier tableId, Schema rowSchema) {
+  OneTableDynamicDestinations(
+      TableIdentifier tableId,
+      Schema rowSchema,
+      @Nullable List<String> partitionFields,
+      @Nullable List<String> sortFields) {
     this.tableIdString = IcebergUtils.tableIdentifierToString(tableId);
     this.rowSchema = rowSchema;
+    this.partitionFields = partitionFields;
+    this.sortFields = sortFields;
   }
 
   @Override
@@ -68,24 +78,41 @@ public String getTableStringIdentifier(ValueInSingleWindow<Row> element) {
 
   @Override
   public IcebergDestination instantiateDestination(String unused) {
+    @Nullable IcebergTableCreateConfig createConfig = null;
+    if (partitionFields != null || sortFields != null) {
+      createConfig =
+          IcebergTableCreateConfig.builder()
+              .setSchema(checkStateNotNull(rowSchema))
+              .setPartitionFields(partitionFields)
+              .setSortFields(sortFields)
+              .build();
+    }
     return IcebergDestination.builder()
         .setTableIdentifier(getTableIdentifier())
-        .setTableCreateConfig(null)
+        .setTableCreateConfig(createConfig)
         .setFileFormat(FileFormat.PARQUET)
         .build();
   }
 
   // Need a public default constructor for custom serialization
   public OneTableDynamicDestinations() {}
 
+  @SuppressWarnings("nullness")
   @Override
   public void writeExternal(ObjectOutput out) throws IOException {
     out.writeUTF(checkStateNotNull(tableIdString));
+    out.writeObject(rowSchema);
+    out.writeObject(partitionFields);
+    out.writeObject(sortFields);
   }
 
+  @SuppressWarnings("nullness")
   @Override
-  public void readExternal(ObjectInput in) throws IOException {
+  public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
     tableIdString = in.readUTF();
     tableId = IcebergUtils.parseTableIdentifier(tableIdString);
+    rowSchema = (Schema) in.readObject();
+    partitionFields = (List<String>) in.readObject();
+    sortFields = (List<String>) in.readObject();
   }
 }

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

@@ -39,6 +39,7 @@
 import org.apache.iceberg.PartitionField;
 import org.apache.iceberg.PartitionKey;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.catalog.Catalog;
@@ -189,6 +190,8 @@ private Table loadOrCreateTable(
       @Nullable IcebergTableCreateConfig createConfig = destination.getTableCreateConfig();
       PartitionSpec partitionSpec =
           createConfig != null ? createConfig.getPartitionSpec() : PartitionSpec.unpartitioned();
+      SortOrder sortOrder =
+          createConfig != null ? createConfig.getSortOrder() : SortOrder.unsorted();
       Map<String, String> tableProperties =
           createConfig != null && createConfig.getTableProperties() != null
               ? createConfig.getTableProperties()
@@ -217,13 +220,19 @@ private Table loadOrCreateTable(
         org.apache.iceberg.Schema tableSchema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema);
         try {
           Table table =
-              catalog.createTable(identifier, tableSchema, partitionSpec, tableProperties);
+              catalog
+                  .buildTable(identifier, tableSchema)
+                  .withPartitionSpec(partitionSpec)
+                  .withSortOrder(sortOrder)
+                  .withProperties(tableProperties)
+                  .create();
           LOG.info(
               "Created Iceberg table '{}' with schema: {}\n"
-                  + ", partition spec: {}, table properties: {}",
+                  + ", partition spec: {}, sort order: {}, table properties: {}",
               identifier,
               tableSchema,
               partitionSpec,
+              sortOrder,
               tableProperties);
           return table;
         } catch (AlreadyExistsException ignored) {

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

@@ -76,7 +76,10 @@
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.DistributionMode;
 import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.NullOrder;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.SortDirection;
+import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.SupportsNamespaces;
@@ -775,4 +778,55 @@ public void process(@Element KV<Integer, Iterable<KV<String, Integer>>> sums) {
             .getCommitted();
     assertEquals(5L, numWaves);
   }
+
+  @Test
+  public void testCreateTableWithPartitionSpecAndSortOrder() {
+    TableIdentifier tableId =
+        TableIdentifier.of(
+            "default", "sorted_partitioned_" + Long.toString(UUID.randomUUID().hashCode(), 16));
+
+    Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA);
+
+    Map<String, String> catalogProps =
+        ImmutableMap.<String, String>builder()
+            .put("type", CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP)
+            .put("warehouse", warehouse.location)
+            .build();
+
+    IcebergCatalogConfig catalog =
+        IcebergCatalogConfig.builder()
+            .setCatalogName("name")
+            .setCatalogProperties(catalogProps)
+            .build();
+
+    testPipeline
+        .apply("Records To Add", Create.of(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1)))
+        .setRowSchema(beamSchema)
+        .apply(
+            "Append To Table",
+            writeTransform(catalog, tableId)
+                .withPartitionFields(ImmutableList.of("bucket(id, 2)"))
+                .withSortOrder(ImmutableList.of("data asc nulls first")));
+
+    testPipeline.run().waitUntilFinish();
+
+    Table table = warehouse.loadTable(tableId);
+
+    // verify partition spec
+    PartitionSpec spec = table.spec();
+    assertTrue(spec.isPartitioned());
+    assertEquals(1, spec.fields().size());
+    assertEquals("id_bucket", spec.fields().get(0).name());
+
+    // verify sort order
+    SortOrder sortOrder = table.sortOrder();
+    assertTrue(sortOrder.isSorted());
+    assertEquals(1, sortOrder.fields().size());
+    assertEquals(SortDirection.ASC, sortOrder.fields().get(0).direction());
+    assertEquals(NullOrder.NULLS_FIRST, sortOrder.fields().get(0).nullOrder());
+
+    // verify data was written correctly
+    List<Record> writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build());
+    assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray()));
+  }
 }