[flink] Preserve complete undo recovery baselines

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/state/WriterState.java

@@ -32,12 +32,30 @@
  * writer is responsible for. During failover recovery, these offsets are used to generate undo logs
  * and rollback to the checkpoint state.
  *
+ * <p>V2 state is a complete but sparse recovery baseline. It stores only positive offsets; a live
+ * bucket absent from the map has an explicit baseline of zero. The state element itself, including
+ * an element with an empty map, proves that the checkpoint was produced with this completeness
+ * guarantee.
+ *
  * <p>This class uses {@link TableBucket} as the key to support partitioned tables. Each bucket is
  * uniquely identified by its table ID, partition ID (if applicable), and bucket ID.
  */
 public class WriterState implements Serializable {
 
     private static final long serialVersionUID = 1L;
+    private static final long NO_TABLE_ID = -1L;
+
+    /** The semantic guarantees carried by the serialized writer state. */
+    public enum StateFormat {
+        /** Legacy state whose bucket map may be sparse. */
+        V1_LEGACY,
+
+        /** Complete state for the table identified by {@link WriterState#getTableId()}. */
+        V2_COMPLETE
+    }
+
+    private final StateFormat stateFormat;
+    private final long tableId;
 
     /**
      * Map from TableBucket to the last successfully written changelog offset.
@@ -50,13 +68,22 @@ public class WriterState implements Serializable {
      */
     private final Map<TableBucket, Long> bucketOffsets;
 
+    /** Creates legacy V1 state whose bucket map must not be treated as complete. */
     public WriterState(Map<TableBucket, Long> bucketOffsets) {
+        this(StateFormat.V1_LEGACY, NO_TABLE_ID, bucketOffsets);
+    }
+
+    private WriterState(
+            StateFormat stateFormat, long tableId, Map<TableBucket, Long> bucketOffsets) {
         if (bucketOffsets == null) {
             throw new IllegalArgumentException("bucketOffsets must not be null");
         }
-        // Validate no null or negative offsets
-        // Note: offset=0 is valid and means the bucket was empty at checkpoint time
-        // (no records have been written to this bucket yet)
+        if (stateFormat == null) {
+            throw new IllegalArgumentException("stateFormat must not be null");
+        }
+        if (stateFormat == StateFormat.V2_COMPLETE && tableId < 0) {
+            throw new IllegalArgumentException("tableId must not be negative: " + tableId);
+        }
         for (Map.Entry<TableBucket, Long> entry : bucketOffsets.entrySet()) {
             if (entry.getKey() == null) {
                 throw new IllegalArgumentException("TableBucket in bucketOffsets must not be null");
@@ -65,14 +92,47 @@ public WriterState(Map<TableBucket, Long> bucketOffsets) {
                 throw new IllegalArgumentException(
                         "Invalid offset for bucket " + entry.getKey() + ": " + entry.getValue());
             }
+            if (stateFormat == StateFormat.V2_COMPLETE && entry.getKey().getTableId() != tableId) {
+                throw new IllegalArgumentException(
+                        "Bucket table ID "
+                                + entry.getKey().getTableId()
+                                + " does not match complete state table ID "
+                                + tableId);
+            }
         }
+        this.stateFormat = stateFormat;
+        this.tableId = tableId;
         this.bucketOffsets = Collections.unmodifiableMap(new HashMap<>(bucketOffsets));
     }
 
+    /** Creates complete V2 state for the given table. */
+    public static WriterState complete(long tableId, Map<TableBucket, Long> bucketOffsets) {
+        return new WriterState(StateFormat.V2_COMPLETE, tableId, bucketOffsets);
+    }
+
+    /** Returns the state format and its completeness guarantee. */
+    public StateFormat getStateFormat() {
+        return stateFormat;
+    }
+
+    /**
+     * Returns the table ID carried by complete V2 state.
+     *
+     * @throws IllegalStateException if this is legacy state
+     */
+    public long getTableId() {
+        if (stateFormat != StateFormat.V2_COMPLETE) {
+            throw new IllegalStateException("Legacy writer state does not carry a table ID");
+        }
+        return tableId;
+    }
+
+    /** Returns the immutable bucket-to-offset map. */
     public Map<TableBucket, Long> getBucketOffsets() {
         return bucketOffsets;
     }
 
+    /** Returns the offset for the bucket, or {@code null} when it is absent. */
     public Long getOffsetForBucket(TableBucket tableBucket) {
         return bucketOffsets.get(tableBucket);
     }
@@ -95,16 +155,25 @@ public boolean equals(Object o) {
             return false;
         }
         WriterState that = (WriterState) o;
-        return Objects.equals(bucketOffsets, that.bucketOffsets);
+        return tableId == that.tableId
+                && stateFormat == that.stateFormat
+                && Objects.equals(bucketOffsets, that.bucketOffsets);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(bucketOffsets);
+        return Objects.hash(stateFormat, tableId, bucketOffsets);
     }
 
     @Override
     public String toString() {
-        return "WriterState{" + "bucketOffsets=" + bucketOffsets + '}';
+        return "WriterState{"
+                + "stateFormat="
+                + stateFormat
+                + ", tableId="
+                + tableId
+                + ", bucketOffsets="
+                + bucketOffsets
+                + '}';
     }
 }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/state/WriterStateSerializer.java

@@ -35,9 +35,10 @@
  * <p>This serializer extends {@link TypeSerializer} for use with Flink's {@code
  * ListStateDescriptor} in Union List State.
  *
- * <p>Serialization format:
+ * <p>V1 serialization format:
  *
  * <ul>
+ *   <li>int: version (1)
  *   <li>int: number of bucket offsets
  *   <li>For each bucket offset:
  *       <ul>
@@ -49,15 +50,28 @@
  *       </ul>
  * </ul>
  *
- * <p>This serializer uses a stable binary format that supports both partitioned and non-partitioned
- * tables via {@link TableBucket}.
+ * <p>V2 serialization format:
+ *
+ * <ul>
+ *   <li>int: version (2)
+ *   <li>long: table ID
+ *   <li>int: number of bucket offsets
+ *   <li>For each bucket offset:
+ *       <ul>
+ *         <li>boolean: has partition ID
+ *         <li>long: partition ID (if has partition ID is true)
+ *         <li>int: bucket ID
+ *         <li>long: offset
+ *       </ul>
+ * </ul>
  */
 public class WriterStateSerializer extends TypeSerializer<WriterState> {
 
     private static final long serialVersionUID = 1L;
 
-    /** The current version of the serialization format. */
-    private static final int CURRENT_VERSION = 1;
+    private static final int V1_VERSION = 1;
+    private static final int V2_VERSION = 2;
+
     // -------------------------------------------------------------------------
     //  TypeSerializer methods
     // -------------------------------------------------------------------------
@@ -99,39 +113,46 @@ public int getLength() {
 
     @Override
     public void serialize(WriterState record, DataOutputView target) throws IOException {
-        target.writeInt(CURRENT_VERSION);
+        if (record.getStateFormat() == WriterState.StateFormat.V1_LEGACY) {
+            serializeV1(record, target);
+        } else if (record.getStateFormat() == WriterState.StateFormat.V2_COMPLETE) {
+            serializeV2(record, target);
+        } else {
+            throw new IOException("Unsupported writer state format: " + record.getStateFormat());
+        }
+    }
+
+    private void serializeV1(WriterState record, DataOutputView target) throws IOException {
+        target.writeInt(V1_VERSION);
         Map<TableBucket, Long> bucketOffsets = record.getBucketOffsets();
         target.writeInt(bucketOffsets.size());
         for (Map.Entry<TableBucket, Long> entry : bucketOffsets.entrySet()) {
             TableBucket bucket = entry.getKey();
             target.writeLong(bucket.getTableId());
-            target.writeBoolean(bucket.getPartitionId() != null);
-            if (bucket.getPartitionId() != null) {
-                target.writeLong(bucket.getPartitionId());
-            }
-            target.writeInt(bucket.getBucket());
-            target.writeLong(entry.getValue());
+            writeBucketOffset(target, bucket, entry.getValue());
+        }
+    }
+
+    private void serializeV2(WriterState record, DataOutputView target) throws IOException {
+        target.writeInt(V2_VERSION);
+        target.writeLong(record.getTableId());
+        Map<TableBucket, Long> bucketOffsets = record.getBucketOffsets();
+        target.writeInt(bucketOffsets.size());
+        for (Map.Entry<TableBucket, Long> entry : bucketOffsets.entrySet()) {
+            writeBucketOffset(target, entry.getKey(), entry.getValue());
         }
     }
 
     @Override
     public WriterState deserialize(DataInputView source) throws IOException {
         int version = source.readInt();
-        if (version != CURRENT_VERSION) {
-            throw new IOException(
-                    "Unsupported version: " + version + ". Expected version: " + CURRENT_VERSION);
-        }
-        int size = source.readInt();
-        Map<TableBucket, Long> bucketOffsets = new HashMap<>(size);
-        for (int i = 0; i < size; i++) {
-            long tableId = source.readLong();
-            boolean hasPartitionId = source.readBoolean();
-            Long partitionId = hasPartitionId ? source.readLong() : null;
-            int bucketId = source.readInt();
-            long offset = source.readLong();
-            bucketOffsets.put(new TableBucket(tableId, partitionId, bucketId), offset);
+        if (version == V1_VERSION) {
+            return deserializeV1(source);
+        } else if (version == V2_VERSION) {
+            return deserializeV2(source);
+        } else {
+            throw new IOException("Unsupported writer state version: " + version);
         }
-        return new WriterState(bucketOffsets);
     }
 
     @Override
@@ -143,28 +164,60 @@ public WriterState deserialize(WriterState reuse, DataInputView source) throws I
     @Override
     public void copy(DataInputView source, DataOutputView target) throws IOException {
         int version = source.readInt();
-        if (version != CURRENT_VERSION) {
-            throw new IOException(
-                    "Unsupported version: " + version + ". Expected version: " + CURRENT_VERSION);
+        WriterState state;
+        if (version == V1_VERSION) {
+            state = deserializeV1(source);
+        } else if (version == V2_VERSION) {
+            state = deserializeV2(source);
+        } else {
+            throw new IOException("Unsupported writer state version: " + version);
         }
-        target.writeInt(version);
-        // Copy bucket offsets size
+        serialize(state, target);
+    }
+
+    private WriterState deserializeV1(DataInputView source) throws IOException {
+        int size = source.readInt();
+        Map<TableBucket, Long> bucketOffsets = readBucketOffsets(source, size, null);
+        return new WriterState(bucketOffsets);
+    }
+
+    private WriterState deserializeV2(DataInputView source) throws IOException {
+        long tableId = source.readLong();
+        validateTableId(tableId);
         int size = source.readInt();
-        target.writeInt(size);
-        // Copy each bucket offset entry
+        Map<TableBucket, Long> bucketOffsets = readBucketOffsets(source, size, tableId);
+        return WriterState.complete(tableId, bucketOffsets);
+    }
+
+    private static Map<TableBucket, Long> readBucketOffsets(
+            DataInputView source, int size, Long tableIdFromStateHeader) throws IOException {
+        Map<TableBucket, Long> bucketOffsets = new HashMap<>(size);
         for (int i = 0; i < size; i++) {
-            // Copy table ID
-            target.writeLong(source.readLong());
-            // Copy has partition ID flag and partition ID if present
+            long tableId =
+                    tableIdFromStateHeader == null ? source.readLong() : tableIdFromStateHeader;
             boolean hasPartitionId = source.readBoolean();
-            target.writeBoolean(hasPartitionId);
-            if (hasPartitionId) {
-                target.writeLong(source.readLong());
-            }
-            // Copy bucket ID
-            target.writeInt(source.readInt());
-            // Copy offset
-            target.writeLong(source.readLong());
+            Long partitionId = hasPartitionId ? source.readLong() : null;
+            int bucketId = source.readInt();
+            long offset = source.readLong();
+            TableBucket bucket = new TableBucket(tableId, partitionId, bucketId);
+            bucketOffsets.put(bucket, offset);
+        }
+        return bucketOffsets;
+    }
+
+    private static void writeBucketOffset(DataOutputView target, TableBucket bucket, long offset)
+            throws IOException {
+        target.writeBoolean(bucket.getPartitionId() != null);
+        if (bucket.getPartitionId() != null) {
+            target.writeLong(bucket.getPartitionId());
+        }
+        target.writeInt(bucket.getBucket());
+        target.writeLong(offset);
+    }
+
+    private static void validateTableId(long tableId) throws IOException {
+        if (tableId < 0) {
+            throw new IOException("Invalid complete writer state table ID: " + tableId);
         }
     }