Flink: Add disk-based insertedRowMap to resolve the OOM while ingesting.

build.gradle

@@ -216,6 +216,7 @@ project(':iceberg-core') {
     implementation "com.fasterxml.jackson.core:jackson-core"
     implementation "com.github.ben-manes.caffeine:caffeine"
     implementation "org.roaringbitmap:RoaringBitmap"
+    implementation "org.rocksdb:rocksdbjni"
     compileOnly("org.apache.hadoop:hadoop-client") {
       exclude group: 'org.apache.avro', module: 'avro'
       exclude group: 'org.slf4j', module: 'slf4j-log4j12'

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

@@ -238,6 +238,8 @@ public void write(Object s, Encoder encoder) throws IOException {
         encoder.writeString((Utf8) s);
       } else if (s instanceof String) {
         encoder.writeString(new Utf8((String) s));
+      } else if (s instanceof CharSequence) {
+        encoder.writeString((CharSequence) s);
       } else if (s == null) {
         throw new IllegalArgumentException("Cannot write null to required string column");
       } else {

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 rowOffset = 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);
+    rowOffset += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long rowOffset() {
+    return rowOffset;
   }
 
   /**
@@ -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 rowOffset = 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);
+    rowOffset += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long rowOffset() {
+    return rowOffset;
   }
 
   /**
@@ -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,145 @@
+/*
+ * 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 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.Function;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.util.StructLikeMapUtil;
+import org.apache.iceberg.util.StructProjection;
+
+public class BaseEqualityDeltaWriter<T> implements EqualityDeltaWriter<T> {
+
+  private final ThreadLocal<PositionDelete<T>> posDelete = ThreadLocal.withInitial(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, StructLike> insertedRowMap;
+
+  private final Function<T, StructLike> asStructLike;
+  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,
+      Map<String, String> tableProperties,
+      Function<T, StructLike> asStructLike) {
+    this.dataWriter = dataWriter;
+    this.equalityWriter = equalityWriter;
+    this.positionWriter = positionWriter;
+
+    this.insertedRowMap = StructLikeMapUtil.load(
+        deleteSchema.asStruct(),
+        PathOffset.schema().asStruct(),
+        tableProperties);
+    this.asStructLike = asStructLike;
+
+    StructProjection projection = StructProjection.create(schema, deleteSchema);
+    this.keyRefFunc = row -> projection.wrap(asStructLike.apply(row));
+    this.keyCopyFunc = row -> StructCopy.copy(keyRefFunc.apply(row));
+  }
+
+  private PositionDelete<T> wrap(StructLike structLike) {
+    CharSequence path = structLike.get(0, CharSequence.class);
+    long offset = structLike.get(1, Long.class);
+    return posDelete.get().set(path, offset, null);
+  }
+
+  @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.
+    StructLike previous = insertedRowMap.put(copiedKey, pathOffset);
+    if (previous != null) {
+      positionWriter.write(wrap(previous), spec, partition);
+    }
+  }
+
+  /**
+   * Retire the old key & position from insertedRowMap cache to position delete file.
+   */
+  private boolean retireOldKey(StructLike key, PartitionSpec spec, StructLike partition) {
+    StructLike previous = insertedRowMap.remove(key);
+    if (previous != null) {
+      positionWriter.write(wrap(previous), 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(asStructLike.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;
@@ -125,7 +124,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);
@@ -141,7 +140,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;
       }
 
@@ -200,28 +199,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.rowOffset());
     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 rowOffset = 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);
+    rowOffset += 1;
+  }
+
+  @Override
+  public CharSequence location() {
+    return location;
+  }
+
+  @Override
+  public long rowOffset() {
+    return rowOffset;
   }
 
   /**