Data: Add GenericFileWriterFactory #9267

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

aokolnychyi merged 3 commits into apache:main from aokolnychyi:generic-file-writer-factory

Dec 12, 2023

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

-Original file line number
+Diff line change
@@ -0,0 +1,196 @@
+    /*
+     * 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 static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT;
+    import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT_DEFAULT;
+    import static org.apache.iceberg.TableProperties.DELETE_DEFAULT_FILE_FORMAT;
+    import java.util.Map;
+    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.avro.DataWriter;
+    import org.apache.iceberg.data.orc.GenericOrcWriter;
+    import org.apache.iceberg.data.parquet.GenericParquetWriter;
+    import org.apache.iceberg.orc.ORC;
+    import org.apache.iceberg.parquet.Parquet;
+    import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+    class GenericFileWriterFactory extends BaseFileWriterFactory<Record> {
+      GenericFileWriterFactory(
+          Table table,
+          FileFormat dataFileFormat,
+          Schema dataSchema,
+          SortOrder dataSortOrder,
+          FileFormat deleteFileFormat,
+          int[] equalityFieldIds,
+          Schema equalityDeleteRowSchema,
+          SortOrder equalityDeleteSortOrder,
+          Schema positionDeleteRowSchema) {
+        super(
+            table,
+            dataFileFormat,
+            dataSchema,
+            dataSortOrder,
+            deleteFileFormat,
+            equalityFieldIds,
+            equalityDeleteRowSchema,
+            equalityDeleteSortOrder,
+            positionDeleteRowSchema);
+      }
+      static Builder builderFor(Table table) {
+        return new Builder(table);
+      }
+      @Override
+      protected void configureDataWrite(Avro.DataWriteBuilder builder) {
+        builder.createWriterFunc(DataWriter::create);
+      }
+      @Override
+      protected void configureEqualityDelete(Avro.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(DataWriter::create);
+      }
+      @Override
+      protected void configurePositionDelete(Avro.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(DataWriter::create);
+      }
+      @Override
+      protected void configureDataWrite(Parquet.DataWriteBuilder builder) {
+        builder.createWriterFunc(GenericParquetWriter::buildWriter);
+      }
+      @Override
+      protected void configureEqualityDelete(Parquet.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(GenericParquetWriter::buildWriter);
+      }
+      @Override
+      protected void configurePositionDelete(Parquet.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(GenericParquetWriter::buildWriter);
+      }
+      @Override
+      protected void configureDataWrite(ORC.DataWriteBuilder builder) {
+        builder.createWriterFunc(GenericOrcWriter::buildWriter);
+      }
+      @Override
+      protected void configureEqualityDelete(ORC.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(GenericOrcWriter::buildWriter);
+      }
+      @Override
+      protected void configurePositionDelete(ORC.DeleteWriteBuilder builder) {
+        builder.createWriterFunc(GenericOrcWriter::buildWriter);
+      }
+      static class Builder {
+        private final Table table;
+        private FileFormat dataFileFormat;
+        private Schema dataSchema;
+        private SortOrder dataSortOrder;
+        private FileFormat deleteFileFormat;
+        private int[] equalityFieldIds;
+        private Schema equalityDeleteRowSchema;
+        private SortOrder equalityDeleteSortOrder;
+        private Schema positionDeleteRowSchema;
+        Builder(Table table) {
+          this.table = table;
+          this.dataSchema = table.schema();
+          Map<String, String> properties = table.properties();
+          String dataFileFormatName =
+              properties.getOrDefault(DEFAULT_FILE_FORMAT, DEFAULT_FILE_FORMAT_DEFAULT);
+          this.dataFileFormat = FileFormat.fromString(dataFileFormatName);
+          String deleteFileFormatName =
+              properties.getOrDefault(DELETE_DEFAULT_FILE_FORMAT, dataFileFormatName);
+          this.deleteFileFormat = FileFormat.fromString(deleteFileFormatName);
+        }
+        Builder dataFileFormat(FileFormat newDataFileFormat) {
+          this.dataFileFormat = newDataFileFormat;
+          return this;
+        }
+        Builder dataSchema(Schema newDataSchema) {
+          this.dataSchema = newDataSchema;
+          return this;
+        }
+        Builder dataSortOrder(SortOrder newDataSortOrder) {
+          this.dataSortOrder = newDataSortOrder;
+          return this;
+        }
+        Builder deleteFileFormat(FileFormat newDeleteFileFormat) {
+          this.deleteFileFormat = newDeleteFileFormat;
+          return this;
+        }
+        Builder equalityFieldIds(int[] newEqualityFieldIds) {
+          this.equalityFieldIds = newEqualityFieldIds;
+          return this;
+        }
+        Builder equalityDeleteRowSchema(Schema newEqualityDeleteRowSchema) {
+          this.equalityDeleteRowSchema = newEqualityDeleteRowSchema;
+          return this;
+        }
+        Builder equalityDeleteSortOrder(SortOrder newEqualityDeleteSortOrder) {
+          this.equalityDeleteSortOrder = newEqualityDeleteSortOrder;
+          return this;
+        }
+        Builder positionDeleteRowSchema(Schema newPositionDeleteRowSchema) {
+          this.positionDeleteRowSchema = newPositionDeleteRowSchema;
+          return this;
+        }
+        GenericFileWriterFactory build() {
+          boolean noEqualityDeleteConf = equalityFieldIds == null && equalityDeleteRowSchema == null;
+          boolean fullEqualityDeleteConf = equalityFieldIds != null && equalityDeleteRowSchema != null;
+          Preconditions.checkArgument(
+              noEqualityDeleteConf || fullEqualityDeleteConf,
+              "Equality field IDs and equality delete row schema must be set together");
+          return new GenericFileWriterFactory(
+              table,
+              dataFileFormat,
+              dataSchema,
+              dataSortOrder,
+              deleteFileFormat,
+              equalityFieldIds,
+              equalityDeleteRowSchema,
+              equalityDeleteSortOrder,
+              positionDeleteRowSchema);
+        }
+      }
+    }

data/src/test/java/org/apache/iceberg/data/FileHelpers.java

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -18,17 +18,11 @@
  
     */

    package org.apache.iceberg.data;

    import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT;

    import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT_DEFAULT;

    import java.io.Closeable;

    import java.io.IOException;

    import java.util.List;

    import java.util.Map;

    import org.apache.iceberg.DataFile;

    import org.apache.iceberg.DataFiles;

    import org.apache.iceberg.DeleteFile;

    import org.apache.iceberg.FileFormat;

    import org.apache.iceberg.Schema;

    import org.apache.iceberg.StructLike;

    import org.apache.iceberg.Table;

    @@ -38,8 +32,8 @@
  
    import org.apache.iceberg.encryption.EncryptedFiles;

    import org.apache.iceberg.encryption.EncryptedOutputFile;

    import org.apache.iceberg.encryption.EncryptionKeyMetadata;

    import org.apache.iceberg.io.FileAppender;

    import org.apache.iceberg.io.FileAppenderFactory;

    import org.apache.iceberg.io.DataWriter;

    import org.apache.iceberg.io.FileWriterFactory;

    import org.apache.iceberg.io.OutputFile;

    import org.apache.iceberg.types.Types;

    import org.apache.iceberg.util.CharSequenceSet;

    @@ -56,11 +50,10 @@ public static Pair<DeleteFile, CharSequenceSet> writeDeleteFile(
  
      public static Pair<DeleteFile, CharSequenceSet> writeDeleteFile(

          Table table, OutputFile out, StructLike partition, List<Pair<CharSequence, Long>> deletes)

          throws IOException {

        FileFormat format = defaultFormat(table.properties());

        FileAppenderFactory<Record> factory = new GenericAppenderFactory(table.schema(), table.spec());

        FileWriterFactory<Record> factory = GenericFileWriterFactory.builderFor(table).build();

        PositionDeleteWriter<Record> writer =

            factory.newPosDeleteWriter(encrypt(out), format, partition);

            factory.newPositionDeleteWriter(encrypt(out), table.spec(), partition);

        PositionDelete<Record> posDelete = PositionDelete.create();

        try (Closeable toClose = writer) {

          for (Pair<CharSequence, Long> delete : deletes) {

    @@ -84,15 +77,16 @@ public static DeleteFile writeDeleteFile(
  
          List<Record> deletes,

          Schema deleteRowSchema)

          throws IOException {

        FileFormat format = defaultFormat(table.properties());

        int[] equalityFieldIds =

            deleteRowSchema.columns().stream().mapToInt(Types.NestedField::fieldId).toArray();

        FileAppenderFactory<Record> factory =

            new GenericAppenderFactory(

                table.schema(), table.spec(), equalityFieldIds, deleteRowSchema, null);

        FileWriterFactory<Record> factory =

            GenericFileWriterFactory.builderFor(table)

                .equalityDeleteRowSchema(deleteRowSchema)

                .equalityFieldIds(equalityFieldIds)

                .build();

        EqualityDeleteWriter<Record> writer =

            factory.newEqDeleteWriter(encrypt(out), format, partition);

            factory.newEqualityDeleteWriter(encrypt(out), table.spec(), partition);

        try (Closeable toClose = writer) {

          writer.write(deletes);

        }

    @@ -102,56 +96,36 @@ public static DeleteFile writeDeleteFile(
  
      public static DataFile writeDataFile(Table table, OutputFile out, List<Record> rows)

          throws IOException {

        FileFormat format = defaultFormat(table.properties());

        GenericAppenderFactory factory = new GenericAppenderFactory(table.schema());

        FileWriterFactory<Record> factory = GenericFileWriterFactory.builderFor(table).build();

        FileAppender<Record> writer = factory.newAppender(out, format);

        DataWriter<Record> writer = factory.newDataWriter(encrypt(out), table.spec(), null);

        try (Closeable toClose = writer) {

          writer.addAll(rows);

          writer.write(rows);

        }

        return DataFiles.builder(table.spec())

            .withFormat(format)

            .withPath(out.location())

            .withFileSizeInBytes(writer.length())

            .withSplitOffsets(writer.splitOffsets())

            .withMetrics(writer.metrics())

            .build();

        return writer.toDataFile();

      }

      public static DataFile writeDataFile(

          Table table, OutputFile out, StructLike partition, List<Record> rows) throws IOException {

        FileFormat format = defaultFormat(table.properties());

        GenericAppenderFactory factory = new GenericAppenderFactory(table.schema(), table.spec());

        FileWriterFactory<Record> factory = GenericFileWriterFactory.builderFor(table).build();

        FileAppender<Record> writer = factory.newAppender(out, format);

        DataWriter<Record> writer = factory.newDataWriter(encrypt(out), table.spec(), partition);

        try (Closeable toClose = writer) {

          writer.addAll(rows);

          writer.write(rows);

        }

        return DataFiles.builder(table.spec())

            .withFormat(format)

            .withPath(out.location())

            .withPartition(partition)

            .withFileSizeInBytes(writer.length())

            .withSplitOffsets(writer.splitOffsets())

            .withMetrics(writer.metrics())

            .build();

        return writer.toDataFile();

      }

      public static DeleteFile writePosDeleteFile(

          Table table, OutputFile out, StructLike partition, List<PositionDelete<?>> deletes)

          throws IOException {

        FileFormat format = defaultFormat(table.properties());

        FileAppenderFactory<Record> factory =

            new GenericAppenderFactory(

                table.schema(),

                table.spec(),

                null, // Equality Fields

                null, // Equality Delete row schema

                table.schema()); // Position Delete row schema (will be wrapped)

        PositionDeleteWriter<?> writer = factory.newPosDeleteWriter(encrypt(out), format, partition);

        FileWriterFactory<Record> factory =

            GenericFileWriterFactory.builderFor(table).positionDeleteRowSchema(table.schema()).build();

        PositionDeleteWriter<?> writer =

            factory.newPositionDeleteWriter(encrypt(out), table.spec(), partition);

        try (Closeable toClose = writer) {

          for (PositionDelete delete : deletes) {

            writer.write(delete);

    @@ -164,9 +138,4 @@ public static DeleteFile writePosDeleteFile(
  
      private static EncryptedOutputFile encrypt(OutputFile out) {

        return EncryptedFiles.encryptedOutput(out, EncryptionKeyMetadata.EMPTY);

      }

      private static FileFormat defaultFormat(Map<String, String> properties) {

        String formatString = properties.getOrDefault(DEFAULT_FILE_FORMAT, DEFAULT_FILE_FORMAT_DEFAULT);

        return FileFormat.fromString(formatString);

      }

    }

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Data: Add GenericFileWriterFactory #9267

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!