[IcebergIO]Create tables with Iceberg Table properties

Author: tarun-google

.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": 3
+    "modification": 5
 }

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

@@ -19,6 +19,7 @@
 
 import com.google.auto.value.AutoValue;
 import java.util.List;
+import java.util.Map;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.iceberg.PartitionSpec;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -40,6 +41,9 @@ public PartitionSpec getPartitionSpec() {
   @Pure
   public abstract @Nullable List<String> getPartitionFields();
 
+  @Pure
+  public abstract @Nullable Map<String, String> getTableProperties();
+
   @Pure
   public static Builder builder() {
     return new AutoValue_IcebergTableCreateConfig.Builder();
@@ -51,6 +55,8 @@ public abstract static class Builder {
 
     public abstract Builder setPartitionFields(@Nullable List<String> partitionFields);
 
+    public abstract Builder setTableProperties(@Nullable Map<String, String> tableProperties);
+
     @Pure
     public abstract IcebergTableCreateConfig build();
   }

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

@@ -124,6 +124,11 @@ public static Builder builder() {
             + "For more information on partition transforms, please visit https://iceberg.apache.org/spec/#partition-transforms.")
     public abstract @Nullable List<String> getPartitionFields();
 
+    @SchemaFieldDescription(
+        "Iceberg table properties to be set on the table when it is created,"
+            + "For more information on table properties, please visit https://iceberg.apache.org/docs/latest/configuration/#table-properties.")
+    public abstract @Nullable Map<String, String> getTableProperties();
+
     @AutoValue.Builder
     public abstract static class Builder {
       public abstract Builder setTable(String table);
@@ -144,6 +149,8 @@ public abstract static class Builder {
 
       public abstract Builder setPartitionFields(List<String> partitionFields);
 
+      public abstract Builder setTableProperties(Map<String, String> tableProperties);
+
       public abstract Configuration build();
     }
 
@@ -209,6 +216,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
                       FileFormat.PARQUET.toString(),
                       rows.getSchema(),
                       configuration.getPartitionFields(),
+                      configuration.getTableProperties(),
                       configuration.getDrop(),
                       configuration.getKeep(),
                       configuration.getOnly()));

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

@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.iceberg;
 
 import java.util.List;
+import java.util.Map;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.util.RowFilter;
 import org.apache.beam.sdk.util.RowStringInterpolator;
@@ -33,17 +34,20 @@ class PortableIcebergDestinations implements DynamicDestinations {
   private final String fileFormat;
 
   private final @Nullable List<String> partitionFields;
+  private final @Nullable Map<String, String> tableProperties;
 
   public PortableIcebergDestinations(
       String destinationTemplate,
       String fileFormat,
       Schema inputSchema,
       @Nullable List<String> partitionFields,
+      @Nullable Map<String, String> tableProperties,
       @Nullable List<String> fieldsToDrop,
       @Nullable List<String> fieldsToKeep,
       @Nullable String onlyField) {
     this.interpolator = new RowStringInterpolator(destinationTemplate, inputSchema);
     this.partitionFields = partitionFields;
+    this.tableProperties = tableProperties;
     RowFilter rf = new RowFilter(inputSchema);
 
     if (fieldsToDrop != null) {
@@ -82,6 +86,7 @@ public IcebergDestination instantiateDestination(String dest) {
             IcebergTableCreateConfig.builder()
                 .setSchema(getDataSchema())
                 .setPartitionFields(partitionFields)
+                .setTableProperties(tableProperties)
                 .build())
         .setFileFormat(FileFormat.fromString(fileFormat))
         .build();

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

@@ -293,6 +293,10 @@ private Table getOrCreateTable(IcebergDestination destination, Schema dataSchema
     @Nullable IcebergTableCreateConfig createConfig = destination.getTableCreateConfig();
     PartitionSpec partitionSpec =
         createConfig != null ? createConfig.getPartitionSpec() : PartitionSpec.unpartitioned();
+    Map<String, String> tableProperties =
+        (createConfig != null && createConfig.getTableProperties() != null)
+            ? createConfig.getTableProperties()
+            : Maps.newHashMap();
 
     synchronized (TABLE_CACHE) {
       // Create namespace if it does not exist yet
@@ -316,7 +320,7 @@ private Table getOrCreateTable(IcebergDestination destination, Schema dataSchema
       } catch (NoSuchTableException e) { // Otherwise, create the table
         org.apache.iceberg.Schema tableSchema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema);
         try {
-          table = catalog.createTable(identifier, tableSchema, partitionSpec);
+          table = catalog.createTable(identifier, tableSchema, partitionSpec, tableProperties);
           LOG.info(
               "Created Iceberg table '{}' with schema: {}\n, partition spec: {}",
               identifier,

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

@@ -543,4 +543,56 @@ public void testWriteCreateTableWithPartitionSpec() {
     Table table = warehouse.loadTable(TableIdentifier.parse(identifier));
     assertEquals(expectedSpec, table.spec());
   }
+
+  @Test
+  public void testWriteCreateTableWithTableProperties() {
+    String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16);
+    Schema schema = Schema.builder().addStringField("str").addInt32Field("int").build();
+
+    // Use real Iceberg table property keys
+    Map<String, String> tableProperties =
+        ImmutableMap.of(
+            "write.format.default", "orc",
+            "commit.retry.num-retries", "5",
+            "read.split.target-size", "134217728");
+
+    Map<String, Object> config =
+        ImmutableMap.of(
+            "table",
+            identifier,
+            "catalog_properties",
+            ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location),
+            "table_properties",
+            tableProperties);
+
+    List<Row> rows = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      Row row = Row.withSchema(schema).addValues("str_" + i, i).build();
+      rows.add(row);
+    }
+
+    PCollection<Row> result =
+        testPipeline
+            .apply("Records To Add", Create.of(rows))
+            .setRowSchema(schema)
+            .apply(Managed.write(Managed.ICEBERG).withConfig(config))
+            .get(SNAPSHOTS_TAG);
+
+    PAssert.that(result)
+        .satisfies(new VerifyOutputs(Collections.singletonList(identifier), "append"));
+    testPipeline.run().waitUntilFinish();
+
+    // Read back and check records are correct
+    Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions());
+    PCollection<Row> readRows =
+        p.apply(Managed.read(Managed.ICEBERG).withConfig(config)).getSinglePCollection();
+    PAssert.that(readRows).containsInAnyOrder(rows);
+    p.run().waitUntilFinish();
+
+    Table table = warehouse.loadTable(TableIdentifier.parse(identifier));
+    // Assert that the table properties are set on the Iceberg table
+    assertEquals("orc", table.properties().get("write.format.default"));
+    assertEquals("5", table.properties().get("commit.retry.num-retries"));
+    assertEquals("134217728", table.properties().get("read.split.target-size"));
+  }
 }

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

@@ -668,6 +668,11 @@ public void testStreamingWrite() throws IOException {
     Map<String, Object> config = new HashMap<>(managedIcebergConfig(tableId()));
     config.put("triggering_frequency_seconds", 4);
     config.put("partition_fields", Arrays.asList("bool_field", "modulo_5"));
+    // Add table properties for testing
+    Map<String, String> tableProperties = new HashMap<>();
+    tableProperties.put("write.format.default", "parquet");
+    tableProperties.put("commit.retry.num-retries", "3");
+    config.put("table_properties", tableProperties);
 
     // create elements from longs in range [0, 1000)
     PCollection<Row> input =
@@ -687,6 +692,8 @@ public void testStreamingWrite() throws IOException {
     List<Record> returnedRecords = readRecords(table);
     assertThat(
         returnedRecords, containsInAnyOrder(inputRows.stream().map(RECORD_FUNC::apply).toArray()));
+    assertEquals("parquet", table.properties().get("write.format.default"));
+    assertEquals("3", table.properties().get("commit.retry.num-retries"));
   }
 
   @Test
@@ -971,4 +978,25 @@ public void runReadBetween(boolean useSnapshotBoundary, boolean streaming) throw
     PAssert.that(rows).containsInAnyOrder(expectedRows);
     pipeline.run().waitUntilFinish();
   }
+
+  @Test
+  public void testWriteWithTableProperties() throws IOException {
+    Map<String, Object> config = new HashMap<>(managedIcebergConfig(tableId()));
+    Map<String, String> tableProperties = new HashMap<>();
+    tableProperties.put("write.format.default", "parquet");
+    tableProperties.put("commit.retry.num-retries", "3");
+    config.put("table_properties", tableProperties);
+    PCollection<Row> input = pipeline.apply(Create.of(inputRows)).setRowSchema(BEAM_SCHEMA);
+    input.apply(Managed.write(ICEBERG).withConfig(config));
+    pipeline.run().waitUntilFinish();
+
+    Table table = catalog.loadTable(TableIdentifier.parse(tableId()));
+    // Read back and check records are correct
+    List<Record> returnedRecords = readRecords(table);
+    assertThat(
+        returnedRecords, containsInAnyOrder(inputRows.stream().map(RECORD_FUNC::apply).toArray()));
+    // Assert that the table properties are set on the Iceberg table
+    assertEquals("parquet", table.properties().get("write.format.default"));
+    assertEquals("3", table.properties().get("commit.retry.num-retries"));
+  }
 }

website/www/site/content/en/documentation/io/managed-io.md

@@ -98,6 +98,8 @@ and Beam SQL is invoked via the Managed API under the hood.
         only (<code style="color: green">str</code>)<br>
         partition_fields (<code>list[<span style="color: green;">str</span>]</code>)<br>
         triggering_frequency_seconds (<code style="color: #f54251">int32</code>)<br>
+        table_properties (<code>map[<span style="color: green;">str</span>, <span style="color: green;
+        ">str</span>]</code>)<br>
       </td>
     </tr>
     <tr>
@@ -431,6 +433,19 @@ For more information on partition transforms, please visit https://iceberg.apach
         For a streaming pipeline, sets the frequency at which snapshots are produced.
       </td>
     </tr>
+    <tr>
+      <td>
+        table_properties
+      </td>
+      <td>
+        <code>map[<span style="color: green;">str</span>, <span style="color: green;">str</span>]</code>
+      </td>
+      <td>
+        Table Properties set while creating Iceberg Table.
+
+        For more information on table properties, please visit https://iceberg.apache.org/docs/latest/configuration/#table-properties
+      </td>
+    </tr>
   </table>
 </div>