[spark] Add create_tag_from_watermark procedure

docs/docs/spark/procedures.md

@@ -78,6 +78,7 @@ Choose a group, then use the page contents to jump to a procedure:
 
 [`create_tag`](./procedures/versions#create_tag),
 [`create_tag_from_timestamp`](./procedures/versions#create_tag_from_timestamp),
+[`create_tag_from_watermark`](./procedures/versions#create_tag_from_watermark),
 [`replace_tag`](./procedures/versions#replace_tag),
 [`rename_tag`](./procedures/versions#rename_tag),
 [`delete_tag`](./procedures/versions#delete_tag),

docs/docs/spark/procedures/versions.md

@@ -69,6 +69,26 @@ CALL sys.create_tag_from_timestamp(
 );
 ```
 
+## create_tag_from_watermark
+
+Create a tag based on given watermark timestamp.
+
+**Arguments**
+
+- `table` (`STRING`, required): the target table identifier.
+- `tag` (`STRING`, required): name of the new tag.
+- `watermark` (`BIGINT`, required): Find the first snapshot (including tagged snapshots) whose watermark is at or after this value.
+- `time_retained` (`STRING`, optional): The maximum time retained for newly created tags.
+
+```sql
+CALL sys.create_tag_from_watermark(
+  `table` => 'default.T',
+  `tag` => 'my_tag',
+  `watermark` => 1724404318750,
+  time_retained => '1 d'
+);
+```
+
 ## replace_tag
 
 Replace an existing tag with new tag info.

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java

@@ -31,6 +31,7 @@
 import org.apache.paimon.spark.procedure.CreateGlobalIndexProcedure;
 import org.apache.paimon.spark.procedure.CreatePolicyProcedure;
 import org.apache.paimon.spark.procedure.CreateTagFromTimestampProcedure;
+import org.apache.paimon.spark.procedure.CreateTagFromWatermarkProcedure;
 import org.apache.paimon.spark.procedure.CreateTagProcedure;
 import org.apache.paimon.spark.procedure.DeleteBranchProcedure;
 import org.apache.paimon.spark.procedure.DeleteTagProcedure;
@@ -106,6 +107,8 @@ private static Map<String, Supplier<ProcedureBuilder>> initProcedureBuilders() {
         procedureBuilders.put("rename_tag", RenameTagProcedure::builder);
         procedureBuilders.put(
                 "create_tag_from_timestamp", CreateTagFromTimestampProcedure::builder);
+        procedureBuilders.put(
+                "create_tag_from_watermark", CreateTagFromWatermarkProcedure::builder);
         procedureBuilders.put("delete_tag", DeleteTagProcedure::builder);
         procedureBuilders.put("expire_tags", ExpireTagsProcedure::builder);
         procedureBuilders.put("create_branch", CreateBranchProcedure::builder);

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/CreateTagFromWatermarkProcedure.java

@@ -0,0 +1,139 @@
+/*
+ * 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.Snapshot;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.utils.SnapshotManager;
+import org.apache.paimon.utils.SnapshotNotExistException;
+import org.apache.paimon.utils.TimeUtils;
+
+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 org.apache.spark.unsafe.types.UTF8String;
+
+import java.time.Duration;
+import java.util.Set;
+
+import static org.apache.spark.sql.types.DataTypes.LongType;
+import static org.apache.spark.sql.types.DataTypes.StringType;
+
+/** The procedure supports creating tags from a snapshot's watermark. */
+public class CreateTagFromWatermarkProcedure extends BaseProcedure {
+
+    private static final ProcedureParameter[] PARAMETERS =
+            new ProcedureParameter[] {
+                ProcedureParameter.required("table", StringType),
+                ProcedureParameter.required("tag", StringType),
+                ProcedureParameter.required("watermark", LongType),
+                ProcedureParameter.optional("time_retained", StringType)
+            };
+
+    private static final StructType OUTPUT_TYPE =
+            new StructType(
+                    new StructField[] {
+                        new StructField("tagName", DataTypes.StringType, true, Metadata.empty()),
+                        new StructField("snapshot", DataTypes.LongType, true, Metadata.empty()),
+                        new StructField("commit_time", DataTypes.LongType, true, Metadata.empty()),
+                        new StructField("watermark", DataTypes.StringType, true, Metadata.empty())
+                    });
+
+    protected CreateTagFromWatermarkProcedure(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 tag = args.getString(1);
+        long watermark = args.getLong(2);
+        Duration timeRetained =
+                args.isNullAt(3) ? null : TimeUtils.parseDuration(args.getString(3));
+
+        return modifyPaimonTable(
+                tableIdent,
+                table -> {
+                    FileStoreTable fileStoreTable = (FileStoreTable) table;
+                    SnapshotManager snapshotManager = fileStoreTable.snapshotManager();
+                    Snapshot snapshot = snapshotManager.laterOrEqualWatermark(watermark);
+
+                    Set<Snapshot> sortedTagsSnapshots = fileStoreTable.tagManager().tags().keySet();
+
+                    // Fall back to the earliest tag whose watermark already covers the target,
+                    // mirroring the Flink create_tag_from_watermark procedure.
+                    for (Snapshot tagSnapshot : sortedTagsSnapshots) {
+                        if (tagSnapshot.watermark() != null
+                                && watermark <= tagSnapshot.watermark()) {
+                            if (snapshot == null
+                                    || snapshot.watermark() == null
+                                    || tagSnapshot.watermark() < snapshot.watermark()) {
+                                snapshot = tagSnapshot;
+                            }
+                            break;
+                        }
+                    }
+
+                    SnapshotNotExistException.checkNotNull(
+                            snapshot,
+                            String.format(
+                                    "Could not find any snapshot whose watermark later than %s.",
+                                    watermark));
+
+                    fileStoreTable.createTag(tag, snapshot.id(), timeRetained);
+
+                    InternalRow outputRow =
+                            newInternalRow(
+                                    UTF8String.fromString(tag),
+                                    snapshot.id(),
+                                    snapshot.timeMillis(),
+                                    UTF8String.fromString(String.valueOf(snapshot.watermark())));
+
+                    return new InternalRow[] {outputRow};
+                });
+    }
+
+    public static ProcedureBuilder builder() {
+        return new BaseProcedure.Builder<CreateTagFromWatermarkProcedure>() {
+            @Override
+            public CreateTagFromWatermarkProcedure doBuild() {
+                return new CreateTagFromWatermarkProcedure(tableCatalog());
+            }
+        };
+    }
+
+    @Override
+    public String description() {
+        return "CreateTagFromWatermarkProcedure";
+    }
+}

paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/CreateTagFromWatermarkProcedureTest.scala

@@ -0,0 +1,89 @@
+/*
+ * 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.data.{BinaryString, GenericRow}
+import org.apache.paimon.manifest.ManifestCommittable
+import org.apache.paimon.spark.PaimonSparkTestBase
+import org.apache.paimon.table.FileStoreTable
+import org.apache.paimon.table.sink.TableCommitImpl
+import org.apache.paimon.utils.SnapshotNotExistException
+
+import org.apache.spark.sql.Row
+
+import scala.collection.JavaConverters._
+
+class CreateTagFromWatermarkProcedureTest extends PaimonSparkTestBase {
+
+  test("Paimon Procedure: create tags from snapshots watermark") {
+    spark.sql(s"""
+                 |CREATE TABLE T (a INT, b STRING)
+                 |TBLPROPERTIES ('primary-key'='a', 'bucket'='1')
+                 |""".stripMargin)
+
+    val table = loadTable("T")
+    // A Spark batch write carries no watermark, so commit snapshots with explicit
+    // watermarks the way a Flink writer would, mirroring the Flink IT case.
+    writeWithWatermark(table, 1L, 1000L, GenericRow.of(1, BinaryString.fromString("a")))
+    writeWithWatermark(table, 2L, 2000L, GenericRow.of(2, BinaryString.fromString("b")))
+    writeWithWatermark(table, 3L, 3000L, GenericRow.of(3, BinaryString.fromString("c")))
+
+    val commitTime1 = table.snapshotManager.snapshot(1).timeMillis
+    val commitTime2 = table.snapshotManager.snapshot(2).timeMillis
+
+    // watermark below snapshot 1 resolves to snapshot 1.
+    checkAnswer(
+      spark.sql(s"""CALL paimon.sys.create_tag_from_watermark(
+                   |table => 'test.T', tag => 'tag1', watermark => 500)""".stripMargin),
+      Row("tag1", 1, commitTime1, "1000") :: Nil
+    )
+
+    // watermark equal to snapshot 2 resolves to snapshot 2.
+    checkAnswer(
+      spark.sql(s"""CALL paimon.sys.create_tag_from_watermark(
+                   |table => 'test.T', tag => 'tag2', watermark => 2000)""".stripMargin),
+      Row("tag2", 2, commitTime2, "2000") :: Nil
+    )
+
+    // watermark later than the latest snapshot throws SnapshotNotExistException.
+    assertThrows[SnapshotNotExistException] {
+      spark.sql(s"""CALL paimon.sys.create_tag_from_watermark(
+                   |table => 'test.T', tag => 'tag3', watermark => ${Long.MaxValue})""".stripMargin)
+    }
+  }
+
+  private def writeWithWatermark(
+      table: FileStoreTable,
+      commitId: Long,
+      watermark: Long,
+      rows: GenericRow*): Unit = {
+    val writeBuilder = table.newStreamWriteBuilder()
+    val writer = writeBuilder.newWrite()
+    val commit = writeBuilder.newCommit().asInstanceOf[TableCommitImpl]
+    try {
+      rows.foreach(writer.write)
+      val committable = new ManifestCommittable(commitId, watermark)
+      writer.prepareCommit(true, commitId).asScala.foreach(committable.addFileCommittable)
+      commit.commit(committable)
+    } finally {
+      writer.close()
+      commit.close()
+    }
+  }
+}