Flink: Refactor to use the BaseEqualityDeltaWriter.

core/src/main/java/org/apache/iceberg/deletes/EqualityDeleteWriter.java

@@ -43,6 +43,7 @@ public class EqualityDeleteWriter<T> implements FileWriter<T, DeleteWriteResult>
   private final int[] equalityFieldIds;
   private final SortOrder sortOrder;
   private DeleteFile deleteFile = null;
+  private long recordCount = 0;
 
   public EqualityDeleteWriter(FileAppender<T> appender, FileFormat format, String location,
                               PartitionSpec spec, StructLike partition, EncryptionKeyMetadata keyMetadata,
@@ -60,6 +61,17 @@ public EqualityDeleteWriter(FileAppender<T> appender, FileFormat format, String
   @Override
   public void write(T row) {
     appender.add(row);
+    recordCount += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long recordCount() {
+    return recordCount;
   }
 
   /**
@@ -69,7 +81,7 @@ public void write(T row) {
    */
   @Deprecated
   public void deleteAll(Iterable<T> rows) {
-    appender.addAll(rows);
+    rows.forEach(this::write);
   }
 
   /**
@@ -79,7 +91,7 @@ public void deleteAll(Iterable<T> rows) {
    */
   @Deprecated
   public void delete(T row) {
-    appender.add(row);
+    write(row);
   }
 
   @Override

core/src/main/java/org/apache/iceberg/deletes/PositionDeleteWriter.java

@@ -43,6 +43,7 @@ public class PositionDeleteWriter<T> implements FileWriter<PositionDelete<T>, De
   private final PositionDelete<T> delete;
   private final CharSequenceSet referencedDataFiles;
   private DeleteFile deleteFile = null;
+  private long recordCount = 0;
 
   public PositionDeleteWriter(FileAppender<StructLike> appender, FileFormat format, String location,
                               PartitionSpec spec, StructLike partition, EncryptionKeyMetadata keyMetadata) {
@@ -60,6 +61,17 @@ public PositionDeleteWriter(FileAppender<StructLike> appender, FileFormat format
   public void write(PositionDelete<T> positionDelete) {
     referencedDataFiles.add(positionDelete.path());
     appender.add(positionDelete);
+    recordCount += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long recordCount() {
+    return recordCount;
   }
 
   /**
@@ -69,7 +81,7 @@ public void write(PositionDelete<T> positionDelete) {
    */
   @Deprecated
   public void delete(CharSequence path, long pos) {
-    delete(path, pos, null);
+    write(delete.set(path, pos, null));
   }
 
   /**
@@ -79,8 +91,7 @@ public void delete(CharSequence path, long pos) {
    */
   @Deprecated
   public void delete(CharSequence path, long pos, T row) {
-    referencedDataFiles.add(path);
-    appender.add(delete.set(path, pos, row));
+    write(delete.set(path, pos, row));
   }
 
   @Override

core/src/main/java/org/apache/iceberg/io/BaseEqualityDeltaWriter.java

@@ -0,0 +1,138 @@
+/*
+ * 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.io;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.function.Function;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.deletes.PositionDelete;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.util.StructLikeMap;
+import org.apache.iceberg.util.StructProjection;
+
+public class BaseEqualityDeltaWriter<T> implements EqualityDeltaWriter<T> {
+
+  private final PositionDelete<T> posDelete = PositionDelete.create();
+
+  private final PartitioningWriter<T, DataWriteResult> dataWriter;
+  private final PartitioningWriter<T, DeleteWriteResult> equalityWriter;
+  private final PartitioningWriter<PositionDelete<T>, DeleteWriteResult> positionWriter;
+
+  private final Map<StructLike, PathOffset> insertedRowMap;
+
+  private final Function<T, StructLike> asStructLikeKey;
+  private final Function<T, StructLike> keyRefFunc;
+  private final Function<T, StructLike> keyCopyFunc;
+
+  private boolean closed = false;
+
+  public BaseEqualityDeltaWriter(
+      PartitioningWriter<T, DataWriteResult> dataWriter,
+      PartitioningWriter<T, DeleteWriteResult> equalityWriter,
+      PartitioningWriter<PositionDelete<T>, DeleteWriteResult> positionWriter,
+      Schema schema,
+      Schema deleteSchema,
+      Function<T, StructLike> asStructLike,
+      Function<T, StructLike> asStructLikeKey) {
+    this.dataWriter = dataWriter;
+    this.equalityWriter = equalityWriter;
+    this.positionWriter = positionWriter;
+
+    this.insertedRowMap = StructLikeMap.create(deleteSchema.asStruct());
+    this.asStructLikeKey = asStructLikeKey;
+
+    StructProjection projection = StructProjection.create(schema, deleteSchema);
+    this.keyRefFunc = row -> projection.wrap(asStructLike.apply(row));
+    this.keyCopyFunc = row -> StructCopy.copy(keyRefFunc.apply(row));
+  }
+
+  @Override
+  public void insert(T row, PartitionSpec spec, StructLike partition) {
+    PathOffset pathOffset = dataWriter.write(row, spec, partition);
+
+    // Create a copied key from this row.
+    StructLike copiedKey = keyCopyFunc.apply(row);
+
+    // Adding a pos-delete to replace the old path-offset.
+    PathOffset previous = insertedRowMap.put(copiedKey, pathOffset);
+    if (previous != null) {
+      posDelete.set(previous.path(), previous.rowOffset(), null);
+      positionWriter.write(posDelete, spec, partition);
+    }
+  }
+
+  /**
+   * Retire the old key & position from insertedRowMap cache to position delete file.
+   */
+  private boolean retireOldKey(StructLike key, PartitionSpec spec, StructLike partition) {
+    PathOffset previous = insertedRowMap.remove(key);
+    if (previous != null) {
+      posDelete.set(previous.path(), previous.rowOffset(), null);
+      positionWriter.write(posDelete, spec, partition);
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  @Override
+  public void delete(T row, PartitionSpec spec, StructLike partition) {
+    if (!retireOldKey(keyRefFunc.apply(row), spec, partition)) {
+      equalityWriter.write(row, spec, partition);
+    }
+  }
+
+  @Override
+  public void deleteKey(T key, PartitionSpec spec, StructLike partition) {
+    if (!retireOldKey(asStructLikeKey.apply(key), spec, partition)) {
+      equalityWriter.write(key, spec, partition);
+    }
+  }
+
+  @Override
+  public WriteResult result() {
+    Preconditions.checkState(closed, "Cannot get result from unclosed writer.");
+
+    DataWriteResult dataWriteResult = dataWriter.result();
+    DeleteWriteResult positionWriteResult = positionWriter.result();
+    DeleteWriteResult equalityWriteResult = equalityWriter.result();
+
+    return WriteResult.builder()
+        .addDataFiles(dataWriteResult.dataFiles())
+        .addDeleteFiles(positionWriteResult.deleteFiles())
+        .addDeleteFiles(equalityWriteResult.deleteFiles())
+        .addReferencedDataFiles(positionWriteResult.referencedDataFiles())
+        .build();
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (!closed) {
+      dataWriter.close();
+      positionWriter.close();
+      equalityWriter.close();
+
+      this.closed = true;
+    }
+  }
+}

core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java

@@ -32,7 +32,6 @@
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.deletes.EqualityDeleteWriter;
 import org.apache.iceberg.encryption.EncryptedOutputFile;
-import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
@@ -130,7 +129,7 @@ public void write(T row) throws IOException {
       PathOffset previous = insertedRowMap.put(copiedKey, pathOffset);
       if (previous != null) {
         // TODO attach the previous row if has a positional-delete row schema in appender factory.
-        posDeleteWriter.delete(previous.path, previous.rowOffset, null);
+        posDeleteWriter.delete(previous.path(), previous.rowOffset(), null);
       }
 
       dataWriter.write(row);
@@ -146,7 +145,7 @@ private boolean internalPosDelete(StructLike key) {
 
       if (previous != null) {
         // TODO attach the previous row if has a positional-delete row schema in appender factory.
-        posDeleteWriter.delete(previous.path, previous.rowOffset, null);
+        posDeleteWriter.delete(previous.path(), previous.rowOffset(), null);
         return true;
       }
 
@@ -205,28 +204,6 @@ public void close() throws IOException {
     }
   }
 
-  private static class PathOffset {
-    private final CharSequence path;
-    private final long rowOffset;
-
-    private PathOffset(CharSequence path, long rowOffset) {
-      this.path = path;
-      this.rowOffset = rowOffset;
-    }
-
-    private static PathOffset of(CharSequence path, long rowOffset) {
-      return new PathOffset(path, rowOffset);
-    }
-
-    @Override
-    public String toString() {
-      return MoreObjects.toStringHelper(this)
-          .add("path", path)
-          .add("row_offset", rowOffset)
-          .toString();
-    }
-  }
-
   private abstract class BaseRollingWriter<W extends Closeable> implements Closeable {
     private static final int ROWS_DIVISOR = 1000;
     private final StructLike partitionKey;

core/src/main/java/org/apache/iceberg/io/ClusteredWriter.java

@@ -64,7 +64,7 @@ abstract class ClusteredWriter<T, R> implements PartitioningWriter<T, R> {
   protected abstract R aggregatedResult();
 
   @Override
-  public void write(T row, PartitionSpec spec, StructLike partition) {
+  public PathOffset write(T row, PartitionSpec spec, StructLike partition) {
     if (!spec.equals(currentSpec)) {
       if (currentSpec != null) {
         closeCurrentWriter();
@@ -85,7 +85,6 @@ public void write(T row, PartitionSpec spec, StructLike partition) {
       // copy the partition key as the key object may be reused
       this.currentPartition = StructCopy.copy(partition);
       this.currentWriter = newWriter(currentSpec, currentPartition);
-
     } else if (partition != currentPartition && partitionComparator.compare(partition, currentPartition) != 0) {
       closeCurrentWriter();
       completedPartitions.add(currentPartition);
@@ -100,7 +99,10 @@ public void write(T row, PartitionSpec spec, StructLike partition) {
       this.currentWriter = newWriter(currentSpec, currentPartition);
     }
 
+    PathOffset pathOffset = PathOffset.of(currentWriter.location(), currentWriter.recordCount());
     currentWriter.write(row);
+
+    return pathOffset;
   }
 
   @Override

core/src/main/java/org/apache/iceberg/io/DataWriter.java

@@ -39,6 +39,7 @@ public class DataWriter<T> implements FileWriter<T, DataWriteResult> {
   private final ByteBuffer keyMetadata;
   private final SortOrder sortOrder;
   private DataFile dataFile = null;
+  private long recordCount = 0;
 
   public DataWriter(FileAppender<T> appender, FileFormat format, String location,
                     PartitionSpec spec, StructLike partition, EncryptionKeyMetadata keyMetadata) {
@@ -59,6 +60,17 @@ public DataWriter(FileAppender<T> appender, FileFormat format, String location,
   @Override
   public void write(T row) {
     appender.add(row);
+    recordCount += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long recordCount() {
+    return recordCount;
   }
 
   /**

core/src/main/java/org/apache/iceberg/io/FanoutEqualityDeleteWriter.java

@@ -0,0 +1,62 @@
+/*
+ * 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.io;
+
+import java.util.List;
+import org.apache.iceberg.DeleteFile;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.CharSequenceSet;
+
+public class FanoutEqualityDeleteWriter<T> extends FanoutWriter<T, DeleteWriteResult> {
+
+  private final FileWriterFactory<T> writerFactory;
+  private final OutputFileFactory fileFactory;
+  private final FileIO io;
+  private final long targetFileSizeInBytes;
+
+  private final List<DeleteFile> deleteFiles;
+
+  public FanoutEqualityDeleteWriter(
+      FileWriterFactory<T> writerFactory, OutputFileFactory fileFactory,
+      FileIO io, long targetFileSizeInBytes) {
+    this.writerFactory = writerFactory;
+    this.fileFactory = fileFactory;
+    this.io = io;
+    this.targetFileSizeInBytes = targetFileSizeInBytes;
+    this.deleteFiles = Lists.newArrayList();
+  }
+
+  @Override
+  protected FileWriter<T, DeleteWriteResult> newWriter(PartitionSpec spec, StructLike partition) {
+    return new RollingEqualityDeleteWriter<>(writerFactory, fileFactory, io, targetFileSizeInBytes, spec, partition);
+  }
+
+  @Override
+  protected void addResult(DeleteWriteResult result) {
+    deleteFiles.addAll(result.deleteFiles());
+  }
+
+  @Override
+  protected DeleteWriteResult aggregatedResult() {
+    return new DeleteWriteResult(deleteFiles, CharSequenceSet.empty());
+  }
+}