[flink] Add restore test for streaming union read pk table (#1674)

caozhen1937 · luoyuxia · web-flow · commit 4c2ddc39087a · 2025-09-11T19:53:00.000+08:00
---------

Co-authored-by: luoyuxia &lt;luoyuxia@alumni.sjtu.edu.cn&gt;
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeRecordRecordEmitter.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeRecordRecordEmitter.java
@@ -57,8 +57,6 @@ public void emitRecord(
             }
 
             ScanRecord scanRecord = recordAndPos.record();
-            // todo: may need a state to mark snapshot phase is finished
-            // just like what we do for HybridSnapshotLogSplitState
             if (scanRecord.logOffset() >= 0) {
                 // record is with a valid offset, means it's in incremental phase,
                 // update the log offset
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitReaderGenerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitReaderGenerator.java
@@ -55,7 +55,10 @@ public void addSplit(SourceSplitBase split, Queue<SourceSplitBase> boundedSplits
         if (split instanceof LakeSnapshotSplit) {
             boundedSplits.add(split);
         } else if (split instanceof LakeSnapshotAndFlussLogSplit) {
-            boundedSplits.add(split);
+            // lake split not finished, add to it
+            if (!((LakeSnapshotAndFlussLogSplit) split).isLakeSplitFinished()) {
+                boundedSplits.add(split);
+            }
         } else {
             throw new UnsupportedOperationException(
                     String.format("The split type of %s is not supported.", split.getClass()));
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitSerializer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitSerializer.java
@@ -78,6 +78,7 @@ public void serialize(DataOutputSerializer out, SourceSplitBase split) throws IO
                             .orElse(LogSplit.NO_STOPPING_OFFSET));
             out.writeLong(lakeSnapshotAndFlussLogSplit.getRecordsToSkip());
             out.writeInt(lakeSnapshotAndFlussLogSplit.getCurrentLakeSplitIndex());
+            out.writeBoolean(lakeSnapshotAndFlussLogSplit.isLakeSplitFinished());
         } else {
             throw new UnsupportedOperationException(
                     "Unsupported split type: " + split.getClass().getName());
@@ -115,14 +116,16 @@ public SourceSplitBase deserialize(
             long stoppingOffset = input.readLong();
             long recordsToSkip = input.readLong();
             int splitIndex = input.readInt();
+            boolean isLakeSplitFinished = input.readBoolean();
             return new LakeSnapshotAndFlussLogSplit(
                     tableBucket,
                     partition,
                     lakeSplits,
                     startingOffset,
                     stoppingOffset,
                     recordsToSkip,
-                    splitIndex);
+                    splitIndex,
+                    isLakeSplitFinished);
         } else {
             throw new UnsupportedOperationException("Unsupported split kind: " + splitKind);
         }
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/split/LakeSnapshotAndFlussLogSplit.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/split/LakeSnapshotAndFlussLogSplit.java
@@ -40,8 +40,11 @@ public class LakeSnapshotAndFlussLogSplit extends SourceSplitBase {
      * lake split via this lake split index.
      */
     private int currentLakeSplitIndex;
-    /** The records to skip when reading a split. */
-    private long recordOffset;
+    /** The records to skip when reading a lake split. */
+    private long recordToSkip;
+
+    /** Whether the lake split has been finished. */
+    private boolean isLakeSplitFinished;
 
     private long startingOffset;
     private final long stoppingOffset;
@@ -52,7 +55,16 @@ public LakeSnapshotAndFlussLogSplit(
             @Nullable List<LakeSplit> snapshotSplits,
             long startingOffset,
             long stoppingOffset) {
-        this(tableBucket, partitionName, snapshotSplits, startingOffset, stoppingOffset, 0, 0);
+        this(
+                tableBucket,
+                partitionName,
+                snapshotSplits,
+                startingOffset,
+                stoppingOffset,
+                0,
+                0,
+                // if lake splits is null, no lake splits, also means LakeSplitFinished
+                snapshotSplits == null);
     }
 
     public LakeSnapshotAndFlussLogSplit(
@@ -62,17 +74,19 @@ public LakeSnapshotAndFlussLogSplit(
             long startingOffset,
             long stoppingOffset,
             long recordsToSkip,
-            int currentLakeSplitIndex) {
+            int currentLakeSplitIndex,
+            boolean isLakeSplitFinished) {
         super(tableBucket, partitionName);
         this.lakeSnapshotSplits = snapshotSplits;
         this.startingOffset = startingOffset;
         this.stoppingOffset = stoppingOffset;
-        this.recordOffset = recordsToSkip;
+        this.recordToSkip = recordsToSkip;
         this.currentLakeSplitIndex = currentLakeSplitIndex;
+        this.isLakeSplitFinished = isLakeSplitFinished;
     }
 
     public LakeSnapshotAndFlussLogSplit updateWithRecordsToSkip(long recordsToSkip) {
-        this.recordOffset = recordsToSkip;
+        this.recordToSkip = recordsToSkip;
         return this;
     }
 
@@ -86,14 +100,23 @@ public LakeSnapshotAndFlussLogSplit updateWithStartingOffset(long startingOffset
         return this;
     }
 
+    public LakeSnapshotAndFlussLogSplit updateWithLakeSplitFinished(boolean isLakeSplitFinished) {
+        this.isLakeSplitFinished = isLakeSplitFinished;
+        return this;
+    }
+
     public long getRecordsToSkip() {
-        return recordOffset;
+        return recordToSkip;
     }
 
     public long getStartingOffset() {
         return startingOffset;
     }
 
+    public boolean isLakeSplitFinished() {
+        return isLakeSplitFinished;
+    }
+
     public Optional<Long> getStoppingOffset() {
         return stoppingOffset >= 0 ? Optional.of(stoppingOffset) : Optional.empty();
     }
@@ -130,8 +153,8 @@ public String toString() {
         return "LakeSnapshotAndFlussLogSplit{"
                 + "lakeSnapshotSplits="
                 + lakeSnapshotSplits
-                + ", recordOffset="
-                + recordOffset
+                + ", recordToSkip="
+                + recordToSkip
                 + ", currentLakeSplitIndex="
                 + currentLakeSplitIndex
                 + ", startingOffset="
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/state/LakeSnapshotAndFlussLogSplitState.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/state/LakeSnapshotAndFlussLogSplitState.java
@@ -30,12 +30,15 @@ public class LakeSnapshotAndFlussLogSplitState extends SourceSplitState {
     private int currentLakeSplitIndex;
     private long nextLogOffset;
 
+    private boolean isLakeSplitFinished;
+
     public LakeSnapshotAndFlussLogSplitState(LakeSnapshotAndFlussLogSplit split) {
         super(split);
         this.recordsToSkip = split.getRecordsToSkip();
         this.split = split;
         this.currentLakeSplitIndex = split.getCurrentLakeSplitIndex();
         this.nextLogOffset = split.getStartingOffset();
+        this.isLakeSplitFinished = split.isLakeSplitFinished();
     }
 
     public void setRecordsToSkip(long recordsToSkip) {
@@ -47,13 +50,16 @@ public void setCurrentLakeSplitIndex(int currentLakeSplitIndex) {
     }
 
     public void setNextLogOffset(long nextOffset) {
+        // if set offset, means lake splits is finished
+        isLakeSplitFinished = true;
         this.nextLogOffset = nextOffset;
     }
 
     @Override
     public SourceSplitBase toSourceSplit() {
         return split.updateWithCurrentLakeSplitIndex(currentLakeSplitIndex)
                 .updateWithRecordsToSkip(recordsToSkip)
-                .updateWithStartingOffset(nextLogOffset);
+                .updateWithStartingOffset(nextLogOffset)
+                .updateWithLakeSplitFinished(isLakeSplitFinished);
     }
 }
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/committer/TieringCommitOperator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/committer/TieringCommitOperator.java
@@ -44,11 +44,13 @@
 
 import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
 import org.apache.flink.runtime.source.event.SourceEventWrapper;
-import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.streaming.api.graph.StreamConfig;
 import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.Output;
 import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
 
 import javax.annotation.Nullable;
 
@@ -109,34 +111,38 @@ public TieringCommitOperator(
         this.flussTableLakeSnapshotCommitter = new FlussTableLakeSnapshotCommitter(flussConf);
         this.collectedTableBucketWriteResults = new HashMap<>();
         this.flussConfig = flussConf;
+        this.operatorEventGateway =
+                parameters
+                        .getOperatorEventDispatcher()
+                        .getOperatorEventGateway(TieringSource.TIERING_SOURCE_OPERATOR_UID);
         this.setup(
                 parameters.getContainingTask(),
                 parameters.getStreamConfig(),
                 parameters.getOutput());
-        operatorEventGateway =
-                parameters
-                        .getOperatorEventDispatcher()
-                        .getOperatorEventGateway(TieringSource.TIERING_SOURCE_OPERATOR_UID);
-    }
-
-    @Override
-    public void open() {
-        flussTableLakeSnapshotCommitter.open();
-        connection = ConnectionFactory.createConnection(flussConfig);
-        admin = connection.getAdmin();
     }
 
     @Override
-    public void initializeState(StateInitializationContext context) throws Exception {
-        super.initializeState(context);
+    public void setup(
+            StreamTask<?, ?> containingTask,
+            StreamConfig config,
+            Output<StreamRecord<CommittableMessage<Committable>>> output) {
+        super.setup(containingTask, config, output);
         int attemptNumber = getRuntimeContext().getAttemptNumber();
         if (attemptNumber > 0) {
+            LOG.info("Send TieringRestoreEvent");
             // attempt number is greater than zero, the job must failover
             operatorEventGateway.sendEventToCoordinator(
                     new SourceEventWrapper(new TieringRestoreEvent()));
         }
     }
 
+    @Override
+    public void open() {
+        flussTableLakeSnapshotCommitter.open();
+        connection = ConnectionFactory.createConnection(flussConfig);
+        admin = connection.getAdmin();
+    }
+
     @Override
     public void processElement(StreamRecord<TableBucketWriteResult<WriteResult>> streamRecord)
             throws Exception {
diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/LakeSplitSerializerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/LakeSplitSerializerTest.java
@@ -90,7 +90,10 @@ void testSerializeAndDeserializeLakeSnapshotAndFlussLogSplit() throws IOExceptio
                         "2025-08-18",
                         Collections.singletonList(LAKE_SPLIT),
                         EARLIEST_OFFSET,
-                        STOPPING_OFFSET);
+                        STOPPING_OFFSET,
+                        2,
+                        1,
+                        true);
 
         DataOutputSerializer output = new DataOutputSerializer(STOPPING_OFFSET);
         serializer.serialize(output, originalSplit);
@@ -105,11 +108,14 @@ void testSerializeAndDeserializeLakeSnapshotAndFlussLogSplit() throws IOExceptio
         assertThat(deserializedSplit instanceof LakeSnapshotAndFlussLogSplit).isTrue();
         LakeSnapshotAndFlussLogSplit result = (LakeSnapshotAndFlussLogSplit) deserializedSplit;
 
-        assertThat(tableBucket).isEqualTo(result.getTableBucket());
-        assertThat("2025-08-18").isEqualTo(result.getPartitionName());
-        assertThat(Collections.singletonList(LAKE_SPLIT)).isEqualTo(result.getLakeSplits());
-        assertThat(EARLIEST_OFFSET).isEqualTo(result.getStartingOffset());
-        assertThat((long) STOPPING_OFFSET).isEqualTo(result.getStoppingOffset().get());
+        assertThat(result.getTableBucket()).isEqualTo(tableBucket);
+        assertThat(result.getPartitionName()).isEqualTo("2025-08-18");
+        assertThat(result.getLakeSplits()).isEqualTo(Collections.singletonList(LAKE_SPLIT));
+        assertThat(result.getStartingOffset()).isEqualTo(EARLIEST_OFFSET);
+        assertThat(result.getStoppingOffset().get()).isEqualTo(STOPPING_OFFSET);
+        assertThat(result.getCurrentLakeSplitIndex()).isEqualTo(1);
+        assertThat(result.getRecordsToSkip()).isEqualTo(2);
+        assertThat(result.isLakeSplitFinished()).isEqualTo(true);
     }
 
     @Test
@@ -149,8 +155,8 @@ public int getVersion() {
 
     private static class TestLakeSplit implements LakeSplit {
 
-        private int bucket;
-        private List<String> partition;
+        private final int bucket;
+        private final List<String> partition;
 
         public TestLakeSplit(int bucket, List<String> partition) {
             this.bucket = bucket;
diff --git a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/flink/FlinkUnionReadPrimaryKeyTableITCase.java b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/flink/FlinkUnionReadPrimaryKeyTableITCase.java

Original file line number	Diff line number	Diff line change
`@@ -57,8 +57,6 @@ public void emitRecord(`
`57`	`57`	`}`
`58`	`58`
`59`	`59`	`ScanRecord scanRecord = recordAndPos.record();`
`60`		`- // todo: may need a state to mark snapshot phase is finished`
`61`		`- // just like what we do for HybridSnapshotLogSplitState`
`62`	`60`	`if (scanRecord.logOffset() >= 0) {`
`63`	`61`	`// record is with a valid offset, means it's in incremental phase,`
`64`	`62`	`// update the log offset`
Original file line number	Diff line number	Diff line change
`@@ -30,12 +30,15 @@ public class LakeSnapshotAndFlussLogSplitState extends SourceSplitState {`
`30`	`30`	`private int currentLakeSplitIndex;`
`31`	`31`	`private long nextLogOffset;`
`32`	`32`
	`33`	`+ private boolean isLakeSplitFinished;`
	`34`	`+`
`33`	`35`	`public LakeSnapshotAndFlussLogSplitState(LakeSnapshotAndFlussLogSplit split) {`
`34`	`36`	`super(split);`
`35`	`37`	`this.recordsToSkip = split.getRecordsToSkip();`
`36`	`38`	`this.split = split;`
`37`	`39`	`this.currentLakeSplitIndex = split.getCurrentLakeSplitIndex();`
`38`	`40`	`this.nextLogOffset = split.getStartingOffset();`
	`41`	`+ this.isLakeSplitFinished = split.isLakeSplitFinished();`
`39`	`42`	`}`
`40`	`43`
`41`	`44`	`public void setRecordsToSkip(long recordsToSkip) {`
`@@ -47,13 +50,16 @@ public void setCurrentLakeSplitIndex(int currentLakeSplitIndex) {`
`47`	`50`	`}`
`48`	`51`
`49`	`52`	`public void setNextLogOffset(long nextOffset) {`
	`53`	`+ // if set offset, means lake splits is finished`
	`54`	`+ isLakeSplitFinished = true;`
`50`	`55`	`this.nextLogOffset = nextOffset;`
`51`	`56`	`}`
`52`	`57`
`53`	`58`	`@Override`
`54`	`59`	`public SourceSplitBase toSourceSplit() {`
`55`	`60`	`return split.updateWithCurrentLakeSplitIndex(currentLakeSplitIndex)`
`56`	`61`	`.updateWithRecordsToSkip(recordsToSkip)`
`57`		`- .updateWithStartingOffset(nextLogOffset);`
	`62`	`+ .updateWithStartingOffset(nextLogOffset)`
	`63`	`+ .updateWithLakeSplitFinished(isLakeSplitFinished);`
`58`	`64`	`}`
`59`	`65`	`}`