AddFiles: Wire a destination-table key through the registration chain

.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

@@ -99,6 +99,7 @@
 * (Python) `Regex.all_matches` now returns every group of the regex, instead of failing with a `TypeError` when the regex has no groups and dropping nested or unmatched trailing groups ([#40562](https://github.com/apache/beam/issues/40562)).
 * (Python) `--dataflow_endpoint` now defaults to `''` instead of `https://dataflow.googleapis.com`. When unset, the Google Cloud client library resolves the Dataflow endpoint and honors `GOOGLE_CLOUD_UNIVERSE_DOMAIN`, so Dataflow jobs can be submitted in non-default universes. Code that reads `GoogleCloudOptions.dataflow_endpoint` should handle an empty value ([#40569](https://github.com/apache/beam/issues/40569)).
 * (Python) The `xgboost` extras now install `xgboost>=2.0.0,<4` and `datatable>=1.1.0,<2`, moving both the lower and upper bounds for both. The previous minimum `xgboost` and `datatable` versions (`1.6.0` and `1.0.0` respectively) are no longer allowed ([#40694](https://github.com/apache/beam/issues/40694)).
+* (Java) [IcebergIO] AddFiles (`IcebergAddFiles` in YAML) now carries the destination table alongside each file through the registration chain, so that one transform can register files into many tables. The coders of its internal `PCollection`s change, which makes this an update incompatible change for streaming pipelines that use it ([#40817](https://github.com/apache/beam/pull/40817)).
 
 ## Deprecations
 

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

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.iceberg;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber;
+
+/**
+ * As much as one manifest file can hold: entries for a single table, all written against a single
+ * partition spec. Data files are batched by this before a manifest is written for each batch.
+ *
+ * <p>A schema type rather than a {@code KV} of the two. The field numbers pin each field's place in
+ * the encoding, so a third field can be added after these two without moving them.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+abstract class ManifestKey {
+
+  static ManifestKey of(String table, int partitionSpecId) {
+    return new AutoValue_ManifestKey(table, partitionSpecId);
+  }
+
+  @SchemaFieldNumber("0")
+  abstract String getTable();
+
+  @SchemaFieldNumber("1")
+  abstract int getPartitionSpecId();
+}

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

@@ -28,6 +28,7 @@
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.values.KV;
 import org.apache.iceberg.FileFormat;
 import org.apache.parquet.hadoop.metadata.ParquetMetadata;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -40,7 +41,8 @@
  * Emits one {@link CollectDistinctSchemas.SchemaGroup} of one file per readable Parquet file: its
  * declared schema and its null-free columns. Unreadable or non-Parquet files contribute nothing.
  */
-class ReadFooterSchema extends DoFn<String, CollectDistinctSchemas.SchemaGroup> {
+class ReadFooterSchema
+    extends DoFn<KV<String, String>, KV<String, CollectDistinctSchemas.SchemaGroup>> {
   private static final Logger LOG = LoggerFactory.getLogger(ReadFooterSchema.class);
 
   static final int DEFAULT_THREAD_POOL_SIZE = 10;
@@ -94,18 +96,21 @@ class ReadFooterSchema extends DoFn<String, CollectDistinctSchemas.SchemaGroup>
 
   /** Counters are updated on the processing thread: metrics touched from the executor are lost. */
   private static class ReadResult {
+    final String table;
     final CollectDistinctSchemas.@Nullable SchemaGroup schema;
     final boolean footerError;
     final Instant timestamp;
     final BoundedWindow window;
     final PaneInfo paneInfo;
 
     ReadResult(
+        String table,
         CollectDistinctSchemas.@Nullable SchemaGroup schema,
         boolean footerError,
         Instant timestamp,
         BoundedWindow window,
         PaneInfo paneInfo) {
+      this.table = table;
       this.schema = schema;
       this.footerError = footerError;
       this.timestamp = timestamp;
@@ -134,14 +139,15 @@ public void teardown() {
 
   @ProcessElement
   public void process(
-      @Element String filePath,
+      @Element KV<String, String> element,
       @Timestamp Instant timestamp,
       BoundedWindow window,
       PaneInfo paneInfo,
-      OutputReceiver<CollectDistinctSchemas.SchemaGroup> output)
+      OutputReceiver<KV<String, CollectDistinctSchemas.SchemaGroup>> output)
       throws Exception {
     numFilesRead.inc();
-    Callable<ReadResult> task = createReadTask(filePath, timestamp, window, paneInfo);
+    Callable<ReadResult> task =
+        createReadTask(element.getKey(), element.getValue(), timestamp, window, paneInfo);
     checkStateNotNull(tasks).submit(task, result -> outputResult(result, output));
   }
 
@@ -153,16 +159,19 @@ public void finishBundle(FinishBundleContext context) throws Exception {
   private static void outputAtFinish(ReadResult result, FinishBundleContext context) {
     count(result);
     if (result.schema != null) {
-      context.output(result.schema, result.timestamp, result.window);
+      context.output(KV.of(result.table, result.schema), result.timestamp, result.window);
     }
   }
 
   private static void outputResult(
-      ReadResult result, OutputReceiver<CollectDistinctSchemas.SchemaGroup> output) {
+      ReadResult result, OutputReceiver<KV<String, CollectDistinctSchemas.SchemaGroup>> output) {
     count(result);
     if (result.schema != null) {
       output.outputWindowedValue(
-          result.schema, result.timestamp, Collections.singleton(result.window), result.paneInfo);
+          KV.of(result.table, result.schema),
+          result.timestamp,
+          Collections.singleton(result.window),
+          result.paneInfo);
     }
   }
 
@@ -176,64 +185,75 @@ private static void count(ReadResult result) {
   }
 
   private Callable<ReadResult> createReadTask(
-      String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
+      String table, String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
     return () -> {
       ParquetMetadata footer;
       if (AddFiles.hasNoExtension(filePath)) {
         try {
           footer = ParquetFooters.read(filePath);
         } catch (Exception e) {
-          return unknownFormat(timestamp, window, paneInfo);
+          return unknownFormat(table, timestamp, window, paneInfo);
         }
       } else {
         FileFormat format;
         try {
           format = AddFiles.inferFormat(filePath);
         } catch (AddFiles.UnknownFormatException e) {
-          return unknownFormat(timestamp, window, paneInfo);
+          return unknownFormat(table, timestamp, window, paneInfo);
         }
         if (!format.equals(FileFormat.PARQUET)) {
           return new ReadResult(
-              dryRunMarker(UNCHECKED_FORMAT_KEY), false, timestamp, window, paneInfo);
+              table, dryRunMarker(UNCHECKED_FORMAT_KEY), false, timestamp, window, paneInfo);
         }
         try {
           footer = ParquetFooters.read(filePath);
         } catch (Exception e) {
-          return unreadableFooter(filePath, e, timestamp, window, paneInfo);
+          return unreadableFooter(table, filePath, e, timestamp, window, paneInfo);
         }
       }
       try {
         List<String> partitionKeysInPath =
             new ArrayList<>(PartitionPaths.partitionKeysInPath(filePath, locationPrefix));
         CollectDistinctSchemas.SchemaGroup group =
             FileSchemas.schemaGroup(footer).withPartitionKeysInPath(partitionKeysInPath);
-        return new ReadResult(group, false, timestamp, window, paneInfo);
+        return new ReadResult(table, group, false, timestamp, window, paneInfo);
       } catch (Exception e) {
-        return unconvertibleSchema(filePath, e, timestamp, window, paneInfo);
+        return unconvertibleSchema(table, filePath, e, timestamp, window, paneInfo);
       }
     };
   }
 
-  private ReadResult unknownFormat(Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
-    return new ReadResult(dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo);
+  private ReadResult unknownFormat(
+      String table, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
+    return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo);
   }
 
   private ReadResult unreadableFooter(
-      String filePath, Exception e, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
+      String table,
+      String filePath,
+      Exception e,
+      Instant timestamp,
+      BoundedWindow window,
+      PaneInfo paneInfo) {
     LOG.warn(
         "Could not read the footer of {}; the file will not contribute to schema inference: {}",
         filePath,
         AddFiles.errorMessage(e));
-    return new ReadResult(dryRunMarker(UNREADABLE_KEY), true, timestamp, window, paneInfo);
+    return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), true, timestamp, window, paneInfo);
   }
 
   private ReadResult unconvertibleSchema(
-      String filePath, Exception e, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
+      String table,
+      String filePath,
+      Exception e,
+      Instant timestamp,
+      BoundedWindow window,
+      PaneInfo paneInfo) {
     LOG.warn(
         "Could not convert the schema of {}; the file will not contribute to schema inference: {}",
         filePath,
         AddFiles.errorMessage(e));
-    return new ReadResult(dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo);
+    return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo);
   }
 
   /** A dry run counts files that contribute no schema; a real run handles them at registration. */

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

@@ -52,19 +52,23 @@
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.iceberg.SchemaEvolutionConfig.UnverifiableFileHandling;
 import org.apache.beam.sdk.metrics.MetricNameFilter;
 import org.apache.beam.sdk.metrics.MetricResult;
 import org.apache.beam.sdk.metrics.MetricsFilter;
 import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
 import org.apache.beam.sdk.testing.ExpectedLogs;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.testing.TestStream;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.PCollection;
@@ -73,6 +77,7 @@
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TimestampedValue;
 import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
@@ -1857,14 +1862,27 @@ private PCollectionTuple convert(SchemaEvolutionConfig config, String... files)
     PCollectionTuple out =
         pipeline
             .apply("Create Input", Create.of(Arrays.asList(files)))
+            .apply(
+                "KeyByTable",
+                WithKeys.<String, String>of(tableId.toString())
+                    .withKeyType(TypeDescriptors.strings()))
             .apply(
                 ParDo.of(
                         new AddFiles.ConvertToDataFile(
-                            catalogConfig, tableId.toString(), null, null, null, null, config))
+                            catalogConfig, null, null, null, null, config))
                     .withOutputTags(
                         AddFiles.ConvertToDataFile.DATA_FILES,
                         TupleTagList.of(AddFiles.ConvertToDataFile.ERRORS)));
     out.get(AddFiles.ConvertToDataFile.ERRORS).setRowSchema(AddFiles.ERROR_SCHEMA);
+    try {
+      out.get(AddFiles.ConvertToDataFile.DATA_FILES)
+          .setCoder(
+              KvCoder.of(
+                  StringUtf8Coder.of(),
+                  SchemaRegistry.createDefault().getSchemaCoder(SerializableDataFile.class)));
+    } catch (NoSuchSchemaException e) {
+      throw new RuntimeException(e);
+    }
     return out;
   }
 

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

@@ -32,6 +32,8 @@
 import java.util.List;
 import org.apache.avro.SchemaBuilder;
 import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.metrics.MetricNameFilter;
 import org.apache.beam.sdk.metrics.MetricQueryResults;
 import org.apache.beam.sdk.metrics.MetricResult;
@@ -40,7 +42,10 @@
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Values;
+import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.hadoop.fs.Path;
 import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Schema;
@@ -341,17 +346,28 @@ private static void writeLegacyRepeated(File file) throws IOException {
   }
 
   private PCollection<CollectDistinctSchemas.SchemaGroup> run(String... paths) {
+    // the read is keyed by destination table; these tests are about the schema, so one key
     return pipeline
         .apply(Create.of(Arrays.asList(paths)))
+        .apply(
+            "KeyByTable",
+            WithKeys.<String, String>of("db.t").withKeyType(TypeDescriptors.strings()))
         .apply(ParDo.of(new ReadFooterSchema(SchemaEvolutionConfig.disabled())))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), CollectDistinctSchemas.groupCoder()))
+        .apply("DropKey", Values.create())
         .setCoder(CollectDistinctSchemas.groupCoder());
   }
 
   private PCollection<CollectDistinctSchemas.SchemaGroup> runWithPrefix(
       String prefix, String... paths) {
     return pipeline
         .apply(Create.of(Arrays.asList(paths)))
+        .apply(
+            "KeyByTable",
+            WithKeys.<String, String>of("db.t").withKeyType(TypeDescriptors.strings()))
         .apply(ParDo.of(new ReadFooterSchema(SchemaEvolutionConfig.disabled(), prefix)))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), CollectDistinctSchemas.groupCoder()))
+        .apply("DropKey", Values.create())
         .setCoder(CollectDistinctSchemas.groupCoder());
   }