Core, Data, Flink: Moving Flink to use the new FormatModel API

core/src/main/java/org/apache/iceberg/formats/FormatModelRegistry.java

@@ -57,7 +57,8 @@ private FormatModelRegistry() {}
   private static final List<String> CLASSES_TO_REGISTER =
       ImmutableList.of(
           "org.apache.iceberg.data.GenericFormatModels",
-          "org.apache.iceberg.arrow.vectorized.ArrowFormatModels");
+          "org.apache.iceberg.arrow.vectorized.ArrowFormatModels",
+          "org.apache.iceberg.flink.data.FlinkFormatModels");
 
   // Format models indexed by file format and object model class
   private static final Map<Pair<FileFormat, Class<?>>, FormatModel<?, ?>> MODELS =

data/src/main/java/org/apache/iceberg/data/RegistryBasedFileWriterFactory.java

@@ -0,0 +1,182 @@
+/*
+ * 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.iceberg.data;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.io.UncheckedIOException;
+import java.util.Map;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.MetricsConfig;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.SortOrder;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.deletes.EqualityDeleteWriter;
+import org.apache.iceberg.deletes.PositionDeleteWriter;
+import org.apache.iceberg.encryption.EncryptedOutputFile;
+import org.apache.iceberg.encryption.EncryptionKeyMetadata;
+import org.apache.iceberg.formats.FileWriterBuilder;
+import org.apache.iceberg.formats.FormatModelRegistry;
+import org.apache.iceberg.io.DataWriter;
+import org.apache.iceberg.io.FileWriterFactory;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+
+/**
+ * A base writer factory to be extended by query engine integrations.
+ *
+ * @param <T> row type
+ */
+public abstract class RegistryBasedFileWriterFactory<T, S>
+    implements FileWriterFactory<T>, Serializable {
+  private final Table table;
+  private final FileFormat dataFileFormat;
+  private final Class<T> inputType;
+  private final Schema dataSchema;
+  private final SortOrder dataSortOrder;
+  private final FileFormat deleteFileFormat;
+  private final int[] equalityFieldIds;
+  private final Schema equalityDeleteRowSchema;
+  private final SortOrder equalityDeleteSortOrder;
+  private final Map<String, String> writerProperties;
+  private final S inputSchema;
+  private final S equalityDeleteInputSchema;
+
+  protected RegistryBasedFileWriterFactory(
+      Table table,
+      FileFormat dataFileFormat,
+      Class<T> inputType,
+      Schema dataSchema,
+      SortOrder dataSortOrder,
+      FileFormat deleteFileFormat,
+      int[] equalityFieldIds,
+      Schema equalityDeleteRowSchema,
+      SortOrder equalityDeleteSortOrder,
+      Map<String, String> writerProperties,
+      S inputSchema,
+      S equalityDeleteInputSchema) {
+    this.table = table;
+    this.dataFileFormat = dataFileFormat;
+    this.inputType = inputType;
+    this.dataSchema = dataSchema;
+    this.dataSortOrder = dataSortOrder;
+    this.deleteFileFormat = deleteFileFormat;
+    this.equalityFieldIds = equalityFieldIds;
+    this.equalityDeleteRowSchema = equalityDeleteRowSchema;
+    this.equalityDeleteSortOrder = equalityDeleteSortOrder;
+    this.writerProperties = writerProperties != null ? writerProperties : ImmutableMap.of();
+    this.inputSchema = inputSchema;
+    this.equalityDeleteInputSchema = equalityDeleteInputSchema;
+  }
+
+  protected S inputSchema() {
+    return inputSchema;
+  }
+
+  protected S equalityDeleteInputSchema() {
+    return equalityDeleteInputSchema;
+  }
+
+  @Override
+  public DataWriter<T> newDataWriter(
+      EncryptedOutputFile file, PartitionSpec spec, StructLike partition) {
+    Preconditions.checkNotNull(dataSchema, "Data schema must not be null");
+    EncryptionKeyMetadata keyMetadata = file.keyMetadata();
+    Map<String, String> properties = table != null ? table.properties() : ImmutableMap.of();
+    MetricsConfig metricsConfig =
+        table != null ? MetricsConfig.forTable(table) : MetricsConfig.getDefault();
+
+    try {
+      FileWriterBuilder<DataWriter<T>, S> builder =
+          FormatModelRegistry.dataWriteBuilder(dataFileFormat, inputType, file);
+      return builder
+          .schema(dataSchema)
+          .engineSchema(inputSchema())
+          .setAll(properties)
+          .setAll(writerProperties)
+          .metricsConfig(metricsConfig)
+          .spec(spec)
+          .partition(partition)
+          .keyMetadata(keyMetadata)
+          .sortOrder(dataSortOrder)
+          .overwrite()
+          .build();
+    } catch (IOException e) {
+      throw new UncheckedIOException("Failed to create new data writer", e);
+    }
+  }
+
+  @Override
+  public EqualityDeleteWriter<T> newEqualityDeleteWriter(
+      EncryptedOutputFile file, PartitionSpec spec, StructLike partition) {
+    Preconditions.checkNotNull(equalityDeleteRowSchema, "Equality delete schema must not be null");
+
+    EncryptionKeyMetadata keyMetadata = file.keyMetadata();
+    Map<String, String> properties = table != null ? table.properties() : ImmutableMap.of();
+    MetricsConfig metricsConfig =
+        table != null ? MetricsConfig.forTable(table) : MetricsConfig.getDefault();
+
+    try {
+      FileWriterBuilder<EqualityDeleteWriter<T>, S> builder =
+          FormatModelRegistry.equalityDeleteWriteBuilder(deleteFileFormat, inputType, file);
+      return builder
+          .setAll(properties)
+          .setAll(writerProperties)
+          .metricsConfig(metricsConfig)
+          .schema(equalityDeleteRowSchema)
+          .engineSchema(equalityDeleteInputSchema())
+          .equalityFieldIds(equalityFieldIds)
+          .spec(spec)
+          .partition(partition)
+          .keyMetadata(keyMetadata)
+          .sortOrder(equalityDeleteSortOrder)
+          .overwrite()
+          .build();
+    } catch (IOException e) {
+      throw new UncheckedIOException("Failed to create new equality delete writer", e);
+    }
+  }
+
+  @Override
+  public PositionDeleteWriter<T> newPositionDeleteWriter(
+      EncryptedOutputFile file, PartitionSpec spec, StructLike partition) {
+    EncryptionKeyMetadata keyMetadata = file.keyMetadata();
+    Map<String, String> properties = table != null ? table.properties() : ImmutableMap.of();
+    MetricsConfig metricsConfig =
+        table != null ? MetricsConfig.forPositionDelete(table) : MetricsConfig.forPositionDelete();
+
+    try {
+      FileWriterBuilder<PositionDeleteWriter<T>, ?> builder =
+          FormatModelRegistry.positionDeleteWriteBuilder(deleteFileFormat, file);
+      return builder
+          .setAll(properties)
+          .setAll(writerProperties)
+          .metricsConfig(metricsConfig)
+          .spec(spec)
+          .partition(partition)
+          .keyMetadata(keyMetadata)
+          .overwrite()
+          .build();
+    } catch (IOException e) {
+      throw new UncheckedIOException("Failed to create new position delete writer", e);
+    }
+  }
+}

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/data/FlinkFormatModels.java

@@ -0,0 +1,58 @@
+/*
+ * 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.iceberg.flink.data;
+
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.avro.AvroFormatModel;
+import org.apache.iceberg.formats.FormatModelRegistry;
+import org.apache.iceberg.orc.ORCFormatModel;
+import org.apache.iceberg.parquet.ParquetFormatModel;
+
+public class FlinkFormatModels {
+  public static void register() {
+    FormatModelRegistry.register(
+        ParquetFormatModel.create(
+            RowData.class,
+            RowType.class,
+            (icebergSchema, fileSchema, engineSchema) ->
+                FlinkParquetWriters.buildWriter(engineSchema, fileSchema),
+            (icebergSchema, fileSchema, engineSchema, idToConstant) ->
+                FlinkParquetReaders.buildReader(icebergSchema, fileSchema, idToConstant)));
+
+    FormatModelRegistry.register(
+        AvroFormatModel.create(
+            RowData.class,
+            RowType.class,
+            (icebergSchema, fileSchema, engineSchema) -> new FlinkAvroWriter(engineSchema),
+            (icebergSchema, fileSchema, engineSchema, idToConstant) ->
+                FlinkPlannedAvroReader.create(icebergSchema, idToConstant)));
+
+    FormatModelRegistry.register(
+        ORCFormatModel.create(
+            RowData.class,
+            RowType.class,
+            (icebergSchema, fileSchema, engineSchema) ->
+                FlinkOrcWriter.buildWriter(engineSchema, icebergSchema),
+            (icebergSchema, fileSchema, engineSchema, idToConstant) ->
+                new FlinkOrcReader(icebergSchema, fileSchema, idToConstant)));
+  }
+
+  private FlinkFormatModels() {}
+}

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/data/FlinkSchemaVisitor.java

@@ -19,6 +19,7 @@
 package org.apache.iceberg.flink.data;
 
 import java.util.List;
+import org.apache.flink.annotation.Internal;
 import org.apache.flink.table.types.logical.ArrayType;
 import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.MapType;
@@ -29,9 +30,10 @@
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 
+@Internal
 abstract class FlinkSchemaVisitor<T> {
 
-  static <T> T visit(RowType flinkType, Schema schema, FlinkSchemaVisitor<T> visitor) {
+  public static <T> T visit(RowType flinkType, Schema schema, FlinkSchemaVisitor<T> visitor) {
     return visit(flinkType, schema.asStruct(), visitor);
   }
 
@@ -94,24 +96,29 @@ private static <T> T visitRecord(
     List<LogicalType> fieldTypes = Lists.newArrayListWithExpectedSize(fieldSize);
     List<Types.NestedField> nestedFields = struct.fields();
 
-    for (int i = 0; i < fieldSize; i++) {
-      Types.NestedField iField = nestedFields.get(i);
-      int fieldIndex = rowType.getFieldIndex(iField.name());
-      Preconditions.checkArgument(
-          fieldIndex >= 0, "NestedField: %s is not found in flink RowType: %s", iField, rowType);
+    visitor.beforeStruct(struct.asStructType());
 
-      LogicalType fieldFlinkType = rowType.getTypeAt(fieldIndex);
+    try {
+      for (int i = 0; i < fieldSize; i++) {
+        Types.NestedField iField = nestedFields.get(i);
+        int fieldIndex = rowType.getFieldIndex(iField.name());
+        Preconditions.checkArgument(
+            fieldIndex >= 0, "NestedField: %s is not found in flink RowType: %s", iField, rowType);
 
-      fieldTypes.add(fieldFlinkType);
+        LogicalType fieldFlinkType = rowType.getTypeAt(fieldIndex);
+        fieldTypes.add(fieldFlinkType);
 
-      visitor.beforeField(iField);
-      try {
-        if (iField.type() != Types.UnknownType.get()) {
-          results.add(visit(fieldFlinkType, iField.type(), visitor));
+        visitor.beforeField(iField);
+        try {
+          if (iField.type() != Types.UnknownType.get()) {
+            results.add(visit(fieldFlinkType, iField.type(), visitor));
+          }
+        } finally {
+          visitor.afterField(iField);
         }
-      } finally {
-        visitor.afterField(iField);
       }
+    } finally {
+      visitor.afterStruct(struct.asStructType());
     }
 
     return visitor.record(struct, results, fieldTypes);
@@ -137,6 +144,10 @@ public void beforeField(Types.NestedField field) {}
 
   public void afterField(Types.NestedField field) {}
 
+  public void beforeStruct(Types.StructType type) {}
+
+  public void afterStruct(Types.StructType type) {}
+
   public void beforeListElement(Types.NestedField elementField) {
     beforeField(elementField);
   }