From 01762b1263f10c2ef6b9aae1397c519d6e3f432b Mon Sep 17 00:00:00 2001
From: Jason Fine <jasonf20@gmail.com>
Date: Mon, 11 Dec 2023 18:37:00 +0200
Subject: [PATCH] MergingSnapshotProducer: Support adding data files at a
 specific sequence number

---
 .../iceberg/MergingSnapshotProducer.java      | 42 +++++++++++++++----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/core/src/main/java/org/apache/iceberg/MergingSnapshotProducer.java b/core/src/main/java/org/apache/iceberg/MergingSnapshotProducer.java
index 6143a4ee9097..632958e242dd 100644
--- a/core/src/main/java/org/apache/iceberg/MergingSnapshotProducer.java
+++ b/core/src/main/java/org/apache/iceberg/MergingSnapshotProducer.java
@@ -33,6 +33,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 import org.apache.iceberg.events.CreateSnapshotEvent;
 import org.apache.iceberg.exceptions.RuntimeIOException;
 import org.apache.iceberg.exceptions.ValidationException;
@@ -83,7 +84,7 @@ abstract class MergingSnapshotProducer<ThisT> extends SnapshotProducer<ThisT> {
   private final boolean snapshotIdInheritanceEnabled;
 
   // update data
-  private final List<DataFile> newDataFiles = Lists.newArrayList();
+  private final List<FileHolder<DataFile>> newDataFiles = Lists.newArrayList();
   private Long newDataFilesDataSequenceNumber;
   private final Map<Integer, List<FileHolder<DeleteFile>>> newDeleteFilesBySpec = Maps.newHashMap();
   private final List<ManifestFile> appendManifests = Lists.newArrayList();
@@ -158,7 +159,8 @@ protected Expression rowFilter() {
   }
 
   protected List<DataFile> addedDataFiles() {
-    return ImmutableList.copyOf(newDataFiles);
+    return ImmutableList.copyOf(
+        newDataFiles.stream().map(FileHolder::file).collect(Collectors.toList()));
   }
 
   protected void failAnyDelete() {
@@ -228,10 +230,20 @@ protected boolean addsDeleteFiles() {
   /** Add a data file to the new snapshot. */
   protected void add(DataFile file) {
     Preconditions.checkNotNull(file, "Invalid data file: null");
-    setDataSpec(file);
-    addedFilesSummary.addedFile(dataSpec(), file);
+    addDataFile(new FileHolder<>(file));
+  }
+
+  /** Add a data file to the new snapshot. */
+  protected void add(DataFile file, long dataSequenceNumber) {
+    Preconditions.checkNotNull(file, "Invalid data file: null");
+    addDataFile(new FileHolder<>(file, dataSequenceNumber));
+  }
+
+  private void addDataFile(FileHolder<DataFile> dataFile) {
+    setDataSpec(dataFile.file());
+    addedFilesSummary.addedFile(dataSpec(), dataFile.file());
     hasNewDataFiles = true;
-    newDataFiles.add(file);
+    newDataFiles.add(dataFile);
   }
 
   /** Add a delete file to the new snapshot. */
@@ -960,9 +972,23 @@ private List<ManifestFile> newDataFilesAsManifests() {
         RollingManifestWriter<DataFile> writer = newRollingManifestWriter(dataSpec());
         try {
           if (newDataFilesDataSequenceNumber == null) {
-            newDataFiles.forEach(writer::add);
+            newDataFiles.forEach(
+                f -> {
+                  if (f.dataSequenceNumber() == null) {
+                    writer.add(f.file());
+                  } else {
+                    writer.add(f.file(), f.dataSequenceNumber);
+                  }
+                });
           } else {
-            newDataFiles.forEach(f -> writer.add(f, newDataFilesDataSequenceNumber));
+            newDataFiles.forEach(
+                f -> {
+                  if (f.dataSequenceNumber() == null) {
+                    writer.add(f.file(), newDataFilesDataSequenceNumber);
+                  } else {
+                    writer.add(f.file(), f.dataSequenceNumber);
+                  }
+                });
           }
         } finally {
           writer.close();
@@ -1132,7 +1158,7 @@ protected ManifestReader<DeleteFile> newManifestReader(ManifestFile manifest) {
     }
   }
 
-  private static class FileHolder<T extends ContentFile<?>>{
+  private static class FileHolder<T extends ContentFile<?>> {
     private final T file;
     private final Long dataSequenceNumber;