Managed Iceberg AddFiles

CHANGES.md

@@ -68,6 +68,7 @@
 * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
 * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)).
 * (Java/Python/Go) DebeziumIO: `readAsJson()` can now include `op` and the remaining Debezium `source` fields (e.g. MySQL `file`, `pos`, `row`, `ts_ms`; PostgreSQL `lsn`, `txId`, `ts_ms`) in `metadata` via `withFullMetadata(true)` / `full_metadata=True` / `debeziumio.FullMetadata()`. Off by default ([#40701](https://github.com/apache/beam/issues/40701)).
+* [IcebergIO] AddFiles is now available through the Managed API with `Managed.write(Managed.ICEBERG_ADD_FILES)` (Java/Python), with URN `beam:schematransform:org.apache.beam:iceberg_add_files:v1` ([#40827](https://github.com/apache/beam/issues/40827)).
 
 ## New Features / Improvements
 

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto

@@ -109,6 +109,8 @@ message ManagedTransforms {
       "beam:schematransform:org.apache.beam:delta_lake_read:v1"];
     DELTA_LAKE_CDC_READ = 14 [(org.apache.beam.model.pipeline.v1.beam_urn) =
       "beam:schematransform:org.apache.beam:delta_lake_cdc_read:v1"];
+    ICEBERG_ADD_FILES = 15 [(org.apache.beam.model.pipeline.v1.beam_urn) =
+      "beam:schematransform:org.apache.beam:iceberg_add_files:v1"];
   }
 }
 

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

@@ -21,6 +21,7 @@
 import static org.apache.beam.sdk.io.iceberg.AddFiles.OUTPUT_TAG;
 import static org.apache.beam.sdk.io.iceberg.AddFilesSchemaTransformProvider.Configuration;
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
 
 import com.google.auto.service.AutoService;
 import com.google.auto.value.AutoValue;
@@ -30,8 +31,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.apache.beam.model.pipeline.v1.ExternalTransforms;
 import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
 import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
@@ -41,6 +45,7 @@
 import org.apache.beam.sdk.transforms.Filter;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -55,7 +60,7 @@ public AddFilesSchemaTransform from(Configuration configuration) {
 
   @Override
   public String identifier() {
-    return "beam:schematransform:iceberg_add_files:v1";
+    return getUrn(ExternalTransforms.ManagedTransforms.Urns.ICEBERG_ADD_FILES);
   }
 
   @DefaultSchema(AutoValueSchema.class)
@@ -302,6 +307,20 @@ public AddFilesSchemaTransform(Configuration configuration) {
       this.configuration = configuration;
     }
 
+    Row getConfigurationRow() {
+      try {
+        // To stay consistent with our SchemaTransform configuration naming conventions,
+        // we sort lexicographically and convert field names to snake_case
+        return SchemaRegistry.createDefault()
+            .getToRowFunction(Configuration.class)
+            .apply(configuration)
+            .sorted()
+            .toSnakeCase();
+      } catch (NoSuchSchemaException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
     @Override
     public PCollectionRowTuple expand(PCollectionRowTuple input) {
       Schema inputSchema = input.getSinglePCollection().getSchema();

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

@@ -23,6 +23,7 @@
 
 import com.google.auto.service.AutoService;
 import java.util.Map;
+import org.apache.beam.sdk.io.iceberg.AddFilesSchemaTransformProvider.AddFilesSchemaTransform;
 import org.apache.beam.sdk.io.iceberg.IcebergCdcReadSchemaTransformProvider.IcebergCdcReadSchemaTransform;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
 import org.apache.beam.sdk.transforms.PTransform;
@@ -87,6 +88,19 @@ public Row toConfigRow(IcebergWriteSchemaTransform transform) {
     }
   }
 
+  static class AddFilesSchemaTransformTranslator
+      extends SchemaTransformPayloadTranslator<AddFilesSchemaTransform> {
+    @Override
+    public SchemaTransformProvider provider() {
+      return new AddFilesSchemaTransformProvider();
+    }
+
+    @Override
+    public Row toConfigRow(AddFilesSchemaTransform transform) {
+      return transform.getConfigurationRow();
+    }
+  }
+
   @AutoService(TransformPayloadTranslatorRegistrar.class)
   public static class WriteRegistrar implements TransformPayloadTranslatorRegistrar {
     @Override
@@ -97,6 +111,7 @@ public static class WriteRegistrar implements TransformPayloadTranslatorRegistra
         getTransformPayloadTranslators() {
       return ImmutableMap.<Class<? extends PTransform>, TransformPayloadTranslator>builder()
           .put(IcebergWriteSchemaTransform.class, new IcebergWriteSchemaTransformTranslator())
+          .put(AddFilesSchemaTransform.class, new AddFilesSchemaTransformTranslator())
           .build();
     }
   }

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

@@ -17,24 +17,50 @@
  */
 package org.apache.beam.sdk.io.iceberg;
 
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.List;
 import org.apache.beam.sdk.io.iceberg.AddFilesSchemaTransformProvider.Configuration;
+import org.apache.beam.sdk.managed.Managed;
+import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.data.IcebergGenerics;
+import org.apache.iceberg.data.Record;
+import org.hamcrest.Matchers;
+import org.junit.ClassRule;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public class AddFilesSchemaTransformProviderTest {
+  @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+
+  @Rule
+  public transient TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default");
+
+  @Rule public transient TestPipeline pipeline = TestPipeline.create();
 
   private static Configuration.Builder base() {
     return Configuration.builder()
@@ -182,4 +208,33 @@ public void testSettingsWithoutOptionsRejectedNamingTheOptionsKey() {
     assertTrue(
         unverifiable.getMessage(), unverifiable.getMessage().contains("schema_evolution_options"));
   }
+
+  @Test
+  public void testAddFilesUsingManagedTransform() throws IOException {
+    TableIdentifier tableId = TableIdentifier.of("default", "managed_add_files");
+    DataFile file =
+        warehouse.writeRecords("data.parquet", TestFixtures.SCHEMA, TestFixtures.FILE1SNAPSHOT1);
+
+    Schema inputSchema = Schema.builder().addStringField("file").build();
+    PCollectionRowTuple output =
+        pipeline
+            .apply(Create.of(Row.withSchema(inputSchema).addValue(file.location()).build()))
+            .setRowSchema(inputSchema)
+            .apply(
+                Managed.write(Managed.ICEBERG_ADD_FILES)
+                    .withConfig(
+                        ImmutableMap.<String, Object>builder()
+                            .put("table", tableId.toString())
+                            .put(
+                                "catalog_properties",
+                                ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location))
+                            .put("error_handling", ImmutableMap.of("output", "errors"))
+                            .build()));
+    PAssert.that(output.get("errors")).empty();
+    pipeline.run().waitUntilFinish();
+
+    Table table = warehouse.loadTable(tableId);
+    List<Record> records = ImmutableList.copyOf(IcebergGenerics.read(table).build());
+    assertThat(records, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray()));
+  }
 }

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

@@ -34,8 +34,11 @@
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.iceberg.AddFilesSchemaTransformProvider.AddFilesSchemaTransform;
 import org.apache.beam.sdk.io.iceberg.IcebergCdcReadSchemaTransformProvider.IcebergCdcReadSchemaTransform;
+import org.apache.beam.sdk.io.iceberg.IcebergSchemaTransformTranslation.AddFilesSchemaTransformTranslator;
 import org.apache.beam.sdk.io.iceberg.IcebergSchemaTransformTranslation.IcebergCdcReadSchemaTransformTranslator;
+import org.apache.beam.sdk.managed.Managed;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.SchemaTranslation;
@@ -74,6 +77,8 @@ public class IcebergSchemaTransformTranslationTest {
       new IcebergReadSchemaTransformProvider();
   static final IcebergCdcReadSchemaTransformProvider READ_CDC_PROVIDER =
       new IcebergCdcReadSchemaTransformProvider();
+  static final AddFilesSchemaTransformProvider ADD_FILES_PROVIDER =
+      new AddFilesSchemaTransformProvider();
 
   private static final Map<String, String> CATALOG_PROPERTIES =
       ImmutableMap.<String, String>builder()
@@ -133,6 +138,15 @@ public class IcebergSchemaTransformTranslationTest {
           .withFieldValue("watermark_column", "event_micros")
           .build();
 
+  private static final Row ADD_FILES_CONFIG_ROW =
+      Row.withSchema(ADD_FILES_PROVIDER.configurationSchema())
+          .withFieldValue("table", "test_table_identifier")
+          .withFieldValue("catalog_properties", CATALOG_PROPERTIES)
+          .withFieldValue("partition_fields", Collections.singletonList("id"))
+          .withFieldValue("manifest_file_size", 100)
+          .withFieldValue("error_handling.output", "errors")
+          .build();
+
   @Test
   public void testReCreateWriteTransformFromRow() {
     IcebergWriteSchemaTransform writeTransform =
@@ -374,6 +388,60 @@ public void testCdcReadTransformProtoTranslation()
     assertEquals(transformConfigRow, readTransformFromSpec.getConfigurationRow());
   }
 
+  @Test
+  public void testAddFilesTransformProtoTranslation()
+      throws InvalidProtocolBufferException, IOException {
+    // First build a pipeline that adds files through the Managed API
+    Pipeline p = Pipeline.create();
+    Schema inputSchema = Schema.builder().addStringField("file").build();
+    p.apply(Create.of(Row.withSchema(inputSchema).addValue("data.parquet").build()))
+        .setRowSchema(inputSchema)
+        .apply(
+            Managed.write(Managed.ICEBERG_ADD_FILES)
+                .withConfig(
+                    ImmutableMap.<String, Object>builder()
+                        .put("table", "test_table_identifier")
+                        .put("catalog_properties", CATALOG_PROPERTIES)
+                        .put("partition_fields", Collections.singletonList("id"))
+                        .put("manifest_file_size", 100)
+                        .put("error_handling", ImmutableMap.of("output", "errors"))
+                        .build()));
+
+    // Then translate the pipeline to a proto and extract AddFilesSchemaTransform proto
+    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
+    List<RunnerApi.PTransform> addFilesTransformProto =
+        pipelineProto.getComponents().getTransformsMap().values().stream()
+            .filter(
+                tr -> {
+                  RunnerApi.FunctionSpec spec = tr.getSpec();
+                  try {
+                    return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
+                        && SchemaTransformPayload.parseFrom(spec.getPayload())
+                            .getIdentifier()
+                            .equals(ADD_FILES_PROVIDER.identifier());
+                  } catch (InvalidProtocolBufferException e) {
+                    throw new RuntimeException(e);
+                  }
+                })
+            .collect(Collectors.toList());
+    assertEquals(1, addFilesTransformProto.size());
+    RunnerApi.FunctionSpec spec = addFilesTransformProto.get(0).getSpec();
+
+    // Check that the proto contains correct values
+    SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
+    Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
+    assertEquals(ADD_FILES_PROVIDER.configurationSchema(), schemaFromSpec);
+    Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
+    assertEquals(ADD_FILES_CONFIG_ROW, rowFromSpec);
+
+    // Use the information in the proto to recreate the AddFilesSchemaTransform
+    AddFilesSchemaTransformTranslator translator = new AddFilesSchemaTransformTranslator();
+    AddFilesSchemaTransform transformFromSpec =
+        translator.fromConfigRow(rowFromSpec, PipelineOptionsFactory.create());
+
+    assertEquals(ADD_FILES_CONFIG_ROW, transformFromSpec.getConfigurationRow());
+  }
+
   private static Record record(long id, String data, long eventMicros) {
     return TestFixtures.createRecord(
         CDC_TRANSLATION_SCHEMA,

sdks/java/managed/src/main/java/org/apache/beam/sdk/managed/Managed.java

@@ -101,6 +101,7 @@ public class Managed {
   public static final String POSTGRES = "postgres";
   public static final String MYSQL = "mysql";
   public static final String SQL_SERVER = "sqlserver";
+  public static final String ICEBERG_ADD_FILES = "iceberg_add_files";
 
   // Supported SchemaTransforms
   public static final Map<String, String> READ_TRANSFORMS =
@@ -124,6 +125,9 @@ DELTA_LAKE_CDC, getUrn(ExternalTransforms.ManagedTransforms.Urns.DELTA_LAKE_CDC_
           .put(POSTGRES, getUrn(ExternalTransforms.ManagedTransforms.Urns.POSTGRES_WRITE))
           .put(MYSQL, getUrn(ExternalTransforms.ManagedTransforms.Urns.MYSQL_WRITE))
           .put(SQL_SERVER, getUrn(ExternalTransforms.ManagedTransforms.Urns.SQL_SERVER_WRITE))
+          .put(
+              ICEBERG_ADD_FILES,
+              getUrn(ExternalTransforms.ManagedTransforms.Urns.ICEBERG_ADD_FILES))
           .build();
 
   /**
@@ -168,6 +172,9 @@ public static ManagedTransform read(String source) {
    *       href="https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/kafka/KafkaIO.html">KafkaIO</a>
    *   <li>{@link Managed#BIGQUERY} : Write to GCP BigQuery tables <a
    *       href="https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.html">BigQueryIO</a>
+   *   <li>{@link Managed#ICEBERG_ADD_FILES} : Register existing data files in Apache Iceberg tables
+   *       using <a
+   *       href="https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/iceberg/AddFiles.html">AddFiles</a>
    * </ul>
    */
   public static ManagedTransform write(String sink) {

sdks/python/apache_beam/transforms/external.py

@@ -88,6 +88,7 @@
     ManagedTransforms.Urns.SQL_SERVER_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET,  # pylint: disable=line-too-long
     ManagedTransforms.Urns.DELTA_LAKE_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,
     ManagedTransforms.Urns.DELTA_LAKE_CDC_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,
+    ManagedTransforms.Urns.ICEBERG_ADD_FILES.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,  # pylint: disable=line-too-long
 }
 
 

sdks/python/apache_beam/transforms/managed.py

@@ -90,13 +90,15 @@
 SQL_SERVER = "sqlserver"
 DELTA = "delta"
 DELTA_CDC = "delta_cdc"
+ICEBERG_ADD_FILES = "iceberg_add_files"
 
 __all__ = [
     "ICEBERG",
     "KAFKA",
     "BIGQUERY",
     "DELTA",
     "DELTA_CDC",
+    "ICEBERG_ADD_FILES",
     "Read",
     "Write",
 ]
@@ -165,7 +167,8 @@ class Write(PTransform):
       BIGQUERY: ManagedTransforms.Urns.BIGQUERY_WRITE.urn,
       POSTGRES: ManagedTransforms.Urns.POSTGRES_WRITE.urn,
       MYSQL: ManagedTransforms.Urns.MYSQL_WRITE.urn,
-      SQL_SERVER: ManagedTransforms.Urns.SQL_SERVER_WRITE.urn
+      SQL_SERVER: ManagedTransforms.Urns.SQL_SERVER_WRITE.urn,
+      ICEBERG_ADD_FILES: ManagedTransforms.Urns.ICEBERG_ADD_FILES.urn,
   }
 
   def __init__(

sdks/python/apache_beam/yaml/standard_io.yaml

@@ -710,7 +710,7 @@
     underlying_provider:
       type: beamJar
       transforms:
-        'IcebergAddFiles': 'beam:schematransform:iceberg_add_files:v1'
+        'IcebergAddFiles': 'beam:schematransform:org.apache.beam:iceberg_add_files:v1'
       config:
         gradle_target: 'sdks:java:io:expansion-service:shadowJar'
 

sdks/standard_expansion_services.yaml

@@ -54,6 +54,7 @@
     - 'beam:schematransform:org.apache.beam:iceberg_cdc_read:v1'
     - 'beam:schematransform:org.apache.beam:delta_lake_read:v1'
     - 'beam:schematransform:org.apache.beam:delta_lake_cdc_read:v1'
+    - 'beam:schematransform:org.apache.beam:iceberg_add_files:v1'
 
 - gradle_target: 'sdks:java:io:messaging-expansion-service:shadowJar'
   destinations: