AddFiles: Dynamic destinations

.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": 4
 }

CHANGES.md

@@ -100,6 +100,7 @@
 * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--bigquery_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) ([#40827](https://github.com/apache/beam/issues/40827)).
+* (Java) IcebergIO: `AddFiles` can register files in several tables in one transform, naming each file's table from its path. Set `path_pattern` (`withPathPattern` in Java) to a glob with `{name}` placeholders, such as `gs://lake/export/{table}/**`, and make `table` a template over what it captures, such as `default.{table}`. The schema pre-pass then runs per table, so a table whose files conflict is refused and its files reported while every other table is still registered. Requires `schema_evolution_options` and `error_handling`; not supported with `dry_run`. Error rows gain a `table` field only when `path_pattern` is set, so an existing pipeline's error output is unchanged ([#40849](https://github.com/apache/beam/issues/40849)).
 
 ## New Features / Improvements
 
@@ -133,6 +134,7 @@
 * (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)).
 * (YAML) `IcebergAddFiles` will be automatically replaced with the compatible managed transform. For older Beam versions, streaming update compatibility can be maintained by specifying the pipeline option `update_compatibility_version` ([#40827](https://github.com/apache/beam/issues/40827)).
+* (Java) [IcebergIO] AddFiles now collapses batch input into the global window. Manifests and snapshots are sized per key and window, so windowing applied upstream of AddFiles only divided `manifest_file_size` into smaller manifests and more snapshots; its outputs are now globally windowed in batch. Streaming pipelines keep the windowing they had ([#40849](https://github.com/apache/beam/issues/40849)).
 * Twister2 runner supported is removed ([#35905](https://github.com/apache/beam/issues/35905))).
 
 ## Deprecations

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

@@ -56,20 +56,25 @@
 import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
 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.Values;
+import org.apache.beam.sdk.transforms.View;
 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.DefaultTrigger;
 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.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
 import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
@@ -194,10 +199,24 @@ public class AddFiles extends PTransform<PCollection<String>, PCollectionRowTupl
   private static final int DEFAULT_MAX_MANIFESTS_PER_SNAPSHOT = 100;
   static final Schema ERROR_SCHEMA =
       Schema.builder().addStringField("file").addStringField("error").build();
+
+  /**
+   * Error rows from the dynamic path, which names the table the file was headed for; empty when
+   * none could be determined. Single-table mode keeps {@link #ERROR_SCHEMA} as it was, so adding
+   * dynamic destinations does not change the shape of an existing pipeline's error output.
+   */
+  static final Schema DYNAMIC_ERROR_SCHEMA =
+      Schema.builder()
+          .addStringField("file")
+          .addStringField("error")
+          .addStringField("table")
+          .build();
+
   private static final long MANIFEST_PREFIX = UUID.randomUUID().getMostSignificantBits();
   private final IcebergCatalogConfig catalogConfig;
   private final String tableIdentifier;
   private @Nullable Duration intervalTrigger;
+  private @Nullable String pathPattern;
   private final int manifestFileSize;
   private final @Nullable String locationPrefix;
   private final @Nullable List<String> partitionFields;
@@ -261,6 +280,59 @@ public AddFiles(
     this.locationPrefix = locationPrefix;
   }
 
+  /**
+   * A rule that routes each file to a destination table, read from the file's own path: a glob with
+   * {@code {name}} placeholders, whose captures fill in the table identifier as a template. So
+   * {@code gs://lake/export/{tbl}/**} with table {@code sales.{tbl}} routes a file under {@code
+   * export/orders/} to {@code sales.orders}.
+   *
+   * <p>Everything that can be settled without the input is settled here: a mistyped pattern, a
+   * template naming something the pattern never captures, or an option this cannot be combined with
+   * is one exception at submit time, not one error row per file.
+   */
+  public AddFiles withPathPattern(String pathPattern) {
+    // these are the only copies; the YAML keys are named because the provider defers to them.
+    // checkArgument, not checkState: PathPattern below throws the same for a bad pattern, and the
+    // provider used to throw it for these.
+    Preconditions.checkArgument(
+        evolution.isEnabled(),
+        "A path pattern needs schema evolution (schema_evolution_options in YAML): the pre-pass is"
+            + " what decides each table.");
+    Preconditions.checkArgument(
+        !evolution.getDryRun(),
+        "A dry run of a path pattern (dry_run in YAML) is not supported yet.");
+    Preconditions.checkArgument(
+        Strings.isNullOrEmpty(locationPrefix),
+        "A path pattern cannot be combined with a location prefix (location_prefix in YAML): one"
+            + " fixed prefix cannot stop short of each file's own table folder, so the folder ends"
+            + " up inside the partition path ('orders/dt=2024-01-15/' where the table is"
+            + " partitioned only by 'dt').");
+    PathPattern compiled = PathPattern.compile(pathPattern);
+    compiled.checkRenderable(tableIdentifier);
+    this.pathPattern = pathPattern;
+    return this;
+  }
+
+  /**
+   * Collapses batch input to one window, for two reasons: the schema pre-pass is one pass over the
+   * whole input, and {@link GroupIntoBatches} sizes per key and window, so upstream windowing only
+   * divides {@code manifestFileSize} into smaller manifests and more snapshots. The default trigger
+   * keeps a triggered input from firing the pre-pass more than once.
+   *
+   * <p>Streaming keeps the caller's windowing: the snapshots output descends from this, and a
+   * global window under the default trigger never closes.
+   *
+   * <p>Applied before the input branches rather than after, so the error output of each branch
+   * shares this windowing. {@code FlattenErrors} rejects a mix.
+   */
+  private static PCollection<String> collapseToOneWindow(PCollection<String> input) {
+    return input.apply(
+        "CollapseToOneWindow",
+        Window.<String>into(new GlobalWindows())
+            .triggering(DefaultTrigger.of())
+            .discardingFiredPanes());
+  }
+
   @Override
   public PCollectionRowTuple expand(PCollection<String> input) {
     if (input.isBounded().equals(UNBOUNDED)) {
@@ -284,35 +356,71 @@ public PCollectionRowTuple expand(PCollection<String> input) {
           "AddFiles configured to build partition metadata after the prefix: '{}'", locationPrefix);
     }
 
-    PCollection<String> paths = input;
-    if (evolution.isEnabled()) {
-      // one commit per window, and for bounded input the window is the whole input
-      PCollection<String> windowed =
-          input.apply("PrePassGlobalWindow", Window.into(new GlobalWindows()));
-      PCollection<List<CollectDistinctSchemas.SchemaGroup>> schemas = distinctSchemas(windowed);
-      CommitSchemaUnion.Settings settings =
-          new CommitSchemaUnion.Settings(
-              evolution,
-              evolution.incompatibleSchemaHandlingFor(input.isBounded()),
-              new CommitSchemaUnion.NewTableSettings(partitionFields, sortFields, tableProps));
-      if (evolution.getDryRun()) {
-        return report(schemas, settings);
+    CommitSchemaUnion.Settings settings =
+        new CommitSchemaUnion.Settings(
+            evolution,
+            evolution.incompatibleSchemaHandlingFor(input.isBounded()),
+            new CommitSchemaUnion.NewTableSettings(partitionFields, sortFields, tableProps));
+
+    // every stage that can reject a file appends here; one Flatten publishes them together
+    List<PCollection<Row>> errorSources = new ArrayList<>();
+
+    PCollection<String> files =
+        input.isBounded().equals(UNBOUNDED) ? input : collapseToOneWindow(input);
+
+    PCollection<KV<String, String>> keyedPaths;
+
+    if (pathPattern != null) {
+      PCollectionTuple routed =
+          files.apply(
+              "RouteFiles",
+              ParDo.of(new RouteFiles(checkStateNotNull(pathPattern), tableIdentifier))
+                  .withOutputTags(RouteFiles.ROUTED, TupleTagList.of(RouteFiles.ERRORS)));
+      errorSources.add(routed.get(RouteFiles.ERRORS).setRowSchema(DYNAMIC_ERROR_SCHEMA));
+
+      PCollection<KV<String, String>> routedPaths =
+          routed
+              .get(RouteFiles.ROUTED)
+              .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
+
+      PCollectionView<Map<String, SchemaCommitVerdict>> verdicts =
+          verdictView(distinctSchemasPerTable(routedPaths), settings);
+
+      PCollectionTuple gated =
+          routedPaths.apply(
+              "TableGate",
+              ParDo.of(new TableGate(verdicts))
+                  .withSideInputs(verdicts)
+                  .withOutputTags(TableGate.RELEASED, TupleTagList.of(TableGate.ERRORS)));
+      errorSources.add(gated.get(TableGate.ERRORS).setRowSchema(DYNAMIC_ERROR_SCHEMA));
+      keyedPaths =
+          gated
+              .get(TableGate.RELEASED)
+              .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
+    } else {
+      PCollection<String> paths = files;
+      if (evolution.isEnabled()) {
+        PCollection<List<CollectDistinctSchemas.SchemaGroup>> schemas = distinctSchemas(files);
+        if (evolution.getDryRun()) {
+          return report(schemas, settings);
+        }
+        PCollection<SchemaCommitVerdict> committed = commitSchema(schemas, settings);
+        paths = files.apply("WaitForSchemaCommit", Wait.on(committed));
       }
-      PCollection<Long> committed = commitSchema(schemas, settings);
-      paths = windowed.apply("WaitForSchemaCommit", Wait.on(committed));
-    }
 
-    PCollection<KV<String, String>> keyedPaths =
-        paths.apply(
-            "KeyByTable",
-            WithKeys.<String, String>of(tableIdentifier).withKeyType(TypeDescriptors.strings()));
+      keyedPaths =
+          paths.apply(
+              "KeyByTable",
+              WithKeys.<String, String>of(tableIdentifier).withKeyType(TypeDescriptors.strings()));
+    }
 
     PCollectionTuple dataFiles =
         keyedPaths.apply(
             "ConvertToDataFiles",
             ParDo.of(
                     new ConvertToDataFile(
                         catalogConfig,
+                        pathPattern != null,
                         locationPrefix,
                         partitionFields,
                         sortFields,
@@ -373,8 +481,15 @@ public void process(
             .apply("CommitManifests", ParDo.of(new CommitManifestFilesDoFn(catalogConfig)))
             .setRowSchema(SnapshotInfo.getSchema());
 
-    return PCollectionRowTuple.of(
-        OUTPUT_TAG, snapshots, ERROR_TAG, dataFiles.get(ERRORS).setRowSchema(ERROR_SCHEMA));
+    errorSources.add(
+        dataFiles
+            .get(ERRORS)
+            .setRowSchema(pathPattern == null ? ERROR_SCHEMA : DYNAMIC_ERROR_SCHEMA));
+
+    PCollection<Row> errors =
+        PCollectionList.of(errorSources).apply("FlattenErrors", Flatten.pCollections());
+
+    return PCollectionRowTuple.of(OUTPUT_TAG, snapshots, ERROR_TAG, errors);
   }
 
   /** The footer read is keyed, so the single-table path supplies its one table as the key. */
@@ -399,8 +514,33 @@ private PCollection<List<CollectDistinctSchemas.SchemaGroup>> distinctSchemas(
             Combine.globally(new CollectDistinctSchemas()).withoutDefaults());
   }
 
+  /** One schema union per table: the same CombineFn, lifted to a per-key combine. */
+  private PCollection<KV<String, List<CollectDistinctSchemas.SchemaGroup>>> distinctSchemasPerTable(
+      PCollection<KV<String, String>> routed) {
+    return fileSchemas(routed)
+        .apply("CollectDistinctSchemasPerTable", Combine.perKey(new CollectDistinctSchemas()));
+  }
+
+  /** The verdicts, as a map the gate can look a table up in. */
+  private PCollectionView<Map<String, SchemaCommitVerdict>> verdictView(
+      PCollection<KV<String, List<CollectDistinctSchemas.SchemaGroup>>> schemas,
+      CommitSchemaUnion.Settings settings) {
+    SchemaCoder<SchemaCommitVerdict> verdictCoder;
+    try {
+      verdictCoder = SchemaRegistry.createDefault().getSchemaCoder(SchemaCommitVerdict.class);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+    return schemas
+        .apply(
+            "CommitSchemaPerTable",
+            ParDo.of(new CommitSchemaPerTable(catalogConfig, settings, committer)))
+        .setCoder(KvCoder.of(StringUtf8Coder.of(), verdictCoder))
+        .apply("VerdictsByTable", View.asMap());
+  }
+
   /** Commits the plan for the window's schemas once; the signal releases the gated paths. */
-  private PCollection<Long> commitSchema(
+  private PCollection<SchemaCommitVerdict> commitSchema(
       PCollection<List<CollectDistinctSchemas.SchemaGroup>> schemas,
       CommitSchemaUnion.Settings settings) {
     return schemas.apply(
@@ -464,6 +604,7 @@ AddFiles withSchemaCommitter(CommitSchemaUnion.Committer committer) {
   static class ConvertToDataFile
       extends DoFn<KV<String, String>, KV<String, SerializableDataFile>> {
     private final IcebergCatalogConfig catalogConfig;
+    private final boolean dynamic;
     public static final TupleTag<Row> ERRORS = new TupleTag<>();
     public static final TupleTag<KV<String, SerializableDataFile>> DATA_FILES = new TupleTag<>();
     private final @Nullable String prefix;
@@ -499,12 +640,14 @@ private TableState stateFor(String tableName) {
 
     public ConvertToDataFile(
         IcebergCatalogConfig catalogConfig,
+        boolean dynamic,
         @Nullable String prefix,
         @Nullable List<String> partitionFields,
         @Nullable List<String> sortFields,
         @Nullable Map<String, String> tableProps) {
       this(
           catalogConfig,
+          dynamic,
           prefix,
           partitionFields,
           sortFields,
@@ -514,12 +657,14 @@ public ConvertToDataFile(
 
     public ConvertToDataFile(
         IcebergCatalogConfig catalogConfig,
+        boolean dynamic,
         @Nullable String prefix,
         @Nullable List<String> partitionFields,
         @Nullable List<String> sortFields,
         @Nullable Map<String, String> tableProps,
         SchemaEvolutionConfig evolution) {
       this.catalogConfig = catalogConfig;
+      this.dynamic = dynamic;
       this.prefix = prefix;
       this.partitionFields = partitionFields;
       this.sortFields = sortFields;
@@ -1035,21 +1180,21 @@ private boolean firstTime(String tableName, String key) {
       return checkStateNotNull(warned).add(tableName + ":" + key);
     }
 
-    private static ProcessResult errorResult(
+    private ProcessResult errorResult(
         String tableName,
         String filePath,
         String message,
         Instant timestamp,
         BoundedWindow window,
         PaneInfo pane) {
-      return new ProcessResult(
-          tableName,
-          null,
-          Row.withSchema(ERROR_SCHEMA).addValues(filePath, message).build(),
-          null,
-          timestamp,
-          window,
-          pane);
+      Row errorRow;
+      if (dynamic) {
+        errorRow =
+            Row.withSchema(DYNAMIC_ERROR_SCHEMA).addValues(filePath, message, tableName).build();
+      } else {
+        errorRow = Row.withSchema(ERROR_SCHEMA).addValues(filePath, message).build();
+      }
+      return new ProcessResult(tableName, null, errorRow, null, timestamp, window, pane);
     }
 
     static <W, T> T transformValue(Transform<W, T> transform, Type type, ByteBuffer bytes) {