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/sink/FlinkFileWriterFactory.java

@@ -25,28 +25,19 @@
 import java.io.Serializable;
 import java.util.Map;
 import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.data.StringData;
 import org.apache.flink.table.types.logical.RowType;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.Table;
-import org.apache.iceberg.avro.Avro;
-import org.apache.iceberg.data.BaseFileWriterFactory;
+import org.apache.iceberg.data.RegistryBasedFileWriterFactory;
 import org.apache.iceberg.flink.FlinkSchemaUtil;
-import org.apache.iceberg.flink.data.FlinkAvroWriter;
-import org.apache.iceberg.flink.data.FlinkOrcWriter;
-import org.apache.iceberg.flink.data.FlinkParquetWriters;
-import org.apache.iceberg.orc.ORC;
-import org.apache.iceberg.parquet.Parquet;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 
-public class FlinkFileWriterFactory extends BaseFileWriterFactory<RowData> implements Serializable {
-  private RowType dataFlinkType;
-  private RowType equalityDeleteFlinkType;
-
-  private FlinkFileWriterFactory(
+public class FlinkFileWriterFactory extends RegistryBasedFileWriterFactory<RowData, RowType>
+    implements Serializable {
+  FlinkFileWriterFactory(
       Table table,
       FileFormat dataFileFormat,
       Schema dataSchema,
@@ -62,85 +53,30 @@ private FlinkFileWriterFactory(
     super(
         table,
         dataFileFormat,
+        RowData.class,
         dataSchema,
         dataSortOrder,
         deleteFileFormat,
         equalityFieldIds,
         equalityDeleteRowSchema,
         equalityDeleteSortOrder,
-        writeProperties);
-
-    this.dataFlinkType = dataFlinkType;
-    this.equalityDeleteFlinkType = equalityDeleteFlinkType;
-  }
-
-  static Builder builderFor(Table table) {
-    return new Builder(table);
-  }
-
-  @Override
-  protected void configureDataWrite(Avro.DataWriteBuilder builder) {
-    builder.createWriterFunc(ignore -> new FlinkAvroWriter(dataFlinkType()));
-  }
-
-  @Override
-  protected void configureEqualityDelete(Avro.DeleteWriteBuilder builder) {
-    builder.createWriterFunc(ignored -> new FlinkAvroWriter(equalityDeleteFlinkType()));
-  }
-
-  @Override
-  protected void configurePositionDelete(Avro.DeleteWriteBuilder builder) {}
-
-  @Override
-  protected void configureDataWrite(Parquet.DataWriteBuilder builder) {
-    builder.createWriterFunc(msgType -> FlinkParquetWriters.buildWriter(dataFlinkType(), msgType));
-  }
-
-  @Override
-  protected void configureEqualityDelete(Parquet.DeleteWriteBuilder builder) {
-    builder.createWriterFunc(
-        msgType -> FlinkParquetWriters.buildWriter(equalityDeleteFlinkType(), msgType));
-  }
-
-  @Override
-  protected void configurePositionDelete(Parquet.DeleteWriteBuilder builder) {
-    builder.transformPaths(path -> StringData.fromString(path.toString()));
-  }
-
-  @Override
-  protected void configureDataWrite(ORC.DataWriteBuilder builder) {
-    builder.createWriterFunc(
-        (iSchema, typDesc) -> FlinkOrcWriter.buildWriter(dataFlinkType(), iSchema));
-  }
-
-  @Override
-  protected void configureEqualityDelete(ORC.DeleteWriteBuilder builder) {
-    builder.createWriterFunc(
-        (iSchema, typDesc) -> FlinkOrcWriter.buildWriter(equalityDeleteFlinkType(), iSchema));
+        writeProperties,
+        dataFlinkType == null ? FlinkSchemaUtil.convert(dataSchema) : dataFlinkType,
+        equalityDeleteInputSchema(equalityDeleteFlinkType, equalityDeleteRowSchema));
   }
 
-  @Override
-  protected void configurePositionDelete(ORC.DeleteWriteBuilder builder) {
-    builder.transformPaths(path -> StringData.fromString(path.toString()));
-  }
-
-  private RowType dataFlinkType() {
-    if (dataFlinkType == null) {
-      Preconditions.checkNotNull(dataSchema(), "Data schema must not be null");
-      this.dataFlinkType = FlinkSchemaUtil.convert(dataSchema());
+  private static RowType equalityDeleteInputSchema(RowType rowType, Schema rowSchema) {
+    if (rowType != null) {
+      return rowType;
+    } else if (rowSchema != null) {
+      return FlinkSchemaUtil.convert(rowSchema);
+    } else {
+      return null;
     }
-
-    return dataFlinkType;
   }
 
-  private RowType equalityDeleteFlinkType() {
-    if (equalityDeleteFlinkType == null) {
-      Preconditions.checkNotNull(
-          equalityDeleteRowSchema(), "Equality delete schema must not be null");
-      this.equalityDeleteFlinkType = FlinkSchemaUtil.convert(equalityDeleteRowSchema());
-    }
-
-    return equalityDeleteFlinkType;
+  static Builder builderFor(Table table) {
+    return new Builder(table);
   }
 
   public static class Builder {