Skip to content

feat(datafusion): support create_tag_from_watermark #1134

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
QuakeWang wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(datafusion): support create_tag_from_watermark #1134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5175997
feat(datafusion): support create_tag_from_watermark
QuakeWang Oct 15, 2026
a9c3a5e
fix(datafusion): reject watermark tags on REST tables
QuakeWang Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 74 additions & 0 deletions crates/integrations/datafusion/src/procedures.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
//! - `CALL sys.rollback_to(table => '...', snapshot_id => ... | tag => '...')`
//! - `CALL sys.rollback_to_timestamp(table => '...', timestamp => ...)`
//! - `CALL sys.create_tag_from_timestamp(table => '...', tag => '...', timestamp => ...)`
//! - `CALL sys.create_tag_from_watermark(table => '...', tag => '...', watermark => ..., time_retained => '...')`
//! - `CALL sys.create_global_index(table => '...', index_column => '...', index_type => 'btree')`
//! - `CALL sys.create_global_index(table => '...', index_column => '...', index_type => 'bitmap')`
//! - `CALL sys.create_global_index(table => '...', index_column => '...', index_type => 'ivf-pq')`
Expand Down Expand Up @@ -134,6 +135,35 @@ async fn later_or_equal_from_all(
}
}

/// Find a snapshot with watermark >= the requested value, including tag-retained
/// snapshots. Matches Java's Spark procedure: compare the first eligible live and
/// tagged snapshots by watermark, then by snapshot ID when watermarks are equal.
async fn later_or_equal_watermark_from_all(
sm: &SnapshotManager,
tm: &TagManager,
watermark: i64,
) -> DFResult<Option<Snapshot>> {
let live = sm
.later_or_equal_watermark(watermark)
.await
.map_err(to_datafusion_error)?;
let tags = tm.list_all().await.map_err(to_datafusion_error)?;
let tag_candidate = tags
.into_iter()
.map(|(_, snapshot)| snapshot)
.filter(|snapshot| {
snapshot
.watermark()
.is_some_and(|value| value != i64::MIN && value >= watermark)
})
// Rust lists tags by name; Java's tag candidates are ordered by snapshot ID.
.min_by_key(Snapshot::id);
Ok(live
.into_iter()
.chain(tag_candidate)
.min_by_key(|snapshot| (snapshot.watermark(), snapshot.id())))
}

/// Find the latest snapshot with commit time <= timestamp_millis,
/// considering both live snapshots and tag-retained snapshots.
async fn earlier_or_equal_from_all(
Expand Down Expand Up @@ -177,6 +207,7 @@ fn declared_parameters(proc_name: &str) -> Option<&'static [&'static str]> {
"rollback_to" => &["table", "snapshot_id", "tag"],
"rollback_to_timestamp" => &["table", "timestamp"],
"create_tag_from_timestamp" => &["table", "tag", "timestamp"],
"create_tag_from_watermark" => &["table", "tag", "watermark", "time_retained"],
"rename_branch" => &["table", "from_branch", "to_branch"],
"create_global_index" => &["table", "index_column", "index_type", "options"],
"drop_global_index" => &[
Expand Down Expand Up @@ -290,6 +321,9 @@ pub async fn execute_call(
"create_tag_from_timestamp" => {
proc_create_tag_from_timestamp(ctx, catalog, catalog_name, &args).await
}
"create_tag_from_watermark" => {
proc_create_tag_from_watermark(ctx, catalog, catalog_name, &args).await
}
"rename_branch" => proc_rename_branch(ctx, catalog, catalog_name, &args).await,
"create_global_index" => proc_create_global_index(ctx, catalog, catalog_name, &args).await,
"drop_global_index" => proc_drop_global_index(ctx, catalog, catalog_name, &args).await,
Expand Down Expand Up @@ -709,6 +743,46 @@ async fn proc_create_tag_from_timestamp(
ok_result(ctx)
}

async fn proc_create_tag_from_watermark(
ctx: &SessionContext,
catalog: &Arc<dyn Catalog>,
catalog_name: &str,
args: &HashMap<String, String>,
) -> DFResult<DataFrame> {
let table = get_table(catalog, catalog_name, args).await?;
let tag_name = require_arg(args, "tag")?;
let watermark_str = require_arg(args, "watermark")?;
let watermark: i64 = watermark_str
.parse()
.map_err(|_| DataFusionError::Plan(format!("Invalid watermark: '{watermark_str}'")))?;
let time_retained = args.get("time_retained").map(String::as_str);

// REST snapshots may exist only in the catalog; storage listings cannot
// enumerate the authoritative history needed for watermark selection.
if table.rest_env().is_some() {
return Err(DataFusionError::NotImplemented(
"create_tag_from_watermark is not supported for tables managed by a REST catalog"
.to_string(),
));
}

let (sm, tm) = managers(&table);
if tm.tag_exists(tag_name).await.map_err(to_datafusion_error)? {
return Err(DataFusionError::Plan(format!(
"Tag '{tag_name}' already exists"
)));
}
let snapshot = later_or_equal_watermark_from_all(&sm, &tm, watermark)

@JingsongLi JingsongLi Oct 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Resolve watermark candidates from authoritative REST metadata

This new entry point accepts REST tables, but managers(&table) creates a filesystem-only SnapshotManager, and later_or_equal_watermark enumerates only physical snapshot files. REST snapshots can legitimately exist only in the catalog: RESTSnapshotCommit::commit sends the snapshot over HTTP without publishing a snapshot file. I reproduced this with two real SQL-written Parquet snapshots (watermarks 1000/2000), retained snapshot-1 on storage, and snapshot-2 served only by the REST catalog. SELECT id FROM paimon.test_db.t1 ORDER BY id returns both rows, but CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'checkpoint', watermark => 2000) fails with No snapshot found with watermark >= 2000; the call makes no /snapshot request. Thus a readable REST table cannot use the advertised checkpoint operation. Please resolve candidates through the catalog-aware metadata path, or explicitly reject/document unsupported REST tables before selection, with this SQL/HTTP regression. Merely switching the helper to table.snapshot_manager() is insufficient because its watermark search still lists storage files.

.await?
.ok_or_else(|| {
DataFusionError::Plan(format!("No snapshot found with watermark >= {watermark}"))
})?;
tm.create_with_retention(tag_name, &snapshot, time_retained)
.await
.map_err(to_datafusion_error)?;
ok_result(ctx)
}

async fn proc_create_lumina_index(
ctx: &SessionContext,
catalog: &Arc<dyn Catalog>,
Expand Down
277 changes: 277 additions & 0 deletions crates/integrations/datafusion/tests/procedures.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,27 @@ async fn setup_table_with_snapshots() -> (tempfile::TempDir, paimon_datafusion::
(tmp, sql_context)
}

async fn setup_table_with_watermarks(
watermarks: [Option<i64>; 3],
) -> (tempfile::TempDir, paimon_datafusion::SQLContext) {
let (tmp, sql_context) = setup_table_with_snapshots().await;
for (index, watermark) in watermarks.into_iter().enumerate() {
let path = tmp
.path()
.join(format!("test_db.db/t1/snapshot/snapshot-{}", index + 1));
let mut snapshot: serde_json::Value =
serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
snapshot["watermark"] = serde_json::json!(watermark);
std::fs::write(path, serde_json::to_vec(&snapshot).unwrap()).unwrap();
}
(tmp, sql_context)
}

fn read_tag(tmp: &Path, name: &str) -> serde_json::Value {
let path = tmp.join(format!("test_db.db/t1/tag/tag-{name}"));
serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap()
}

#[tokio::test]
async fn test_create_tag() {
let (_tmp, sql_context) = setup_table_with_snapshots().await;
Expand Down Expand Up @@ -2476,6 +2497,262 @@ async fn test_create_tag_from_timestamp() {
assert_eq!(count, 1);
}

#[tokio::test]
async fn test_create_tag_from_watermark_selects_first_live_snapshot() {
for (watermark, expected_id) in [(500, 1), (1000, 1), (1500, 2), (2000, 2)] {
let (tmp, sql_context) =
setup_table_with_watermarks([Some(1000), Some(2000), Some(2000)]).await;
let tag_name = format!("wm_{watermark}");
exec(
&sql_context,
&format!(
"CALL sys.create_tag_from_watermark(table => 'paimon.test_db.t1', \
tag => '{tag_name}', watermark => {watermark})"
),
)
.await;
let tag = read_tag(tmp.path(), &tag_name);
assert_eq!(tag["id"], expected_id);
assert!(tag.get("tagTimeRetained").is_none());
assert!(tag.get("tagCreateTime").is_none());
assert_eq!(
row_count(
&sql_context,
&format!("SELECT * FROM paimon.test_db.t1 VERSION AS OF '{tag_name}'"),
)
.await,
expected_id as usize
);
}
}

#[tokio::test]
async fn test_create_tag_from_watermark_uses_retained_snapshots() {
let (tmp, sql_context) =
setup_table_with_watermarks([Some(1000), Some(2000), Some(3000)]).await;
let snapshot_dir = tmp.path().join("test_db.db/t1/snapshot");
// Tag names deliberately put the later snapshot first in Rust's tag listing.
for (name, id) in [("z_earlier", 1), ("a_later", 2)] {
exec(
&sql_context,
&format!(
"CALL sys.create_tag(table => 'test_db.t1', tag => '{name}', snapshot_id => {id})"
),
)
.await;
std::fs::remove_file(snapshot_dir.join(format!("snapshot-{id}"))).unwrap();
}

for (watermark, expected_id) in [(999, 1), (1000, 1), (1500, 2), (2000, 2), (2500, 3)] {
let tag_name = format!("wm_{watermark}");
exec(
&sql_context,
&format!(
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', \
tag => '{tag_name}', watermark => {watermark}, time_retained => '7d')"
),
)
.await;
let tag = read_tag(tmp.path(), &tag_name);
assert_eq!(tag["id"], expected_id);
assert_eq!(
tag["tagTimeRetained"].as_number().unwrap().to_string(),
"604800.000000000"
);
assert_eq!(tag["tagCreateTime"].as_array().unwrap().len(), 7);
assert_eq!(
row_count(
&sql_context,
&format!("SELECT * FROM paimon.test_db.t1 VERSION AS OF '{tag_name}'"),
)
.await,
expected_id as usize
);
}

// All snapshot files can be gone while tag metadata still supplies a candidate.
std::fs::remove_file(snapshot_dir.join("snapshot-3")).unwrap();
exec(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'tags_only', watermark => 1500)",
)
.await;
assert_eq!(read_tag(tmp.path(), "tags_only")["id"], 2);
}

#[tokio::test]
async fn test_create_tag_from_watermark_equal_watermarks_prefer_earlier_snapshot() {
for (retained_id, expire_first) in [(1, true), (2, false)] {
let (tmp, sql_context) =
setup_table_with_watermarks([Some(1000), Some(1000), Some(2000)]).await;
exec(
&sql_context,
&format!(
"CALL sys.create_tag(table => 'test_db.t1', tag => 'keep', snapshot_id => {retained_id})"
),
)
.await;
if expire_first {
std::fs::remove_file(tmp.path().join("test_db.db/t1/snapshot/snapshot-1")).unwrap();
}
exec(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'equal', watermark => 1000)",
)
.await;
assert_eq!(read_tag(tmp.path(), "equal")["id"], 1);
}
}

#[tokio::test]
async fn test_create_tag_from_watermark_skips_missing_live_and_tag_watermarks() {
let (tmp, sql_context) = setup_table_with_watermarks([None, Some(i64::MIN), Some(2000)]).await;
for id in [1, 2] {
exec(
&sql_context,
&format!(
"CALL sys.create_tag(table => 'test_db.t1', tag => 'missing_{id}', snapshot_id => {id})"
),
)
.await;
}
exec(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'live', watermark => 0)",
)
.await;
assert_eq!(read_tag(tmp.path(), "live")["id"], 3);
for id in [1, 2] {
std::fs::remove_file(
tmp.path()
.join(format!("test_db.db/t1/snapshot/snapshot-{id}")),
)
.unwrap();
}
for (index, watermark) in [i64::MIN, 0, 2000].into_iter().enumerate() {
let name = format!("valid_{index}");
exec(
&sql_context,
&format!(
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => '{name}', watermark => {watermark})"
),
)
.await;
assert_eq!(read_tag(tmp.path(), &name)["id"], 3);
}
}

#[tokio::test]
async fn test_create_tag_from_watermark_without_candidate_does_not_write() {
for (watermarks, target) in [
([None, None, None], 0),
([None, Some(i64::MIN), None], i64::MIN),
([Some(1000), Some(2000), Some(3000)], 3001),
] {
let (tmp, sql_context) = setup_table_with_watermarks(watermarks).await;
exec(
&sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'keep', snapshot_id => 2)",
)
.await;
let keep_path = tmp.path().join("test_db.db/t1/tag/tag-keep");
let original = std::fs::read(&keep_path).unwrap();
assert_sql_error(
&sql_context,
&format!(
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'missing', watermark => {target})"
),
&format!("No snapshot found with watermark >= {target}"),
)
.await;
assert!(!tmp.path().join("test_db.db/t1/tag/tag-missing").exists());
assert_eq!(std::fs::read(keep_path).unwrap(), original);
assert_eq!(
row_count(&sql_context, "SELECT * FROM paimon.test_db.`t1$tags`").await,
1
);
}

let (tmp, sql_context) = setup_sql_context().await;
exec(&sql_context, "CREATE TABLE paimon.test_db.t1 (id INT)").await;
assert_sql_error(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'empty', watermark => 0)",
"No snapshot found with watermark >= 0",
)
.await;
assert!(!tmp.path().join("test_db.db/t1/tag").exists());
}

#[tokio::test]
async fn test_create_tag_from_watermark_rejects_invalid_arguments_without_writes() {
let (tmp, sql_context) =
setup_table_with_watermarks([Some(1000), Some(2000), Some(3000)]).await;
exec(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'keep', watermark => 1000)",
)
.await;
let keep_path = tmp.path().join("test_db.db/t1/tag/tag-keep");
let original = std::fs::read(&keep_path).unwrap();

for (args, expected_error) in [
("tag => 'invalid', watermark => 1000", "Missing required argument: 'table'"),
("table => 'test_db.t1', watermark => 1000", "Missing required argument: 'tag'"),
("table => 'test_db.t1', tag => 'invalid'", "Missing required argument: 'watermark'"),
("table => 'test_db.t1', tag => 'invalid', watermark => 'abc'", "Invalid watermark"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1.5", "Invalid watermark"),
("table => 'test_db.t1', tag => 'invalid', watermark => 9223372036854775808", "Invalid watermark"),
("table => 'test_db.t1', tag => 'invalid', watermark => -9223372036854775809", "Invalid watermark"),
("table => 'test_db.t1', tag => 'invalid', watermak => 1000", "Argument watermak is unknown"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1000, time_retain => '1d'", "Argument time_retain is unknown"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1000, watermark => 2000", "Procedure argument watermark is duplicated"),
("table => 'test_db.t1', tag => '', watermark => 1000", "Invalid tag name"),
("table => 'test_db.t1', tag => '../invalid', watermark => 1000", "Invalid tag name"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1000, time_retained => ''", "Invalid tag retention"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1000, time_retained => 'invalid'", "Invalid tag retention"),
("table => 'test_db.t1', tag => 'invalid', watermark => 1000, time_retained => '-1d'", "Invalid tag retention"),
] {
assert_sql_error(
&sql_context,
&format!("CALL sys.create_tag_from_watermark({args})"),
expected_error,
)
.await;
assert_eq!(std::fs::read(&keep_path).unwrap(), original);
}
assert!(!tmp.path().join("test_db.db/t1/tag/tag-invalid").exists());
assert_eq!(
row_count(&sql_context, "SELECT * FROM paimon.test_db.`t1$tags`").await,
1
);
}

#[tokio::test]
async fn test_create_tag_from_watermark_preserves_existing_tag() {
let (tmp, sql_context) =
setup_table_with_watermarks([Some(1000), Some(2000), Some(3000)]).await;
exec(
&sql_context,
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'keep', watermark => 1000, time_retained => '1d')",
)
.await;
let keep_path = tmp.path().join("test_db.db/t1/tag/tag-keep");
let original = std::fs::read(&keep_path).unwrap();
for (watermark, retention) in [(2000, "2d"), (2000, "invalid"), (4000, "1d")] {
assert_sql_error(
&sql_context,
&format!(
"CALL sys.create_tag_from_watermark(table => 'test_db.t1', tag => 'keep', \
watermark => {watermark}, time_retained => '{retention}')"
),
"Tag 'keep' already exists",
)
.await;
assert_eq!(std::fs::read(&keep_path).unwrap(), original);
}
}

#[tokio::test]
async fn test_rollback_cleans_newer_tags() {
let (_tmp, sql_context) = setup_table_with_snapshots().await;
Expand Down
Loading
Loading