CP to 2.77: Addfiles wiring provider (#40150)

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

CHANGES.md

@@ -69,6 +69,7 @@
 * BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed (Java) ([#39597](https://github.com/apache/beam/issues/39597)) .
 * SolaceIO now supports reading and writing binary and text content data payload (Java) ([#39875](https://github.com/apache/beam/issues/39875)).
 * ClickHouseIO: support writing `Decimal(P, S)` / `Decimal32/64/128/256` columns (Java) ([#39840](https://github.com/apache/beam/issues/39840)).
+* [IcebergIO] AddFiles (`IcebergAddFiles` in YAML) can evolve the table schema before registering files, with `schema_evolution_options`, `required_columns`, `incompatible_schema_handling` and `unverifiable_file_handling` (Java/YAML, batch only) ([#40144](https://github.com/apache/beam/issues/40144)).
 
 ## New Features / Improvements
 

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

@@ -38,11 +38,13 @@
 import java.util.UUID;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.io.iceberg.SchemaEvolutionConfig.IncompatibleSchemaHandling;
 import org.apache.beam.sdk.io.iceberg.SchemaEvolutionConfig.UnverifiableFileHandling;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.schemas.Schema;
@@ -51,13 +53,17 @@
 import org.apache.beam.sdk.state.StateSpec;
 import org.apache.beam.sdk.state.StateSpecs;
 import org.apache.beam.sdk.state.ValueState;
+import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupIntoBatches;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Wait;
 import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
@@ -112,8 +118,38 @@
 import org.slf4j.LoggerFactory;
 
 /**
- * A transform that takes in a stream of file paths, converts them to Iceberg {@link DataFile}s with
- * partition metadata and metrics, then commits them to an Iceberg {@link Table}.
+ * Registers existing Parquet, ORC or Avro files in an Iceberg table without rewriting them: each
+ * path becomes a {@link DataFile} with partition metadata and column stats, batched into manifests
+ * and committed as snapshots.
+ *
+ * <p>Outputs: {@code snapshots} (one row per commit), {@code errors} (one row per file that could
+ * not be registered: {@code file}, {@code error}).
+ *
+ * <p><b>Schema evolution.</b> With a {@link SchemaEvolutionConfig} whose options are set, a
+ * pre-pass reads every Parquet footer, classifies the change each distinct file schema needs on the
+ * table (add a column, relax a required column, promote a type), commits the allowed changes in one
+ * transaction, and only then registers the files. Manifest entries are immutable, so this ordering
+ * is what guarantees that every registered file carries stats for every column it has. Files whose
+ * schema needs a change that is not allowed, or that conflicts with the table or with another file,
+ * are incompatible: by default the pipeline fails before committing anything, or routes them to
+ * {@code errors} (see {@link SchemaEvolutionConfig.IncompatibleSchemaHandling}). The per-file
+ * checks read Parquet footers: an ORC or Avro file, or a pinned column the footer has no null count
+ * for, cannot be verified and goes to {@code errors} unless {@link
+ * SchemaEvolutionConfig.UnverifiableFileHandling#ACCEPT} registers it on trust. Schema evolution
+ * currently requires bounded input; unbounded input with options set is rejected at construction.
+ *
+ * <pre>{@code
+ * SchemaEvolutionConfig evolution =
+ *     SchemaEvolutionConfig.builder()
+ *         .setOptions(EnumSet.of(ALLOW_FIELD_ADDITION, ALLOW_TYPE_PROMOTION))
+ *         .setRequiredColumns(Collections.singleton("id"))
+ *         .build();
+ * paths.apply(new AddFiles(catalog, "db.sales", null, null, null, null, null, null, evolution));
+ * }</pre>
+ *
+ * <p>Without options the table schema is never changed and files register as-is: columns the table
+ * does not have get no stats and are not readable, and a nested column the table does not know can
+ * make that file, and any scan that includes it, fail in an Iceberg reader.
  */
 public class AddFiles extends PTransform<PCollection<String>, PCollectionRowTuple> {
   static final String OUTPUT_TAG = "snapshots";
@@ -142,6 +178,8 @@ public class AddFiles extends PTransform<PCollection<String>, PCollectionRowTupl
   private final @Nullable List<String> partitionFields;
   private final @Nullable List<String> sortFields;
   private final @Nullable Map<String, String> tableProps;
+  private final SchemaEvolutionConfig evolution;
+  private CommitSchemaUnion.Committer committer = CommitSchemaUnion.DEFAULT_COMMITTER;
 
   public AddFiles(
       IcebergCatalogConfig catalogConfig,
@@ -152,6 +190,40 @@ public AddFiles(
       @Nullable Map<String, String> tableProps,
       @Nullable Integer manifestFileSize,
       @Nullable Duration intervalTrigger) {
+    this(
+        catalogConfig,
+        tableIdentifier,
+        locationPrefix,
+        partitionFields,
+        sortFields,
+        tableProps,
+        manifestFileSize,
+        intervalTrigger,
+        null);
+  }
+
+  /**
+   * @param locationPrefix when set on a partitioned table, the partition is read from the path
+   *     after this prefix instead of from the file's column stats
+   * @param partitionFields partition spec applied when the table is created by this transform
+   * @param sortFields sort order applied when the table is created by this transform
+   * @param tableProps table properties applied when the table is created by this transform
+   * @param manifestFileSize data files per manifest
+   * @param intervalTrigger streaming only: how often manifests are committed
+   * @param evolution schema evolution settings; null or no options means the schema is never
+   *     changed
+   */
+  public AddFiles(
+      IcebergCatalogConfig catalogConfig,
+      String tableIdentifier,
+      @Nullable String locationPrefix,
+      @Nullable List<String> partitionFields,
+      @Nullable List<String> sortFields,
+      @Nullable Map<String, String> tableProps,
+      @Nullable Integer manifestFileSize,
+      @Nullable Duration intervalTrigger,
+      @Nullable SchemaEvolutionConfig evolution) {
+    this.evolution = evolution != null ? evolution : SchemaEvolutionConfig.disabled();
     this.catalogConfig = catalogConfig;
     this.tableIdentifier = tableIdentifier;
     this.partitionFields = partitionFields;
@@ -166,6 +238,10 @@ public AddFiles(
   @Override
   public PCollectionRowTuple expand(PCollection<String> input) {
     if (input.isBounded().equals(UNBOUNDED)) {
+      Preconditions.checkArgument(
+          !evolution.isEnabled(),
+          "Schema evolution is not yet supported for unbounded input: run a batch pipeline or"
+              + " remove the schema evolution options.");
       intervalTrigger = intervalTrigger != null ? intervalTrigger : DEFAULT_TRIGGER_INTERVAL;
       LOG.info(
           "AddFiles configured to generate a new manifest after accumulating {} files, or after {} seconds.",
@@ -182,8 +258,13 @@ public PCollectionRowTuple expand(PCollection<String> input) {
           "AddFiles configured to build partition metadata after the prefix: '{}'", locationPrefix);
     }
 
+    PCollection<String> paths = input;
+    if (evolution.isEnabled()) {
+      paths = gateOnSchemaCommit(input);
+    }
+
     PCollectionTuple dataFiles =
-        input.apply(
+        paths.apply(
             "ConvertToDataFiles",
             ParDo.of(
                     new ConvertToDataFile(
@@ -192,7 +273,8 @@ public PCollectionRowTuple expand(PCollection<String> input) {
                         locationPrefix,
                         partitionFields,
                         sortFields,
-                        tableProps))
+                        tableProps,
+                        evolution))
                 .withOutputTags(DATA_FILES, TupleTagList.of(ERRORS)));
     SchemaCoder<SerializableDataFile> sdfCoder;
     try {
@@ -238,6 +320,47 @@ public PCollectionRowTuple expand(PCollection<String> input) {
         OUTPUT_TAG, snapshots, ERROR_TAG, dataFiles.get(ERRORS).setRowSchema(ERROR_SCHEMA));
   }
 
+  /**
+   * Holds every path until the schema commit has landed. The pre-pass window is the unit of commit,
+   * and for bounded input that unit is the whole input: paths are rewindowed into the global window
+   * first, so one commit covers everything and the Wait.on signal lines up with the main input
+   * whatever windowing the caller applied upstream.
+   */
+  private PCollection<String> gateOnSchemaCommit(PCollection<String> input) {
+    PCollection<String> windowed =
+        input.apply("PrePassGlobalWindow", Window.into(new GlobalWindows()));
+    CommitSchemaUnion.TableCreation creation =
+        new CommitSchemaUnion.TableCreation(partitionFields, sortFields, tableProps);
+    // unset handling follows the mode (fail the job in batch, route in streaming); the pre-pass
+    // only runs on bounded input, see expand
+    IncompatibleSchemaHandling onIncompatible =
+        evolution.incompatibleSchemaHandlingFor(input.isBounded());
+    PCollection<Long> signal =
+        windowed
+            .apply("ReadFooterSchema", ParDo.of(new ReadFooterSchema()))
+            .setCoder(CollectDistinctSchemas.groupCoder())
+            .apply(
+                "CollectDistinctSchemas",
+                Combine.globally(new CollectDistinctSchemas()).withoutDefaults())
+            .apply(
+                "CommitSchemaOnce",
+                ParDo.of(
+                    new CommitSchemaOnce(
+                        catalogConfig,
+                        tableIdentifier,
+                        evolution,
+                        onIncompatible,
+                        creation,
+                        committer)));
+    return windowed.apply("WaitForSchemaCommit", Wait.on(signal));
+  }
+
+  /** Test hook: how the schema pre-pass commits its transaction. */
+  AddFiles withSchemaCommitter(CommitSchemaUnion.Committer committer) {
+    this.committer = committer;
+    return this;
+  }
+
   /**
    * Reads incoming file paths, extracts Iceberg metadata, and converts them into {@link
    * SerializableDataFile} objects.
@@ -272,6 +395,7 @@ static class ConvertToDataFile extends DoFn<String, SerializableDataFile> {
     private transient @MonotonicNonNull BoundedAsyncTasks<ProcessResult> tasks;
     private transient volatile @MonotonicNonNull Table table;
     private transient @MonotonicNonNull Set<String> warned;
+    private final AtomicBoolean refreshedThisBundle = new AtomicBoolean();
 
     // Number of parallel threads processing incoming files
     private static final int THREAD_POOL_SIZE = 10;
@@ -392,8 +516,8 @@ public void setup() {
     /** Clears anything left behind if the runner reuses this instance after a failed bundle. */
     @StartBundle
     public void startBundle() {
-
       checkStateNotNull(tasks).cancelAll();
+      refreshedThisBundle.set(false);
     }
 
     @Teardown
@@ -610,17 +734,21 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta
 
     /**
      * The pre-pass commits the schema before paths reach this stage, so the cached table normally
-     * covers every file. If not, refresh once (a commit may have landed since the table was cached)
-     * and report the remaining delta. Never changes the schema.
+     * covers every file. If not, refresh once per bundle (a commit may have landed since the table
+     * was cached) and report the remaining delta. Never changes the schema.
      */
     private @Nullable String uncoveredReason(org.apache.iceberg.Schema fileSchema) {
       Table table = checkStateNotNull(this.table);
       SchemaDelta delta = SchemaDelta.classify(table, fileSchema);
       if (delta.isEmpty()) {
         return null;
       }
-      synchronized (this) {
-        table.refresh();
+      // a commit can land after the table was cached; one refresh per bundle is enough to see it,
+      // and a bundle full of routed files must not load the table once per file
+      if (refreshedThisBundle.compareAndSet(false, true)) {
+        synchronized (this) {
+          table.refresh();
+        }
       }
       delta = SchemaDelta.classify(table, fileSchema);
       if (delta.isEmpty()) {