Support write mode for Parquet and Avro using table property

core/src/main/java/org/apache/iceberg/TableProperties.java

@@ -56,9 +56,15 @@ private TableProperties() {}
   public static final String PARQUET_COMPRESSION = "write.parquet.compression-codec";
   public static final String PARQUET_COMPRESSION_DEFAULT = "gzip";
 
+  public static final String PARQUET_WRITE_MODE = "write.parquet.write-mode";
+  public static final String PARQUET_WRITE_MODE_DEFAULT = "overwrite";
+
   public static final String AVRO_COMPRESSION = "write.avro.compression-codec";
   public static final String AVRO_COMPRESSION_DEFAULT = "gzip";
 
+  public static final String AVRO_WRITE_MODE = "write.avro.write-mode";
+  public static final String AVRO_WRITE_MODE_DEFAULT = "overwrite";
+
   public static final String SPLIT_SIZE = "read.split.target-size";
   public static final long SPLIT_SIZE_DEFAULT = 134217728; // 128 MB
 

core/src/main/java/org/apache/iceberg/avro/Avro.java

@@ -34,12 +34,15 @@
 import org.apache.avro.io.DatumWriter;
 import org.apache.avro.specific.SpecificData;
 import org.apache.iceberg.SchemaParser;
+import org.apache.iceberg.Table;
 import org.apache.iceberg.io.FileAppender;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.OutputFile;
 
 import static org.apache.iceberg.TableProperties.AVRO_COMPRESSION;
 import static org.apache.iceberg.TableProperties.AVRO_COMPRESSION_DEFAULT;
+import static org.apache.iceberg.TableProperties.AVRO_WRITE_MODE;
+import static org.apache.iceberg.TableProperties.AVRO_WRITE_MODE_DEFAULT;
 
 public class Avro {
   private Avro() {
@@ -73,6 +76,11 @@ public CodecFactory get() {
     DEFAULT_MODEL.addLogicalTypeConversion(new UUIDConversion());
   }
 
+  public enum WriteMode {
+    CREATE,
+    OVERWRITE
+  }
+
   public static WriteBuilder write(OutputFile file) {
     return new WriteBuilder(file);
   }
@@ -89,6 +97,12 @@ private WriteBuilder(OutputFile file) {
       this.file = file;
     }
 
+    public WriteBuilder forTable(Table table) {
+      schema(table.schema());
+      setAll(table.properties());
+      return this;
+    }
+
     public WriteBuilder schema(org.apache.iceberg.Schema newSchema) {
       this.schema = newSchema;
       return this;
@@ -133,6 +147,15 @@ private CodecFactory codec() {
       }
     }
 
+    private WriteMode writeMode() {
+      String writeMode = config.getOrDefault(AVRO_WRITE_MODE, AVRO_WRITE_MODE_DEFAULT);
+      try {
+        return WriteMode.valueOf(writeMode.toUpperCase(Locale.ENGLISH));
+      } catch (IllegalArgumentException e) {
+        throw new IllegalArgumentException("Unsupported write mode: " + writeMode);
+      }
+    }
+
     public <D> FileAppender<D> build() throws IOException {
       Preconditions.checkNotNull(schema, "Schema is required");
       Preconditions.checkNotNull(name, "Table name is required and cannot be null");
@@ -141,7 +164,7 @@ public <D> FileAppender<D> build() throws IOException {
       meta("iceberg.schema", SchemaParser.toJson(schema));
 
       return new AvroFileAppender<>(
-          AvroSchemaUtil.convert(schema, name), file, createWriterFunc, codec(), metadata);
+          AvroSchemaUtil.convert(schema, name), file, createWriterFunc, codec(), metadata, writeMode());
     }
   }
 

core/src/main/java/org/apache/iceberg/avro/AvroFileAppender.java

@@ -39,8 +39,9 @@ class AvroFileAppender<D> implements FileAppender<D> {
 
   AvroFileAppender(Schema schema, OutputFile file,
                    Function<Schema, DatumWriter<?>> createWriterFunc,
-                   CodecFactory codec, Map<String, String> metadata) throws IOException {
-    this.stream = file.create();
+                   CodecFactory codec, Map<String, String> metadata,
+                   Avro.WriteMode writeMode) throws IOException {
+    this.stream = writeMode == Avro.WriteMode.CREATE ? file.create() : file.createOrOverwrite();
     this.writer = newAvroWriter(schema, stream, createWriterFunc, codec, metadata);
   }
 
@@ -92,7 +93,6 @@ private static <D> DataFileWriter<D> newAvroWriter(
       writer.setMeta(entry.getKey(), entry.getValue());
     }
 
-    // TODO: support overwrite
     return writer.create(schema, stream);
   }
 }

core/src/test/java/org/apache/iceberg/avro/TestAvroWrite.java

@@ -0,0 +1,116 @@
+/*
+ * 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.avro;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import java.io.File;
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.avro.generic.GenericData;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.exceptions.AlreadyExistsException;
+import org.apache.iceberg.io.FileAppender;
+import org.apache.iceberg.types.Types;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.rules.TemporaryFolder;
+
+public class TestAvroWrite {
+
+  @Rule
+  public TemporaryFolder temp = new TemporaryFolder();
+
+  @Rule
+  public ExpectedException thrown = ExpectedException.none();
+
+  @Test
+  public void testWriteCreate() throws IOException {
+    Map<String, String> properties = ImmutableMap.of(TableProperties.AVRO_WRITE_MODE, "create");
+
+    File file = write(10, properties, null);
+
+    thrown.expect(AlreadyExistsException.class);
+    write(5, properties, file);
+  }
+
+  @Test
+  public void testWriteOverwrite() throws IOException {
+    Map<String, String> properties = ImmutableMap.of(TableProperties.AVRO_WRITE_MODE, "overwrite");
+
+    File file = write(10, properties, null);
+    Assert.assertEquals(10, Lists.newArrayList(Avro.read(Files.localInput(file))
+        .project(schema())
+        .build()
+    ).size());
+
+    write(5, properties, file);
+    Assert.assertEquals(5, Lists.newArrayList(Avro.read(Files.localInput(file))
+        .project(schema())
+        .build()
+    ).size());
+  }
+
+  @Test
+  public void testUnsupportedWriteMode() throws IOException {
+    thrown.expect(IllegalArgumentException.class);
+    write(10, ImmutableMap.of(TableProperties.AVRO_WRITE_MODE, "abc"), null);
+  }
+
+  private Schema schema() {
+    return new Schema(
+        Types.NestedField.required(0, "col_int", Types.IntegerType.get())
+    );
+  }
+
+  private File write(int recordsNumber, Map<String, String> properties, File file) throws IOException {
+    Schema schema = schema();
+    org.apache.avro.Schema avroSchema = AvroSchemaUtil.convert(schema, "table");
+
+    List<GenericData.Record> records = IntStream.rangeClosed(1, recordsNumber)
+        .mapToObj(index -> {
+          GenericData.Record record = new GenericData.Record(avroSchema);
+          record.put("col_int", index);
+          return record;
+        })
+        .collect(Collectors.toList());
+
+    File outputFile = file;
+    if (outputFile == null) {
+      outputFile = temp.newFile();
+      Assert.assertTrue("File should have been deleted", outputFile.delete());
+    }
+
+    try (FileAppender<GenericData.Record> appender = Avro.write(Files.localOutput(outputFile))
+        .schema(schema)
+        .setAll(properties)
+        .build()) {
+      appender.addAll(records);
+    }
+
+    return outputFile;
+  }
+}

parquet/src/main/java/org/apache/iceberg/parquet/Parquet.java

@@ -64,6 +64,8 @@
 import static org.apache.iceberg.TableProperties.PARQUET_PAGE_SIZE_BYTES_DEFAULT;
 import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES;
 import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT;
+import static org.apache.iceberg.TableProperties.PARQUET_WRITE_MODE;
+import static org.apache.iceberg.TableProperties.PARQUET_WRITE_MODE_DEFAULT;
 
 public class Parquet {
   private Parquet() {
@@ -159,6 +161,15 @@ private CompressionCodecName codec() {
       }
     }
 
+    private ParquetFileWriter.Mode writeMode() {
+      String writeMode = config.getOrDefault(PARQUET_WRITE_MODE, PARQUET_WRITE_MODE_DEFAULT);
+      try {
+        return ParquetFileWriter.Mode.valueOf(writeMode.toUpperCase(Locale.ENGLISH));
+      } catch (IllegalArgumentException e) {
+        throw new IllegalArgumentException("Unsupported write mode: " + writeMode);
+      }
+    }
+
     public <D> FileAppender<D> build() throws IOException {
       Preconditions.checkNotNull(schema, "Schema is required");
       Preconditions.checkNotNull(name, "Table name is required and cannot be null");
@@ -201,7 +212,7 @@ public <D> FileAppender<D> build() throws IOException {
 
         return new org.apache.iceberg.parquet.ParquetWriter<>(
             conf, file, schema, rowGroupSize, metadata, createWriterFunc, codec(),
-            parquetProperties, metricsConfig);
+            parquetProperties, metricsConfig, writeMode());
       } else {
         return new ParquetWriteAdapter<>(new ParquetWriteBuilder<D>(ParquetIO.file(file))
             .withWriterVersion(writerVersion)
@@ -210,7 +221,7 @@ conf, file, schema, rowGroupSize, metadata, createWriterFunc, codec(),
             .setKeyValueMetadata(metadata)
             .setWriteSupport(getWriteSupport(type))
             .withCompressionCodec(codec())
-            .withWriteMode(ParquetFileWriter.Mode.OVERWRITE) // TODO: support modes
+            .withWriteMode(writeMode())
             .withRowGroupSize(rowGroupSize)
             .withPageSize(pageSize)
             .withDictionaryPageSize(dictionaryPageSize)

parquet/src/main/java/org/apache/iceberg/parquet/ParquetWriter.java

@@ -83,7 +83,8 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
                 Function<MessageType, ParquetValueWriter<?>> createWriterFunc,
                 CompressionCodecName codec,
                 ParquetProperties properties,
-                MetricsConfig metricsConfig) {
+                MetricsConfig metricsConfig,
+                ParquetFileWriter.Mode writeMode) {
     this.output = output;
     this.targetRowGroupSize = rowGroupSize;
     this.props = properties;
@@ -95,7 +96,7 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
 
     try {
       this.writer = new ParquetFileWriter(ParquetIO.file(output, conf), parquetSchema,
-          ParquetFileWriter.Mode.OVERWRITE, rowGroupSize, 0);
+          writeMode, rowGroupSize, 0);
     } catch (IOException e) {
       throw new RuntimeIOException(e, "Failed to create Parquet file");
     }

parquet/src/test/java/org/apache/iceberg/parquet/BaseParquetWritingTest.java

@@ -45,23 +45,28 @@ public abstract class BaseParquetWritingTest {
   public TemporaryFolder temp = new TemporaryFolder();
 
   File writeRecords(Schema schema, GenericData.Record... records) throws IOException {
-    return writeRecords(schema, Collections.emptyMap(), null, records);
+    return writeRecords(schema, Collections.emptyMap(), null, null, records);
   }
 
-  File writeRecords(
-      Schema schema, Map<String, String> properties,
-      Function<MessageType, ParquetValueWriter<?>> createWriterFunc,
-      GenericData.Record... records) throws IOException {
-    File tmpFolder = temp.newFolder("parquet");
-    String filename = UUID.randomUUID().toString();
-    File file = new File(tmpFolder, FileFormat.PARQUET.addExtension(filename));
-    try (FileAppender<GenericData.Record> writer = Parquet.write(localOutput(file))
+  File writeRecords(Schema schema,
+                    Map<String, String> properties,
+                    Function<MessageType, ParquetValueWriter<?>> createWriterFunc,
+                    File file,
+                    GenericData.Record... records) throws IOException {
+    File outputFile = file;
+    if (outputFile == null) {
+      File tmpFolder = temp.newFolder("parquet");
+      String filename = UUID.randomUUID().toString();
+      outputFile = new File(tmpFolder, FileFormat.PARQUET.addExtension(filename));
+    }
+
+    try (FileAppender<GenericData.Record> writer = Parquet.write(localOutput(outputFile))
         .schema(schema)
         .setAll(properties)
         .createWriterFunc(createWriterFunc)
         .build()) {
       writer.addAll(Lists.newArrayList(records));
     }
-    return file;
+    return outputFile;
   }
 }