AddFiles: tighten file schemas

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

@@ -17,87 +17,219 @@
  */
 package org.apache.beam.sdk.io.iceberg;
 
+import com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.TreeMap;
+import java.util.TreeSet;
+import org.apache.beam.sdk.coders.AtomicCoder;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderRegistry;
-import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.ListCoder;
 import org.apache.beam.sdk.coders.MapCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber;
 import org.apache.beam.sdk.transforms.Combine;
-import org.apache.beam.sdk.values.KV;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
- * Collects the distinct schemas among canonical file schema JSONs (see {@link FileSchemas}), with
- * the number of files per schema, most common first (ties broken by JSON). The commit side applies
- * schemas in this order, so the schema covering the most files wins a conflict.
- *
- * <p>Inputs are compared as strings, so they must already be canonical.
+ * 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.
  */
 class CollectDistinctSchemas
-    extends Combine.CombineFn<String, Map<String, Long>, List<KV<String, Long>>> {
+    extends Combine.CombineFn<
+        CollectDistinctSchemas.SchemaGroup,
+        Map<String, CollectDistinctSchemas.Group>,
+        List<CollectDistinctSchemas.SchemaGroup>> {
+
+  /** Mutable accumulator counterpart of {@link SchemaGroup}. */
+  static final class Group {
+    long files;
+    TreeSet<String> nullFreeColumns;
+
+    Group(long files, TreeSet<String> nullFreeColumns) {
+      this.files = files;
+      this.nullFreeColumns = nullFreeColumns;
+    }
+
+    @Override
+    public boolean equals(@Nullable Object other) {
+      if (!(other instanceof Group)) {
+        return false;
+      }
+      Group that = (Group) other;
+      return files == that.files && nullFreeColumns.equals(that.nullFreeColumns);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(files, nullFreeColumns);
+    }
+  }
+
+  /**
+   * A schema, how many files carry it, and the columns all of them proved free of nulls.
+   * ReadFooterSchema emits one per file ({@code files} = 1); this combiner merges them.
+   */
+  @DefaultSchema(AutoValueSchema.class)
+  @AutoValue
+  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);
+    }
+
+    static SchemaCoder<SchemaGroup> getCoder() {
+      if (coder == null) {
+        try {
+          coder = SchemaRegistry.createDefault().getSchemaCoder(SchemaGroup.class);
+        } catch (NoSuchSchemaException e) {
+          throw new RuntimeException(e);
+        }
+      }
+      return coder;
+    }
+
+    @SchemaFieldNumber("0")
+    abstract String getSchemaJson();
+
+    @SchemaFieldNumber("1")
+    abstract long getFiles();
+
+    @SchemaFieldNumber("2")
+    abstract List<String> getNullFreeColumns();
+
+    @Override
+    public final String toString() {
+      return getFiles()
+          + " file(s), null-free in "
+          + getNullFreeColumns()
+          + ", schema "
+          + getSchemaJson();
+    }
+  }
 
   @Override
-  public Map<String, Long> createAccumulator() {
+  public Map<String, Group> createAccumulator() {
     return new TreeMap<>();
   }
 
   @Override
-  public Map<String, Long> addInput(Map<String, Long> accumulator, String schemaJson) {
-    add(accumulator, schemaJson, 1L);
+  public Map<String, Group> addInput(Map<String, Group> accumulator, SchemaGroup file) {
+    add(accumulator, file.getSchemaJson(), file.getFiles(), file.getNullFreeColumns());
     return accumulator;
   }
 
   @Override
-  public Map<String, Long> mergeAccumulators(Iterable<Map<String, Long>> accumulators) {
-    Map<String, Long> merged = createAccumulator();
-    for (Map<String, Long> accumulator : accumulators) {
-      for (Map.Entry<String, Long> entry : accumulator.entrySet()) {
-        add(merged, entry.getKey(), entry.getValue());
+  public Map<String, Group> mergeAccumulators(Iterable<Map<String, Group>> accumulators) {
+    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);
       }
     }
     return merged;
   }
 
   @Override
-  public List<KV<String, Long>> extractOutput(Map<String, Long> accumulator) {
-    List<KV<String, Long>> schemas = new ArrayList<>();
-    for (Map.Entry<String, Long> entry : accumulator.entrySet()) {
-      schemas.add(KV.of(entry.getKey(), entry.getValue()));
+  public List<SchemaGroup> extractOutput(Map<String, Group> accumulator) {
+    List<SchemaGroup> schemas = new ArrayList<>();
+    for (Map.Entry<String, Group> entry : accumulator.entrySet()) {
+      schemas.add(
+          SchemaGroup.of(
+              entry.getKey(),
+              entry.getValue().files,
+              new ArrayList<>(entry.getValue().nullFreeColumns)));
     }
     schemas.sort(
         (a, b) -> {
-          int byCount = Long.compare(b.getValue(), a.getValue());
+          int byCount = Long.compare(b.getFiles(), a.getFiles());
           if (byCount != 0) {
             return byCount;
           }
-          return a.getKey().compareTo(b.getKey());
+          return a.getSchemaJson().compareTo(b.getSchemaJson());
         });
     return schemas;
   }
 
   @Override
-  public Coder<Map<String, Long>> getAccumulatorCoder(
-      CoderRegistry registry, Coder<String> inputCoder) {
-    return MapCoder.of(StringUtf8Coder.of(), VarLongCoder.of());
+  public Coder<Map<String, Group>> getAccumulatorCoder(
+      CoderRegistry registry, Coder<SchemaGroup> inputCoder) {
+    return MapCoder.of(StringUtf8Coder.of(), GroupCoder.INSTANCE);
   }
 
   @Override
-  public Coder<List<KV<String, Long>>> getDefaultOutputCoder(
-      CoderRegistry registry, Coder<String> inputCoder) {
-    return ListCoder.of(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
+  public Coder<List<SchemaGroup>> getDefaultOutputCoder(
+      CoderRegistry registry, Coder<SchemaGroup> inputCoder) {
+    return outputCoder();
+  }
+
+  static Coder<SchemaGroup> groupCoder() {
+    return SchemaGroup.getCoder();
+  }
+
+  static Coder<List<SchemaGroup>> outputCoder() {
+    return ListCoder.of(SchemaGroup.getCoder());
+  }
+
+  private static final Coder<List<String>> COLUMNS_CODER = ListCoder.of(StringUtf8Coder.of());
+
+  /** Sorted columns, so the encoding is deterministic. */
+  private static class GroupCoder extends AtomicCoder<Group> {
+    static final GroupCoder INSTANCE = new GroupCoder();
+
+    private GroupCoder() {}
+
+    @Override
+    public void encode(Group value, OutputStream out) throws IOException {
+      VarLongCoder.of().encode(value.files, out);
+      COLUMNS_CODER.encode(new ArrayList<>(value.nullFreeColumns), 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)));
+    }
   }
 
-  private static void add(Map<String, Long> accumulator, String schemaJson, long count) {
-    Long existing = accumulator.get(schemaJson);
+  private static void add(
+      Map<String, Group> accumulator,
+      String schemaJson,
+      long files,
+      Iterable<String> nullFreeColumns) {
+    Group existing = accumulator.get(schemaJson);
     if (existing == null) {
-      accumulator.put(schemaJson, count);
-    } else {
-      accumulator.put(schemaJson, existing + count);
+      TreeSet<String> copy = new TreeSet<>();
+      for (String column : nullFreeColumns) {
+        copy.add(column);
+      }
+      accumulator.put(schemaJson, new Group(files, copy));
+      return;
+    }
+    existing.files += files;
+    TreeSet<String> stillNullFree = new TreeSet<>();
+    for (String column : nullFreeColumns) {
+      if (existing.nullFreeColumns.contains(column)) {
+        stillNullFree.add(column);
+      }
     }
+    existing.nullFreeColumns = stillNullFree;
   }
 }