[core] Support rollback to tag from snapshot id

Author: JingsongLi

paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java

@@ -85,8 +85,6 @@ abstract class AbstractFileStoreTable implements FileStoreTable {
 
     private static final long serialVersionUID = 1L;
 
-    private static final String WATERMARK_PREFIX = "watermark-";
-
     protected final FileIO fileIO;
     protected final Path path;
     protected final TableSchema tableSchema;
@@ -502,14 +500,22 @@ public void rollbackTo(long snapshotId) {
         try {
             snapshotManager.rollback(Instant.snapshot(snapshotId));
         } catch (UnsupportedOperationException e) {
-            Snapshot snapshot;
             try {
-                snapshot = snapshotManager.tryGetSnapshot(snapshotId);
+                Snapshot snapshot = snapshotManager.tryGetSnapshot(snapshotId);
+                rollbackHelper().cleanLargerThan(snapshot);
             } catch (FileNotFoundException ex) {
+                // try to get snapshot from tag
+                TagManager tagManager = tagManager();
+                SortedMap<Snapshot, List<String>> tags = tagManager.tags();
+                for (Map.Entry<Snapshot, List<String>> entry : tags.entrySet()) {
+                    if (entry.getKey().id() == snapshotId) {
+                        rollbackTo(entry.getValue().get(0));
+                        return;
+                    }
+                }
                 throw new IllegalArgumentException(
                         String.format("Rollback snapshot '%s' doesn't exist.", snapshotId), ex);
             }
-            rollbackHelper().cleanLargerThan(snapshot);
         }
     }
 

paimon-core/src/test/java/org/apache/paimon/table/SimpleTableTestBase.java

@@ -931,6 +931,35 @@ public void testRollbackToTag(boolean expire) throws Exception {
         assertRollbackTo(table, singletonList(1L), 1, 1, singletonList("test1"));
     }
 
+    @Test
+    public void testRollbackToTagFromSnapshotId() throws Exception {
+        int commitTimes = ThreadLocalRandom.current().nextInt(100) + 5;
+        FileStoreTable table = prepareRollbackTable(commitTimes);
+
+        table.createTag("test", 1);
+
+        // expire snapshots
+        Options options = new Options();
+        options.set(CoreOptions.SNAPSHOT_NUM_RETAINED_MIN, 5);
+        options.set(CoreOptions.SNAPSHOT_NUM_RETAINED_MAX, 5);
+        options.set(SNAPSHOT_EXPIRE_LIMIT, Integer.MAX_VALUE);
+        options.set(CHANGELOG_NUM_RETAINED_MIN, 5);
+        options.set(CHANGELOG_NUM_RETAINED_MAX, 5);
+        table.copy(options.toMap()).newCommit("").expireSnapshots();
+
+        table.rollbackTo(1);
+        ReadBuilder readBuilder = table.newReadBuilder();
+        List<String> result =
+                getResult(
+                        readBuilder.newRead(),
+                        readBuilder.newScan().plan().splits(),
+                        BATCH_ROW_TO_STRING);
+        assertThat(result)
+                .containsExactlyInAnyOrder("0|0|0|binary|varbinary|mapKey:mapVal|multiset");
+
+        assertRollbackTo(table, singletonList(1L), 1, 1, singletonList("test"));
+    }
+
     private FileStoreTable prepareRollbackTable(int commitTimes) throws Exception {
         FileStoreTable table = createFileStoreTable();
         return prepareRollbackTable(commitTimes, table);