From fe1900f6fbc59862a221ac0a4f5ae909d76714b6 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sat, 26 Sep 2026 18:35:10 +0800 Subject: [PATCH 1/4] [spark] Add rollback_to_as_latest procedure Spark's rollback drops every snapshot and tag created after the target, so a user who only wants to restore an earlier state loses the newer history. Flink already exposes rollback_to_as_latest, which rolls the target forward as a new latest snapshot non-destructively (core rollbackToAsLatest, #8450), but Spark has no equivalent. Add RollbackToAsLatestProcedure, registered as rollback_to_as_latest, taking exactly one of tag / snapshot_id. It mirrors the Flink procedure: resolve the target snapshot (creating a protection tag for the snapshot_id path), commit.rollbackToAsLatest(tag), and clean that tag up if the commit did not land. Returns previous / rolled_back / current snapshot ids. Tests: RollbackToAsLatestProcedureTest writes three snapshots, rolls back to snapshot 1 as latest and asserts snapshots 2 and 3 survive (non-destructive) and the data is snapshot 1's, then rolls forward to snapshot 3. Written with Claude Code; verification is mine. --- .../apache/paimon/spark/SparkProcedures.java | 2 + .../RollbackToAsLatestProcedure.java | 233 ++++++++++++++++++ .../RollbackToAsLatestProcedureTest.scala | 50 ++++ 3 files changed, 285 insertions(+) create mode 100644 paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java create mode 100644 paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java index f534773959f6..18bb32837533 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java @@ -67,6 +67,7 @@ import org.apache.paimon.spark.procedure.RevokePermissionProcedure; import org.apache.paimon.spark.procedure.RewriteFileIndexProcedure; import org.apache.paimon.spark.procedure.RollbackProcedure; +import org.apache.paimon.spark.procedure.RollbackToAsLatestProcedure; import org.apache.paimon.spark.procedure.RollbackToTimestampProcedure; import org.apache.paimon.spark.procedure.RollbackToWatermarkProcedure; import org.apache.paimon.spark.procedure.TriggerTagAutomaticCreationProcedure; @@ -98,6 +99,7 @@ private static Map> initProcedureBuilders() { ImmutableMap.Builder> procedureBuilders = ImmutableMap.builder(); procedureBuilders.put("rollback", RollbackProcedure::builder); + procedureBuilders.put("rollback_to_as_latest", RollbackToAsLatestProcedure::builder); procedureBuilders.put("rollback_to_timestamp", RollbackToTimestampProcedure::builder); procedureBuilders.put("rollback_to_watermark", RollbackToWatermarkProcedure::builder); procedureBuilders.put("purge_files", PurgeFilesProcedure::builder); diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java new file mode 100644 index 000000000000..11c21efd1cd7 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java @@ -0,0 +1,233 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.spark.procedure; + +import org.apache.paimon.FileStore; +import org.apache.paimon.Snapshot; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.sink.TableCommitImpl; +import org.apache.paimon.tag.Tag; +import org.apache.paimon.utils.Preconditions; +import org.apache.paimon.utils.SnapshotManager; +import org.apache.paimon.utils.StringUtils; +import org.apache.paimon.utils.TagManager; + +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.SortedMap; +import java.util.UUID; + +import static org.apache.spark.sql.types.DataTypes.LongType; +import static org.apache.spark.sql.types.DataTypes.StringType; + +/** + * Rollback to a snapshot or tag as the latest snapshot, without dropping the snapshots and tags + * created after it (unlike {@link RollbackProcedure}). Mirrors Flink's {@code + * rollback_to_as_latest} procedure. + */ +public class RollbackToAsLatestProcedure extends BaseProcedure { + + private static final String ROLLBACK_TO_AS_LATEST_TAG_PREFIX = "rollback-to-as-latest-"; + + private static final ProcedureParameter[] PARAMETERS = + new ProcedureParameter[] { + ProcedureParameter.required("table", StringType), + ProcedureParameter.optional("tag", StringType), + ProcedureParameter.optional("snapshot_id", LongType) + }; + + private static final StructType OUTPUT_TYPE = + new StructType( + new StructField[] { + new StructField( + "previous_snapshot_id", + DataTypes.LongType, + false, + Metadata.empty()), + new StructField( + "rolled_back_snapshot_id", + DataTypes.LongType, + false, + Metadata.empty()), + new StructField( + "current_snapshot_id", DataTypes.LongType, false, Metadata.empty()) + }); + + protected RollbackToAsLatestProcedure(TableCatalog tableCatalog) { + super(tableCatalog); + } + + @Override + public ProcedureParameter[] parameters() { + return PARAMETERS; + } + + @Override + public StructType outputType() { + return OUTPUT_TYPE; + } + + @Override + public InternalRow[] call(InternalRow args) { + Identifier tableIdent = toIdentifier(args.getString(0), PARAMETERS[0].name()); + String tagName = args.isNullAt(1) ? null : args.getString(1); + Long snapshotId = args.isNullAt(2) ? null : args.getLong(2); + + return modifyPaimonTable( + tableIdent, + table -> { + FileStoreTable fileStoreTable = (FileStoreTable) table; + FileStore store = fileStoreTable.store(); + Snapshot latestSnapshot = store.snapshotManager().latestSnapshot(); + Preconditions.checkNotNull( + latestSnapshot, "Latest snapshot is null, can not roll back."); + + boolean hasTag = !StringUtils.isNullOrWhitespaceOnly(tagName); + boolean hasSnapshot = snapshotId != null; + Preconditions.checkArgument( + hasTag != hasSnapshot, + "Must specify exactly one of tag and snapshot_id."); + + TagManager tagManager = store.newTagManager(); + Tag targetTag; + Snapshot targetSnapshot; + if (hasTag) { + targetTag = tagManager.getOrThrow(tagName); + targetSnapshot = targetTag.trimToSnapshot(); + } else { + targetTag = null; + targetSnapshot = findSnapshot(store, tagManager, snapshotId); + } + + String createdRollbackTag = null; + String commitUser = ROLLBACK_TO_AS_LATEST_TAG_PREFIX + UUID.randomUUID(); + try { + if (!hasTag) { + createdRollbackTag = + createRollbackToAsLatestTag(tagManager, targetSnapshot); + targetTag = tagManager.getOrThrow(createdRollbackTag); + } + try (TableCommitImpl commit = fileStoreTable.newCommit(commitUser)) { + boolean success = commit.rollbackToAsLatest(targetTag); + Preconditions.checkState( + success, + "Failed to roll back to snapshot %s as latest.", + targetSnapshot.id()); + } + } catch (Exception e) { + try { + deleteCreatedRollbackTagIfNotCommitted( + store, + tagManager, + createdRollbackTag, + latestSnapshot.id() + 1, + commitUser); + } catch (Exception cleanupException) { + e.addSuppressed(cleanupException); + } + throw new RuntimeException( + String.format( + "Failed to roll back to snapshot %s as latest.", + targetSnapshot.id()), + e); + } + + InternalRow outputRow = + newInternalRow( + latestSnapshot.id(), + targetSnapshot.id(), + store.snapshotManager().latestSnapshotId()); + return new InternalRow[] {outputRow}; + }); + } + + private String createRollbackToAsLatestTag(TagManager tagManager, Snapshot targetSnapshot) { + String tagName = + ROLLBACK_TO_AS_LATEST_TAG_PREFIX + targetSnapshot.id() + "-" + UUID.randomUUID(); + tagManager.createTag(targetSnapshot, tagName, null, Collections.emptyList(), false); + return tagName; + } + + private void deleteCreatedRollbackTagIfNotCommitted( + FileStore store, + TagManager tagManager, + String createdRollbackTag, + long rollbackSnapshotId, + String commitUser) { + if (createdRollbackTag == null + || rollbackSnapshotCommitted( + store.snapshotManager(), rollbackSnapshotId, commitUser)) { + return; + } + + tagManager.deleteTag( + createdRollbackTag, + store.newTagDeletion(), + store.snapshotManager(), + Collections.emptyList()); + } + + private boolean rollbackSnapshotCommitted( + SnapshotManager snapshotManager, long rollbackSnapshotId, String commitUser) { + return snapshotManager.snapshotExists(rollbackSnapshotId) + && commitUser.equals(snapshotManager.snapshot(rollbackSnapshotId).commitUser()); + } + + private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { + SnapshotManager snapshotManager = store.snapshotManager(); + if (snapshotManager.snapshotExists(snapshotId)) { + return snapshotManager.snapshot(snapshotId); + } + + SortedMap> tags = tagManager.tags(); + for (Map.Entry> entry : tags.entrySet()) { + if (entry.getKey().id() == snapshotId) { + return entry.getKey(); + } else if (entry.getKey().id() > snapshotId) { + break; + } + } + + throw new IllegalArgumentException( + String.format("Snapshot '%s' to roll back to doesn't exist.", snapshotId)); + } + + public static ProcedureBuilder builder() { + return new BaseProcedure.Builder() { + @Override + public RollbackToAsLatestProcedure doBuild() { + return new RollbackToAsLatestProcedure(tableCatalog()); + } + }; + } + + @Override + public String description() { + return "RollbackToAsLatestProcedure"; + } +} diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala new file mode 100644 index 000000000000..cec04b925195 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.spark.procedure + +import org.apache.paimon.spark.PaimonSparkTestBase + +import org.apache.spark.sql.Row + +class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { + + test("Paimon Procedure: rollback to a snapshot as the latest snapshot") { + spark.sql("CREATE TABLE T (id INT, name STRING)") + spark.sql("INSERT INTO T VALUES (1, 'a')") + spark.sql("INSERT INTO T VALUES (2, 'b')") + spark.sql("INSERT INTO T VALUES (3, 'c')") + + val snapshotManager = loadTable("T").snapshotManager + + // Non-destructive: roll snapshot 1 forward as the new latest (snapshot 4); + // snapshots 2 and 3 stay, unlike the destructive `rollback`. + checkAnswer( + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)"), + Row(3L, 1L, 4L) :: Nil) + assert(snapshotManager.snapshotExists(2)) + assert(snapshotManager.snapshotExists(3)) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Nil) + + // Roll forward to snapshot 3 as the latest (snapshot 5). + checkAnswer( + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 3)"), + Row(4L, 3L, 5L) :: Nil) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil) + } +} From 123273ede19aee51002c5c4019df0d6f36a1d12b Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sat, 26 Sep 2026 22:49:28 +0800 Subject: [PATCH 2/4] [spark] Fix rollback_to_as_latest cleanup deleting the protection tag after an interleaved commit Data-safety fix in the failure-path cleanup of rollback_to_as_latest. The procedure reads the latest snapshot, creates a protection tag for the rollback target, then commits the rollback. On any exception (e.g. a post-commit callback throwing) it decided whether the rollback had committed by probing snapshot `latestSnapshot.id() + 1`. But FileStoreCommitImpl.rollbackToAsLatest re-reads the latest snapshot, so a writer committing between the two reads pushes the rollback to a higher id: 1. procedure reads latest = 3 2. another writer commits snapshot 4 3. the rollback commits snapshot 5 4. a post-commit callback throws 5. cleanup probes snapshot 4, sees a different commitUser, and deletes the protection tag though snapshot 5 committed -- expiring older snapshots then drops the restored files (FileNotFoundException on read). Fix: establish the real outcome instead of inferring it from the stale id. Cleanup now scans every snapshot committed after the initial read for one authored by this rollback's unique commit user, and only deletes the tag when it positively confirms the rollback did not commit; if it committed -- or the outcome is uncertain -- the tag is retained. The mirrored Flink procedure had the identical flaw and is fixed the same way. Tests: a regression test drives two rollbacks so snapshots 4 and 5 carry distinct commit users, then asserts the cleanup check finds the rollback at 5 rather than probing the stale id 4 (the old point-probe returned false). Written with Claude Code; verification is mine. --- .../RollbackToAsLatestProcedure.java | 44 +++++++++++++---- .../RollbackToAsLatestProcedure.java | 48 +++++++++++++++---- .../RollbackToAsLatestProcedureTest.scala | 30 ++++++++++++ 3 files changed, 105 insertions(+), 17 deletions(-) diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java index 46df2916c2bd..729a95eb235f 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java @@ -113,7 +113,7 @@ Row[] call(ProcedureContext procedureContext, String tableId, String tagName, Lo } catch (Exception e) { try { deleteCreatedRollbackTagIfNotCommitted( - store, tagManager, createdRollbackTag, latestSnapshot.id() + 1, commitUser); + store, tagManager, createdRollbackTag, latestSnapshot.id(), commitUser); } catch (Exception cleanupException) { e.addSuppressed(cleanupException); } @@ -145,11 +145,27 @@ private void deleteCreatedRollbackTagIfNotCommitted( FileStore store, TagManager tagManager, String createdRollbackTag, - long rollbackSnapshotId, + long previousLatestSnapshotId, String commitUser) { - if (createdRollbackTag == null - || rollbackSnapshotCommitted( - store.snapshotManager(), rollbackSnapshotId, commitUser)) { + if (createdRollbackTag == null) { + return; + } + + // Determine the real commit outcome instead of assuming the rollback landed at + // previousLatestSnapshotId + 1: FileStoreCommitImpl.rollbackToAsLatest re-reads the latest + // snapshot, so another writer committing between our read and the rollback pushes the + // rollback snapshot to a higher id. Only delete the protection tag when we positively + // confirm the rollback did not commit; if it committed (or the outcome is uncertain), keep + // the tag so snapshot expiration cannot drop the file group the rollback restored. + boolean committed; + try { + committed = + rollbackSnapshotCommitted( + store.snapshotManager(), previousLatestSnapshotId, commitUser); + } catch (Exception e) { + return; + } + if (committed) { return; } @@ -161,9 +177,21 @@ private void deleteCreatedRollbackTagIfNotCommitted( } private boolean rollbackSnapshotCommitted( - SnapshotManager snapshotManager, long rollbackSnapshotId, String commitUser) { - return snapshotManager.snapshotExists(rollbackSnapshotId) - && commitUser.equals(snapshotManager.snapshot(rollbackSnapshotId).commitUser()); + SnapshotManager snapshotManager, long previousLatestSnapshotId, String commitUser) { + Long latestSnapshotId = snapshotManager.latestSnapshotId(); + if (latestSnapshotId == null) { + return false; + } + // The rollback snapshot is authored by our unique commitUser. Scan every snapshot committed + // after our initial read (there may be several if other writers interleaved) for one it + // authored rather than probing a single, possibly-stale id. + for (long id = latestSnapshotId; id > previousLatestSnapshotId; id--) { + if (snapshotManager.snapshotExists(id) + && commitUser.equals(snapshotManager.snapshot(id).commitUser())) { + return true; + } + } + return false; } private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java index 11c21efd1cd7..98cfae7b118c 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java @@ -20,6 +20,7 @@ import org.apache.paimon.FileStore; import org.apache.paimon.Snapshot; +import org.apache.paimon.annotation.VisibleForTesting; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.tag.Tag; @@ -145,7 +146,7 @@ public InternalRow[] call(InternalRow args) { store, tagManager, createdRollbackTag, - latestSnapshot.id() + 1, + latestSnapshot.id(), commitUser); } catch (Exception cleanupException) { e.addSuppressed(cleanupException); @@ -177,11 +178,27 @@ private void deleteCreatedRollbackTagIfNotCommitted( FileStore store, TagManager tagManager, String createdRollbackTag, - long rollbackSnapshotId, + long previousLatestSnapshotId, String commitUser) { - if (createdRollbackTag == null - || rollbackSnapshotCommitted( - store.snapshotManager(), rollbackSnapshotId, commitUser)) { + if (createdRollbackTag == null) { + return; + } + + // Determine the real commit outcome instead of assuming the rollback landed at + // previousLatestSnapshotId + 1: FileStoreCommitImpl.rollbackToAsLatest re-reads the latest + // snapshot, so another writer committing between our read and the rollback pushes the + // rollback snapshot to a higher id. Only delete the protection tag when we positively + // confirm the rollback did not commit; if it committed (or the outcome is uncertain), keep + // the tag so snapshot expiration cannot drop the file group the rollback restored. + boolean committed; + try { + committed = + rollbackSnapshotCommitted( + store.snapshotManager(), previousLatestSnapshotId, commitUser); + } catch (Exception e) { + return; + } + if (committed) { return; } @@ -192,10 +209,23 @@ private void deleteCreatedRollbackTagIfNotCommitted( Collections.emptyList()); } - private boolean rollbackSnapshotCommitted( - SnapshotManager snapshotManager, long rollbackSnapshotId, String commitUser) { - return snapshotManager.snapshotExists(rollbackSnapshotId) - && commitUser.equals(snapshotManager.snapshot(rollbackSnapshotId).commitUser()); + @VisibleForTesting + static boolean rollbackSnapshotCommitted( + SnapshotManager snapshotManager, long previousLatestSnapshotId, String commitUser) { + Long latestSnapshotId = snapshotManager.latestSnapshotId(); + if (latestSnapshotId == null) { + return false; + } + // The rollback snapshot is authored by our unique commitUser. Scan every snapshot committed + // after our initial read (there may be several if other writers interleaved) for one it + // authored rather than probing a single, possibly-stale id. + for (long id = latestSnapshotId; id > previousLatestSnapshotId; id--) { + if (snapshotManager.snapshotExists(id) + && commitUser.equals(snapshotManager.snapshot(id).commitUser())) { + return true; + } + } + return false; } private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala index cec04b925195..2bfb33841ee9 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala @@ -47,4 +47,34 @@ class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { Row(4L, 3L, 5L) :: Nil) checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil) } + + test("Paimon Procedure: rollback cleanup detects the commit after an interleaved snapshot") { + spark.sql("CREATE TABLE T (id INT, name STRING)") + spark.sql("INSERT INTO T VALUES (1, 'a')") + spark.sql("INSERT INTO T VALUES (2, 'b')") + spark.sql("INSERT INTO T VALUES (3, 'c')") + + // Each rollback commits under its own unique commit user, so snapshots 4 and 5 + // carry distinct users -- the shape the failure-path cleanup must handle when + // another writer commits between the procedure's snapshot read and the rollback. + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)") + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 3)") + + val snapshotManager = loadTable("T").snapshotManager + val user4 = snapshotManager.snapshot(4).commitUser + val user5 = snapshotManager.snapshot(5).commitUser + assert(user4 != user5) + + // The rollback that committed snapshot 5 began when 3 was the latest, but snapshot 4 + // (a different user) landed first. Cleanup must still find the rollback at 5 rather + // than probing the stale id 4 and wrongly concluding it "did not commit" -- which + // previously deleted the protection tag and let expiration drop the restored files. + assert(RollbackToAsLatestProcedure.rollbackSnapshotCommitted(snapshotManager, 3L, user5)) + assert(RollbackToAsLatestProcedure.rollbackSnapshotCommitted(snapshotManager, 4L, user5)) + // A rollback that never committed (no snapshot after 3 authored by it) is reported as + // such, so a genuine cleanup still removes its leftover helper tag. + assert( + !RollbackToAsLatestProcedure + .rollbackSnapshotCommitted(snapshotManager, 3L, "never-committed")) + } } From 5a084fdf2298c8c6b3f4af8d31f76df23906cb60 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Wed, 30 Sep 2026 19:39:24 +0800 Subject: [PATCH 3/4] Keep the Spark protection tag unless the commit returns false, and drop the Flink change The failure cleanup previously inferred the rollback outcome from a snapshot id, which is unsafe when another writer commits between the procedure's snapshot read and the core rollback. Track whether cleanup is safe from the call boundary instead: only delete the created protection tag when we never entered the commit or the commit explicitly returned false; if the commit threw after publishing (e.g. a post-commit callback), keep the tag so expiration cannot drop the restored files. This mirrors the Flink-side approach in #10204, so the Flink change is dropped here and left to that PR. --- docs/docs/spark/procedures.md | 1 + docs/docs/spark/procedures/versions.md | 23 +++ .../RollbackToAsLatestProcedure.java | 44 +----- .../RollbackToAsLatestProcedure.java | 75 ++-------- .../RollbackToAsLatestProcedureTest.scala | 138 ++++++++++++++---- 5 files changed, 158 insertions(+), 123 deletions(-) diff --git a/docs/docs/spark/procedures.md b/docs/docs/spark/procedures.md index 2048c49b0718..947de28348bb 100644 --- a/docs/docs/spark/procedures.md +++ b/docs/docs/spark/procedures.md @@ -89,6 +89,7 @@ Choose a group, then use the page contents to jump to a procedure: [`fast_forward`](./procedures/versions#fast_forward), [`merge_branch`](./procedures/versions#merge_branch), [`rollback`](./procedures/versions#rollback), +[`rollback_to_as_latest`](./procedures/versions#rollback_to_as_latest), [`rollback_to_timestamp`](./procedures/versions#rollback_to_timestamp), [`rollback_to_watermark`](./procedures/versions#rollback_to_watermark) diff --git a/docs/docs/spark/procedures/versions.md b/docs/docs/spark/procedures/versions.md index d9b96729722c..f4d4a76f086b 100644 --- a/docs/docs/spark/procedures/versions.md +++ b/docs/docs/spark/procedures/versions.md @@ -240,6 +240,29 @@ CALL sys.rollback(table => 'default.T', tag => 'tag1'); CALL sys.rollback(table => 'default.T', snapshot => 2); ``` +## rollback_to_as_latest + +Roll back to a retained snapshot or tag and materialize it as the latest snapshot, without deleting +later snapshots or tags. Specify exactly one of `snapshot_id` or `tag`. Batch and time-travel reads +are correct; for deletion-vector tables, a rollback whose only difference is a deletion-vector +change is not guaranteed to be observed by streaming overwrite readers. + +When rolling back to a snapshot, the procedure creates a tag named +`rollback-to-as-latest--` that protects the restored data files from snapshot +expiration. + +**Arguments** + +- `table` (`STRING`, required): the target table identifier. +- `tag` (`STRING`, optional): tag that will roll back to. +- `snapshot_id` (`BIGINT`, optional): snapshot that will roll back to. + +```sql +CALL sys.rollback_to_as_latest(table => 'default.T', snapshot_id => 10); + +CALL sys.rollback_to_as_latest(table => 'default.T', tag => 'tag1'); +``` + ## rollback_to_timestamp Rollback to the snapshot which earlier or equal than timestamp. diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java index 729a95eb235f..46df2916c2bd 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RollbackToAsLatestProcedure.java @@ -113,7 +113,7 @@ Row[] call(ProcedureContext procedureContext, String tableId, String tagName, Lo } catch (Exception e) { try { deleteCreatedRollbackTagIfNotCommitted( - store, tagManager, createdRollbackTag, latestSnapshot.id(), commitUser); + store, tagManager, createdRollbackTag, latestSnapshot.id() + 1, commitUser); } catch (Exception cleanupException) { e.addSuppressed(cleanupException); } @@ -145,27 +145,11 @@ private void deleteCreatedRollbackTagIfNotCommitted( FileStore store, TagManager tagManager, String createdRollbackTag, - long previousLatestSnapshotId, + long rollbackSnapshotId, String commitUser) { - if (createdRollbackTag == null) { - return; - } - - // Determine the real commit outcome instead of assuming the rollback landed at - // previousLatestSnapshotId + 1: FileStoreCommitImpl.rollbackToAsLatest re-reads the latest - // snapshot, so another writer committing between our read and the rollback pushes the - // rollback snapshot to a higher id. Only delete the protection tag when we positively - // confirm the rollback did not commit; if it committed (or the outcome is uncertain), keep - // the tag so snapshot expiration cannot drop the file group the rollback restored. - boolean committed; - try { - committed = - rollbackSnapshotCommitted( - store.snapshotManager(), previousLatestSnapshotId, commitUser); - } catch (Exception e) { - return; - } - if (committed) { + if (createdRollbackTag == null + || rollbackSnapshotCommitted( + store.snapshotManager(), rollbackSnapshotId, commitUser)) { return; } @@ -177,21 +161,9 @@ private void deleteCreatedRollbackTagIfNotCommitted( } private boolean rollbackSnapshotCommitted( - SnapshotManager snapshotManager, long previousLatestSnapshotId, String commitUser) { - Long latestSnapshotId = snapshotManager.latestSnapshotId(); - if (latestSnapshotId == null) { - return false; - } - // The rollback snapshot is authored by our unique commitUser. Scan every snapshot committed - // after our initial read (there may be several if other writers interleaved) for one it - // authored rather than probing a single, possibly-stale id. - for (long id = latestSnapshotId; id > previousLatestSnapshotId; id--) { - if (snapshotManager.snapshotExists(id) - && commitUser.equals(snapshotManager.snapshot(id).commitUser())) { - return true; - } - } - return false; + SnapshotManager snapshotManager, long rollbackSnapshotId, String commitUser) { + return snapshotManager.snapshotExists(rollbackSnapshotId) + && commitUser.equals(snapshotManager.snapshot(rollbackSnapshotId).commitUser()); } private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java index 98cfae7b118c..fd28d4c3c5c9 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java @@ -20,7 +20,6 @@ import org.apache.paimon.FileStore; import org.apache.paimon.Snapshot; -import org.apache.paimon.annotation.VisibleForTesting; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.tag.Tag; @@ -126,6 +125,7 @@ public InternalRow[] call(InternalRow args) { } String createdRollbackTag = null; + boolean canDeleteCreatedTag = true; String commitUser = ROLLBACK_TO_AS_LATEST_TAG_PREFIX + UUID.randomUUID(); try { if (!hasTag) { @@ -134,7 +134,13 @@ public InternalRow[] call(InternalRow args) { targetTag = tagManager.getOrThrow(createdRollbackTag); } try (TableCommitImpl commit = fileStoreTable.newCommit(commitUser)) { + // The core rollback reads latest again, so another writer may change + // the rollback snapshot ID. An exception may also come after a + // successful commit (for example, from a callback). Keep the + // protection tag unless the commit returns false. + canDeleteCreatedTag = false; boolean success = commit.rollbackToAsLatest(targetTag); + canDeleteCreatedTag = !success; Preconditions.checkState( success, "Failed to roll back to snapshot %s as latest.", @@ -142,12 +148,13 @@ public InternalRow[] call(InternalRow args) { } } catch (Exception e) { try { - deleteCreatedRollbackTagIfNotCommitted( - store, - tagManager, - createdRollbackTag, - latestSnapshot.id(), - commitUser); + if (createdRollbackTag != null && canDeleteCreatedTag) { + tagManager.deleteTag( + createdRollbackTag, + store.newTagDeletion(), + store.snapshotManager(), + Collections.emptyList()); + } } catch (Exception cleanupException) { e.addSuppressed(cleanupException); } @@ -174,60 +181,6 @@ private String createRollbackToAsLatestTag(TagManager tagManager, Snapshot targe return tagName; } - private void deleteCreatedRollbackTagIfNotCommitted( - FileStore store, - TagManager tagManager, - String createdRollbackTag, - long previousLatestSnapshotId, - String commitUser) { - if (createdRollbackTag == null) { - return; - } - - // Determine the real commit outcome instead of assuming the rollback landed at - // previousLatestSnapshotId + 1: FileStoreCommitImpl.rollbackToAsLatest re-reads the latest - // snapshot, so another writer committing between our read and the rollback pushes the - // rollback snapshot to a higher id. Only delete the protection tag when we positively - // confirm the rollback did not commit; if it committed (or the outcome is uncertain), keep - // the tag so snapshot expiration cannot drop the file group the rollback restored. - boolean committed; - try { - committed = - rollbackSnapshotCommitted( - store.snapshotManager(), previousLatestSnapshotId, commitUser); - } catch (Exception e) { - return; - } - if (committed) { - return; - } - - tagManager.deleteTag( - createdRollbackTag, - store.newTagDeletion(), - store.snapshotManager(), - Collections.emptyList()); - } - - @VisibleForTesting - static boolean rollbackSnapshotCommitted( - SnapshotManager snapshotManager, long previousLatestSnapshotId, String commitUser) { - Long latestSnapshotId = snapshotManager.latestSnapshotId(); - if (latestSnapshotId == null) { - return false; - } - // The rollback snapshot is authored by our unique commitUser. Scan every snapshot committed - // after our initial read (there may be several if other writers interleaved) for one it - // authored rather than probing a single, possibly-stale id. - for (long id = latestSnapshotId; id > previousLatestSnapshotId; id--) { - if (snapshotManager.snapshotExists(id) - && commitUser.equals(snapshotManager.snapshot(id).commitUser())) { - return true; - } - } - return false; - } - private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { SnapshotManager snapshotManager = store.snapshotManager(); if (snapshotManager.snapshotExists(snapshotId)) { diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala index 2bfb33841ee9..0c39beaa8fe6 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala @@ -18,9 +18,14 @@ package org.apache.paimon.spark.procedure +import org.apache.paimon.data.{BinaryString, GenericRow} +import org.apache.paimon.manifest.ManifestCommittable import org.apache.paimon.spark.PaimonSparkTestBase +import org.apache.paimon.table.{ExpireSnapshotsImpl, FileStoreTable} +import org.apache.paimon.table.sink.CommitCallback import org.apache.spark.sql.Row +import org.assertj.core.api.Assertions.assertThatThrownBy class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { @@ -48,33 +53,114 @@ class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil) } - test("Paimon Procedure: rollback cleanup detects the commit after an interleaved snapshot") { - spark.sql("CREATE TABLE T (id INT, name STRING)") - spark.sql("INSERT INTO T VALUES (1, 'a')") - spark.sql("INSERT INTO T VALUES (2, 'b')") - spark.sql("INSERT INTO T VALUES (3, 'c')") + test("Paimon Procedure: rollback keeps the protection tag after a post-commit failure") { + Seq(false, true).foreach { + concurrentCommit => + withTable("T") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + // The restored file must not also belong to the snapshots retained before the rollback. + spark.sql("INSERT OVERWRITE T VALUES (2, 'replacement')") + spark.sql("INSERT INTO T VALUES (3, 'retained')") - // Each rollback commits under its own unique commit user, so snapshots 4 and 5 - // carry distinct users -- the shape the failure-path cleanup must handle when - // another writer commits between the procedure's snapshot read and the rollback. - spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)") - spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 3)") + FailingRollbackCallback.concurrentCommit = concurrentCommit + FailingRollbackCallback.failRollbackCommit = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected post-commit callback failure") + } finally { + FailingRollbackCallback.reset() + } - val snapshotManager = loadTable("T").snapshotManager - val user4 = snapshotManager.snapshot(4).commitUser - val user5 = snapshotManager.snapshot(5).commitUser - assert(user4 != user5) - - // The rollback that committed snapshot 5 began when 3 was the latest, but snapshot 4 - // (a different user) landed first. Cleanup must still find the rollback at 5 rather - // than probing the stale id 4 and wrongly concluding it "did not commit" -- which - // previously deleted the protection tag and let expiration drop the restored files. - assert(RollbackToAsLatestProcedure.rollbackSnapshotCommitted(snapshotManager, 3L, user5)) - assert(RollbackToAsLatestProcedure.rollbackSnapshotCommitted(snapshotManager, 4L, user5)) - // A rollback that never committed (no snapshot after 3 authored by it) is reported as - // such, so a genuine cleanup still removes its leftover helper tag. - assert( - !RollbackToAsLatestProcedure - .rollbackSnapshotCommitted(snapshotManager, 3L, "never-committed")) + val table = loadTable("T") + val latest = table.snapshotManager().latestSnapshot() + assert(latest.id() == (if (concurrentCommit) 5L else 4L)) + assert(latest.commitUser().startsWith("rollback-to-as-latest-")) + val tags = table.tagManager().allTagNames() + assert(tags.size() == 1 && tags.get(0).startsWith("rollback-to-as-latest-1-")) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + val expire = table.newExpireSnapshots().asInstanceOf[ExpireSnapshotsImpl] + expire.expireUntil(1, latest.id()) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + } + } + } + + test("Paimon Procedure: rollback removes the protection tag when it never started") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + + FailingRollbackCallback.failCommitCreation = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected commit creation failure") + } finally { + FailingRollbackCallback.reset() + } + + val table = loadTable("T") + assert(table.tagManager().allTagNames().isEmpty) + assert(table.snapshotManager().latestSnapshotId() == 1L) + } + + private def createTableWithFailingCallback(): Unit = { + val callback = classOf[FailingRollbackCallback].getName + spark.sql( + s"CREATE TABLE T (id INT, name STRING) TBLPROPERTIES ('commit.callbacks' = '$callback')") + } + + private def rollbackToSnapshot1(): Unit = { + spark + .sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)") + .collect() + } +} + +/** Injects rollback failures, optionally after another writer commits a snapshot first. */ +class FailingRollbackCallback extends CommitCallback { + + override def setTable(table: FileStoreTable): Unit = { + if (FailingRollbackCallback.failCommitCreation) { + throw new RuntimeException("Injected commit creation failure") + } + if (FailingRollbackCallback.concurrentCommit) { + // The procedure already read latest; the core rollback has not read it yet. + FailingRollbackCallback.concurrentCommit = false + val builder = table.newBatchWriteBuilder() + val write = builder.newWrite() + val commit = builder.newCommit() + try { + write.write(GenericRow.of(Int.box(4), BinaryString.fromString("concurrent"))) + commit.commit(write.prepareCommit()) + } finally { + write.close() + commit.close() + } + } + } + + override def call(context: CommitCallback.Context): Unit = { + if ( + FailingRollbackCallback.failRollbackCommit && + context.snapshot.commitUser().startsWith("rollback-to-as-latest-") + ) { + throw new RuntimeException("Injected post-commit callback failure") + } + } + + override def retry(committable: ManifestCommittable): Unit = {} + + override def close(): Unit = {} +} + +object FailingRollbackCallback { + @volatile var failRollbackCommit = false + @volatile var concurrentCommit = false + @volatile var failCommitCreation = false + + def reset(): Unit = { + failRollbackCommit = false + concurrentCommit = false + failCommitCreation = false } } From ab2a2d2f9d093c1a2a51570598ad3c83981c3164 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Fri, 2 Oct 2026 22:44:53 +0800 Subject: [PATCH 4/4] [spark] Refresh the Spark cache when rollback_to_as_latest fails after publishing --- .../RollbackToAsLatestProcedure.java | 28 ++++++++++++++++++- .../RollbackToAsLatestProcedureTest.scala | 20 +++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java index fd28d4c3c5c9..69366598eee0 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java @@ -20,6 +20,7 @@ import org.apache.paimon.FileStore; import org.apache.paimon.Snapshot; +import org.apache.paimon.spark.SparkTable; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.tag.Tag; @@ -41,6 +42,7 @@ import java.util.Map; import java.util.SortedMap; import java.util.UUID; +import java.util.function.Function; import static org.apache.spark.sql.types.DataTypes.LongType; import static org.apache.spark.sql.types.DataTypes.StringType; @@ -98,7 +100,7 @@ public InternalRow[] call(InternalRow args) { String tagName = args.isNullAt(1) ? null : args.getString(1); Long snapshotId = args.isNullAt(2) ? null : args.getLong(2); - return modifyPaimonTable( + return modifyPaimonTableRefreshingCacheOnFailure( tableIdent, table -> { FileStoreTable fileStoreTable = (FileStoreTable) table; @@ -174,6 +176,30 @@ public InternalRow[] call(InternalRow args) { }); } + /** + * Like {@link #modifyPaimonTable} but also refreshes Spark's cached plans when {@code func} + * throws. {@code rollback_to_as_latest} can publish the rollback snapshot and then fail (for + * example a post-commit callback throws), so the table state is already durable; the shared + * success-only refresh would otherwise leave {@code CACHE TABLE} serving the pre-rollback data. + * The original failure is preserved; a refresh error is only added as suppressed. + */ + private InternalRow[] modifyPaimonTableRefreshingCacheOnFailure( + Identifier ident, Function func) { + SparkTable sparkTable = loadSparkTable(ident); + try { + InternalRow[] result = func.apply(sparkTable.getTable()); + refreshSparkCache(ident, sparkTable); + return result; + } catch (RuntimeException e) { + try { + refreshSparkCache(ident, sparkTable); + } catch (RuntimeException refreshError) { + e.addSuppressed(refreshError); + } + throw e; + } + } + private String createRollbackToAsLatestTag(TagManager tagManager, Snapshot targetSnapshot) { String tagName = ROLLBACK_TO_AS_LATEST_TAG_PREFIX + targetSnapshot.id() + "-" + UUID.randomUUID(); diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala index 0c39beaa8fe6..e8b01d611b40 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala @@ -103,6 +103,26 @@ class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { assert(table.snapshotManager().latestSnapshotId() == 1L) } + test("Paimon Procedure: rollback refreshes a cached table after a post-commit failure") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + spark.sql("INSERT OVERWRITE T VALUES (2, 'replacement')") + spark.sql("CACHE TABLE T") + checkAnswer(spark.sql("SELECT * FROM T"), Row(2, "replacement") :: Nil) + + FailingRollbackCallback.failRollbackCommit = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected post-commit callback failure") + } finally { + FailingRollbackCallback.reset() + } + + // Snapshot 3 (the rollback) is durable despite the callback failure, so the cached read must + // reflect snapshot 1's data rather than the stale pre-rollback replacement row. + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + } + private def createTableWithFailingCallback(): Unit = { val callback = classOf[FailingRollbackCallback].getName spark.sql(