AddFiles: Parse partition keys in filepath for requirement check

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

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

@@ -211,7 +211,8 @@ public AddFiles(
 
   /**
    * @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
+   *     after this prefix instead of from the file's column stats. A table partitioned by {@code
+   *     dt} then accepts files under {@code dt=2024-01-15/} that have no {@code dt} column
    * @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
@@ -343,7 +344,7 @@ public PCollectionRowTuple expand(PCollection<String> input) {
   private PCollection<List<CollectDistinctSchemas.SchemaGroup>> distinctSchemas(
       PCollection<String> windowed) {
     return windowed
-        .apply("ReadFooterSchema", ParDo.of(new ReadFooterSchema(evolution)))
+        .apply("ReadFooterSchema", ParDo.of(new ReadFooterSchema(evolution, locationPrefix)))
         .setCoder(CollectDistinctSchemas.groupCoder())
         .apply(
             "CollectDistinctSchemas",
@@ -811,21 +812,37 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta
       } catch (Exception e) {
         return Verdict.error(UNREADABLE_SCHEMA_ERROR + errorMessage(e));
       }
-      @Nullable String uncovered = uncoveredReason(fileSchema);
+      Set<Integer> identityPartitionColumnsFromPath = identityPartitionColumnsFromPath(filePath);
+      @Nullable String uncovered = uncoveredReason(fileSchema, identityPartitionColumnsFromPath);
       if (uncovered != null) {
         return Verdict.error(uncovered);
       }
-      return checkPins(filePath, fileSchema, parquetFooter);
+      return checkPins(filePath, fileSchema, parquetFooter, identityPartitionColumnsFromPath);
+    }
+
+    /**
+     * Identity partition columns this file may lack because its path gives them a value: {@code dt}
+     * for a file under {@code dt=2024-01-15/} when the table is partitioned by {@code dt}. Empty
+     * without a prefix.
+     */
+    private Set<Integer> identityPartitionColumnsFromPath(String filePath) {
+      Table table = checkStateNotNull(this.table);
+      if (Strings.isNullOrEmpty(prefix) || table.spec().isUnpartitioned()) {
+        return Collections.emptySet();
+      }
+      return PartitionPaths.identityColumnsFilledByKeys(
+          table.spec(), PartitionPaths.partitionKeysInPath(filePath, prefix));
     }
 
     /**
      * The pre-pass commits the schema before paths reach this stage, so the cached table normally
      * 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) {
+    private @Nullable String uncoveredReason(
+        org.apache.iceberg.Schema fileSchema, Set<Integer> identityPartitionColumnsFromPath) {
       Table table = checkStateNotNull(this.table);
-      SchemaDelta delta = SchemaDelta.classify(table, fileSchema);
+      SchemaDelta delta = SchemaDelta.classify(table, fileSchema, identityPartitionColumnsFromPath);
       if (delta.isEmpty()) {
         return null;
       }
@@ -836,7 +853,7 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta
           table.refresh();
         }
       }
-      delta = SchemaDelta.classify(table, fileSchema);
+      delta = SchemaDelta.classify(table, fileSchema, identityPartitionColumnsFromPath);
       if (delta.isEmpty()) {
         return null;
       }
@@ -857,14 +874,23 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta
      * goes on, so a pin the footer does count nulls for still fails the file.
      */
     private Verdict checkPins(
-        String filePath, org.apache.iceberg.Schema fileSchema, ParquetMetadata footer) {
+        String filePath,
+        org.apache.iceberg.Schema fileSchema,
+        ParquetMetadata footer,
+        Set<Integer> identityPartitionColumnsFromPath) {
       Table table = checkStateNotNull(this.table);
       List<String> unproven = new ArrayList<>();
       for (String pinned : evolution.getRequiredColumns()) {
-        if (table.schema().findField(pinned) == null) {
+        org.apache.iceberg.types.Types.@Nullable NestedField tableField =
+            table.schema().findField(pinned);
+        if (tableField == null) {
           continue;
         }
         if (fileSchema.findField(pinned) == null) {
+          if (identityPartitionColumnsFromPath.contains(tableField.fieldId())) {
+            // the path gives it a value, e.g. dt=2024-01-15, so it is never null
+            continue;
+          }
           return Verdict.error(PINNED_COLUMN_ERROR + pinned + " is absent from the file");
         }
         @Nullable Long nulls = FileSchemas.nullCount(footer, fileSchema, pinned);
@@ -969,10 +995,7 @@ private String getPartitionFromFilePath(String filePath) {
       if (checkStateNotNull(table).spec().isUnpartitioned()) {
         return "";
       }
-      String partitionPath = filePath.substring(checkStateNotNull(prefix).length());
-      int lastSlashIndex = partitionPath.lastIndexOf('/');
-
-      return lastSlashIndex > 0 ? partitionPath.substring(0, lastSlashIndex) : "";
+      return PartitionPaths.relative(filePath, checkStateNotNull(prefix));
     }
 
     /**

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

@@ -22,6 +22,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -47,9 +48,10 @@
 /**
  * One output entry per distinct schema: its file count and the columns EVERY file carrying it
  * proved free of nulls; one file with a null in "name" forces "name" to relax, however many clean
- * files sit next to it. Entries come out most common first (ties broken by the JSON text) because
- * the commit side applies schemas in that order and the most common schema should win a conflict.
- * Schemas are compared as strings, so inputs must already be canonical.
+ * files sit next to it. Partition keys in path are kept the same way: {@code dt} stays only if
+ * every file sits under a {@code dt=...} folder. Entries come out most common first (ties broken by
+ * the JSON text) because the commit side applies schemas in that order and the most common schema
+ * should win a conflict. Schemas are compared as strings, so inputs must already be canonical.
  */
 class CollectDistinctSchemas
     extends Combine.CombineFn<
@@ -61,10 +63,12 @@ class CollectDistinctSchemas
   static final class Group {
     long files;
     TreeSet<String> nullFreeColumns;
+    TreeSet<String> partitionKeysInPath;
 
-    Group(long files, TreeSet<String> nullFreeColumns) {
+    Group(long files, TreeSet<String> nullFreeColumns, TreeSet<String> partitionKeysInPath) {
       this.files = files;
       this.nullFreeColumns = nullFreeColumns;
+      this.partitionKeysInPath = partitionKeysInPath;
     }
 
     @Override
@@ -73,12 +77,14 @@ public boolean equals(@Nullable Object other) {
         return false;
       }
       Group that = (Group) other;
-      return files == that.files && nullFreeColumns.equals(that.nullFreeColumns);
+      return files == that.files
+          && nullFreeColumns.equals(that.nullFreeColumns)
+          && partitionKeysInPath.equals(that.partitionKeysInPath);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(files, nullFreeColumns);
+      return Objects.hash(files, nullFreeColumns, partitionKeysInPath);
     }
   }
 
@@ -92,7 +98,20 @@ abstract static class SchemaGroup {
     private static @MonotonicNonNull SchemaCoder<SchemaGroup> coder;
 
     static SchemaGroup of(String schemaJson, long files, List<String> nullFreeColumns) {
-      return new AutoValue_CollectDistinctSchemas_SchemaGroup(schemaJson, files, nullFreeColumns);
+      return of(schemaJson, files, nullFreeColumns, Collections.emptyList());
+    }
+
+    static SchemaGroup of(
+        String schemaJson,
+        long files,
+        List<String> nullFreeColumns,
+        List<String> partitionKeysInPath) {
+      return new AutoValue_CollectDistinctSchemas_SchemaGroup(
+          schemaJson, files, nullFreeColumns, partitionKeysInPath);
+    }
+
+    SchemaGroup withPartitionKeysInPath(List<String> partitionKeysInPath) {
+      return of(getSchemaJson(), getFiles(), getNullFreeColumns(), partitionKeysInPath);
     }
 
     static SchemaCoder<SchemaGroup> getCoder() {
@@ -115,11 +134,21 @@ static SchemaCoder<SchemaGroup> getCoder() {
     @SchemaFieldNumber("2")
     abstract List<String> getNullFreeColumns();
 
+    /**
+     * Folder keys every file of the group has a value for, e.g. {@code [dt]} when every file sits
+     * under a {@code dt=...} folder. A table partitioned by {@code dt} then needs no {@code dt}
+     * column in these files.
+     */
+    @SchemaFieldNumber("3")
+    abstract List<String> getPartitionKeysInPath();
+
     @Override
     public final String toString() {
       return getFiles()
           + " file(s), null-free in "
           + getNullFreeColumns()
+          + ", partition keys in path "
+          + getPartitionKeysInPath()
           + ", schema "
           + getSchemaJson();
     }
@@ -132,7 +161,12 @@ public Map<String, Group> createAccumulator() {
 
   @Override
   public Map<String, Group> addInput(Map<String, Group> accumulator, SchemaGroup file) {
-    add(accumulator, file.getSchemaJson(), file.getFiles(), file.getNullFreeColumns());
+    add(
+        accumulator,
+        file.getSchemaJson(),
+        file.getFiles(),
+        file.getNullFreeColumns(),
+        file.getPartitionKeysInPath());
     return accumulator;
   }
 
@@ -141,7 +175,8 @@ public Map<String, Group> mergeAccumulators(Iterable<Map<String, Group>> accumul
     Map<String, Group> merged = createAccumulator();
     for (Map<String, Group> accumulator : accumulators) {
       for (Map.Entry<String, Group> entry : accumulator.entrySet()) {
-        add(merged, entry.getKey(), entry.getValue().files, entry.getValue().nullFreeColumns);
+        Group group = entry.getValue();
+        add(merged, entry.getKey(), group.files, group.nullFreeColumns, group.partitionKeysInPath);
       }
     }
     return merged;
@@ -155,7 +190,8 @@ public List<SchemaGroup> extractOutput(Map<String, Group> accumulator) {
           SchemaGroup.of(
               entry.getKey(),
               entry.getValue().files,
-              new ArrayList<>(entry.getValue().nullFreeColumns)));
+              new ArrayList<>(entry.getValue().nullFreeColumns),
+              new ArrayList<>(entry.getValue().partitionKeysInPath)));
     }
     schemas.sort(
         (a, b) -> {
@@ -200,36 +236,50 @@ private GroupCoder() {}
     public void encode(Group value, OutputStream out) throws IOException {
       VarLongCoder.of().encode(value.files, out);
       COLUMNS_CODER.encode(new ArrayList<>(value.nullFreeColumns), out);
+      COLUMNS_CODER.encode(new ArrayList<>(value.partitionKeysInPath), out);
     }
 
     @Override
     public Group decode(InputStream in) throws IOException {
       long files = VarLongCoder.of().decode(in);
-      return new Group(files, new TreeSet<>(COLUMNS_CODER.decode(in)));
+      TreeSet<String> nullFreeColumns = new TreeSet<>(COLUMNS_CODER.decode(in));
+      TreeSet<String> partitionKeysInPath = new TreeSet<>(COLUMNS_CODER.decode(in));
+      return new Group(files, nullFreeColumns, partitionKeysInPath);
     }
   }
 
   private static void add(
       Map<String, Group> accumulator,
       String schemaJson,
       long files,
-      Iterable<String> nullFreeColumns) {
+      Iterable<String> nullFreeColumns,
+      Iterable<String> partitionKeysInPath) {
     Group existing = accumulator.get(schemaJson);
     if (existing == null) {
-      TreeSet<String> copy = new TreeSet<>();
-      for (String column : nullFreeColumns) {
-        copy.add(column);
-      }
-      accumulator.put(schemaJson, new Group(files, copy));
+      accumulator.put(
+          schemaJson, new Group(files, copy(nullFreeColumns), copy(partitionKeysInPath)));
       return;
     }
     existing.files += files;
-    TreeSet<String> stillNullFree = new TreeSet<>();
-    for (String column : nullFreeColumns) {
-      if (existing.nullFreeColumns.contains(column)) {
-        stillNullFree.add(column);
+    existing.nullFreeColumns = intersect(existing.nullFreeColumns, nullFreeColumns);
+    existing.partitionKeysInPath = intersect(existing.partitionKeysInPath, partitionKeysInPath);
+  }
+
+  private static TreeSet<String> copy(Iterable<String> values) {
+    TreeSet<String> copy = new TreeSet<>();
+    for (String value : values) {
+      copy.add(value);
+    }
+    return copy;
+  }
+
+  private static TreeSet<String> intersect(TreeSet<String> kept, Iterable<String> values) {
+    TreeSet<String> stillKept = new TreeSet<>();
+    for (String value : values) {
+      if (kept.contains(value)) {
+        stillKept.add(value);
       }
     }
-    existing.nullFreeColumns = stillNullFree;
+    return stillKept;
   }
 }

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

@@ -0,0 +1,91 @@
+/*
+ * 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 java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeSet;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Partition values in a file's path. With prefix {@code gs://lake/orders/}, the file {@code
+ * gs://lake/orders/dt=2024-01-15/f.parquet} has {@code dt = 2024-01-15}. If the table is
+ * partitioned by {@code dt} as is (not by a transform such as {@code bucket}), readers take {@code
+ * dt} from that value, so the file needs no {@code dt} column.
+ */
+final class PartitionPaths {
+  /**
+   * Hive writes {@code dt=__HIVE_DEFAULT_PARTITION__} when dt is null; Iceberg reads it as null.
+   */
+  static final String HIVE_NULL = "__HIVE_DEFAULT_PARTITION__";
+
+  private PartitionPaths() {}
+
+  /**
+   * With prefix {@code gs://lake/orders/}: {@code gs://lake/orders/dt=1/region=EU/f.parquet} gives
+   * {@code dt=1/region=EU}, and {@code gs://lake/orders/f.parquet} gives "".
+   */
+  static String relative(String filePath, String prefix) {
+    String partitionPath = filePath.substring(prefix.length());
+    int lastSlashIndex = partitionPath.lastIndexOf('/');
+    return lastSlashIndex > 0 ? partitionPath.substring(0, lastSlashIndex) : "";
+  }
+
+  /**
+   * The partition keys the file's folders give a value: {@code dt=1/region=EU/f.parquet} gives
+   * {@code [dt, region]}, {@code dt=__HIVE_DEFAULT_PARTITION__/region=EU/f.parquet} gives {@code
+   * [region]}, and {@code 2024/f.parquet} gives none. Also none without a prefix or for a file
+   * outside it. Needs no table, unlike Iceberg's parser, so the pre-pass can use it.
+   */
+  static Set<String> partitionKeysInPath(String filePath, @Nullable String prefix) {
+    Set<String> keys = new TreeSet<>();
+    if (prefix == null || prefix.isEmpty() || !filePath.startsWith(prefix)) {
+      return keys;
+    }
+    String relative = relative(filePath, prefix);
+    if (relative.isEmpty()) {
+      return keys;
+    }
+    for (String folder : relative.split("/", -1)) {
+      String[] parts = folder.split("=", 2);
+      if (parts.length == 2 && !parts[0].isEmpty() && !HIVE_NULL.equals(parts[1])) {
+        keys.add(parts[0]);
+      }
+    }
+    return keys;
+  }
+
+  /**
+   * Ids of the columns that identity partitions named by these keys fill: {@code dt} for key {@code
+   * dt} in a table partitioned by {@code dt}. Other transforms fill nothing: a file under {@code
+   * id_bucket=3/} in a table partitioned by {@code bucket(id, 4)} is in bucket 3, which does not
+   * say what its {@code id} is.
+   */
+  static Set<Integer> identityColumnsFilledByKeys(PartitionSpec spec, Collection<String> keys) {
+    Set<Integer> ids = new HashSet<>();
+    for (PartitionField field : spec.fields()) {
+      if (field.transform().isIdentity() && keys.contains(field.name())) {
+        ids.add(field.sourceId());
+      }
+    }
+    return ids;
+  }
+}