[IcebergIO]Create tables with Iceberg Table properties

Author: tarun-google

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

@@ -84,12 +84,19 @@ public org.apache.iceberg.catalog.Catalog catalog() {
   }
 
   public void createTable(
-      String tableIdentifier, Schema tableSchema, @Nullable List<String> partitionFields) {
+      String tableIdentifier,
+      Schema tableSchema,
+      @Nullable List<String> partitionFields,
+      @Nullable Map<String, String> tableProperties) {
     TableIdentifier icebergIdentifier = TableIdentifier.parse(tableIdentifier);
     org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(tableSchema);
     PartitionSpec icebergSpec = PartitionUtils.toPartitionSpec(partitionFields, tableSchema);
     try {
-      catalog().createTable(icebergIdentifier, icebergSchema, icebergSpec);
+      if (tableProperties == null) {
+        catalog().createTable(icebergIdentifier, icebergSchema, icebergSpec);
+      } else {
+        catalog().createTable(icebergIdentifier, icebergSchema, icebergSpec, tableProperties);
+      }
       LOG.info(
           "Created table '{}' with schema: {}\n, partition spec: {}",
           icebergIdentifier,

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

@@ -293,11 +293,9 @@ private Table getOrCreateTable(IcebergDestination destination, Schema dataSchema
     @Nullable IcebergTableCreateConfig createConfig = destination.getTableCreateConfig();
     PartitionSpec partitionSpec =
         createConfig != null ? createConfig.getPartitionSpec() : PartitionSpec.unpartitioned();
-    // Ensure table properties are not null, Iceberg's createTable does not accept null values.
+    @Nullable
     Map<String, String> tableProperties =
-        (createConfig != null && createConfig.getTableProperties() != null)
-            ? createConfig.getTableProperties()
-            : Maps.newHashMap();
+        createConfig != null ? createConfig.getTableProperties() : null;
 
     synchronized (TABLE_CACHE) {
       // Create namespace if it does not exist yet
@@ -321,7 +319,11 @@ 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, tableProperties);
+          if (tableProperties != null) {
+            table = catalog.createTable(identifier, tableSchema, partitionSpec, tableProperties);
+          } else {
+            table = catalog.createTable(identifier, tableSchema, partitionSpec);
+          }
           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

@@ -545,7 +545,7 @@ public void testWriteCreateTableWithPartitionSpec() {
   }
 
   @Test
-  public void testWriteCreateTableWithTableProperties() {
+  public void testWriteCreateTableWithTablePropertiesSpec() {
     String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16);
     Schema schema = Schema.builder().addStringField("str").addInt32Field("int").build();
 
@@ -595,4 +595,57 @@ public void testWriteCreateTableWithTableProperties() {
     assertEquals("5", table.properties().get("commit.retry.num-retries"));
     assertEquals("134217728", table.properties().get("read.split.target-size"));
   }
+
+  @Test
+  public void testWriteCreateTableWithTableProperties() {
+    String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16);
+    Schema schema = Schema.builder().addStringField("str").addInt32Field("int").build();
+    org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema);
+    PartitionSpec spec = PartitionSpec.unpartitioned();
+    Map<String, String> tableProperties =
+        ImmutableMap.of(
+            "write.format.default", "orc",
+            "commit.retry.num-retries", "5",
+            "read.split.target-size", "134217728");
+
+    // Create the table with properties
+    warehouse.createTable(TableIdentifier.parse(identifier), icebergSchema, spec, 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);
+    }
+
+    Map<String, Object> config =
+        ImmutableMap.of(
+            "table",
+            identifier,
+            "catalog_properties",
+            ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location));
+
+    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/TestDataWarehouse.java

@@ -163,13 +163,22 @@ public DataFile writeRecords(
   }
 
   public Table createTable(TableIdentifier tableId, Schema schema) {
-    return createTable(tableId, schema, null);
+    return createTable(tableId, schema, null, null);
   }
 
   public Table createTable(
       TableIdentifier tableId, Schema schema, @Nullable PartitionSpec partitionSpec) {
     someTableHasBeenCreated = true;
-    return catalog.createTable(tableId, schema, partitionSpec);
+    return catalog.createTable(tableId, schema, partitionSpec, null);
+  }
+
+  public Table createTable(
+      TableIdentifier tableId,
+      Schema schema,
+      @Nullable PartitionSpec partitionSpec,
+      @Nullable Map<String, String> tableProperties) {
+    someTableHasBeenCreated = true;
+    return catalog.createTable(tableId, schema, partitionSpec, tableProperties);
   }
 
   public Catalog.TableBuilder buildTable(TableIdentifier tableId, Schema schema) {