AddFiles: read side of the schema pre-pass

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

@@ -0,0 +1,103 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+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.transforms.Combine;
+import org.apache.beam.sdk.values.KV;
+
+/**
+ * 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.
+ */
+class CollectDistinctSchemas
+    extends Combine.CombineFn<String, Map<String, Long>, List<KV<String, Long>>> {
+
+  @Override
+  public Map<String, Long> createAccumulator() {
+    return new TreeMap<>();
+  }
+
+  @Override
+  public Map<String, Long> addInput(Map<String, Long> accumulator, String schemaJson) {
+    add(accumulator, schemaJson, 1L);
+    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());
+      }
+    }
+    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()));
+    }
+    schemas.sort(
+        (a, b) -> {
+          int byCount = Long.compare(b.getValue(), a.getValue());
+          if (byCount != 0) {
+            return byCount;
+          }
+          return a.getKey().compareTo(b.getKey());
+        });
+    return schemas;
+  }
+
+  @Override
+  public Coder<Map<String, Long>> getAccumulatorCoder(
+      CoderRegistry registry, Coder<String> inputCoder) {
+    return MapCoder.of(StringUtf8Coder.of(), VarLongCoder.of());
+  }
+
+  @Override
+  public Coder<List<KV<String, Long>>> getDefaultOutputCoder(
+      CoderRegistry registry, Coder<String> inputCoder) {
+    return ListCoder.of(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
+  }
+
+  private static void add(Map<String, Long> accumulator, String schemaJson, long count) {
+    Long existing = accumulator.get(schemaJson);
+    if (existing == null) {
+      accumulator.put(schemaJson, count);
+    } else {
+      accumulator.put(schemaJson, existing + count);
+    }
+  }
+}

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

@@ -0,0 +1,98 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.SchemaParser;
+import org.apache.iceberg.parquet.ParquetSchemaUtil;
+import org.apache.iceberg.types.Type;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+
+/**
+ * Derives the schema a file contributes to schema inference. The canonical form sorts struct fields
+ * by name at every level and renumbers ids in deterministic order, so files that differ only in
+ * column order produce identical JSON. Ids are positional and meaningless: the commit side
+ * reconciles columns by name.
+ */
+final class FileSchemas {
+  private FileSchemas() {}
+
+  static String canonicalJson(ParquetMetadata footer) {
+    Schema converted = ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema());
+    return SchemaParser.toJson(canonical(converted));
+  }
+
+  static Schema canonical(Schema schema) {
+    Type sorted = TypeUtil.visit(schema.asStruct(), new SortFields());
+    int[] nextId = {0};
+    return TypeUtil.assignFreshIds(new Schema(sorted.asStructType().fields()), () -> ++nextId[0]);
+  }
+
+  /**
+   * Rebuilds every struct with its fields sorted by name; every other attribute (optionality, doc,
+   * defaults) is preserved. Iceberg owns the traversal, so nested types this code has never heard
+   * of (variant, and whatever comes next) are visited rather than silently passed through.
+   */
+  private static class SortFields extends TypeUtil.SchemaVisitor<Type> {
+    @Override
+    public Type struct(Types.StructType struct, List<Type> fieldTypes) {
+      List<Types.NestedField> rebuilt = new ArrayList<>();
+      for (int i = 0; i < struct.fields().size(); i++) {
+        Types.NestedField field = struct.fields().get(i);
+        rebuilt.add(Types.NestedField.from(field).ofType(fieldTypes.get(i)).build());
+      }
+      rebuilt.sort((a, b) -> a.name().compareTo(b.name()));
+      return Types.StructType.of(rebuilt);
+    }
+
+    @Override
+    public Type field(Types.NestedField field, Type fieldType) {
+      return fieldType;
+    }
+
+    @Override
+    public Type list(Types.ListType list, Type elementType) {
+      if (list.isElementOptional()) {
+        return Types.ListType.ofOptional(list.elementId(), elementType);
+      }
+      return Types.ListType.ofRequired(list.elementId(), elementType);
+    }
+
+    @Override
+    public Type map(Types.MapType map, Type keyType, Type valueType) {
+      if (map.isValueOptional()) {
+        return Types.MapType.ofOptional(map.keyId(), map.valueId(), keyType, valueType);
+      }
+      return Types.MapType.ofRequired(map.keyId(), map.valueId(), keyType, valueType);
+    }
+
+    @Override
+    public Type variant(Types.VariantType variant) {
+      return variant;
+    }
+
+    @Override
+    public Type primitive(Type.PrimitiveType primitive) {
+      return primitive;
+    }
+  }
+}

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

@@ -0,0 +1,181 @@
+/*
+ * 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 static org.apache.beam.sdk.metrics.Metrics.counter;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.util.Collections;
+import java.util.concurrent.Callable;
+import org.apache.beam.sdk.metrics.Counter;
+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.iceberg.FileFormat;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Emits the canonical schema (see {@link FileSchemas}) of every readable Parquet file as JSON.
+ * Unreadable or non-Parquet files contribute nothing.
+ */
+class ReadFooterSchema extends DoFn<String, String> {
+  private static final Logger LOG = LoggerFactory.getLogger(ReadFooterSchema.class);
+
+  static final int DEFAULT_THREAD_POOL_SIZE = 10;
+  static final int DEFAULT_MAX_IN_FLIGHT_TASKS = 100;
+  static final String FILES_READ_COUNTER = "numFilesRead";
+  static final String SCHEMAS_EMITTED_COUNTER = "numSchemasEmitted";
+  static final String FOOTER_READ_ERRORS_COUNTER = "numFooterReadErrors";
+  private static final Counter numFilesRead = counter(ReadFooterSchema.class, FILES_READ_COUNTER);
+  private static final Counter numSchemasEmitted =
+      counter(ReadFooterSchema.class, SCHEMAS_EMITTED_COUNTER);
+  private static final Counter numFooterReadErrors =
+      counter(ReadFooterSchema.class, FOOTER_READ_ERRORS_COUNTER);
+
+  private final int threadPoolSize;
+  private final int maxInFlightTasks;
+  private transient @MonotonicNonNull BoundedAsyncTasks<ReadResult> tasks;
+
+  ReadFooterSchema() {
+    this(DEFAULT_THREAD_POOL_SIZE, DEFAULT_MAX_IN_FLIGHT_TASKS);
+  }
+
+  ReadFooterSchema(int threadPoolSize, int maxInFlightTasks) {
+    this.threadPoolSize = threadPoolSize;
+    this.maxInFlightTasks = maxInFlightTasks;
+  }
+
+  /**
+   * {@code schemaJson} is null when the file contributes no schema. Counters are updated when the
+   * result is delivered, on the processing thread: metrics touched from the executor are lost.
+   */
+  private static class ReadResult {
+    final @Nullable String schemaJson;
+    final boolean footerError;
+    final Instant timestamp;
+    final BoundedWindow window;
+    final PaneInfo paneInfo;
+
+    ReadResult(
+        @Nullable String schemaJson,
+        boolean footerError,
+        Instant timestamp,
+        BoundedWindow window,
+        PaneInfo paneInfo) {
+      this.schemaJson = schemaJson;
+      this.footerError = footerError;
+      this.timestamp = timestamp;
+      this.window = window;
+      this.paneInfo = paneInfo;
+    }
+  }
+
+  @Setup
+  public void setup() {
+    tasks = new BoundedAsyncTasks<>(threadPoolSize, maxInFlightTasks);
+  }
+
+  /** Clears anything left behind if the runner reuses this instance after a failed bundle. */
+  @StartBundle
+  public void startBundle() {
+    checkStateNotNull(tasks).cancelAll();
+  }
+
+  @Teardown
+  public void teardown() {
+    if (tasks != null) {
+      tasks.shutdown();
+    }
+  }
+
+  @ProcessElement
+  public void process(
+      @Element String filePath,
+      @Timestamp Instant timestamp,
+      BoundedWindow window,
+      PaneInfo paneInfo,
+      OutputReceiver<String> output)
+      throws Exception {
+    numFilesRead.inc();
+    Callable<ReadResult> task = createReadTask(filePath, timestamp, window, paneInfo);
+    checkStateNotNull(tasks).submit(task, result -> outputResult(result, output));
+  }
+
+  @FinishBundle
+  public void finishBundle(FinishBundleContext context) throws Exception {
+    checkStateNotNull(tasks).awaitAll(result -> outputAtFinish(result, context));
+  }
+
+  private static void outputAtFinish(ReadResult result, FinishBundleContext context) {
+    count(result);
+    if (result.schemaJson != null) {
+      context.output(result.schemaJson, result.timestamp, result.window);
+    }
+  }
+
+  private static void outputResult(ReadResult result, OutputReceiver<String> output) {
+    count(result);
+    if (result.schemaJson != null) {
+      output.outputWindowedValue(
+          result.schemaJson,
+          result.timestamp,
+          Collections.singleton(result.window),
+          result.paneInfo);
+    }
+  }
+
+  private static void count(ReadResult result) {
+    if (result.schemaJson != null) {
+      numSchemasEmitted.inc();
+    }
+    if (result.footerError) {
+      numFooterReadErrors.inc();
+    }
+  }
+
+  private static Callable<ReadResult> createReadTask(
+      String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
+    return () -> {
+      FileFormat format;
+      try {
+        format = AddFiles.inferFormat(filePath);
+      } catch (AddFiles.UnknownFormatException e) {
+        return new ReadResult(null, false, timestamp, window, paneInfo);
+      }
+      if (!format.equals(FileFormat.PARQUET)) {
+        return new ReadResult(null, false, timestamp, window, paneInfo);
+      }
+      try {
+        ParquetMetadata footer = ParquetFooters.read(filePath);
+        return new ReadResult(
+            FileSchemas.canonicalJson(footer), false, timestamp, window, paneInfo);
+      } catch (Exception e) {
+        LOG.warn(
+            "Could not read the footer of {}; the file will not contribute to schema inference: {}",
+            filePath,
+            AddFiles.errorMessage(e));
+        return new ReadResult(null, true, timestamp, window, paneInfo);
+      }
+    };
+  }
+}