[lake/iceberg] Add IT for Iceberg files compaction for primary key table (#1658)

---------

Co-authored-by: maxcwang <[email protected]>

Author: beryllw

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/maintenance/IcebergRewriteDataFiles.java

@@ -26,6 +26,7 @@
 import org.apache.iceberg.ContentScanTask;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.Snapshot;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.data.IcebergGenericReader;
 import org.apache.iceberg.data.Record;
@@ -82,10 +83,15 @@ public IcebergRewriteDataFiles targetSizeInBytes(long targetSize) {
         return this;
     }
 
-    private List<CombinedScanTask> planRewriteFileGroups() throws IOException {
+    private List<CombinedScanTask> planRewriteFileGroups(long snapshotId) throws IOException {
         List<FileScanTask> fileScanTasks = new ArrayList<>();
         try (CloseableIterable<FileScanTask> tasks =
-                table.newScan().includeColumnStats().filter(filter).ignoreResiduals().planFiles()) {
+                table.newScan()
+                        .useSnapshot(snapshotId)
+                        .includeColumnStats()
+                        .filter(filter)
+                        .ignoreResiduals()
+                        .planFiles()) {
             tasks.forEach(fileScanTasks::add);
         }
 
@@ -137,7 +143,12 @@ private Comparator<FileScanTask> sortFileScanTask(int sortFiledId) {
     public RewriteDataFileResult execute() {
         try {
             // plan the file groups to be rewrite
-            List<CombinedScanTask> tasksToRewrite = planRewriteFileGroups();
+            Snapshot snapshot = table.currentSnapshot();
+            // if no snapshot, just return
+            if (snapshot == null) {
+                return null;
+            }
+            List<CombinedScanTask> tasksToRewrite = planRewriteFileGroups(snapshot.snapshotId());
             if (tasksToRewrite.isEmpty()) {
                 return null;
             }
@@ -152,7 +163,8 @@ public RewriteDataFileResult execute() {
                                 .collect(Collectors.toList()));
             }
             LOG.info("Finish rewriting files from {} to {}.", deletedDataFiles, addedDataFiles);
-            return new RewriteDataFileResult(deletedDataFiles, addedDataFiles);
+            return new RewriteDataFileResult(
+                    snapshot.snapshotId(), deletedDataFiles, addedDataFiles);
         } catch (Exception e) {
             throw new RuntimeException(
                     String.format("Fail to compact bucket %s of table %s.", bucket, table.name()),

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/maintenance/RewriteDataFileResult.java

@@ -28,10 +28,13 @@ public class RewriteDataFileResult implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
+    private final long snapshotId;
     private final List<DataFile> deletedDataFiles;
     private final List<DataFile> addedDataFiles;
 
-    public RewriteDataFileResult(List<DataFile> deletedDataFiles, List<DataFile> addedDataFiles) {
+    public RewriteDataFileResult(
+            long snapshotId, List<DataFile> deletedDataFiles, List<DataFile> addedDataFiles) {
+        this.snapshotId = snapshotId;
         this.deletedDataFiles = deletedDataFiles;
         this.addedDataFiles = addedDataFiles;
     }
@@ -44,10 +47,16 @@ public List<DataFile> addedDataFiles() {
         return addedDataFiles;
     }
 
+    public long snapshotId() {
+        return snapshotId;
+    }
+
     @Override
     public String toString() {
         return "RewriteDataFileResult{"
-                + "deletedDataFiles="
+                + "snapshotId="
+                + snapshotId
+                + ", deletedDataFiles="
                 + deletedDataFiles
                 + ", addedDataFiles="
                 + addedDataFiles

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/IcebergLakeCommitter.java

@@ -28,6 +28,7 @@
 
 import org.apache.iceberg.AppendFiles;
 import org.apache.iceberg.CatalogUtil;
+import org.apache.iceberg.ContentFile;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.DeleteFile;
 import org.apache.iceberg.RewriteFiles;
@@ -162,19 +163,28 @@ private Long commitRewrite(
             Map<String, String> snapshotProperties) {
         icebergTable.refresh();
         RewriteFiles rewriteFiles = icebergTable.newRewrite();
-        for (RewriteDataFileResult rewriteDataFileResult : rewriteDataFileResults) {
-            rewriteDataFileResult.addedDataFiles().forEach(rewriteFiles::addFile);
-            rewriteDataFileResult.deletedDataFiles().forEach(rewriteFiles::deleteFile);
-        }
         try {
+            if (rewriteDataFileResults.stream()
+                            .map(RewriteDataFileResult::snapshotId)
+                            .distinct()
+                            .count()
+                    > 1) {
+                throw new IllegalArgumentException(
+                        "Rewrite data file results must have same snapshot id.");
+            }
+            rewriteFiles.validateFromSnapshot(rewriteDataFileResults.get(0).snapshotId());
+            for (RewriteDataFileResult rewriteDataFileResult : rewriteDataFileResults) {
+                rewriteDataFileResult.addedDataFiles().forEach(rewriteFiles::addFile);
+                rewriteDataFileResult.deletedDataFiles().forEach(rewriteFiles::deleteFile);
+            }
             return commit(rewriteFiles, snapshotProperties);
         } catch (Exception e) {
             List<String> rewriteAddedDataFiles =
                     rewriteDataFileResults.stream()
                             .flatMap(
                                     rewriteDataFileResult ->
                                             rewriteDataFileResult.addedDataFiles().stream())
-                            .map(dataFile -> dataFile.path().toString())
+                            .map(ContentFile::location)
                             .collect(Collectors.toList());
             LOG.error(
                     "Failed to commit rewrite files to iceberg, delete rewrite added files {}.",

fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/maintenance/IcebergRewriteITCase.java

@@ -19,27 +19,47 @@
 package org.apache.fluss.lake.iceberg.maintenance;
 
 import org.apache.fluss.lake.iceberg.testutils.FlinkIcebergTieringTestBase;
+import org.apache.fluss.metadata.Schema;
 import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.types.DataTypes;
 
 import org.apache.flink.core.execution.JobClient;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.iceberg.data.Record;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.Iterator;
 import java.util.List;
 
 import static org.apache.fluss.testutils.DataTestUtils.row;
+import static org.assertj.core.api.Assertions.assertThat;
 
 /** Integration test for Iceberg compaction. */
 class IcebergRewriteITCase extends FlinkIcebergTieringTestBase {
     protected static final String DEFAULT_DB = "fluss";
 
     private static StreamExecutionEnvironment execEnv;
 
+    private static final Schema pkSchema =
+            Schema.newBuilder()
+                    .column("f_int", DataTypes.INT())
+                    .column("f_string", DataTypes.STRING())
+                    .primaryKey("f_int")
+                    .build();
+
+    private static final Schema logSchema =
+            Schema.newBuilder()
+                    .column("f_int", DataTypes.INT())
+                    .column("f_string", DataTypes.STRING())
+                    .build();
+
     @BeforeAll
     protected static void beforeAll() {
         FlinkIcebergTieringTestBase.beforeAll();
@@ -48,42 +68,148 @@ protected static void beforeAll() {
         execEnv.enableCheckpointing(1000);
     }
 
+    @Test
+    void testPkTableCompaction() throws Exception {
+        JobClient jobClient = buildTieringJob(execEnv);
+        try {
+            TablePath t1 = TablePath.of(DEFAULT_DB, "pk_table_1");
+            long t1Id = createPkTable(t1, 1, true, pkSchema);
+            TableBucket t1Bucket = new TableBucket(t1Id, 0);
+            List<InternalRow> flussRows = new ArrayList<>();
+
+            List<InternalRow> rows = Collections.singletonList(row(1, "v1"));
+            writeIcebergTableRecords(t1, t1Bucket, 1, false, rows);
+            flussRows.addAll(rows);
+
+            rows = Collections.singletonList(row(2, "v1"));
+            writeIcebergTableRecords(t1, t1Bucket, 2, false, rows);
+            flussRows.addAll(rows);
+
+            // add pos-delete
+            rows = Arrays.asList(row(3, "v1"), row(3, "v2"));
+            writeIcebergTableRecords(t1, t1Bucket, 5, false, rows);
+            // one UPDATE_BEFORE and one UPDATE_AFTER
+            checkFileStatusInIcebergTable(t1, 3, true);
+            flussRows.add(rows.get(1));
+
+            // trigger compaction
+            rows = Collections.singletonList(row(4, "v1"));
+            writeIcebergTableRecords(t1, t1Bucket, 6, false, rows);
+            checkFileStatusInIcebergTable(t1, 2, false);
+            flussRows.addAll(rows);
+
+            checkRecords(getIcebergRecords(t1), flussRows);
+        } finally {
+            jobClient.cancel().get();
+        }
+    }
+
+    private void checkRecords(List<Record> actualRows, List<InternalRow> expectedRows) {
+        // check records size
+        assertThat(actualRows.size()).isEqualTo(expectedRows.size());
+
+        // check records content
+        Iterator<Record> actualIterator =
+                actualRows.stream()
+                        .sorted(Comparator.comparingInt((Record r) -> (int) r.get(0)))
+                        .iterator();
+        Iterator<InternalRow> expectedIterator =
+                expectedRows.stream().sorted(Comparator.comparingInt(r -> r.getInt(0))).iterator();
+        while (actualIterator.hasNext() && expectedIterator.hasNext()) {
+            Record record = actualIterator.next();
+            InternalRow row = expectedIterator.next();
+            assertThat(record.get(0)).isEqualTo(row.getInt(0));
+            assertThat(record.get(1)).isEqualTo(row.getString(1).toString());
+        }
+        assertThat(actualIterator.hasNext()).isFalse();
+        assertThat(expectedIterator.hasNext()).isFalse();
+    }
+
+    @Test
+    void testPkTableCompactionWithConflict() throws Exception {
+        JobClient jobClient = buildTieringJob(execEnv);
+        try {
+            TablePath t1 = TablePath.of(DEFAULT_DB, "pk_table_2");
+            long t1Id = createPkTable(t1, 1, true, pkSchema);
+            TableBucket t1Bucket = new TableBucket(t1Id, 0);
+            List<InternalRow> flussRows = new ArrayList<>();
+
+            List<InternalRow> rows = Collections.singletonList(row(1, "v1"));
+            flussRows.addAll(writeIcebergTableRecords(t1, t1Bucket, 1, false, rows));
+            checkFileStatusInIcebergTable(t1, 1, false);
+
+            rows = Collections.singletonList(row(2, "v1"));
+            flussRows.addAll(writeIcebergTableRecords(t1, t1Bucket, 2, false, rows));
+
+            rows = Collections.singletonList(row(3, "v1"));
+            flussRows.addAll(writeIcebergTableRecords(t1, t1Bucket, 3, false, rows));
+
+            // add pos-delete and trigger compaction
+            rows = Arrays.asList(row(4, "v1"), row(4, "v2"));
+            flussRows.add(writeIcebergTableRecords(t1, t1Bucket, 6, false, rows).get(1));
+            // rewritten files should fail to commit due to conflict, add check here
+            checkRecords(getIcebergRecords(t1), flussRows);
+            // 4 data file and 1 delete file
+            checkFileStatusInIcebergTable(t1, 4, true);
+
+            // previous compaction conflicts won't prevent further compaction, and check iceberg
+            // records
+            rows = Collections.singletonList(row(5, "v1"));
+            flussRows.addAll(writeIcebergTableRecords(t1, t1Bucket, 7, false, rows));
+            checkRecords(getIcebergRecords(t1), flussRows);
+            checkFileStatusInIcebergTable(t1, 2, false);
+        } finally {
+            jobClient.cancel().get();
+        }
+    }
+
     @Test
     void testLogTableCompaction() throws Exception {
         JobClient jobClient = buildTieringJob(execEnv);
         try {
             TablePath t1 = TablePath.of(DEFAULT_DB, "log_table");
-            long t1Id = createLogTable(t1, true);
+            long t1Id = createLogTable(t1, 1, true, logSchema);
             TableBucket t1Bucket = new TableBucket(t1Id, 0);
 
             int i = 0;
             List<InternalRow> flussRows = new ArrayList<>();
-            flussRows.addAll(writeLogTableRecords(t1, t1Bucket, ++i));
+            flussRows.addAll(
+                    writeIcebergTableRecords(
+                            t1, t1Bucket, ++i, true, Collections.singletonList(row(1, "v1"))));
 
-            flussRows.addAll(writeLogTableRecords(t1, t1Bucket, ++i));
+            flussRows.addAll(
+                    writeIcebergTableRecords(
+                            t1, t1Bucket, ++i, true, Collections.singletonList(row(1, "v1"))));
 
-            flussRows.addAll(writeLogTableRecords(t1, t1Bucket, ++i));
-            checkFileCountInIcebergTable(t1, 3);
+            flussRows.addAll(
+                    writeIcebergTableRecords(
+                            t1, t1Bucket, ++i, true, Collections.singletonList(row(1, "v1"))));
+            checkFileStatusInIcebergTable(t1, 3, false);
 
             // Write should trigger compaction now since the current data file count is greater or
             // equal MIN_FILES_TO_COMPACT
-            flussRows.addAll(writeLogTableRecords(t1, t1Bucket, ++i));
+            flussRows.addAll(
+                    writeIcebergTableRecords(
+                            t1, t1Bucket, ++i, true, Collections.singletonList(row(1, "v1"))));
             // Should only have two files now, one file it for newly written, one file is for target
             // compacted file
-            checkFileCountInIcebergTable(t1, 2);
+            checkFileStatusInIcebergTable(t1, 2, false);
 
             // check data in iceberg to make sure compaction won't lose data or duplicate data
-            checkDataInIcebergAppendOnlyTable(t1, flussRows, 0);
+            checkRecords(getIcebergRecords(t1), flussRows);
         } finally {
             jobClient.cancel().get();
         }
     }
 
-    private List<InternalRow> writeLogTableRecords(
-            TablePath tablePath, TableBucket tableBucket, long expectedLogEndOffset)
+    private List<InternalRow> writeIcebergTableRecords(
+            TablePath tablePath,
+            TableBucket tableBucket,
+            long expectedLogEndOffset,
+            boolean append,
+            List<InternalRow> rows)
             throws Exception {
-        List<InternalRow> rows = Arrays.asList(row(1, "v1"));
-        writeRows(tablePath, rows, true);
+        writeRows(tablePath, rows, append);
         assertReplicaStatus(tableBucket, expectedLogEndOffset);
         return rows;
     }