Flink Sink: Add WriteObserver plugin interface for per-record metadata

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergCommittable.java

@@ -20,7 +20,10 @@
 
 import java.io.Serializable;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
+import javax.annotation.Nullable;
 import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
 
 /**
@@ -36,12 +39,23 @@ class IcebergCommittable implements Serializable {
   private final String jobId;
   private final String operatorId;
   private final long checkpointId;
+  private final Map<String, String> observerMetadata;
 
   IcebergCommittable(byte[] manifest, String jobId, String operatorId, long checkpointId) {
+    this(manifest, jobId, operatorId, checkpointId, null);
+  }
+
+  IcebergCommittable(
+      byte[] manifest,
+      String jobId,
+      String operatorId,
+      long checkpointId,
+      @Nullable Map<String, String> observerMetadata) {
     this.manifest = manifest;
     this.jobId = jobId;
     this.operatorId = operatorId;
     this.checkpointId = checkpointId;
+    this.observerMetadata = observerMetadata != null ? observerMetadata : Collections.emptyMap();
   }
 
   byte[] manifest() {
@@ -60,6 +74,10 @@ Long checkpointId() {
     return checkpointId;
   }
 
+  Map<String, String> observerMetadata() {
+    return observerMetadata;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergCommittableSerializer.java

@@ -20,9 +20,11 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.Map;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
 import org.apache.flink.core.memory.DataInputDeserializer;
 import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 
 /**
  * This serializer is used for serializing the {@link IcebergCommittable} objects between the Writer
@@ -31,7 +33,7 @@
  * <p>In both cases only the respective part is serialized.
  */
 public class IcebergCommittableSerializer implements SimpleVersionedSerializer<IcebergCommittable> {
-  private static final int VERSION = 1;
+  private static final int VERSION = 2;
 
   @Override
   public int getVersion() {
@@ -47,22 +49,62 @@ public byte[] serialize(IcebergCommittable committable) throws IOException {
     view.writeLong(committable.checkpointId());
     view.writeInt(committable.manifest().length);
     view.write(committable.manifest());
+
+    Map<String, String> metadata = committable.observerMetadata();
+    if (metadata != null && !metadata.isEmpty()) {
+      view.writeBoolean(true);
+      view.writeInt(metadata.size());
+      for (Map.Entry<String, String> entry : metadata.entrySet()) {
+        view.writeUTF(entry.getKey());
+        view.writeUTF(entry.getValue());
+      }
+    } else {
+      view.writeBoolean(false);
+    }
+
     return out.toByteArray();
   }
 
   @Override
   public IcebergCommittable deserialize(int version, byte[] serialized) throws IOException {
     if (version == 1) {
-      DataInputDeserializer view = new DataInputDeserializer(serialized);
-      String jobId = view.readUTF();
-      String operatorId = view.readUTF();
-      long checkpointId = view.readLong();
-      int manifestLen = view.readInt();
-      byte[] manifestBuf;
-      manifestBuf = new byte[manifestLen];
-      view.read(manifestBuf);
-      return new IcebergCommittable(manifestBuf, jobId, operatorId, checkpointId);
+      return deserializeV1(serialized);
+    } else if (version == 2) {
+      return deserializeV2(serialized);
     }
     throw new IOException("Unrecognized version or corrupt state: " + version);
   }
+
+  private IcebergCommittable deserializeV1(byte[] serialized) throws IOException {
+    DataInputDeserializer view = new DataInputDeserializer(serialized);
+    String jobId = view.readUTF();
+    String operatorId = view.readUTF();
+    long checkpointId = view.readLong();
+    int manifestLen = view.readInt();
+    byte[] manifestBuf = new byte[manifestLen];
+    view.read(manifestBuf);
+    return new IcebergCommittable(manifestBuf, jobId, operatorId, checkpointId);
+  }
+
+  private IcebergCommittable deserializeV2(byte[] serialized) throws IOException {
+    DataInputDeserializer view = new DataInputDeserializer(serialized);
+    String jobId = view.readUTF();
+    String operatorId = view.readUTF();
+    long checkpointId = view.readLong();
+    int manifestLen = view.readInt();
+    byte[] manifestBuf = new byte[manifestLen];
+    view.read(manifestBuf);
+
+    Map<String, String> metadata = null;
+    boolean hasMetadata = view.readBoolean();
+    if (hasMetadata) {
+      int mapSize = view.readInt();
+      metadata = Maps.newHashMapWithExpectedSize(mapSize);
+      for (int i = 0; i < mapSize; i++) {
+        metadata.put(view.readUTF(), view.readUTF());
+      }
+    }
+
+    return new IcebergCommittable(manifestBuf, jobId, operatorId, checkpointId, metadata);
+  }
 }

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergCommitter.java

@@ -26,6 +26,7 @@
 import java.util.NavigableMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 import org.apache.flink.api.connector.sink2.Committer;
 import org.apache.flink.core.io.SimpleVersionedSerialization;
 import org.apache.iceberg.AppendFiles;
@@ -157,22 +158,31 @@ private void commitPendingRequests(
     long checkpointId = commitRequestMap.lastKey();
     List<ManifestFile> manifests = Lists.newArrayList();
     NavigableMap<Long, WriteResult> pendingResults = Maps.newTreeMap();
+
+    // Merge observer metadata from all committables
+    Map<String, String> mergedObserverMetadata = Maps.newHashMap();
     for (Map.Entry<Long, CommitRequest<IcebergCommittable>> e : commitRequestMap.entrySet()) {
-      if (Arrays.equals(EMPTY_MANIFEST_DATA, e.getValue().getCommittable().manifest())) {
+      IcebergCommittable committable = e.getValue().getCommittable();
+      if (!committable.observerMetadata().isEmpty()) {
+        mergedObserverMetadata.putAll(committable.observerMetadata());
+      }
+      if (Arrays.equals(EMPTY_MANIFEST_DATA, committable.manifest())) {
         pendingResults.put(e.getKey(), EMPTY_WRITE_RESULT);
       } else {
         DeltaManifests deltaManifests =
             SimpleVersionedSerialization.readVersionAndDeSerialize(
-                DeltaManifestsSerializer.INSTANCE, e.getValue().getCommittable().manifest());
+                DeltaManifestsSerializer.INSTANCE, committable.manifest());
         pendingResults.put(
             e.getKey(),
             FlinkManifestUtil.readCompletedFiles(deltaManifests, table.io(), table.specs()));
         manifests.addAll(deltaManifests.manifests());
       }
     }
 
+    Map<String, String> observerMetadata =
+        mergedObserverMetadata.isEmpty() ? null : mergedObserverMetadata;
     CommitSummary summary = new CommitSummary(pendingResults);
-    commitPendingResult(pendingResults, summary, newFlinkJobId, operatorId);
+    commitPendingResult(pendingResults, summary, newFlinkJobId, operatorId, observerMetadata);
     if (committerMetrics != null) {
       committerMetrics.updateCommitSummary(summary);
     }
@@ -195,14 +205,15 @@ private void commitPendingResult(
       NavigableMap<Long, WriteResult> pendingResults,
       CommitSummary summary,
       String newFlinkJobId,
-      String operatorId) {
+      String operatorId,
+      @Nullable Map<String, String> observerMetadata) {
     long totalFiles = summary.dataFilesCount() + summary.deleteFilesCount();
     continuousEmptyCheckpoints = totalFiles == 0 ? continuousEmptyCheckpoints + 1 : 0;
     if (totalFiles != 0 || continuousEmptyCheckpoints % maxContinuousEmptyCommits == 0) {
       if (replacePartitions) {
-        replacePartitions(pendingResults, summary, newFlinkJobId, operatorId);
+        replacePartitions(pendingResults, summary, newFlinkJobId, operatorId, observerMetadata);
       } else {
-        commitDeltaTxn(pendingResults, summary, newFlinkJobId, operatorId);
+        commitDeltaTxn(pendingResults, summary, newFlinkJobId, operatorId, observerMetadata);
       }
       continuousEmptyCheckpoints = 0;
     } else {
@@ -215,7 +226,8 @@ private void replacePartitions(
       NavigableMap<Long, WriteResult> pendingResults,
       CommitSummary summary,
       String newFlinkJobId,
-      String operatorId) {
+      String operatorId,
+      @Nullable Map<String, String> observerMetadata) {
     long checkpointId = pendingResults.lastKey();
     Preconditions.checkState(
         summary.deleteFilesCount() == 0, "Cannot overwrite partitions with delete files.");
@@ -229,14 +241,16 @@ private void replacePartitions(
     String description = "dynamic partition overwrite";
 
     logCommitSummary(summary, description);
-    commitOperation(dynamicOverwrite, description, newFlinkJobId, operatorId, checkpointId);
+    commitOperation(
+        dynamicOverwrite, description, newFlinkJobId, operatorId, checkpointId, observerMetadata);
   }
 
   private void commitDeltaTxn(
       NavigableMap<Long, WriteResult> pendingResults,
       CommitSummary summary,
       String newFlinkJobId,
-      String operatorId) {
+      String operatorId,
+      @Nullable Map<String, String> observerMetadata) {
     long checkpointId = pendingResults.lastKey();
     if (summary.deleteFilesCount() == 0) {
       // To be compatible with iceberg format V1.
@@ -250,7 +264,8 @@ private void commitDeltaTxn(
       String description = "append";
       logCommitSummary(summary, description);
       // fail all commits as really its only one
-      commitOperation(appendFiles, description, newFlinkJobId, operatorId, checkpointId);
+      commitOperation(
+          appendFiles, description, newFlinkJobId, operatorId, checkpointId, observerMetadata);
     } else {
       // To be compatible with iceberg format V2.
       for (Map.Entry<Long, WriteResult> e : pendingResults.entrySet()) {
@@ -274,7 +289,8 @@ private void commitDeltaTxn(
 
         String description = "rowDelta";
         logCommitSummary(summary, description);
-        commitOperation(rowDelta, description, newFlinkJobId, operatorId, e.getKey());
+        commitOperation(
+            rowDelta, description, newFlinkJobId, operatorId, e.getKey(), observerMetadata);
       }
     }
   }
@@ -284,9 +300,13 @@ private void commitOperation(
       String description,
       String newFlinkJobId,
       String operatorId,
-      long checkpointId) {
+      long checkpointId,
+      @Nullable Map<String, String> observerMetadata) {
 
     snapshotProperties.forEach(operation::set);
+    if (observerMetadata != null) {
+      observerMetadata.forEach(operation::set);
+    }
     // custom snapshot metadata properties will be overridden if they conflict with internal ones
     // used by the sink.
     operation.set(SinkUtil.MAX_COMMITTED_CHECKPOINT_ID, Long.toString(checkpointId));

flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java

@@ -29,6 +29,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.flink.annotation.Experimental;
 import org.apache.flink.api.common.SupportsConcurrentExecutionAttempts;
@@ -174,6 +175,8 @@ public class IcebergSink
   private final transient List<MaintenanceTaskBuilder<?>> maintenanceTasks;
   private final transient FlinkMaintenanceConfig flinkMaintenanceConfig;
 
+  @Nullable private final WriteObserver writeObserver;
+
   private IcebergSink(
       TableLoader tableLoader,
       Table table,
@@ -188,7 +191,8 @@ private IcebergSink(
       boolean overwriteMode,
       List<MaintenanceTaskBuilder<?>> maintenanceTasks,
       FlinkMaintenanceConfig flinkMaintenanceConfig,
-      Set<String> equalityFieldColumns) {
+      Set<String> equalityFieldColumns,
+      @Nullable WriteObserver writeObserver) {
     this.tableLoader = tableLoader;
     this.snapshotProperties = snapshotProperties;
     this.uidSuffix = uidSuffix;
@@ -212,6 +216,7 @@ private IcebergSink(
     this.maintenanceTasks = maintenanceTasks;
     this.flinkMaintenanceConfig = flinkMaintenanceConfig;
     this.equalityFieldColumns = equalityFieldColumns;
+    this.writeObserver = writeObserver;
   }
 
   @Override
@@ -232,7 +237,8 @@ public SinkWriter<RowData> createWriter(WriterInitContext context) {
         taskWriterFactory,
         metrics,
         context.getTaskInfo().getIndexOfThisSubtask(),
-        context.getTaskInfo().getAttemptNumber());
+        context.getTaskInfo().getAttemptNumber(),
+        writeObserver);
   }
 
   @Override
@@ -349,6 +355,7 @@ public static class Builder implements IcebergSinkBuilder<Builder> {
     private ReadableConfig readableConfig = new Configuration();
     private List<String> equalityFieldColumns = null;
     private final List<MaintenanceTaskBuilder<?>> maintenanceTasks = Lists.newArrayList();
+    @Nullable private WriteObserver writeObserver;
 
     private Builder() {}
 
@@ -716,6 +723,11 @@ public Builder toBranch(String branch) {
       return this;
     }
 
+    public Builder writeObserver(WriteObserver observer) {
+      this.writeObserver = observer;
+      return this;
+    }
+
     IcebergSink build() {
 
       Preconditions.checkArgument(
@@ -806,7 +818,8 @@ IcebergSink build() {
           overwriteMode,
           maintenanceTasks,
           flinkMaintenanceConfig,
-          equalityFieldColumnsSet);
+          equalityFieldColumnsSet,
+          writeObserver);
     }
 
     /**