From ee6738df50e90951075c26dc3c68ba7728e18606 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Tue, 7 Jul 2026 21:26:57 -0700 Subject: [PATCH 01/13] datafusion: support reading branch tables --- .../integration_tests/tests/append_tables.rs | 9 + crates/integration_tests/tests/read_tables.rs | 2 + crates/integrations/datafusion/src/catalog.rs | 40 +++-- crates/integrations/datafusion/src/delete.rs | 6 +- .../integrations/datafusion/src/merge_into.rs | 2 + .../datafusion/src/physical_plan/sink.rs | 2 +- .../datafusion/src/sql_context.rs | 115 +++++++++--- .../datafusion/src/system_tables/manifests.rs | 4 +- .../datafusion/src/system_tables/mod.rs | 129 +++++++------ .../datafusion/src/system_tables/snapshots.rs | 7 +- .../src/system_tables/table_indexes.rs | 4 +- .../datafusion/src/system_tables/tags.rs | 7 +- .../integrations/datafusion/src/table/mod.rs | 6 + crates/integrations/datafusion/src/update.rs | 6 +- .../datafusion/tests/read_tables.rs | 2 + .../datafusion/tests/sql_context_tests.rs | 170 ++++++++++++++++++ crates/paimon-rest-server/tests/e2e.rs | 14 +- crates/paimon/src/catalog/mod.rs | 135 ++++++++++++++ crates/paimon/src/table/mod.rs | 77 +++++++- crates/paimon/src/table/table_commit.rs | 2 +- crates/paimon/src/table/table_scan.rs | 11 +- crates/paimon/src/table/time_travel.rs | 35 ++-- crates/paimon/src/table/write_builder.rs | 28 ++- docs/src/sql.md | 15 +- 24 files changed, 683 insertions(+), 145 deletions(-) diff --git a/crates/integration_tests/tests/append_tables.rs b/crates/integration_tests/tests/append_tables.rs index 6d22d5ea..09cdaf71 100644 --- a/crates/integration_tests/tests/append_tables.rs +++ b/crates/integration_tests/tests/append_tables.rs @@ -132,6 +132,7 @@ async fn write_commit_read(table: &Table, batches: Vec) -> Vec(catalog: &C) { .add_matched_batch(data_evolution_update_batch()) .expect("Failed to add data-evolution update batch"); wb.new_commit() + .unwrap() .commit( update .prepare_commit() @@ -220,6 +221,7 @@ async fn append_data_evolution_batch(table: &paimon::Table, batch: RecordBatch) .await .expect("Failed to write fixture batch"); wb.new_commit() + .unwrap() .commit( write .prepare_commit() diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index e74d3baa..dc788fd1 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -364,13 +364,13 @@ impl SchemaProvider for PaimonSchemaProvider { } } - let (base, system_name) = system_tables::split_object_name(name); - if let Some(system_name) = system_name { + let object = system_tables::parse_object_name_for_datafusion(name)?; + if let Some(system_name) = object.system_table().map(str::to_string) { return await_with_runtime(system_tables::load( Arc::clone(&self.catalog), self.database.clone(), - base.to_string(), - system_name.to_string(), + object, + system_name, )) .await; } @@ -378,10 +378,17 @@ impl SchemaProvider for PaimonSchemaProvider { let catalog = Arc::clone(&self.catalog); let dynamic_options = Arc::clone(&self.dynamic_options); let blob_reader_registry = self.blob_reader_registry.clone(); - let identifier = Identifier::new(self.database.clone(), base); + let identifier = Identifier::new(self.database.clone(), object.table().to_string()); + let branch = object.branch().map(str::to_string); await_with_runtime(async move { match catalog.get_table(&identifier).await { - Ok(table) => { + Ok(mut table) => { + if let Some(branch) = branch.as_deref() { + table = table + .copy_with_branch(branch) + .await + .map_err(to_datafusion_error)?; + } let opts = dynamic_options.read().unwrap().clone(); let provider = if opts.is_empty() { PaimonTableProvider::try_new_with_blob_reader_registry( @@ -419,19 +426,32 @@ impl SchemaProvider for PaimonSchemaProvider { } } - let (base, system_name) = system_tables::split_object_name(name); - if let Some(system_name) = system_name { + let object = match system_tables::parse_object_name_for_datafusion(name) { + Ok(object) => object, + Err(e) => { + log::error!("failed to parse Paimon object name '{name}': {e}"); + return false; + } + }; + if let Some(system_name) = object.system_table() { if !system_tables::is_registered(system_name) { return false; } } let catalog = Arc::clone(&self.catalog); - let identifier = Identifier::new(self.database.clone(), base.to_string()); + let identifier = Identifier::new(self.database.clone(), object.table().to_string()); + let branch = object.branch().map(str::to_string); block_on_with_runtime( async move { match catalog.get_table(&identifier).await { - Ok(_) => true, + Ok(table) => { + if let Some(branch) = branch.as_deref() { + table.copy_with_branch(branch).await.is_ok() + } else { + true + } + } Err(paimon::Error::TableNotExist { .. }) => false, Err(e) => { log::error!("failed to check table '{}': {e}", identifier); diff --git a/crates/integrations/datafusion/src/delete.rs b/crates/integrations/datafusion/src/delete.rs index 832acd92..84084c8d 100644 --- a/crates/integrations/datafusion/src/delete.rs +++ b/crates/integrations/datafusion/src/delete.rs @@ -115,6 +115,7 @@ async fn execute_data_evolution_delete_once( let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; if !messages.is_empty() { wb.new_commit() + .map_err(to_datafusion_error)? .commit(messages) .await .map_err(to_datafusion_error)?; @@ -165,7 +166,10 @@ async fn execute_cow_delete_once( let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; if !messages.is_empty() { - let commit = table.new_write_builder().new_commit(); + let commit = table + .new_write_builder() + .new_commit() + .map_err(to_datafusion_error)?; commit.commit(messages).await.map_err(to_datafusion_error)?; } diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index 203538d4..f7ec4546 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -394,6 +394,7 @@ async fn execute_cow_merge_once( if !all_messages.is_empty() { wb.new_commit() + .map_err(to_datafusion_error)? .commit(all_messages) .await .map_err(to_datafusion_error)?; @@ -769,6 +770,7 @@ async fn execute_merge_into_once( // 6. Commit all messages atomically if !all_messages.is_empty() { wb.new_commit() + .map_err(to_datafusion_error)? .commit(all_messages) .await .map_err(to_datafusion_error)?; diff --git a/crates/integrations/datafusion/src/physical_plan/sink.rs b/crates/integrations/datafusion/src/physical_plan/sink.rs index f021ba35..1abd76ab 100644 --- a/crates/integrations/datafusion/src/physical_plan/sink.rs +++ b/crates/integrations/datafusion/src/physical_plan/sink.rs @@ -92,7 +92,7 @@ impl DataSink for PaimonDataSink { } let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?; - let commit = wb.new_commit(); + let commit = wb.new_commit().map_err(to_datafusion_error)?; if self.overwrite { commit diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index babd650c..cec783e8 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -57,7 +57,7 @@ use datafusion::sql::sqlparser::ast::{ use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; use futures::StreamExt; -use paimon::catalog::{Catalog, Identifier}; +use paimon::catalog::{parse_object_name, Catalog, Identifier, DEFAULT_MAIN_BRANCH}; use paimon::spec::{ ArrayType as PaimonArrayType, BigIntType, BinaryType, BlobType, BooleanType, CharType, DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType, @@ -510,10 +510,8 @@ impl SQLContext { let (catalog, _catalog_name, identifier) = self.resolve_table_name_from_ref(&table_ref)?; - let paimon_table = catalog - .get_table(&identifier) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let (paimon_table, base_identifier, system_name) = + Self::load_table_for_read(&catalog, &identifier).await?; // Merge dynamic options with time-travel options let mut options = self.dynamic_options.read().unwrap().clone(); @@ -523,10 +521,22 @@ impl SQLContext { .copy_with_time_travel(options) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; - let provider = Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( - table_with_options, - self.blob_reader_registry.clone(), - )?); + let provider: Arc = if let Some(system_name) = system_name { + crate::system_tables::provider_for_table( + Arc::clone(&catalog), + base_identifier, + table_with_options, + &system_name, + )? + .ok_or_else(|| { + DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}")) + })? + } else { + Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( + table_with_options, + self.blob_reader_registry.clone(), + )?) + }; let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); self.register_temp_table(uuid_name.as_str(), provider)?; @@ -540,10 +550,8 @@ impl SQLContext { let (catalog, _catalog_name, identifier) = self.resolve_table_name_from_ref(&table_ref)?; - let paimon_table = catalog - .get_table(&identifier) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let (paimon_table, base_identifier, system_name) = + Self::load_table_for_read(&catalog, &identifier).await?; let millis = Self::parse_timestamp_to_millis(&info.timestamp)?; @@ -555,10 +563,22 @@ impl SQLContext { .copy_with_time_travel(options) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; - let provider = Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( - table_with_options, - self.blob_reader_registry.clone(), - )?); + let provider: Arc = if let Some(system_name) = system_name { + crate::system_tables::provider_for_table( + Arc::clone(&catalog), + base_identifier, + table_with_options, + &system_name, + )? + .ok_or_else(|| { + DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}")) + })? + } else { + Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( + table_with_options, + self.blob_reader_registry.clone(), + )?) + }; let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); self.register_temp_table(uuid_name.as_str(), provider)?; @@ -649,6 +669,37 @@ impl SQLContext { } } + async fn load_table_for_read( + catalog: &Arc, + identifier: &Identifier, + ) -> DFResult<(paimon::table::Table, Identifier, Option)> { + let parsed = identifier + .parsed_object_name() + .map_err(to_datafusion_error)?; + let base_identifier = Identifier::new( + identifier.database().to_string(), + parsed.table().to_string(), + ); + let mut table = catalog + .get_table(&base_identifier) + .await + .map_err(to_datafusion_error)?; + let system_table = parsed.system_table().map(str::to_string); + if let Some(branch) = parsed.branch() { + let is_branches_table = system_table + .as_deref() + .is_some_and(|name| name.eq_ignore_ascii_case("branches")); + if is_branches_table { + return Ok((table, base_identifier, system_table)); + } + table = table + .copy_with_branch(branch) + .await + .map_err(to_datafusion_error)?; + } + Ok((table, base_identifier, system_table)) + } + async fn handle_create_table( &self, catalog: &Arc, @@ -899,6 +950,7 @@ impl SQLContext { operations: &[AlterTableOperation], if_exists: bool, ) -> DFResult { + Self::ensure_main_branch_write_target(name, "ALTER TABLE")?; let identifier = self.resolve_table_name(name)?; let mut changes = Vec::new(); @@ -1029,6 +1081,7 @@ impl SQLContext { ))) } }; + Self::ensure_main_branch_write_target(&table_name, "MERGE INTO")?; let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?; let table = catalog @@ -1049,6 +1102,7 @@ impl SQLContext { ))) } }; + Self::ensure_main_branch_write_target(&table_name, "UPDATE")?; let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?; let table = catalog @@ -1076,6 +1130,7 @@ impl SQLContext { ))) } }; + Self::ensure_main_branch_write_target(&table_name, "DELETE")?; let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?; let table = catalog @@ -1097,6 +1152,7 @@ impl SQLContext { ))) } }; + Self::ensure_main_branch_write_target(&table_name, "INSERT OVERWRITE")?; let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?; let table = catalog .get_table(&identifier) @@ -1216,7 +1272,7 @@ impl SQLContext { } let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?; - let commit = wb.new_commit(); + let commit = wb.new_commit().map_err(to_datafusion_error)?; let overwrite_partitions = if static_partitions.is_empty() { None @@ -1241,6 +1297,7 @@ impl SQLContext { let target = truncate.table_names.first().ok_or_else(|| { DataFusionError::Plan("TRUNCATE TABLE requires a table name".to_string()) })?; + Self::ensure_main_branch_write_target(&target.name, "TRUNCATE TABLE")?; let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&target.name)?; let table = match catalog.get_table(&identifier).await { Ok(t) => t, @@ -1251,7 +1308,7 @@ impl SQLContext { }; let wb = table.new_write_builder(); - let commit = wb.new_commit(); + let commit = wb.new_commit().map_err(to_datafusion_error)?; if let Some(partitions) = &truncate.partitions { if partitions.is_empty() { @@ -1336,7 +1393,7 @@ impl SQLContext { )?; let wb = table.new_write_builder(); - let commit = wb.new_commit(); + let commit = wb.new_commit().map_err(to_datafusion_error)?; commit .truncate_partitions(partition_values) .await @@ -1425,6 +1482,24 @@ impl SQLContext { } } + fn ensure_main_branch_write_target(name: &ObjectName, operation: &str) -> DFResult<()> { + let object = name + .0 + .last() + .and_then(|part| part.as_ident()) + .map(|ident| ident.value.as_str()) + .ok_or_else(|| DataFusionError::Plan(format!("Invalid table reference: {name}")))?; + let parsed = parse_object_name(object).map_err(to_datafusion_error)?; + if let Some(branch) = parsed.branch() { + if branch != DEFAULT_MAIN_BRANCH { + return Err(DataFusionError::NotImplemented(format!( + "{operation} on Paimon branch '{branch}' is not supported" + ))); + } + } + Ok(()) + } + /// Resolve an ObjectName to just the Identifier (for backward compat in handle_alter_table). fn resolve_table_name(&self, name: &ObjectName) -> DFResult { let (_catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(name)?; diff --git a/crates/integrations/datafusion/src/system_tables/manifests.rs b/crates/integrations/datafusion/src/system_tables/manifests.rs index 4e438ae2..9d9f9886 100644 --- a/crates/integrations/datafusion/src/system_tables/manifests.rs +++ b/crates/integrations/datafusion/src/system_tables/manifests.rs @@ -29,7 +29,7 @@ use datafusion::error::Result as DFResult; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; use paimon::spec::{BinaryRow, DataField, ManifestFileMeta, ManifestList}; -use paimon::table::{SnapshotManager, Table}; +use paimon::table::Table; use super::row_string_cast::format_row_as_java_cast_string; use crate::error::to_datafusion_error; @@ -160,7 +160,7 @@ impl TableProvider for ManifestsTable { async fn collect_manifests(table: &Table) -> paimon::Result> { let file_io = table.file_io(); - let sm = SnapshotManager::new(file_io.clone(), table.location().to_string()); + let sm = table.snapshot_manager(); let snapshot = match sm.get_latest_snapshot().await? { Some(s) => s, None => return Ok(Vec::new()), diff --git a/crates/integrations/datafusion/src/system_tables/mod.rs b/crates/integrations/datafusion/src/system_tables/mod.rs index c46d3cd7..392db7da 100644 --- a/crates/integrations/datafusion/src/system_tables/mod.rs +++ b/crates/integrations/datafusion/src/system_tables/mod.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use datafusion::datasource::TableProvider; use datafusion::error::{DataFusionError, Result as DFResult}; -use paimon::catalog::{Catalog, Identifier, SYSTEM_BRANCH_PREFIX, SYSTEM_TABLE_SPLITTER}; +use paimon::catalog::{parse_object_name, Catalog, Identifier, ParsedObjectName}; use paimon::table::Table; use crate::error::to_datafusion_error; @@ -74,36 +74,16 @@ const SYSTEM_TABLE_NAMES: &[&str] = &[ "tags", ]; -/// Parse a Paimon object name into `(base_table, optional system_table_name)`. +/// Parse a Paimon object name into table, branch, and optional system table. /// /// Mirrors Java [Identifier.splitObjectName](https://github.com/apache/paimon/blob/release-1.3/paimon-api/src/main/java/org/apache/paimon/catalog/Identifier.java). /// -/// - `t` → `("t", None)` -/// - `t$options` → `("t", Some("options"))` -/// - `t$branch_main` → `("t", None)` (branch reference, not a system table) -/// - `t$branch_main$options` → `("t", Some("options"))` (branch + system table) -pub(crate) fn split_object_name(name: &str) -> (&str, Option<&str>) { - let mut parts = name.splitn(3, SYSTEM_TABLE_SPLITTER); - let base = parts.next().unwrap_or(name); - match (parts.next(), parts.next()) { - (None, _) => (base, None), - (Some(second), None) => { - if second.starts_with(SYSTEM_BRANCH_PREFIX) { - (base, None) - } else { - (base, Some(second)) - } - } - (Some(second), Some(third)) => { - if second.starts_with(SYSTEM_BRANCH_PREFIX) { - (base, Some(third)) - } else { - // `$` is legal in table names, so `t$foo$bar` falls through as - // plain `t` and errors later as "table not found". - (base, None) - } - } - } +/// - `t` → table `t` +/// - `t$options` → table `t`, system table `options` +/// - `t$branch_b1` → table `t`, branch `b1` +/// - `t$branch_b1$options` → table `t`, branch `b1`, system table `options` +pub(crate) fn parse_object_name_for_datafusion(name: &str) -> DFResult { + parse_object_name(name).map_err(to_datafusion_error) } /// Returns true if `name` is a recognised Paimon system table suffix. @@ -121,6 +101,27 @@ fn wrap_to_system_table(name: &str, base_table: Table) -> Option, + identifier: Identifier, + table: Table, + system_name: &str, +) -> DFResult>> { + if !is_registered(system_name) { + return Ok(None); + } + // Fail closed: system tables expose file metadata the client can't authorize. + paimon::spec::CoreOptions::new(table.schema().options()) + .ensure_read_authorized() + .map_err(to_datafusion_error)?; + if system_name.eq_ignore_ascii_case("partitions") { + return partitions::build(catalog, identifier, table).map(Some); + } + wrap_to_system_table(system_name, table) + .expect("is_registered guarantees a builder") + .map(Some) +} + /// Loads `$` from the catalog and wraps it as a system /// table provider. /// @@ -130,29 +131,29 @@ fn wrap_to_system_table(name: &str, base_table: Table) -> Option, database: String, - base: String, + object: ParsedObjectName, system_name: String, ) -> DFResult>> { if !is_registered(&system_name) { return Ok(None); } - let identifier = Identifier::new(database, base.clone()); + let identifier = Identifier::new(database, object.table().to_string()); match catalog.get_table(&identifier).await { - Ok(table) => { - // Fail closed: system tables expose file metadata the client can't authorize. - paimon::spec::CoreOptions::new(table.schema().options()) - .ensure_read_authorized() - .map_err(to_datafusion_error)?; - if system_name.eq_ignore_ascii_case("partitions") { - return partitions::build(catalog, identifier, table).map(Some); + Ok(mut table) => { + if let Some(branch) = object.branch() { + if !system_name.eq_ignore_ascii_case("branches") { + table = table + .copy_with_branch(branch) + .await + .map_err(to_datafusion_error)?; + } } - wrap_to_system_table(&system_name, table) - .expect("is_registered guarantees a builder") - .map(Some) + provider_for_table(catalog, identifier, table, &system_name) } Err(paimon::Error::TableNotExist { .. }) => Err(DataFusionError::Plan(format!( "Cannot read system table `${system_name}`: \ - base table `{base}` does not exist" + base table `{}` does not exist", + object.table() ))), Err(e) => Err(to_datafusion_error(e)), } @@ -160,7 +161,7 @@ pub(crate) async fn load( #[cfg(test)] mod tests { - use super::{is_registered, split_object_name, SYSTEM_TABLE_NAMES, TABLES}; + use super::{is_registered, parse_object_name_for_datafusion, SYSTEM_TABLE_NAMES, TABLES}; /// Guards against the two registries drifting: anything in `TABLES` must /// also be in `SYSTEM_TABLE_NAMES`, and the only name allowed to be in @@ -214,40 +215,54 @@ mod tests { #[test] fn plain_table_name() { - assert_eq!(split_object_name("orders"), ("orders", None)); + let parsed = parse_object_name_for_datafusion("orders").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), None); + assert_eq!(parsed.system_table(), None); } #[test] fn system_table_only() { - assert_eq!( - split_object_name("orders$options"), - ("orders", Some("options")) - ); + let parsed = parse_object_name_for_datafusion("orders$options").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), None); + assert_eq!(parsed.system_table(), Some("options")); } #[test] fn branch_reference_is_not_a_system_table() { - assert_eq!(split_object_name("orders$branch_main"), ("orders", None)); + let parsed = parse_object_name_for_datafusion("orders$branch_main").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), Some("main")); + assert_eq!(parsed.system_table(), None); + } + + #[test] + fn branch_reference_does_not_drop_branch_name() { + let parsed = parse_object_name_for_datafusion("orders$branch_b1").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), Some("b1")); + assert_eq!(parsed.system_table(), None); } #[test] fn branch_plus_system_table() { - assert_eq!( - split_object_name("orders$branch_main$options"), - ("orders", Some("options")) - ); + let parsed = parse_object_name_for_datafusion("orders$branch_main$options").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), Some("main")); + assert_eq!(parsed.system_table(), Some("options")); } #[test] fn three_parts_without_branch_prefix_is_not_a_system_table() { - assert_eq!(split_object_name("orders$foo$bar"), ("orders", None)); + assert!(parse_object_name_for_datafusion("orders$foo$bar").is_err()); } #[test] fn system_table_name_preserves_case() { - assert_eq!( - split_object_name("orders$Options"), - ("orders", Some("Options")) - ); + let parsed = parse_object_name_for_datafusion("orders$Options").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), None); + assert_eq!(parsed.system_table(), Some("Options")); } } diff --git a/crates/integrations/datafusion/src/system_tables/snapshots.rs b/crates/integrations/datafusion/src/system_tables/snapshots.rs index 17f3ca89..da3428f5 100644 --- a/crates/integrations/datafusion/src/system_tables/snapshots.rs +++ b/crates/integrations/datafusion/src/system_tables/snapshots.rs @@ -28,7 +28,7 @@ use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::Result as DFResult; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; -use paimon::table::{SnapshotManager, Table}; +use paimon::table::Table; use crate::error::to_datafusion_error; @@ -86,10 +86,7 @@ impl TableProvider for SnapshotsTable { _filters: &[Expr], _limit: Option, ) -> DFResult> { - let sm = SnapshotManager::new( - self.table.file_io().clone(), - self.table.location().to_string(), - ); + let sm = self.table.snapshot_manager(); let snapshots = crate::runtime::await_with_runtime(async move { sm.list_all().await }) .await .map_err(to_datafusion_error)?; diff --git a/crates/integrations/datafusion/src/system_tables/table_indexes.rs b/crates/integrations/datafusion/src/system_tables/table_indexes.rs index 43a8fd02..fd19af53 100644 --- a/crates/integrations/datafusion/src/system_tables/table_indexes.rs +++ b/crates/integrations/datafusion/src/system_tables/table_indexes.rs @@ -34,7 +34,7 @@ use datafusion::physical_plan::ExecutionPlan; use paimon::spec::{ BinaryRow, DataField, DeletionVectorMeta, FileKind, IndexManifest, IndexManifestEntry, }; -use paimon::table::{SnapshotManager, Table}; +use paimon::table::Table; use super::row_string_cast::format_row_as_java_cast_string; use crate::error::to_datafusion_error; @@ -187,7 +187,7 @@ impl TableProvider for TableIndexesTable { async fn collect_index_entries(table: &Table) -> paimon::Result> { let file_io = table.file_io(); - let sm = SnapshotManager::new(file_io.clone(), table.location().to_string()); + let sm = table.snapshot_manager(); let snapshot = match sm.get_latest_snapshot().await? { Some(s) => s, None => return Ok(Vec::new()), diff --git a/crates/integrations/datafusion/src/system_tables/tags.rs b/crates/integrations/datafusion/src/system_tables/tags.rs index 41f51ca3..5b9a22b4 100644 --- a/crates/integrations/datafusion/src/system_tables/tags.rs +++ b/crates/integrations/datafusion/src/system_tables/tags.rs @@ -30,7 +30,7 @@ use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::Result as DFResult; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; -use paimon::table::{Table, TagManager}; +use paimon::table::Table; use crate::error::to_datafusion_error; @@ -85,10 +85,7 @@ impl TableProvider for TagsTable { _filters: &[Expr], _limit: Option, ) -> DFResult> { - let tm = TagManager::new( - self.table.file_io().clone(), - self.table.location().to_string(), - ); + let tm = self.table.tag_manager(); let tags = crate::runtime::await_with_runtime(async move { tm.list_all().await }) .await .map_err(to_datafusion_error)?; diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index ba0a119a..e625adde 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -402,6 +402,12 @@ impl TableProvider for PaimonTableProvider { input: Arc, insert_op: InsertOp, ) -> DFResult> { + if !self.table.is_main_branch() { + return Err(datafusion::error::DataFusionError::NotImplemented(format!( + "Writing to Paimon branch '{}' is not supported", + self.table.branch() + ))); + } let overwrite = match insert_op { InsertOp::Append => false, InsertOp::Overwrite => true, diff --git a/crates/integrations/datafusion/src/update.rs b/crates/integrations/datafusion/src/update.rs index a987aaa8..06b78568 100644 --- a/crates/integrations/datafusion/src/update.rs +++ b/crates/integrations/datafusion/src/update.rs @@ -158,6 +158,7 @@ async fn execute_update_once( .map_err(to_datafusion_error)?; if !messages.is_empty() { wb.new_commit() + .map_err(to_datafusion_error)? .commit(messages) .await .map_err(to_datafusion_error)?; @@ -220,7 +221,10 @@ async fn execute_cow_update_once( let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; if !messages.is_empty() { - let commit = table.new_write_builder().new_commit(); + let commit = table + .new_write_builder() + .new_commit() + .map_err(to_datafusion_error)?; commit.commit(messages).await.map_err(to_datafusion_error)?; } diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index f2b8894e..2061fb29 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1803,6 +1803,7 @@ mod vector_search_tests { .expect("Failed to prepare commit"); write_builder .new_commit() + .unwrap() .commit(messages) .await .expect("Failed to commit vector data"); @@ -1886,6 +1887,7 @@ mod vector_search_tests { .expect("Failed to prepare commit"); write_builder .new_commit() + .unwrap() .commit(messages) .await .expect("Failed to commit vector data"); diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 1d5ea349..d0deb111 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -19,6 +19,7 @@ use std::sync::Arc; +use datafusion::arrow::array::Int64Array; use datafusion::catalog::CatalogProvider; use datafusion::datasource::MemTable; use paimon::catalog::Identifier; @@ -27,6 +28,7 @@ use paimon::spec::{ LocalZonedTimestampType, MapType, MultisetType, SchemaChange, TimeType, VarBinaryType, VarCharType, VectorType, }; +use paimon::table::{BranchManager, SnapshotManager, TagManager}; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::{PaimonCatalogProvider, SQLContext}; use tempfile::TempDir; @@ -46,6 +48,53 @@ async fn create_sql_context(catalog: Arc) -> SQLContext { ctx } +async fn collect_ids(sql_context: &SQLContext, sql: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut ids = Vec::new(); + for batch in batches { + let id_array = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("id column"); + for row in 0..batch.num_rows() { + ids.push(id_array.value(row)); + } + } + ids.sort_unstable(); + ids +} + +async fn collect_i64_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column_by_name(column) + .and_then(|c| c.as_any().downcast_ref::()) + .expect(column); + for row in 0..batch.num_rows() { + values.push(array.value(row)); + } + } + values.sort_unstable(); + values +} + +async fn assert_sql_error_contains(sql_context: &SQLContext, sql: &str, expected: &str) { + let err = match sql_context.sql(sql).await { + Ok(df) => df + .collect() + .await + .expect_err("SQL should fail but succeeded") + .to_string(), + Err(err) => err.to_string(), + }; + assert!( + err.contains(expected), + "expected error containing '{expected}', got: {err}" + ); +} + #[tokio::test] async fn test_show_tables_is_enabled() { let (_tmp, catalog) = create_test_env(); @@ -60,6 +109,127 @@ async fn test_show_tables_is_enabled() { .expect("SHOW TABLES should execute"); } +#[tokio::test] +async fn test_select_branch_table_reads_branch_snapshot() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + + sql_context + .sql("CREATE TABLE paimon.default.branch_orders (id INT, name STRING)") + .await + .unwrap() + .collect() + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.default.branch_orders VALUES (1, 'branch')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let identifier = Identifier::new("default", "branch_orders"); + let table = catalog.get_table(&identifier).await.unwrap(); + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let tag_manager = TagManager::new(table.file_io().clone(), table.location().to_string()); + tag_manager.create("branch_base", &snapshot).await.unwrap(); + let branch_manager = BranchManager::new(table.file_io().clone(), table.location().to_string()); + branch_manager + .create_branch_from_tag("b1", "branch_base") + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.default.branch_orders VALUES (2, 'main')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!( + collect_ids(&sql_context, "SELECT id FROM paimon.default.branch_orders").await, + vec![1, 2] + ); + assert_eq!( + collect_ids( + &sql_context, + "SELECT id FROM paimon.default.branch_orders$branch_b1" + ) + .await, + vec![1] + ); + assert_eq!( + collect_i64_column( + &sql_context, + "SELECT snapshot_id FROM paimon.default.branch_orders$snapshots", + "snapshot_id" + ) + .await, + vec![1, 2] + ); + assert_eq!( + collect_i64_column( + &sql_context, + "SELECT snapshot_id FROM paimon.default.branch_orders$branch_b1$snapshots", + "snapshot_id" + ) + .await, + vec![1] + ); + assert_eq!( + collect_i64_column( + &sql_context, + "SELECT record_count FROM paimon.default.branch_orders$files VERSION AS OF 'branch_base'", + "record_count" + ) + .await, + vec![1] + ); + assert_eq!( + collect_i64_column( + &sql_context, + "SELECT record_count FROM paimon.default.branch_orders$branch_b1$files VERSION AS OF 'branch_base'", + "record_count" + ) + .await, + vec![1] + ); + + let branch_table = table.copy_with_branch("b1").await.unwrap(); + let write_builder = branch_table.new_write_builder(); + assert!(write_builder.new_write().is_err()); + assert!(write_builder.new_update(vec!["name".to_string()]).is_err()); + assert!(write_builder.new_delete().is_err()); + assert!(write_builder.new_commit().is_err()); + + assert_sql_error_contains( + &sql_context, + "INSERT INTO paimon.default.branch_orders$branch_b1 VALUES (3, 'blocked')", + "Writing to Paimon branch 'b1' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "UPDATE paimon.default.branch_orders$branch_b1 SET name = 'blocked' WHERE id = 1", + "UPDATE on Paimon branch 'b1' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "DELETE FROM paimon.default.branch_orders$branch_b1 WHERE id = 1", + "DELETE on Paimon branch 'b1' is not supported", + ) + .await; +} + // ======================= CREATE / DROP SCHEMA ======================= #[tokio::test] diff --git a/crates/paimon-rest-server/tests/e2e.rs b/crates/paimon-rest-server/tests/e2e.rs index 6edf274b..390c8332 100644 --- a/crates/paimon-rest-server/tests/e2e.rs +++ b/crates/paimon-rest-server/tests/e2e.rs @@ -242,7 +242,12 @@ async fn test_write_commit_read_roundtrip() { writer.write_arrow_batch(&sample_batch()).await.unwrap(); let messages = writer.prepare_commit().await.unwrap(); assert!(!messages.is_empty(), "expected at least one commit message"); - write_builder.new_commit().commit(messages).await.unwrap(); + write_builder + .new_commit() + .unwrap() + .commit(messages) + .await + .unwrap(); // Read back. let table = cat.get_table(&ident).await.unwrap(); @@ -352,7 +357,12 @@ async fn test_list_partitions() { .unwrap(); let messages = writer.prepare_commit().await.unwrap(); assert!(!messages.is_empty(), "expected at least one commit message"); - write_builder.new_commit().commit(messages).await.unwrap(); + write_builder + .new_commit() + .unwrap() + .commit(messages) + .await + .unwrap(); // The partitions endpoint must serve the two partitions (no 404 fallback); // the client receives them directly from the server. diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 166d59a4..3276ae89 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -68,6 +68,35 @@ pub struct Identifier { object: String, } +/// Parsed form of a Paimon object name. +/// +/// Mirrors Java `Identifier.splitObjectName`: `table$branch_b1$snapshots` +/// resolves to table `table`, branch `b1`, and system table `snapshots`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ParsedObjectName { + table: String, + branch: Option, + system_table: Option, +} + +impl ParsedObjectName { + pub fn table(&self) -> &str { + &self.table + } + + pub fn branch(&self) -> Option<&str> { + self.branch.as_deref() + } + + pub fn branch_or_default(&self) -> &str { + self.branch.as_deref().unwrap_or(DEFAULT_MAIN_BRANCH) + } + + pub fn system_table(&self) -> Option<&str> { + self.system_table.as_deref() + } +} + impl Identifier { /// Create an identifier from database and object name. pub fn new(database: impl Into, object: impl Into) -> Self { @@ -112,6 +141,75 @@ impl Identifier { format!("{}.{}", self.database, self.object) } } + + /// Parse the object name into table, branch, and system-table components. + pub fn parsed_object_name(&self) -> Result { + parse_object_name(&self.object) + } + + pub fn table_name(&self) -> Result { + Ok(self.parsed_object_name()?.table) + } + + pub fn branch_name(&self) -> Result> { + Ok(self.parsed_object_name()?.branch) + } + + pub fn branch_name_or_default(&self) -> Result { + Ok(self + .branch_name()? + .unwrap_or_else(|| DEFAULT_MAIN_BRANCH.to_string())) + } + + pub fn system_table_name(&self) -> Result> { + Ok(self.parsed_object_name()?.system_table) + } +} + +/// Parse a Paimon object name into table, optional branch, and optional system table. +pub fn parse_object_name(object: &str) -> Result { + let parts: Vec<&str> = object.split(SYSTEM_TABLE_SPLITTER).collect(); + let invalid = || Error::IdentifierInvalid { + message: format!("Invalid object name: {object}"), + }; + let branch_from = |part: &str| { + let branch = part + .strip_prefix(SYSTEM_BRANCH_PREFIX) + .ok_or_else(invalid)?; + if branch.trim().is_empty() { + return Err(Error::IdentifierInvalid { + message: format!("Branch name cannot be empty in object name: {object}"), + }); + } + validate_identifier_name("branch", branch)?; + Ok(branch.to_string()) + }; + + match parts.as_slice() { + [table] => Ok(ParsedObjectName { + table: (*table).to_string(), + branch: None, + system_table: None, + }), + [table, second] if second.starts_with(SYSTEM_BRANCH_PREFIX) => Ok(ParsedObjectName { + table: (*table).to_string(), + branch: Some(branch_from(second)?), + system_table: None, + }), + [table, system_table] => Ok(ParsedObjectName { + table: (*table).to_string(), + branch: None, + system_table: Some((*system_table).to_string()), + }), + [table, branch, system_table] if branch.starts_with(SYSTEM_BRANCH_PREFIX) => { + Ok(ParsedObjectName { + table: (*table).to_string(), + branch: Some(branch_from(branch)?), + system_table: Some((*system_table).to_string()), + }) + } + _ => Err(invalid()), + } } fn validate_identifier_name(kind: &str, name: &str) -> Result<()> { @@ -300,6 +398,43 @@ pub trait Catalog: Send + Sync { mod tests { use super::*; + #[test] + fn parses_plain_object_name() { + let parsed = parse_object_name("orders").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), None); + assert_eq!(parsed.branch_or_default(), DEFAULT_MAIN_BRANCH); + assert_eq!(parsed.system_table(), None); + } + + #[test] + fn parses_branch_object_name() { + let parsed = parse_object_name("orders$branch_b1").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), Some("b1")); + assert_eq!(parsed.system_table(), None); + } + + #[test] + fn parses_branch_system_table_object_name() { + let parsed = parse_object_name("orders$branch_b1$snapshots").unwrap(); + assert_eq!(parsed.table(), "orders"); + assert_eq!(parsed.branch(), Some("b1")); + assert_eq!(parsed.system_table(), Some("snapshots")); + } + + #[test] + fn rejects_invalid_three_part_object_name() { + assert!(parse_object_name("orders$foo$bar").is_err()); + } + + #[test] + fn rejects_path_unsafe_branch_name() { + assert!(parse_object_name("orders$branch_../../other").is_err()); + assert!(parse_object_name("orders$branch_nested/name").is_err()); + assert!(parse_object_name("orders$branch_..").is_err()); + } + #[test] fn test_identifier_validate_should_reject_path_control_names() { for (database, object) in [ diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 70eef2db..67c325f5 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -106,7 +106,7 @@ pub use vector_search_builder::{BatchVectorSearchBuilder, VectorSearchBuilder}; pub use vindex_index_build_builder::VindexIndexBuildBuilder; pub use write_builder::WriteBuilder; -use crate::catalog::Identifier; +use crate::catalog::{Identifier, DEFAULT_MAIN_BRANCH}; use crate::io::FileIO; use crate::spec::{CoreOptions, DataField, Snapshot, TableSchema}; use std::collections::HashMap; @@ -119,6 +119,7 @@ pub struct Table { location: String, schema: TableSchema, schema_manager: SchemaManager, + branch: String, rest_env: Option, /// True when this table copy was switched to a historical schema by /// [`Table::copy_with_time_travel`]. Such a copy is read-only. @@ -139,12 +140,14 @@ impl Table { rest_env: Option, ) -> Self { let schema_manager = SchemaManager::new(file_io.clone(), location.clone()); + let branch = DEFAULT_MAIN_BRANCH.to_string(); Self { file_io, identifier, location, schema, schema_manager, + branch, rest_env, time_traveled: false, travel_snapshot: None, @@ -176,6 +179,32 @@ impl Table { &self.schema_manager } + pub fn branch(&self) -> &str { + &self.branch + } + + pub fn is_main_branch(&self) -> bool { + self.branch == DEFAULT_MAIN_BRANCH + } + + pub fn snapshot_manager(&self) -> SnapshotManager { + let manager = SnapshotManager::new(self.file_io.clone(), self.location.clone()); + if self.is_main_branch() { + manager + } else { + manager.with_branch(&self.branch) + } + } + + pub fn tag_manager(&self) -> TagManager { + let manager = TagManager::new(self.file_io.clone(), self.location.clone()); + if self.is_main_branch() { + manager + } else { + manager.with_branch(&self.branch) + } + } + /// Get the REST environment, if this table was loaded from a REST catalog. pub fn rest_env(&self) -> Option<&RESTEnv> { self.rest_env.as_ref() @@ -257,6 +286,7 @@ impl Table { location: self.location.clone(), schema: self.schema.copy_with_options(extra), schema_manager: self.schema_manager.clone(), + branch: self.branch.clone(), rest_env: self.rest_env.clone(), time_traveled: self.time_traveled, travel_snapshot: if selector_changed { @@ -287,9 +317,12 @@ impl Table { CoreOptions::new(table.schema().options()).validate_scan_options()?; // travel_to_snapshot returns Ok(None) without IO when the merged // options contain no selector. - if let Ok(Some(snapshot)) = - time_travel::travel_to_snapshot(&table.file_io, &table.location, table.schema.options()) - .await + if let Ok(Some(snapshot)) = time_travel::travel_to_snapshot( + &table.snapshot_manager(), + &table.tag_manager(), + table.schema.options(), + ) + .await { if snapshot.schema_id() != table.schema.id() { let snapshot_schema = table.schema_manager.schema(snapshot.schema_id()).await?; @@ -302,6 +335,42 @@ impl Table { Ok(table) } + pub async fn copy_with_branch(&self, branch_name: &str) -> Result { + let branch = if branch_name.trim().is_empty() { + return Err(crate::Error::DataInvalid { + message: "Branch name cannot be empty.".to_string(), + source: None, + }); + } else { + branch_name.to_string() + }; + let schema_manager = if branch == DEFAULT_MAIN_BRANCH { + SchemaManager::new(self.file_io.clone(), self.location.clone()) + } else { + SchemaManager::new(self.file_io.clone(), self.location.clone()).with_branch(&branch) + }; + let schema = schema_manager + .latest() + .await? + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("Branch '{branch}' does not exist."), + source: None, + })?; + let mut options = schema.options().clone(); + options.insert("branch".to_string(), branch.clone()); + Ok(Self { + file_io: self.file_io.clone(), + identifier: self.identifier.clone(), + location: self.location.clone(), + schema: schema.copy_with_replaced_options(options), + schema_manager, + branch, + rest_env: self.rest_env.clone(), + time_traveled: false, + travel_snapshot: None, + }) + } + /// Whether this table copy reads a historical snapshot with its /// historical schema (see [`Table::copy_with_time_travel`]). pub fn is_time_traveled(&self) -> bool { diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 453d8aa7..ae6fa989 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -73,7 +73,7 @@ pub struct TableCommit { } impl TableCommit { - pub fn new(table: Table, commit_user: String) -> Self { + pub(crate) fn new(table: Table, commit_user: String) -> Self { let snapshot_manager = SnapshotManager::new(table.file_io.clone(), table.location.clone()); let snapshot_commit = if let Some(env) = &table.rest_env { env.snapshot_commit() diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 67184a2c..0d9a3cd6 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -47,7 +47,6 @@ use crate::table::source::{ DataSplitBuilder, DeletionFile, PartitionBucket, Plan, RowRange, }; use crate::table::ScanTrace; -use crate::table::SnapshotManager; use futures::{StreamExt, TryStreamExt}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -768,20 +767,16 @@ impl<'a> TableScan<'a> { }); } - let file_io = self.table.file_io(); - let table_path = self.table.location(); - match super::time_travel::travel_to_snapshot( - file_io, - table_path, + &self.table.snapshot_manager(), + &self.table.tag_manager(), self.table.schema().options(), ) .await? { Some(snapshot) => Ok(Some(snapshot)), None => { - let snapshot_manager = - SnapshotManager::new(file_io.clone(), table_path.to_string()); + let snapshot_manager = self.table.snapshot_manager(); snapshot_manager.get_latest_snapshot().await } } diff --git a/crates/paimon/src/table/time_travel.rs b/crates/paimon/src/table/time_travel.rs index ebb67ae8..a70f6c10 100644 --- a/crates/paimon/src/table/time_travel.rs +++ b/crates/paimon/src/table/time_travel.rs @@ -17,7 +17,6 @@ //! Snapshot resolution for time travel, mirroring Java `TimeTravelUtil`. -use crate::io::FileIO; use crate::spec::{CoreOptions, Snapshot, TimeTravelSelector}; use crate::table::SnapshotManager; use crate::table::TagManager; @@ -31,12 +30,11 @@ use std::collections::HashMap; /// match any snapshot — callers that need Java `tryTravelToSnapshot`'s silent /// fallback (keep the current schema on failure) handle the `Err` themselves. pub(crate) async fn travel_to_snapshot( - file_io: &FileIO, - table_path: &str, + snapshot_manager: &SnapshotManager, + tag_manager: &TagManager, options: &HashMap, ) -> crate::Result> { let core_options = CoreOptions::new(options); - let snapshot_manager = SnapshotManager::new(file_io.clone(), table_path.to_string()); match core_options.try_time_travel_selector()? { Some(TimeTravelSelector::TimestampMillis(ts)) => { @@ -53,9 +51,8 @@ pub(crate) async fn travel_to_snapshot( option_name, }) => { // `scan.version` is ambiguous by design: tag first, then snapshot id. - let tag_manager = TagManager::new(file_io.clone(), table_path.to_string()); if tag_manager.tag_exists(v).await? { - resolve_tag(&tag_manager, v).await.map(Some) + resolve_tag(tag_manager, v).await.map(Some) } else if let Ok(id) = v.parse::() { snapshot_manager.get_snapshot(id).await.map(Some) } else { @@ -83,9 +80,8 @@ pub(crate) async fn travel_to_snapshot( option_name: _, }) => { // An explicit tag name: resolve strictly by tag, never as a snapshot id. - let tag_manager = TagManager::new(file_io.clone(), table_path.to_string()); if tag_manager.tag_exists(v).await? { - resolve_tag(&tag_manager, v).await.map(Some) + resolve_tag(tag_manager, v).await.map(Some) } else { Err(Error::DataInvalid { message: format!("Tag '{v}' doesn't exist."), @@ -415,7 +411,9 @@ mod tests { async fn test_invalid_snapshot_id_error_names_original_option() { let (file_io, table_path) = setup_evolved_table().await; let opts = options(&[("scan.snapshot-id", "abc")]); - let err = super::travel_to_snapshot(&file_io, &table_path, &opts) + let sm = SnapshotManager::new(file_io.clone(), table_path.clone()); + let tm = TagManager::new(file_io.clone(), table_path.clone()); + let err = super::travel_to_snapshot(&sm, &tm, &opts) .await .expect_err("non-numeric snapshot-id must fail"); match err { @@ -545,13 +543,9 @@ mod tests { // scan.snapshot-id must only accept a numeric snapshot id; it must not // fall back to resolving a tag of the same name. - let err = super::travel_to_snapshot( - &file_io, - &table_path, - &options(&[("scan.snapshot-id", "v1-tag")]), - ) - .await - .expect_err("scan.snapshot-id must not resolve a tag name"); + let err = super::travel_to_snapshot(&sm, &tm, &options(&[("scan.snapshot-id", "v1-tag")])) + .await + .expect_err("scan.snapshot-id must not resolve a tag name"); assert!( matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("scan.snapshot-id")), @@ -563,10 +557,11 @@ mod tests { async fn test_tag_name_selector_rejects_snapshot_id() { let (file_io, table_path) = setup_evolved_table().await; // No tag named "1" exists, but snapshot 1 does. - let err = - super::travel_to_snapshot(&file_io, &table_path, &options(&[("scan.tag-name", "1")])) - .await - .expect_err("scan.tag-name must not resolve a snapshot id"); + let sm = SnapshotManager::new(file_io.clone(), table_path.clone()); + let tm = TagManager::new(file_io.clone(), table_path.clone()); + let err = super::travel_to_snapshot(&sm, &tm, &options(&[("scan.tag-name", "1")])) + .await + .expect_err("scan.tag-name must not resolve a snapshot id"); assert!( matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("Tag '1'")), diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 9671e3ba..edf061b1 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -71,8 +71,12 @@ impl<'a> WriteBuilder<'a> { } /// Create a new TableCommit for committing write results. - pub fn new_commit(&self) -> TableCommit { - TableCommit::new(self.table.clone(), self.commit_user.clone()) + pub fn new_commit(&self) -> crate::Result { + self.ensure_main_branch_write()?; + Ok(TableCommit::new( + self.table.clone(), + self.commit_user.clone(), + )) } /// Create a new TableWrite for writing Arrow data. @@ -80,6 +84,7 @@ impl<'a> WriteBuilder<'a> { /// For primary-key tables, sequence numbers are lazily scanned per partition /// when the first writer for that partition is created. pub fn new_write(&self) -> crate::Result { + self.ensure_main_branch_write()?; // A table with a time-travel selector reads a pinned snapshot (and may // carry that snapshot's historical schema), so writing through the // same copy would be inconsistent with what its reads observe — even @@ -108,13 +113,28 @@ impl<'a> WriteBuilder<'a> { /// Create a new TableUpdate for data-evolution row-id updates. pub fn new_update(&self, update_columns: Vec) -> crate::Result { + self.ensure_main_branch_write()?; TableUpdate::new(self.table, update_columns) } /// Create a new writer for data-evolution row-id deletes. pub fn new_delete(&self) -> crate::Result { + self.ensure_main_branch_write()?; DataEvolutionDeleteWriter::new(self.table) } + + fn ensure_main_branch_write(&self) -> crate::Result<()> { + if self.table.is_main_branch() { + Ok(()) + } else { + Err(crate::Error::Unsupported { + message: format!( + "Writing to Paimon branch '{}' is not supported", + self.table.branch() + ), + }) + } + } } fn validate_commit_user(commit_user: &str) -> crate::Result<()> { @@ -305,7 +325,7 @@ mod tests { messages[0].new_files[0].file_name ); - wb.new_commit().commit(messages).await.unwrap(); + wb.new_commit().unwrap().commit(messages).await.unwrap(); let snapshot_manager = crate::table::SnapshotManager::new(file_io.clone(), table_path.to_string()); @@ -339,7 +359,7 @@ mod tests { "Overwrite-aware writer must not produce input changelog files" ); - wb.new_commit().commit(messages).await.unwrap(); + wb.new_commit().unwrap().commit(messages).await.unwrap(); let snapshot_manager = crate::table::SnapshotManager::new(file_io.clone(), table_path.to_string()); diff --git a/docs/src/sql.md b/docs/src/sql.md index 934eb279..3b42f463 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1466,12 +1466,23 @@ WHERE r.source = 'total'; ### Branch References -System tables support branch syntax: +Read a table branch with Java-compatible `$branch_` syntax: ```sql -SELECT * FROM paimon.default.my_table$branch_main$options; +SELECT * FROM paimon.default.my_table$branch_b1; ``` +System tables support the same branch syntax: + +```sql +SELECT * FROM paimon.default.my_table$branch_b1$options; +SELECT * FROM paimon.default.my_table$branch_b1$snapshots; +``` + +Branch references are read-only in DataFusion. `INSERT`, `UPDATE`, `DELETE`, +`MERGE INTO`, `TRUNCATE TABLE`, and `ALTER TABLE` against a branch reference are +rejected. + ## Table Options Set via `WITH ('key' = 'value')` at table creation time, or dynamically via `SET`. From 3d002ad6634ed2278359d85ba1255b6505566e7e Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 8 Jul 2026 00:46:45 -0700 Subject: [PATCH 02/13] datafusion: address branch review findings --- crates/integration_tests/tests/append_tables.rs | 9 --------- crates/integration_tests/tests/read_tables.rs | 2 -- crates/integrations/datafusion/src/catalog.rs | 6 ++++++ crates/integrations/datafusion/src/delete.rs | 4 ++-- crates/integrations/datafusion/src/merge_into.rs | 4 ++-- .../datafusion/src/physical_plan/sink.rs | 2 +- crates/integrations/datafusion/src/sql_context.rs | 6 +++--- .../datafusion/src/system_tables/manifests.rs | 9 ++++++--- .../datafusion/src/system_tables/table_indexes.rs | 9 ++++++--- crates/integrations/datafusion/src/table/mod.rs | 2 +- crates/integrations/datafusion/src/update.rs | 4 ++-- .../integrations/datafusion/tests/read_tables.rs | 2 -- .../datafusion/tests/sql_context_tests.rs | 8 +++++++- crates/paimon-rest-server/tests/e2e.rs | 14 ++------------ crates/paimon/src/catalog/mod.rs | 6 +++++- crates/paimon/src/table/mod.rs | 13 +++++++++++-- crates/paimon/src/table/write_builder.rs | 14 ++++++++++---- 17 files changed, 64 insertions(+), 50 deletions(-) diff --git a/crates/integration_tests/tests/append_tables.rs b/crates/integration_tests/tests/append_tables.rs index 09cdaf71..6d22d5ea 100644 --- a/crates/integration_tests/tests/append_tables.rs +++ b/crates/integration_tests/tests/append_tables.rs @@ -132,7 +132,6 @@ async fn write_commit_read(table: &Table, batches: Vec) -> Vec(catalog: &C) { .add_matched_batch(data_evolution_update_batch()) .expect("Failed to add data-evolution update batch"); wb.new_commit() - .unwrap() .commit( update .prepare_commit() @@ -221,7 +220,6 @@ async fn append_data_evolution_batch(table: &paimon::Table, batch: RecordBatch) .await .expect("Failed to write fixture batch"); wb.new_commit() - .unwrap() .commit( write .prepare_commit() diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index dc788fd1..7e32bae5 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -442,11 +442,17 @@ impl SchemaProvider for PaimonSchemaProvider { let catalog = Arc::clone(&self.catalog); let identifier = Identifier::new(self.database.clone(), object.table().to_string()); let branch = object.branch().map(str::to_string); + let is_branches_table = object + .system_table() + .is_some_and(|name| name.eq_ignore_ascii_case("branches")); block_on_with_runtime( async move { match catalog.get_table(&identifier).await { Ok(table) => { if let Some(branch) = branch.as_deref() { + if is_branches_table { + return true; + } table.copy_with_branch(branch).await.is_ok() } else { true diff --git a/crates/integrations/datafusion/src/delete.rs b/crates/integrations/datafusion/src/delete.rs index 84084c8d..11a14c02 100644 --- a/crates/integrations/datafusion/src/delete.rs +++ b/crates/integrations/datafusion/src/delete.rs @@ -114,7 +114,7 @@ async fn execute_data_evolution_delete_once( let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; if !messages.is_empty() { - wb.new_commit() + wb.try_new_commit() .map_err(to_datafusion_error)? .commit(messages) .await @@ -168,7 +168,7 @@ async fn execute_cow_delete_once( if !messages.is_empty() { let commit = table .new_write_builder() - .new_commit() + .try_new_commit() .map_err(to_datafusion_error)?; commit.commit(messages).await.map_err(to_datafusion_error)?; } diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index f7ec4546..0171e7f2 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -393,7 +393,7 @@ async fn execute_cow_merge_once( all_messages.extend(insert_messages); if !all_messages.is_empty() { - wb.new_commit() + wb.try_new_commit() .map_err(to_datafusion_error)? .commit(all_messages) .await @@ -769,7 +769,7 @@ async fn execute_merge_into_once( // 6. Commit all messages atomically if !all_messages.is_empty() { - wb.new_commit() + wb.try_new_commit() .map_err(to_datafusion_error)? .commit(all_messages) .await diff --git a/crates/integrations/datafusion/src/physical_plan/sink.rs b/crates/integrations/datafusion/src/physical_plan/sink.rs index 1abd76ab..3d40e2ad 100644 --- a/crates/integrations/datafusion/src/physical_plan/sink.rs +++ b/crates/integrations/datafusion/src/physical_plan/sink.rs @@ -92,7 +92,7 @@ impl DataSink for PaimonDataSink { } let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?; - let commit = wb.new_commit().map_err(to_datafusion_error)?; + let commit = wb.try_new_commit().map_err(to_datafusion_error)?; if self.overwrite { commit diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index cec783e8..27ecaa43 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -1272,7 +1272,7 @@ impl SQLContext { } let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?; - let commit = wb.new_commit().map_err(to_datafusion_error)?; + let commit = wb.try_new_commit().map_err(to_datafusion_error)?; let overwrite_partitions = if static_partitions.is_empty() { None @@ -1308,7 +1308,7 @@ impl SQLContext { }; let wb = table.new_write_builder(); - let commit = wb.new_commit().map_err(to_datafusion_error)?; + let commit = wb.try_new_commit().map_err(to_datafusion_error)?; if let Some(partitions) = &truncate.partitions { if partitions.is_empty() { @@ -1393,7 +1393,7 @@ impl SQLContext { )?; let wb = table.new_write_builder(); - let commit = wb.new_commit().map_err(to_datafusion_error)?; + let commit = wb.try_new_commit().map_err(to_datafusion_error)?; commit .truncate_partitions(partition_values) .await diff --git a/crates/integrations/datafusion/src/system_tables/manifests.rs b/crates/integrations/datafusion/src/system_tables/manifests.rs index 9d9f9886..e3b52081 100644 --- a/crates/integrations/datafusion/src/system_tables/manifests.rs +++ b/crates/integrations/datafusion/src/system_tables/manifests.rs @@ -161,9 +161,12 @@ impl TableProvider for ManifestsTable { async fn collect_manifests(table: &Table) -> paimon::Result> { let file_io = table.file_io(); let sm = table.snapshot_manager(); - let snapshot = match sm.get_latest_snapshot().await? { - Some(s) => s, - None => return Ok(Vec::new()), + let snapshot = match table.travel_snapshot().cloned() { + Some(snapshot) => snapshot, + None => match sm.get_latest_snapshot().await? { + Some(snapshot) => snapshot, + None => return Ok(Vec::new()), + }, }; let base_path = sm.manifest_path(snapshot.base_manifest_list()); diff --git a/crates/integrations/datafusion/src/system_tables/table_indexes.rs b/crates/integrations/datafusion/src/system_tables/table_indexes.rs index fd19af53..ecb843f4 100644 --- a/crates/integrations/datafusion/src/system_tables/table_indexes.rs +++ b/crates/integrations/datafusion/src/system_tables/table_indexes.rs @@ -188,9 +188,12 @@ impl TableProvider for TableIndexesTable { async fn collect_index_entries(table: &Table) -> paimon::Result> { let file_io = table.file_io(); let sm = table.snapshot_manager(); - let snapshot = match sm.get_latest_snapshot().await? { - Some(s) => s, - None => return Ok(Vec::new()), + let snapshot = match table.travel_snapshot().cloned() { + Some(snapshot) => snapshot, + None => match sm.get_latest_snapshot().await? { + Some(snapshot) => snapshot, + None => return Ok(Vec::new()), + }, }; let Some(index_manifest_name) = snapshot.index_manifest() else { return Ok(Vec::new()); diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index e625adde..f9294b3f 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -402,7 +402,7 @@ impl TableProvider for PaimonTableProvider { input: Arc, insert_op: InsertOp, ) -> DFResult> { - if !self.table.is_main_branch() { + if self.table.is_branch_reference() { return Err(datafusion::error::DataFusionError::NotImplemented(format!( "Writing to Paimon branch '{}' is not supported", self.table.branch() diff --git a/crates/integrations/datafusion/src/update.rs b/crates/integrations/datafusion/src/update.rs index 06b78568..e78037a5 100644 --- a/crates/integrations/datafusion/src/update.rs +++ b/crates/integrations/datafusion/src/update.rs @@ -157,7 +157,7 @@ async fn execute_update_once( .await .map_err(to_datafusion_error)?; if !messages.is_empty() { - wb.new_commit() + wb.try_new_commit() .map_err(to_datafusion_error)? .commit(messages) .await @@ -223,7 +223,7 @@ async fn execute_cow_update_once( if !messages.is_empty() { let commit = table .new_write_builder() - .new_commit() + .try_new_commit() .map_err(to_datafusion_error)?; commit.commit(messages).await.map_err(to_datafusion_error)?; } diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index 2061fb29..f2b8894e 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1803,7 +1803,6 @@ mod vector_search_tests { .expect("Failed to prepare commit"); write_builder .new_commit() - .unwrap() .commit(messages) .await .expect("Failed to commit vector data"); @@ -1887,7 +1886,6 @@ mod vector_search_tests { .expect("Failed to prepare commit"); write_builder .new_commit() - .unwrap() .commit(messages) .await .expect("Failed to commit vector data"); diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index d0deb111..0f48d6c3 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -208,7 +208,7 @@ async fn test_select_branch_table_reads_branch_snapshot() { assert!(write_builder.new_write().is_err()); assert!(write_builder.new_update(vec!["name".to_string()]).is_err()); assert!(write_builder.new_delete().is_err()); - assert!(write_builder.new_commit().is_err()); + assert!(write_builder.try_new_commit().is_err()); assert_sql_error_contains( &sql_context, @@ -216,6 +216,12 @@ async fn test_select_branch_table_reads_branch_snapshot() { "Writing to Paimon branch 'b1' is not supported", ) .await; + assert_sql_error_contains( + &sql_context, + "INSERT INTO paimon.default.branch_orders$branch_main VALUES (3, 'blocked')", + "Writing to Paimon branch 'main' is not supported", + ) + .await; assert_sql_error_contains( &sql_context, "UPDATE paimon.default.branch_orders$branch_b1 SET name = 'blocked' WHERE id = 1", diff --git a/crates/paimon-rest-server/tests/e2e.rs b/crates/paimon-rest-server/tests/e2e.rs index 390c8332..6edf274b 100644 --- a/crates/paimon-rest-server/tests/e2e.rs +++ b/crates/paimon-rest-server/tests/e2e.rs @@ -242,12 +242,7 @@ async fn test_write_commit_read_roundtrip() { writer.write_arrow_batch(&sample_batch()).await.unwrap(); let messages = writer.prepare_commit().await.unwrap(); assert!(!messages.is_empty(), "expected at least one commit message"); - write_builder - .new_commit() - .unwrap() - .commit(messages) - .await - .unwrap(); + write_builder.new_commit().commit(messages).await.unwrap(); // Read back. let table = cat.get_table(&ident).await.unwrap(); @@ -357,12 +352,7 @@ async fn test_list_partitions() { .unwrap(); let messages = writer.prepare_commit().await.unwrap(); assert!(!messages.is_empty(), "expected at least one commit message"); - write_builder - .new_commit() - .unwrap() - .commit(messages) - .await - .unwrap(); + write_builder.new_commit().commit(messages).await.unwrap(); // The partitions endpoint must serve the two partitions (no 404 fallback); // the client receives them directly from the server. diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 3276ae89..6d5db12b 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -181,7 +181,7 @@ pub fn parse_object_name(object: &str) -> Result { message: format!("Branch name cannot be empty in object name: {object}"), }); } - validate_identifier_name("branch", branch)?; + validate_branch_name(branch)?; Ok(branch.to_string()) }; @@ -212,6 +212,10 @@ pub fn parse_object_name(object: &str) -> Result { } } +pub(crate) fn validate_branch_name(name: &str) -> Result<()> { + validate_identifier_name("branch", name) +} + fn validate_identifier_name(kind: &str, name: &str) -> Result<()> { let invalid = if name.trim().is_empty() { Some("cannot be empty or whitespace") diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 67c325f5..6d75c9b4 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -106,7 +106,7 @@ pub use vector_search_builder::{BatchVectorSearchBuilder, VectorSearchBuilder}; pub use vindex_index_build_builder::VindexIndexBuildBuilder; pub use write_builder::WriteBuilder; -use crate::catalog::{Identifier, DEFAULT_MAIN_BRANCH}; +use crate::catalog::{validate_branch_name, Identifier, DEFAULT_MAIN_BRANCH}; use crate::io::FileIO; use crate::spec::{CoreOptions, DataField, Snapshot, TableSchema}; use std::collections::HashMap; @@ -120,6 +120,7 @@ pub struct Table { schema: TableSchema, schema_manager: SchemaManager, branch: String, + branch_reference: bool, rest_env: Option, /// True when this table copy was switched to a historical schema by /// [`Table::copy_with_time_travel`]. Such a copy is read-only. @@ -148,6 +149,7 @@ impl Table { schema, schema_manager, branch, + branch_reference: false, rest_env, time_traveled: false, travel_snapshot: None, @@ -187,6 +189,10 @@ impl Table { self.branch == DEFAULT_MAIN_BRANCH } + pub fn is_branch_reference(&self) -> bool { + self.branch_reference + } + pub fn snapshot_manager(&self) -> SnapshotManager { let manager = SnapshotManager::new(self.file_io.clone(), self.location.clone()); if self.is_main_branch() { @@ -287,6 +293,7 @@ impl Table { schema: self.schema.copy_with_options(extra), schema_manager: self.schema_manager.clone(), branch: self.branch.clone(), + branch_reference: self.branch_reference, rest_env: self.rest_env.clone(), time_traveled: self.time_traveled, travel_snapshot: if selector_changed { @@ -342,6 +349,7 @@ impl Table { source: None, }); } else { + validate_branch_name(branch_name)?; branch_name.to_string() }; let schema_manager = if branch == DEFAULT_MAIN_BRANCH { @@ -365,6 +373,7 @@ impl Table { schema: schema.copy_with_replaced_options(options), schema_manager, branch, + branch_reference: true, rest_env: self.rest_env.clone(), time_traveled: false, travel_snapshot: None, @@ -387,7 +396,7 @@ impl Table { /// The snapshot resolved by [`Table::copy_with_time_travel`] from this /// copy's options, if any. Lets scans skip re-resolving the selector. - pub(crate) fn travel_snapshot(&self) -> Option<&Snapshot> { + pub fn travel_snapshot(&self) -> Option<&Snapshot> { self.travel_snapshot.as_ref() } } diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index edf061b1..546ce342 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -71,7 +71,13 @@ impl<'a> WriteBuilder<'a> { } /// Create a new TableCommit for committing write results. - pub fn new_commit(&self) -> crate::Result { + pub fn new_commit(&self) -> TableCommit { + self.try_new_commit() + .expect("Cannot create TableCommit for a branch table") + } + + /// Try to create a new TableCommit for committing write results. + pub fn try_new_commit(&self) -> crate::Result { self.ensure_main_branch_write()?; Ok(TableCommit::new( self.table.clone(), @@ -124,7 +130,7 @@ impl<'a> WriteBuilder<'a> { } fn ensure_main_branch_write(&self) -> crate::Result<()> { - if self.table.is_main_branch() { + if !self.table.is_branch_reference() { Ok(()) } else { Err(crate::Error::Unsupported { @@ -325,7 +331,7 @@ mod tests { messages[0].new_files[0].file_name ); - wb.new_commit().unwrap().commit(messages).await.unwrap(); + wb.new_commit().commit(messages).await.unwrap(); let snapshot_manager = crate::table::SnapshotManager::new(file_io.clone(), table_path.to_string()); @@ -359,7 +365,7 @@ mod tests { "Overwrite-aware writer must not produce input changelog files" ); - wb.new_commit().unwrap().commit(messages).await.unwrap(); + wb.new_commit().commit(messages).await.unwrap(); let snapshot_manager = crate::table::SnapshotManager::new(file_io.clone(), table_path.to_string()); From 847ad3a978dfaa901b3904d2eb525e4d42336c00 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 8 Jul 2026 19:27:55 -0700 Subject: [PATCH 03/13] fix(core): reject branch references in commit paths --- .../table/btree_global_index_build_builder.rs | 2 + .../table/btree_global_index_drop_builder.rs | 2 + .../src/table/lumina_index_build_builder.rs | 2 + crates/paimon/src/table/mod.rs | 13 ++++++ crates/paimon/src/table/table_commit.rs | 14 ++++++ .../src/table/vindex_index_build_builder.rs | 2 + crates/paimon/src/table/write_builder.rs | 46 +++++++++++++++---- 7 files changed, 71 insertions(+), 10 deletions(-) diff --git a/crates/paimon/src/table/btree_global_index_build_builder.rs b/crates/paimon/src/table/btree_global_index_build_builder.rs index 73b9e7c1..cc87c70f 100644 --- a/crates/paimon/src/table/btree_global_index_build_builder.rs +++ b/crates/paimon/src/table/btree_global_index_build_builder.rs @@ -68,6 +68,8 @@ impl<'a> BTreeGlobalIndexBuildBuilder<'a> { } pub async fn execute(&self) -> Result { + self.table.ensure_not_branch_reference_for_write()?; + let index_type = normalize_sorted_global_index_type(&self.index_type).ok_or_else(|| { Error::Unsupported { message: format!( diff --git a/crates/paimon/src/table/btree_global_index_drop_builder.rs b/crates/paimon/src/table/btree_global_index_drop_builder.rs index a1c497e8..b56a1e2c 100644 --- a/crates/paimon/src/table/btree_global_index_drop_builder.rs +++ b/crates/paimon/src/table/btree_global_index_drop_builder.rs @@ -47,6 +47,8 @@ impl<'a> BTreeGlobalIndexDropBuilder<'a> { } pub async fn execute(&self) -> Result { + self.table.ensure_not_branch_reference_for_write()?; + let index_type = normalize_sorted_global_index_type(&self.index_type).ok_or_else(|| { Error::Unsupported { message: format!( diff --git a/crates/paimon/src/table/lumina_index_build_builder.rs b/crates/paimon/src/table/lumina_index_build_builder.rs index 5ead6de2..4a2980a6 100644 --- a/crates/paimon/src/table/lumina_index_build_builder.rs +++ b/crates/paimon/src/table/lumina_index_build_builder.rs @@ -71,6 +71,8 @@ impl<'a> LuminaIndexBuildBuilder<'a> { } pub async fn execute(&self) -> Result { + self.table.ensure_not_branch_reference_for_write()?; + if !is_lumina_index_type(&self.index_type) { return Err(Error::DataInvalid { message: format!("Unsupported Lumina index type: {}", self.index_type), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 312c73b5..0951c5a3 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -194,6 +194,19 @@ impl Table { self.branch_reference } + pub(crate) fn ensure_not_branch_reference_for_write(&self) -> Result<()> { + if self.is_branch_reference() { + Err(crate::Error::Unsupported { + message: format!( + "Writing to Paimon branch '{}' is not supported", + self.branch() + ), + }) + } else { + Ok(()) + } + } + pub fn snapshot_manager(&self) -> SnapshotManager { let manager = SnapshotManager::new(self.file_io.clone(), self.location.clone()); if self.is_main_branch() { diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 065f0177..5274888c 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -128,6 +128,8 @@ impl TableCommit { commit_messages: Vec, commit_identifier: i64, ) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + if commit_messages.is_empty() { return Ok(()); } @@ -168,6 +170,8 @@ impl TableCommit { expected_snapshot_id: i64, commit_identifier: i64, ) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + if commit_messages.is_empty() { return Ok(()); } @@ -217,6 +221,8 @@ impl TableCommit { static_partitions: Option>>, commit_identifier: i64, ) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + if commit_messages.is_empty() && static_partitions.is_none() { return Ok(()); } @@ -450,6 +456,8 @@ impl TableCommit { partitions: Vec>>, commit_identifier: i64, ) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + if partitions.is_empty() { return Ok(()); } @@ -489,6 +497,8 @@ impl TableCommit { partitions: Vec>>, commit_identifier: i64, ) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + if partitions.is_empty() { return Err(crate::Error::DataInvalid { message: "Partitions list cannot be empty.".to_string(), @@ -507,6 +517,8 @@ impl TableCommit { /// Truncate the entire table with a caller-provided commit identifier. pub async fn truncate_table_with_identifier(&self, commit_identifier: i64) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + self.try_commit( CommitEntriesPlan::Overwrite { partition_filter: None, @@ -529,6 +541,8 @@ impl TableCommit { /// files or storage errors are ignored so abort cleanup never masks the /// original write failure. pub async fn abort(&self, commit_messages: &[CommitMessage]) -> Result<()> { + self.table.ensure_not_branch_reference_for_write()?; + for message in commit_messages { let bucket_path = self.bucket_path(&message.partition, message.bucket)?; for file in message diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs b/crates/paimon/src/table/vindex_index_build_builder.rs index 7adb48ad..5803af82 100644 --- a/crates/paimon/src/table/vindex_index_build_builder.rs +++ b/crates/paimon/src/table/vindex_index_build_builder.rs @@ -62,6 +62,8 @@ impl<'a> VindexIndexBuildBuilder<'a> { } pub async fn execute(&self) -> Result { + self.table.ensure_not_branch_reference_for_write()?; + if !is_vindex_index_type(&self.index_type) { return Err(Error::DataInvalid { message: format!("Unsupported vindex index type: {}", self.index_type), diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 546ce342..0a0644fc 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -130,16 +130,7 @@ impl<'a> WriteBuilder<'a> { } fn ensure_main_branch_write(&self) -> crate::Result<()> { - if !self.table.is_branch_reference() { - Ok(()) - } else { - Err(crate::Error::Unsupported { - message: format!( - "Writing to Paimon branch '{}' is not supported", - self.table.branch() - ), - }) - } + self.table.ensure_not_branch_reference_for_write() } } @@ -220,6 +211,13 @@ mod tests { ) } + fn as_main_branch_reference(table: Table) -> Table { + Table { + branch_reference: true, + ..table + } + } + fn input_changelog_pk_table(file_io: &FileIO, table_path: &str) -> Table { let schema = Schema::builder() .column("id", DataType::Int(IntType::new())) @@ -343,6 +341,34 @@ mod tests { assert_eq!(snapshot.commit_user(), "my-commit-user"); } + #[tokio::test] + async fn test_branch_reference_rejects_write_and_index_builders() { + let table = as_main_branch_reference(test_postpone_pk_table( + &test_file_io(), + "memory:/test_branch_reference_writes", + )); + + let write_err = table.new_write_builder().try_new_commit().err().unwrap(); + assert!( + matches!(write_err, crate::Error::Unsupported { ref message } + if message == "Writing to Paimon branch 'main' is not supported"), + "Expected branch write rejection, got: {write_err:?}" + ); + + let index_err = table + .new_btree_global_index_build_builder() + .with_index_column("value") + .execute() + .await + .err() + .unwrap(); + assert!( + matches!(index_err, crate::Error::Unsupported { ref message } + if message == "Writing to Paimon branch 'main' is not supported"), + "Expected branch index-build rejection, got: {index_err:?}" + ); + } + #[tokio::test] async fn test_with_overwrite_marks_new_write_as_overwrite_aware() { let file_io = test_file_io(); From e49afcbaa0d54635bf2ca1f3f428f546b8cc56e0 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 8 Jul 2026 19:29:56 -0700 Subject: [PATCH 04/13] fix(datafusion): reject main branch write references --- .../datafusion/src/sql_context.rs | 10 ++--- .../datafusion/tests/sql_context_tests.rs | 40 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 7aeb3dfd..3d555c7d 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -57,7 +57,7 @@ use datafusion::sql::sqlparser::ast::{ use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; use futures::StreamExt; -use paimon::catalog::{parse_object_name, Catalog, Identifier, DEFAULT_MAIN_BRANCH}; +use paimon::catalog::{parse_object_name, Catalog, Identifier}; use paimon::spec::{ ArrayType as PaimonArrayType, BigIntType, BinaryType, BlobType, BooleanType, CharType, DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType, @@ -1492,11 +1492,9 @@ impl SQLContext { .ok_or_else(|| DataFusionError::Plan(format!("Invalid table reference: {name}")))?; let parsed = parse_object_name(object).map_err(to_datafusion_error)?; if let Some(branch) = parsed.branch() { - if branch != DEFAULT_MAIN_BRANCH { - return Err(DataFusionError::NotImplemented(format!( - "{operation} on Paimon branch '{branch}' is not supported" - ))); - } + return Err(DataFusionError::NotImplemented(format!( + "{operation} on Paimon branch '{branch}' is not supported" + ))); } Ok(()) } diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 0f48d6c3..1fcf916c 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -228,12 +228,52 @@ async fn test_select_branch_table_reads_branch_snapshot() { "UPDATE on Paimon branch 'b1' is not supported", ) .await; + assert_sql_error_contains( + &sql_context, + "UPDATE paimon.default.branch_orders$branch_main SET name = 'blocked' WHERE id = 1", + "UPDATE on Paimon branch 'main' is not supported", + ) + .await; assert_sql_error_contains( &sql_context, "DELETE FROM paimon.default.branch_orders$branch_b1 WHERE id = 1", "DELETE on Paimon branch 'b1' is not supported", ) .await; + assert_sql_error_contains( + &sql_context, + "DELETE FROM paimon.default.branch_orders$branch_main WHERE id = 1", + "DELETE on Paimon branch 'main' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "MERGE INTO paimon.default.branch_orders$branch_main AS target \ + USING (SELECT 1 AS id, 'blocked' AS name) AS source \ + ON target.id = source.id \ + WHEN MATCHED THEN UPDATE SET name = source.name", + "MERGE INTO on Paimon branch 'main' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "TRUNCATE TABLE paimon.default.branch_orders$branch_main", + "TRUNCATE TABLE on Paimon branch 'main' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "ALTER TABLE paimon.default.branch_orders$branch_main ADD COLUMN blocked INT", + "ALTER TABLE on Paimon branch 'main' is not supported", + ) + .await; + assert_sql_error_contains( + &sql_context, + "INSERT OVERWRITE paimon.default.branch_orders$branch_main \ + PARTITION (id = 1) SELECT 'blocked'", + "INSERT OVERWRITE on Paimon branch 'main' is not supported", + ) + .await; } // ======================= CREATE / DROP SCHEMA ======================= From eadcc9b53fc2a4e92032a5b5f1223864a5ef423b Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 02:31:07 -0700 Subject: [PATCH 05/13] fix(datafusion): respect branch partitions system table --- .../src/system_tables/partitions.rs | 24 +++-- .../datafusion/tests/sql_context_tests.rs | 89 +++++++++++++++++++ .../paimon/src/catalog/partition_listing.rs | 11 +-- 3 files changed, 112 insertions(+), 12 deletions(-) diff --git a/crates/integrations/datafusion/src/system_tables/partitions.rs b/crates/integrations/datafusion/src/system_tables/partitions.rs index f43501b0..6561d6cb 100644 --- a/crates/integrations/datafusion/src/system_tables/partitions.rs +++ b/crates/integrations/datafusion/src/system_tables/partitions.rs @@ -31,7 +31,7 @@ use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; -use paimon::catalog::{Catalog, Identifier}; +use paimon::catalog::{list_partitions_from_file_system, Catalog, Identifier}; use paimon::spec::{CoreOptions, Partition}; use paimon::table::Table; @@ -50,6 +50,7 @@ pub(super) fn build( Ok(Arc::new(PartitionsTable { catalog, identifier, + table, partition_keys, default_partition_name, })) @@ -87,6 +88,7 @@ fn partitions_schema() -> SchemaRef { struct PartitionsTable { catalog: Arc, identifier: Identifier, + table: Table, partition_keys: Vec, default_partition_name: String, } @@ -116,12 +118,20 @@ impl TableProvider for PartitionsTable { _filters: &[Expr], _limit: Option, ) -> DFResult> { - let catalog = self.catalog.clone(); - let identifier = self.identifier.clone(); - let partitions = crate::runtime::await_with_runtime(async move { - catalog.list_partitions(&identifier).await - }) - .await + let partitions = if self.table.is_branch_reference() { + let table = self.table.clone(); + crate::runtime::await_with_runtime(async move { + list_partitions_from_file_system(&table).await + }) + .await + } else { + let catalog = self.catalog.clone(); + let identifier = self.identifier.clone(); + crate::runtime::await_with_runtime( + async move { catalog.list_partitions(&identifier).await }, + ) + .await + } .map_err(to_datafusion_error)?; let mut rows: Vec<(String, Partition)> = partitions diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 99603d3d..1d6b9937 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -80,6 +80,24 @@ async fn collect_i64_column(sql_context: &SQLContext, sql: &str, column: &str) - values } +async fn collect_string_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column_by_name(column) + .and_then(|c| c.as_any().downcast_ref::()) + .expect(column); + for row in 0..batch.num_rows() { + if !array.is_null(row) { + values.push(array.value(row).to_string()); + } + } + } + values.sort_unstable(); + values +} + async fn assert_sql_error_contains(sql_context: &SQLContext, sql: &str, expected: &str) { let err = match sql_context.sql(sql).await { Ok(df) => df @@ -276,6 +294,77 @@ async fn test_select_branch_table_reads_branch_snapshot() { .await; } +#[tokio::test] +async fn test_branch_partitions_system_table_reads_branch_snapshot() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + + sql_context + .sql( + "CREATE TABLE paimon.default.branch_partition_orders \ + (id INT, name STRING) PARTITIONED BY (id)", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.default.branch_partition_orders VALUES (1, 'branch')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let identifier = Identifier::new("default", "branch_partition_orders"); + let table = catalog.get_table(&identifier).await.unwrap(); + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let tag_manager = TagManager::new(table.file_io().clone(), table.location().to_string()); + tag_manager + .create("partition_branch_base", &snapshot) + .await + .unwrap(); + let branch_manager = BranchManager::new(table.file_io().clone(), table.location().to_string()); + branch_manager + .create_branch_from_tag("b1", "partition_branch_base") + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.default.branch_partition_orders VALUES (2, 'main')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!( + collect_string_column( + &sql_context, + "SELECT \"partition\" FROM paimon.default.branch_partition_orders$partitions", + "partition", + ) + .await, + vec!["id=1".to_string(), "id=2".to_string()] + ); + assert_eq!( + collect_string_column( + &sql_context, + "SELECT \"partition\" FROM paimon.default.branch_partition_orders$branch_b1$partitions", + "partition", + ) + .await, + vec!["id=1".to_string()] + ); +} + // ======================= CREATE / DROP SCHEMA ======================= #[tokio::test] diff --git a/crates/paimon/src/catalog/partition_listing.rs b/crates/paimon/src/catalog/partition_listing.rs index 7652abae..46822149 100644 --- a/crates/paimon/src/catalog/partition_listing.rs +++ b/crates/paimon/src/catalog/partition_listing.rs @@ -33,14 +33,15 @@ use crate::Result; /// matching the shape catalogs would otherwise return from a metastore. pub async fn list_partitions_from_file_system(table: &Table) -> Result> { let file_io = table.file_io(); - let sm = SnapshotManager::new(file_io.clone(), table.location().to_string()); - let snapshot = match sm.get_latest_snapshot().await? { + let snapshot_sm = table.snapshot_manager(); + let manifest_sm = SnapshotManager::new(file_io.clone(), table.location().to_string()); + let snapshot = match snapshot_sm.get_latest_snapshot().await? { Some(s) => s, None => return Ok(Vec::new()), }; - let base_path = sm.manifest_path(snapshot.base_manifest_list()); - let delta_path = sm.manifest_path(snapshot.delta_manifest_list()); + let base_path = manifest_sm.manifest_path(snapshot.base_manifest_list()); + let delta_path = manifest_sm.manifest_path(snapshot.delta_manifest_list()); let (base_metas, delta_metas) = futures::try_join!( ManifestList::read(file_io, &base_path), ManifestList::read(file_io, &delta_path), @@ -48,7 +49,7 @@ pub async fn list_partitions_from_file_system(table: &Table) -> Result Date: Thu, 9 Jul 2026 03:10:04 -0700 Subject: [PATCH 06/13] fix(datafusion): read branch system metadata from root manifests --- .../datafusion/src/system_tables/manifests.rs | 12 ++++--- .../src/system_tables/partitions.rs | 31 ++++++++++--------- .../src/system_tables/table_indexes.rs | 8 +++-- .../datafusion/tests/sql_context_tests.rs | 27 ++++++++++++++++ .../paimon/src/catalog/partition_listing.rs | 9 ++++-- 5 files changed, 61 insertions(+), 26 deletions(-) diff --git a/crates/integrations/datafusion/src/system_tables/manifests.rs b/crates/integrations/datafusion/src/system_tables/manifests.rs index e3b52081..6f8748f7 100644 --- a/crates/integrations/datafusion/src/system_tables/manifests.rs +++ b/crates/integrations/datafusion/src/system_tables/manifests.rs @@ -160,20 +160,22 @@ impl TableProvider for ManifestsTable { async fn collect_manifests(table: &Table) -> paimon::Result> { let file_io = table.file_io(); - let sm = table.snapshot_manager(); + let snapshot_sm = table.snapshot_manager(); + let manifest_sm = + paimon::table::SnapshotManager::new(file_io.clone(), table.location().to_string()); let snapshot = match table.travel_snapshot().cloned() { Some(snapshot) => snapshot, - None => match sm.get_latest_snapshot().await? { + None => match snapshot_sm.get_latest_snapshot().await? { Some(snapshot) => snapshot, None => return Ok(Vec::new()), }, }; - let base_path = sm.manifest_path(snapshot.base_manifest_list()); - let delta_path = sm.manifest_path(snapshot.delta_manifest_list()); + let base_path = manifest_sm.manifest_path(snapshot.base_manifest_list()); + let delta_path = manifest_sm.manifest_path(snapshot.delta_manifest_list()); let changelog_path = snapshot .changelog_manifest_list() - .map(|c| sm.manifest_path(c)); + .map(|c| manifest_sm.manifest_path(c)); let base_fut = ManifestList::read(file_io, &base_path); let delta_fut = ManifestList::read(file_io, &delta_path); let changelog_fut = async { diff --git a/crates/integrations/datafusion/src/system_tables/partitions.rs b/crates/integrations/datafusion/src/system_tables/partitions.rs index 6561d6cb..769f3dbd 100644 --- a/crates/integrations/datafusion/src/system_tables/partitions.rs +++ b/crates/integrations/datafusion/src/system_tables/partitions.rs @@ -118,21 +118,22 @@ impl TableProvider for PartitionsTable { _filters: &[Expr], _limit: Option, ) -> DFResult> { - let partitions = if self.table.is_branch_reference() { - let table = self.table.clone(); - crate::runtime::await_with_runtime(async move { - list_partitions_from_file_system(&table).await - }) - .await - } else { - let catalog = self.catalog.clone(); - let identifier = self.identifier.clone(); - crate::runtime::await_with_runtime( - async move { catalog.list_partitions(&identifier).await }, - ) - .await - } - .map_err(to_datafusion_error)?; + let partitions = + if self.table.is_branch_reference() || self.table.travel_snapshot().is_some() { + let table = self.table.clone(); + crate::runtime::await_with_runtime(async move { + list_partitions_from_file_system(&table).await + }) + .await + } else { + let catalog = self.catalog.clone(); + let identifier = self.identifier.clone(); + crate::runtime::await_with_runtime(async move { + catalog.list_partitions(&identifier).await + }) + .await + } + .map_err(to_datafusion_error)?; let mut rows: Vec<(String, Partition)> = partitions .into_iter() diff --git a/crates/integrations/datafusion/src/system_tables/table_indexes.rs b/crates/integrations/datafusion/src/system_tables/table_indexes.rs index ecb843f4..b48a8a1c 100644 --- a/crates/integrations/datafusion/src/system_tables/table_indexes.rs +++ b/crates/integrations/datafusion/src/system_tables/table_indexes.rs @@ -187,10 +187,12 @@ impl TableProvider for TableIndexesTable { async fn collect_index_entries(table: &Table) -> paimon::Result> { let file_io = table.file_io(); - let sm = table.snapshot_manager(); + let snapshot_sm = table.snapshot_manager(); + let manifest_sm = + paimon::table::SnapshotManager::new(file_io.clone(), table.location().to_string()); let snapshot = match table.travel_snapshot().cloned() { Some(snapshot) => snapshot, - None => match sm.get_latest_snapshot().await? { + None => match snapshot_sm.get_latest_snapshot().await? { Some(snapshot) => snapshot, None => return Ok(Vec::new()), }, @@ -199,7 +201,7 @@ async fn collect_index_entries(table: &Table) -> paimon::Result Result s, - None => return Ok(Vec::new()), + let snapshot = match table.travel_snapshot().cloned() { + Some(snapshot) => snapshot, + None => match snapshot_sm.get_latest_snapshot().await? { + Some(snapshot) => snapshot, + None => return Ok(Vec::new()), + }, }; let base_path = manifest_sm.manifest_path(snapshot.base_manifest_list()); From 86c25d2b1a06d15749101fc881b392787186a5fc Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 03:11:02 -0700 Subject: [PATCH 07/13] fix(datafusion): reject reserved table suffixes on create --- .../datafusion/src/sql_context.rs | 9 +++++++ .../datafusion/tests/sql_context_tests.rs | 24 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index a7eacd06..a9506ddb 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -723,6 +723,15 @@ impl SQLContext { } let identifier = self.resolve_table_name(&ct.name)?; + let parsed = identifier + .parsed_object_name() + .map_err(to_datafusion_error)?; + if parsed.branch().is_some() || parsed.system_table().is_some() { + return Err(DataFusionError::Plan(format!( + "Cannot create Paimon table '{}': '$branch_' and '$' suffixes are reserved for branch and system-table reads", + identifier.object() + ))); + } let mut builder = paimon::spec::Schema::builder(); let table_options = extract_options(&ct.table_options)?; diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 3fcdd202..54c2192d 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -492,6 +492,30 @@ async fn test_create_table() { assert_eq!(schema.primary_keys(), &["id"]); } +#[tokio::test] +async fn test_create_table_rejects_reserved_branch_and_system_suffixes() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + + catalog + .create_database("mydb", false, Default::default()) + .await + .unwrap(); + + assert_sql_error_contains( + &sql_context, + "CREATE TABLE paimon.mydb.t$branch_b1 (id INT)", + "suffixes are reserved for branch and system-table reads", + ) + .await; + assert_sql_error_contains( + &sql_context, + "CREATE TABLE paimon.mydb.t$snapshots (id INT)", + "suffixes are reserved for branch and system-table reads", + ) + .await; +} + #[tokio::test] async fn test_create_table_with_blob_type() { let (_tmp, catalog) = create_test_env(); From e5526dacd4a863f3a3574f0e49adb7830c78443f Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 03:32:13 -0700 Subject: [PATCH 08/13] Revert "fix(datafusion): reject reserved table suffixes on create" This reverts commit 86c25d2b1a06d15749101fc881b392787186a5fc. --- .../datafusion/src/sql_context.rs | 9 ------- .../datafusion/tests/sql_context_tests.rs | 24 ------------------- 2 files changed, 33 deletions(-) diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index a9506ddb..a7eacd06 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -723,15 +723,6 @@ impl SQLContext { } let identifier = self.resolve_table_name(&ct.name)?; - let parsed = identifier - .parsed_object_name() - .map_err(to_datafusion_error)?; - if parsed.branch().is_some() || parsed.system_table().is_some() { - return Err(DataFusionError::Plan(format!( - "Cannot create Paimon table '{}': '$branch_' and '$' suffixes are reserved for branch and system-table reads", - identifier.object() - ))); - } let mut builder = paimon::spec::Schema::builder(); let table_options = extract_options(&ct.table_options)?; diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 54c2192d..3fcdd202 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -492,30 +492,6 @@ async fn test_create_table() { assert_eq!(schema.primary_keys(), &["id"]); } -#[tokio::test] -async fn test_create_table_rejects_reserved_branch_and_system_suffixes() { - let (_tmp, catalog) = create_test_env(); - let sql_context = create_sql_context(catalog.clone()).await; - - catalog - .create_database("mydb", false, Default::default()) - .await - .unwrap(); - - assert_sql_error_contains( - &sql_context, - "CREATE TABLE paimon.mydb.t$branch_b1 (id INT)", - "suffixes are reserved for branch and system-table reads", - ) - .await; - assert_sql_error_contains( - &sql_context, - "CREATE TABLE paimon.mydb.t$snapshots (id INT)", - "suffixes are reserved for branch and system-table reads", - ) - .await; -} - #[tokio::test] async fn test_create_table_with_blob_type() { let (_tmp, catalog) = create_test_env(); From 0398cbb02fd536e05eaef9b7e4016a5c2bb7f2e0 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 03:34:26 -0700 Subject: [PATCH 09/13] fix(core): keep branch commit creation non-panicking --- crates/paimon/src/table/table_commit.rs | 2 +- crates/paimon/src/table/write_builder.rs | 16 ++++++++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 5274888c..dad551f4 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -73,7 +73,7 @@ pub struct TableCommit { } impl TableCommit { - pub(crate) fn new(table: Table, commit_user: String) -> Self { + pub fn new(table: Table, commit_user: String) -> Self { let snapshot_manager = SnapshotManager::new(table.file_io.clone(), table.location.clone()); let snapshot_commit = if let Some(env) = &table.rest_env { env.snapshot_commit() diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 6c80f39f..3297da93 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -72,8 +72,7 @@ impl<'a> WriteBuilder<'a> { /// Create a new TableCommit for committing write results. pub fn new_commit(&self) -> TableCommit { - self.try_new_commit() - .expect("Cannot create TableCommit for a branch table") + TableCommit::new(self.table.clone(), self.commit_user.clone()) } /// Try to create a new TableCommit for committing write results. @@ -355,6 +354,19 @@ mod tests { "Expected branch write rejection, got: {write_err:?}" ); + let commit_err = table + .new_write_builder() + .new_commit() + .commit(Vec::new()) + .await + .err() + .unwrap(); + assert!( + matches!(commit_err, crate::Error::Unsupported { ref message } + if message == "Writing to Paimon branch 'main' is not supported"), + "Expected branch commit rejection, got: {commit_err:?}" + ); + let index_err = table .new_btree_global_index_build_builder() .with_index_column("value") From 10bb3fde71c56da23ab9b0d17d62c9c502d90315 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 21:12:24 -0700 Subject: [PATCH 10/13] fix(core): keep branch manifest paths at table root --- crates/paimon/src/table/snapshot_manager.rs | 57 +++++++++++++++++++-- 1 file changed, 54 insertions(+), 3 deletions(-) diff --git a/crates/paimon/src/table/snapshot_manager.rs b/crates/paimon/src/table/snapshot_manager.rs index 73effe69..d7d8fbbd 100644 --- a/crates/paimon/src/table/snapshot_manager.rs +++ b/crates/paimon/src/table/snapshot_manager.rs @@ -18,6 +18,7 @@ //! Snapshot manager for reading snapshot metadata using FileIO. //! //! Reference:[org.apache.paimon.utils.SnapshotManager](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/utils/SnapshotManager.java). +use crate::catalog::DEFAULT_MAIN_BRANCH; use crate::io::FileIO; use crate::spec::Snapshot; use futures::future::try_join_all; @@ -35,6 +36,7 @@ const EARLIEST_HINT: &str = "EARLIEST"; pub struct SnapshotManager { file_io: FileIO, table_path: String, + branch: String, } impl SnapshotManager { @@ -43,6 +45,7 @@ impl SnapshotManager { Self { file_io, table_path, + branch: DEFAULT_MAIN_BRANCH.to_string(), } } @@ -52,13 +55,26 @@ impl SnapshotManager { /// Path to the snapshot directory (e.g. `table_path/snapshot`). pub fn snapshot_dir(&self) -> String { - format!("{}/{}", self.table_path, SNAPSHOT_DIR) + let branch_path = if self.branch == DEFAULT_MAIN_BRANCH { + self.table_path.clone() + } else { + format!("{}/branch/branch-{}", self.table_path, self.branch) + }; + format!("{branch_path}/{SNAPSHOT_DIR}") } /// Create a SnapshotManager for a branch of this table. pub fn with_branch(&self, branch_name: &str) -> Self { - let branch_path = format!("{}/branch/branch-{}", self.table_path, branch_name); - Self::new(self.file_io.clone(), branch_path) + let branch = if branch_name.trim().is_empty() { + DEFAULT_MAIN_BRANCH + } else { + branch_name + }; + Self { + file_io: self.file_io.clone(), + table_path: self.table_path.clone(), + branch: branch.to_string(), + } } /// Path to the LATEST hint file. @@ -503,4 +519,39 @@ mod tests { let ids: Vec = snaps.iter().map(|s| s.id()).collect(); assert_eq!(ids, vec![1, 2, 3]); } + + #[test] + fn test_branch_scopes_snapshot_paths_only() { + let sm = SnapshotManager::new(test_file_io(), "memory:/test_branch_paths".to_string()); + let branch_sm = sm.with_branch("b1"); + + assert_eq!( + branch_sm.snapshot_path(1), + "memory:/test_branch_paths/branch/branch-b1/snapshot/snapshot-1" + ); + assert_eq!( + branch_sm.latest_hint_path(), + "memory:/test_branch_paths/branch/branch-b1/snapshot/LATEST" + ); + assert_eq!( + branch_sm.manifest_path("manifest-list-1"), + "memory:/test_branch_paths/manifest/manifest-list-1" + ); + + let other_branch_sm = branch_sm.with_branch("b2"); + assert_eq!( + other_branch_sm.snapshot_dir(), + "memory:/test_branch_paths/branch/branch-b2/snapshot" + ); + assert_eq!( + other_branch_sm.manifest_dir(), + "memory:/test_branch_paths/manifest" + ); + assert_eq!( + other_branch_sm + .with_branch(DEFAULT_MAIN_BRANCH) + .snapshot_dir(), + "memory:/test_branch_paths/snapshot" + ); + } } From de10281b40417e2e45225c5bcb6435a025e82f24 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 21:12:33 -0700 Subject: [PATCH 11/13] fix(datafusion): preserve branch partition metadata --- .../src/system_tables/partitions.rs | 39 ++-- .../datafusion/tests/sql_context_tests.rs | 196 +++++++++++++++++- 2 files changed, 215 insertions(+), 20 deletions(-) diff --git a/crates/integrations/datafusion/src/system_tables/partitions.rs b/crates/integrations/datafusion/src/system_tables/partitions.rs index 769f3dbd..97cf51b8 100644 --- a/crates/integrations/datafusion/src/system_tables/partitions.rs +++ b/crates/integrations/datafusion/src/system_tables/partitions.rs @@ -118,22 +118,31 @@ impl TableProvider for PartitionsTable { _filters: &[Expr], _limit: Option, ) -> DFResult> { - let partitions = - if self.table.is_branch_reference() || self.table.travel_snapshot().is_some() { - let table = self.table.clone(); - crate::runtime::await_with_runtime(async move { - list_partitions_from_file_system(&table).await - }) - .await + let table = self.table.clone(); + let partitions = if table.travel_snapshot().is_some() { + crate::runtime::await_with_runtime(async move { + list_partitions_from_file_system(&table).await + }) + .await + } else { + let catalog = self.catalog.clone(); + let identifier = if table.is_main_branch() { + self.identifier.clone() } else { - let catalog = self.catalog.clone(); - let identifier = self.identifier.clone(); - crate::runtime::await_with_runtime(async move { - catalog.list_partitions(&identifier).await - }) - .await - } - .map_err(to_datafusion_error)?; + Identifier::new( + self.identifier.database(), + format!("{}$branch_{}", self.identifier.object(), table.branch()), + ) + }; + crate::runtime::await_with_runtime(async move { + match catalog.list_partitions(&identifier).await { + Ok(partitions) => Ok(partitions), + Err(_) => list_partitions_from_file_system(&table).await, + } + }) + .await + } + .map_err(to_datafusion_error)?; let mut rows: Vec<(String, Partition)> = partitions .into_iter() diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 3fcdd202..d8934ee3 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -17,12 +17,14 @@ //! SQL context integration tests for paimon-datafusion. -use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use async_trait::async_trait; use datafusion::arrow::array::{Array, Int64Array}; use datafusion::catalog::CatalogProvider; use datafusion::datasource::MemTable; -use paimon::catalog::Identifier; +use paimon::catalog::{list_partitions_from_file_system, Identifier}; use paimon::spec::{ ArrayType, BinaryType, BlobType, CharType, DataType, FloatType, IntType, LocalZonedTimestampType, MapType, MultisetType, SchemaChange, TimeType, VarBinaryType, @@ -48,6 +50,142 @@ async fn create_sql_context(catalog: Arc) -> SQLContext { ctx } +struct PartitionCatalog { + inner: Arc, + fail_list_partitions: AtomicBool, + partition_identifiers: Mutex>, +} + +impl PartitionCatalog { + fn new(inner: Arc) -> Self { + Self { + inner, + fail_list_partitions: AtomicBool::new(false), + partition_identifiers: Mutex::new(Vec::new()), + } + } + + fn set_fail_list_partitions(&self, fail: bool) { + self.fail_list_partitions.store(fail, Ordering::SeqCst); + } + + fn take_partition_identifiers(&self) -> Vec { + std::mem::take(&mut *self.partition_identifiers.lock().unwrap()) + } +} + +#[async_trait] +impl Catalog for PartitionCatalog { + async fn list_databases(&self) -> paimon::Result> { + self.inner.list_databases().await + } + + async fn create_database( + &self, + name: &str, + ignore_if_exists: bool, + properties: std::collections::HashMap, + ) -> paimon::Result<()> { + self.inner + .create_database(name, ignore_if_exists, properties) + .await + } + + async fn get_database(&self, name: &str) -> paimon::Result { + self.inner.get_database(name).await + } + + async fn drop_database( + &self, + name: &str, + ignore_if_not_exists: bool, + cascade: bool, + ) -> paimon::Result<()> { + self.inner + .drop_database(name, ignore_if_not_exists, cascade) + .await + } + + async fn get_table(&self, identifier: &Identifier) -> paimon::Result { + self.inner.get_table(identifier).await + } + + async fn list_tables(&self, database_name: &str) -> paimon::Result> { + self.inner.list_tables(database_name).await + } + + async fn create_table( + &self, + identifier: &Identifier, + creation: paimon::spec::Schema, + ignore_if_exists: bool, + ) -> paimon::Result<()> { + self.inner + .create_table(identifier, creation, ignore_if_exists) + .await + } + + async fn drop_table( + &self, + identifier: &Identifier, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .drop_table(identifier, ignore_if_not_exists) + .await + } + + async fn rename_table( + &self, + from: &Identifier, + to: &Identifier, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .rename_table(from, to, ignore_if_not_exists) + .await + } + + async fn alter_table( + &self, + identifier: &Identifier, + changes: Vec, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .alter_table(identifier, changes, ignore_if_not_exists) + .await + } + + async fn list_partitions( + &self, + identifier: &Identifier, + ) -> paimon::Result> { + self.partition_identifiers + .lock() + .unwrap() + .push(identifier.clone()); + + let Some(branch) = identifier.branch_name()? else { + return self.inner.list_partitions(identifier).await; + }; + if self.fail_list_partitions.load(Ordering::SeqCst) { + return Err(paimon::Error::Unsupported { + message: "injected list_partitions failure".to_string(), + }); + } + + let base = Identifier::new(identifier.database(), identifier.table_name()?); + let table = self.inner.get_table(&base).await?; + let table = table.copy_with_branch(&branch).await?; + let mut partitions = list_partitions_from_file_system(&table).await?; + for partition in &mut partitions { + partition.created_by = Some("catalog".to_string()); + } + Ok(partitions) + } +} + async fn collect_ids(sql_context: &SQLContext, sql: &str) -> Vec { let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); let mut ids = Vec::new(); @@ -303,8 +441,13 @@ async fn test_select_branch_table_reads_branch_snapshot() { #[tokio::test] async fn test_branch_partitions_system_table_reads_branch_snapshot() { - let (_tmp, catalog) = create_test_env(); - let sql_context = create_sql_context(catalog.clone()).await; + let (_tmp, file_catalog) = create_test_env(); + let catalog = Arc::new(PartitionCatalog::new(file_catalog.clone())); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); sql_context .sql( @@ -325,7 +468,7 @@ async fn test_branch_partitions_system_table_reads_branch_snapshot() { .unwrap(); let identifier = Identifier::new("default", "branch_partition_orders"); - let table = catalog.get_table(&identifier).await.unwrap(); + let table = file_catalog.get_table(&identifier).await.unwrap(); let snapshot_manager = SnapshotManager::new(table.file_io().clone(), table.location().to_string()); let snapshot = snapshot_manager @@ -361,6 +504,40 @@ async fn test_branch_partitions_system_table_reads_branch_snapshot() { .await, vec!["id=1".to_string(), "id=2".to_string()] ); + catalog.take_partition_identifiers(); + + assert_eq!( + collect_string_column( + &sql_context, + "SELECT created_by FROM paimon.default.branch_partition_orders$branch_b1$partitions", + "created_by", + ) + .await, + vec!["catalog".to_string()] + ); + assert_eq!( + catalog.take_partition_identifiers(), + vec![Identifier::new( + "default", + "branch_partition_orders$branch_b1" + )] + ); + + assert_eq!( + collect_string_column( + &sql_context, + "SELECT \"partition\" FROM paimon.default.branch_partition_orders$branch_main$partitions", + "partition", + ) + .await, + vec!["id=1".to_string(), "id=2".to_string()] + ); + assert_eq!( + catalog.take_partition_identifiers(), + vec![Identifier::new("default", "branch_partition_orders")] + ); + + catalog.set_fail_list_partitions(true); assert_eq!( collect_string_column( &sql_context, @@ -370,6 +547,14 @@ async fn test_branch_partitions_system_table_reads_branch_snapshot() { .await, vec!["id=1".to_string()] ); + assert_eq!( + catalog.take_partition_identifiers(), + vec![Identifier::new( + "default", + "branch_partition_orders$branch_b1" + )] + ); + assert_eq!( collect_string_column( &sql_context, @@ -390,6 +575,7 @@ async fn test_branch_partitions_system_table_reads_branch_snapshot() { .await, vec!["id=1".to_string()] ); + assert!(catalog.take_partition_identifiers().is_empty()); } // ======================= CREATE / DROP SCHEMA ======================= From 7e4e9ed522111bda18218a75a031ba128620fe48 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 21:36:07 -0700 Subject: [PATCH 12/13] fix(core): delegate fallible commits for format tables --- crates/paimon/src/table/format_write_builder.rs | 5 +++++ crates/paimon/src/table/write_builder.rs | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/crates/paimon/src/table/format_write_builder.rs b/crates/paimon/src/table/format_write_builder.rs index 3608e822..ffd12af8 100644 --- a/crates/paimon/src/table/format_write_builder.rs +++ b/crates/paimon/src/table/format_write_builder.rs @@ -56,6 +56,11 @@ impl<'a> FormatWriteBuilder<'a> { TableCommit::new(self.table.clone(), self.commit_user.clone()) } + pub(crate) fn try_new_commit(&self) -> crate::Result { + self.table.ensure_not_branch_reference_for_write()?; + Ok(self.new_commit()) + } + pub(crate) fn new_write(&self) -> crate::Result { Err(crate::Error::Unsupported { message: "Writing format tables is not supported by the Rust client yet".to_string(), diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 10cc6b09..52ba5775 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -83,6 +83,14 @@ impl<'a> WriteBuilder<'a> { } } + /// Try to create a new TableCommit for committing write results. + pub fn try_new_commit(&self) -> crate::Result { + match &self.0 { + WriteBuilderKind::Paimon(builder) => builder.try_new_commit(), + WriteBuilderKind::Format(builder) => builder.try_new_commit(), + } + } + /// Create a new TableWrite for writing Arrow data. pub fn new_write(&self) -> crate::Result { match &self.0 { From fc2f5f9f056e520152518b93910aebad0c087625 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Thu, 9 Jul 2026 23:24:45 -0700 Subject: [PATCH 13/13] fix(datafusion): make read functions branch-aware --- .../integrations/datafusion/src/blob_view.rs | 6 +- .../datafusion/src/full_text_search.rs | 6 +- .../datafusion/src/hybrid_search.rs | 6 +- crates/integrations/datafusion/src/lib.rs | 1 + .../datafusion/src/sql_context.rs | 36 +--- .../datafusion/src/table_loader.rs | 69 +++++++ .../datafusion/src/vector_search.rs | 6 +- .../datafusion/tests/blob_tests.rs | 70 +++++++ .../datafusion/tests/read_tables.rs | 178 +++++++++++++++--- .../src/table/full_text_search_builder.rs | 12 +- .../paimon/src/table/vector_search_builder.rs | 12 +- 11 files changed, 312 insertions(+), 90 deletions(-) create mode 100644 crates/integrations/datafusion/src/table_loader.rs diff --git a/crates/integrations/datafusion/src/blob_view.rs b/crates/integrations/datafusion/src/blob_view.rs index 345da1ff..a91d4db1 100644 --- a/crates/integrations/datafusion/src/blob_view.rs +++ b/crates/integrations/datafusion/src/blob_view.rs @@ -36,6 +36,7 @@ use paimon::spec::BlobViewStruct; use crate::error::to_datafusion_error; use crate::runtime::block_on_with_runtime; use crate::table_function_args::parse_table_identifier; +use crate::table_loader::load_data_table_for_read; const FUNCTION_NAME: &str = "blob_view"; @@ -68,10 +69,9 @@ impl BlobViewFunc { let identifier = parse_table_identifier(FUNCTION_NAME, table_name, &self.default_database)?; let catalog = Arc::clone(&self.catalog); let table = block_on_with_runtime( - async move { catalog.get_table(&identifier).await }, + async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await }, "blob_view: catalog access thread panicked", - ) - .map_err(to_datafusion_error)?; + )?; let field = table .schema() diff --git a/crates/integrations/datafusion/src/full_text_search.rs b/crates/integrations/datafusion/src/full_text_search.rs index d5f74d36..27c8db1b 100644 --- a/crates/integrations/datafusion/src/full_text_search.rs +++ b/crates/integrations/datafusion/src/full_text_search.rs @@ -46,6 +46,7 @@ use crate::table::{PaimonScanBuilder, PaimonTableProvider}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; +use crate::table_loader::load_data_table_for_read; const FUNCTION_NAME: &str = "full_text_search"; @@ -110,10 +111,9 @@ impl TableFunctionImpl for FullTextSearchFunction { let catalog = Arc::clone(&self.catalog); let table = block_on_with_runtime( - async move { catalog.get_table(&identifier).await }, + async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await }, "full_text_search: catalog access thread panicked", - ) - .map_err(to_datafusion_error)?; + )?; let inner = PaimonTableProvider::try_new(table)?; diff --git a/crates/integrations/datafusion/src/hybrid_search.rs b/crates/integrations/datafusion/src/hybrid_search.rs index e6083668..2d5f997c 100644 --- a/crates/integrations/datafusion/src/hybrid_search.rs +++ b/crates/integrations/datafusion/src/hybrid_search.rs @@ -51,6 +51,7 @@ use crate::table::{PaimonScanBuilder, PaimonTableProvider}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; +use crate::table_loader::load_data_table_for_read; const FUNCTION_NAME: &str = "hybrid_search"; @@ -123,10 +124,9 @@ impl TableFunctionImpl for HybridSearchFunction { parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?; let catalog = Arc::clone(&self.catalog); let table = block_on_with_runtime( - async move { catalog.get_table(&identifier).await }, + async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await }, "hybrid_search: catalog access thread panicked", - ) - .map_err(to_datafusion_error)?; + )?; Ok(Arc::new(HybridSearchTableProvider { inner: PaimonTableProvider::try_new(table)?, diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 675fda40..9c1d9143 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -55,6 +55,7 @@ mod sql_context; mod system_tables; mod table; mod table_function_args; +mod table_loader; mod update; mod variant_functions; mod variant_pushdown; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index a7eacd06..1513ea9d 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -67,6 +67,7 @@ use paimon::spec::{ }; use crate::error::to_datafusion_error; +use crate::table_loader::load_table_for_read; use crate::{BlobReaderRegistry, DynamicOptions}; /// A SQL context that supports registering multiple Paimon catalogs and executing SQL. @@ -511,7 +512,7 @@ impl SQLContext { self.resolve_table_name_from_ref(&table_ref)?; let (paimon_table, base_identifier, system_name) = - Self::load_table_for_read(&catalog, &identifier).await?; + load_table_for_read(&catalog, &identifier).await?; // Merge dynamic options with time-travel options let mut options = self.dynamic_options.read().unwrap().clone(); @@ -551,7 +552,7 @@ impl SQLContext { self.resolve_table_name_from_ref(&table_ref)?; let (paimon_table, base_identifier, system_name) = - Self::load_table_for_read(&catalog, &identifier).await?; + load_table_for_read(&catalog, &identifier).await?; let millis = Self::parse_timestamp_to_millis(&info.timestamp)?; @@ -669,37 +670,6 @@ impl SQLContext { } } - async fn load_table_for_read( - catalog: &Arc, - identifier: &Identifier, - ) -> DFResult<(paimon::table::Table, Identifier, Option)> { - let parsed = identifier - .parsed_object_name() - .map_err(to_datafusion_error)?; - let base_identifier = Identifier::new( - identifier.database().to_string(), - parsed.table().to_string(), - ); - let mut table = catalog - .get_table(&base_identifier) - .await - .map_err(to_datafusion_error)?; - let system_table = parsed.system_table().map(str::to_string); - if let Some(branch) = parsed.branch() { - let is_branches_table = system_table - .as_deref() - .is_some_and(|name| name.eq_ignore_ascii_case("branches")); - if is_branches_table { - return Ok((table, base_identifier, system_table)); - } - table = table - .copy_with_branch(branch) - .await - .map_err(to_datafusion_error)?; - } - Ok((table, base_identifier, system_table)) - } - async fn handle_create_table( &self, catalog: &Arc, diff --git a/crates/integrations/datafusion/src/table_loader.rs b/crates/integrations/datafusion/src/table_loader.rs new file mode 100644 index 00000000..f1d61f9a --- /dev/null +++ b/crates/integrations/datafusion/src/table_loader.rs @@ -0,0 +1,69 @@ +// 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. + +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result as DFResult}; +use paimon::catalog::{Catalog, Identifier}; +use paimon::table::Table; + +use crate::error::to_datafusion_error; + +pub(crate) async fn load_table_for_read( + catalog: &Arc, + identifier: &Identifier, +) -> DFResult<(Table, Identifier, Option)> { + let parsed = identifier + .parsed_object_name() + .map_err(to_datafusion_error)?; + let base_identifier = Identifier::new( + identifier.database().to_string(), + parsed.table().to_string(), + ); + let mut table = catalog + .get_table(&base_identifier) + .await + .map_err(to_datafusion_error)?; + let system_table = parsed.system_table().map(str::to_string); + if let Some(branch) = parsed.branch() { + let is_branches_table = system_table + .as_deref() + .is_some_and(|name| name.eq_ignore_ascii_case("branches")); + if is_branches_table { + return Ok((table, base_identifier, system_table)); + } + table = table + .copy_with_branch(branch) + .await + .map_err(to_datafusion_error)?; + } + Ok((table, base_identifier, system_table)) +} + +pub(crate) async fn load_data_table_for_read( + catalog: &Arc, + identifier: &Identifier, + caller: &str, +) -> DFResult { + let (table, _, system_table) = load_table_for_read(catalog, identifier).await?; + if system_table.is_some() { + return Err(DataFusionError::Plan(format!( + "{caller} requires a data table" + ))); + } + Ok(table) +} diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index 694f21bb..d101cce6 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -37,6 +37,7 @@ use crate::table::{PaimonScanBuilder, PaimonTableProvider}; use crate::table_function_args::{ extract_int_literal, extract_string_literal, parse_table_identifier, }; +use crate::table_loader::load_data_table_for_read; const FUNCTION_NAME: &str = "vector_search"; @@ -96,10 +97,9 @@ impl TableFunctionImpl for VectorSearchFunction { let catalog = Arc::clone(&self.catalog); let table = block_on_with_runtime( - async move { catalog.get_table(&identifier).await }, + async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await }, "vector_search: catalog access thread panicked", - ) - .map_err(to_datafusion_error)?; + )?; let inner = PaimonTableProvider::try_new(table)?; let query_vector_json = diff --git a/crates/integrations/datafusion/tests/blob_tests.rs b/crates/integrations/datafusion/tests/blob_tests.rs index 10d46248..84fad62a 100644 --- a/crates/integrations/datafusion/tests/blob_tests.rs +++ b/crates/integrations/datafusion/tests/blob_tests.rs @@ -23,7 +23,10 @@ mod common; use arrow_array::{Array, BinaryArray, Int32Array, RecordBatch, StringArray}; use common::{create_sql_context, create_test_env, exec}; +use paimon::catalog::Identifier; use paimon::spec::{BlobDescriptor, BlobViewStruct}; +use paimon::table::BranchManager; +use paimon::Catalog; use paimon_datafusion::SQLContext; // ======================= Helpers ======================= @@ -872,6 +875,73 @@ async fn test_blob_view_without_rest_env_preserves_reference() { drop(tmp); } +#[tokio::test] +async fn test_blob_view_preserves_branch_reference() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + sql_context + .sql("CREATE SCHEMA paimon.test_db") + .await + .unwrap(); + sql_context + .sql( + "CREATE TABLE paimon.test_db.src (\ + id INT, \ + picture BLOB\ + ) WITH (\ + 'data-evolution.enabled' = 'true', \ + 'row-tracking.enabled' = 'true'\ + )", + ) + .await + .unwrap(); + exec( + &sql_context, + "INSERT INTO paimon.test_db.src (id, picture) VALUES (1, X'616C696365')", + ) + .await; + + let table = catalog + .get_table(&Identifier::new("test_db", "src")) + .await + .expect("load table"); + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .expect("load latest snapshot") + .expect("latest snapshot"); + table + .tag_manager() + .create("branch-source", &snapshot) + .await + .expect("create tag"); + BranchManager::new(table.file_io().clone(), table.location().to_string()) + .create_branch_from_tag("b1", "branch-source") + .await + .expect("create branch"); + + let batches = sql_context + .sql( + "SELECT blob_view('test_db.src$branch_b1', 'picture', \"_ROW_ID\") AS picture \ + FROM paimon.test_db.src", + ) + .await + .expect("blob_view SQL should parse") + .collect() + .await + .expect("branch blob_view should execute"); + let value = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("binary blob view") + .value(0); + let view = BlobViewStruct::deserialize(value).expect("deserialize blob view"); + assert_eq!(view.identifier().full_name(), "test_db.src$branch_b1"); + assert_eq!(view.row_id(), 0); +} + #[tokio::test] async fn test_blob_view_resolve_disabled_preserves_reference() { let (_tmp, sql_context) = setup( diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index cb9a68f9..55db55b5 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1399,6 +1399,8 @@ mod fulltext_tests { use std::sync::Arc; use datafusion::arrow::array::{Int32Array, StringArray}; + use paimon::catalog::Identifier; + use paimon::table::BranchManager; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::{register_full_text_search, SQLContext}; @@ -1420,19 +1422,18 @@ mod fulltext_tests { (tmp, warehouse) } - async fn create_fulltext_context() -> (SQLContext, tempfile::TempDir) { + async fn create_fulltext_context() -> (SQLContext, Arc, tempfile::TempDir) { let (tmp, warehouse) = extract_test_warehouse(); let mut options = Options::new(); options.set(CatalogOptions::WAREHOUSE, warehouse); - let catalog = FileSystemCatalog::new(options).expect("Failed to create catalog"); - let catalog: Arc = Arc::new(catalog); + let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog")); let mut ctx = SQLContext::new(); ctx.register_catalog("paimon", catalog.clone()) .await .expect("Failed to register catalog"); - register_full_text_search(ctx.ctx(), catalog, "default"); - (ctx, tmp) + register_full_text_search(ctx.ctx(), catalog.clone(), "default"); + (ctx, catalog, tmp) } fn extract_id_content_rows( @@ -1459,7 +1460,7 @@ mod fulltext_tests { /// Search for 'paimon' — rows 0, 2, 4 mention "paimon". #[tokio::test] async fn test_full_text_search_paimon() { - let (ctx, _tmp) = create_fulltext_context().await; + let (ctx, _catalog, _tmp) = create_fulltext_context().await; let batches = ctx .sql("SELECT id, content FROM full_text_search('paimon.default.test_tantivy_fulltext', 'content', 'paimon', 10)") .await @@ -1477,10 +1478,49 @@ mod fulltext_tests { ); } + #[tokio::test] + async fn test_full_text_search_branch() { + let (ctx, catalog, _tmp) = create_fulltext_context().await; + let identifier = Identifier::new("default", "test_tantivy_fulltext"); + let table = catalog.get_table(&identifier).await.expect("load table"); + let snapshot_manager = table.snapshot_manager(); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .expect("load latest snapshot") + .expect("latest snapshot"); + table + .tag_manager() + .create("branch-source", &snapshot) + .await + .expect("create tag"); + BranchManager::new(table.file_io().clone(), table.location().to_string()) + .create_branch_from_tag("b1", "branch-source") + .await + .expect("create branch"); + table + .file_io() + .delete_dir(&snapshot_manager.snapshot_dir()) + .await + .expect("delete main snapshots"); + + let batches = ctx + .sql("SELECT id, content FROM full_text_search('paimon.default.test_tantivy_fulltext$branch_b1', 'content', 'paimon', 10)") + .await + .expect("SQL should parse") + .collect() + .await + .expect("branch query should execute"); + + let rows = extract_id_content_rows(&batches); + let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); + assert_eq!(ids, vec![0, 2, 4]); + } + /// Search for 'tantivy' — only row 1. #[tokio::test] async fn test_full_text_search_tantivy() { - let (ctx, _tmp) = create_fulltext_context().await; + let (ctx, _catalog, _tmp) = create_fulltext_context().await; let batches = ctx .sql("SELECT id, content FROM full_text_search('paimon.default.test_tantivy_fulltext', 'content', 'tantivy', 10)") .await @@ -1497,7 +1537,7 @@ mod fulltext_tests { /// Search for 'search' — rows 1, 3 mention "full-text search". #[tokio::test] async fn test_full_text_search_search() { - let (ctx, _tmp) = create_fulltext_context().await; + let (ctx, _catalog, _tmp) = create_fulltext_context().await; let batches = ctx .sql("SELECT id, content FROM full_text_search('paimon.default.test_tantivy_fulltext', 'content', 'search', 10)") .await @@ -1526,6 +1566,7 @@ mod vector_search_tests { use datafusion::datasource::MemTable; use paimon::catalog::Identifier; use paimon::spec::{ArrayType, DataType, FloatType, IntType, Schema}; + use paimon::table::BranchManager; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::{register_vector_search, SQLContext}; @@ -1547,26 +1588,29 @@ mod vector_search_tests { (tmp, warehouse) } - async fn create_vector_search_context(archive_name: &str) -> (SQLContext, tempfile::TempDir) { + async fn create_vector_search_context( + archive_name: &str, + ) -> (SQLContext, Arc, tempfile::TempDir) { let (tmp, warehouse) = extract_test_warehouse(archive_name); let mut options = Options::new(); options.set(CatalogOptions::WAREHOUSE, warehouse); - let catalog = FileSystemCatalog::new(options).expect("Failed to create catalog"); - let catalog: Arc = Arc::new(catalog); + let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog")); let mut ctx = SQLContext::new(); ctx.register_catalog("paimon", catalog.clone()) .await .expect("Failed to register catalog"); - register_vector_search(ctx.ctx(), catalog, "default"); - (ctx, tmp) + register_vector_search(ctx.ctx(), catalog.clone(), "default"); + (ctx, catalog, tmp) } - async fn create_lumina_vector_search_context() -> (SQLContext, tempfile::TempDir) { + async fn create_lumina_vector_search_context( + ) -> (SQLContext, Arc, tempfile::TempDir) { create_vector_search_context("test_lumina_vector.tar.gz").await } - async fn create_java_vindex_vector_search_context() -> (SQLContext, tempfile::TempDir) { + async fn create_java_vindex_vector_search_context( + ) -> (SQLContext, Arc, tempfile::TempDir) { create_vector_search_context("test_java_vindex_vector.tar.gz").await } @@ -1748,7 +1792,7 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_top3() { - let (ctx, _tmp) = create_lumina_vector_search_context().await; + let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await; let batches = ctx .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)") .await @@ -1762,9 +1806,48 @@ mod vector_search_tests { assert!(ids.contains(&0), "exact match [1,0,0,0] should be in top 3"); } + #[tokio::test] + async fn test_vector_search_branch() { + let (ctx, catalog, _tmp) = create_java_vindex_vector_search_context().await; + let identifier = Identifier::new("default", "test_java_vindex_vector"); + let table = catalog.get_table(&identifier).await.expect("load table"); + let snapshot_manager = table.snapshot_manager(); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .expect("load latest snapshot") + .expect("latest snapshot"); + table + .tag_manager() + .create("branch-source", &snapshot) + .await + .expect("create tag"); + BranchManager::new(table.file_io().clone(), table.location().to_string()) + .create_branch_from_tag("b1", "branch-source") + .await + .expect("create branch"); + table + .file_io() + .delete_dir(&snapshot_manager.snapshot_dir()) + .await + .expect("delete main snapshots"); + + let batches = ctx + .sql("SELECT id FROM vector_search('paimon.default.test_java_vindex_vector$branch_b1', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)") + .await + .expect("SQL should parse") + .collect() + .await + .expect("branch query should execute"); + + let ids = extract_ids(&batches); + assert_eq!(ids.len(), 3); + assert!(ids.contains(&0)); + } + #[tokio::test] async fn test_vector_search_top6_returns_all() { - let (ctx, _tmp) = create_lumina_vector_search_context().await; + let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await; let batches = ctx .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 6)") .await @@ -1779,7 +1862,7 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_without_matching_index_returns_empty() { - let (ctx, _tmp) = create_lumina_vector_search_context().await; + let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await; let batches = ctx .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'missing_embedding', '[1.0]', 10)") .await @@ -1797,7 +1880,7 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_java_vindex_table() { - let (ctx, _tmp) = create_java_vindex_vector_search_context().await; + let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await; let batches = ctx .sql("SELECT id FROM vector_search('paimon.default.test_java_vindex_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)") .await @@ -1812,7 +1895,7 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_lateral_join_uses_query_vectors() { - let (ctx, _tmp) = create_java_vindex_vector_search_context().await; + let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await; let query_batch = build_vector_batch( vec![10, 20], vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]], @@ -2018,6 +2101,8 @@ mod hybrid_search_tests { use std::sync::Arc; use datafusion::arrow::array::Int32Array; + use paimon::catalog::Identifier; + use paimon::table::BranchManager; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::SQLContext; @@ -2039,18 +2124,18 @@ mod hybrid_search_tests { (tmp, warehouse) } - async fn create_hybrid_search_context() -> (SQLContext, tempfile::TempDir) { + async fn create_hybrid_search_context( + ) -> (SQLContext, Arc, tempfile::TempDir) { let (tmp, warehouse) = extract_test_warehouse("test_java_vindex_vector.tar.gz"); let mut options = Options::new(); options.set(CatalogOptions::WAREHOUSE, warehouse); - let catalog = FileSystemCatalog::new(options).expect("Failed to create catalog"); - let catalog: Arc = Arc::new(catalog); + let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog")); let mut ctx = SQLContext::new(); ctx.register_catalog("paimon", catalog.clone()) .await .expect("Failed to register catalog"); - (ctx, tmp) + (ctx, catalog, tmp) } fn extract_ids(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec { @@ -2070,7 +2155,7 @@ mod hybrid_search_tests { #[tokio::test] async fn test_hybrid_search_multiple_vector_routes_spark_shape() { - let (ctx, _tmp) = create_hybrid_search_context().await; + let (ctx, _catalog, _tmp) = create_hybrid_search_context().await; let batches = ctx .sql( "SELECT id FROM hybrid_search( \ @@ -2097,4 +2182,47 @@ mod hybrid_search_tests { assert_eq!(extract_ids(&batches), vec![0, 1, 2]); } + + #[tokio::test] + async fn test_hybrid_search_branch() { + let (ctx, catalog, _tmp) = create_hybrid_search_context().await; + let identifier = Identifier::new("default", "test_java_vindex_vector"); + let table = catalog.get_table(&identifier).await.expect("load table"); + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .expect("load latest snapshot") + .expect("latest snapshot"); + table + .tag_manager() + .create("branch-source", &snapshot) + .await + .expect("create tag"); + BranchManager::new(table.file_io().clone(), table.location().to_string()) + .create_branch_from_tag("b1", "branch-source") + .await + .expect("create branch"); + + let batches = ctx + .sql( + "SELECT id FROM hybrid_search( \ + 'paimon.default.test_java_vindex_vector$branch_b1', \ + array(named_struct( \ + 'field', 'embedding', \ + 'query_vector', array(1.0, 0.0, 0.0, 0.0), \ + 'limit', 3, \ + 'weight', 1.0)), \ + array(), \ + 3, \ + 'rrf')", + ) + .await + .expect("hybrid_search SQL should parse") + .collect() + .await + .expect("branch query should execute"); + + assert_eq!(extract_ids(&batches), vec![0, 1, 2]); + } } diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 042a9e88..30e1d358 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -28,7 +28,6 @@ use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, RowRangeIndex, }; -use crate::table::snapshot_manager::SnapshotManager; use crate::table::{find_field_id_by_name, merge_row_ranges, RowRange, Table}; use crate::tantivy::full_text_search::{FullTextSearch, SearchResult}; use crate::tantivy::reader::TantivyFullTextReader; @@ -120,10 +119,7 @@ impl<'a> FullTextSearchBuilder<'a> { let search = FullTextSearch::new(query_text.to_string(), limit, text_column.to_string())?; - let snapshot_manager = SnapshotManager::new( - self.table.file_io().clone(), - self.table.location().to_string(), - ); + let snapshot_manager = self.table.snapshot_manager(); let snapshot = match snapshot_manager.get_latest_snapshot().await? { Some(s) => s, @@ -132,11 +128,7 @@ impl<'a> FullTextSearchBuilder<'a> { let index_entries = match snapshot.index_manifest() { Some(index_manifest_name) => { - let manifest_path = format!( - "{}/manifest/{}", - self.table.location().trim_end_matches('/'), - index_manifest_name - ); + let manifest_path = snapshot_manager.manifest_path(index_manifest_name); IndexManifest::read(self.table.file_io(), &manifest_path).await? } None => Vec::new(), diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 8d75c492..c439bc67 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -26,7 +26,6 @@ use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, RowRangeIndex, }; -use crate::table::snapshot_manager::SnapshotManager; use crate::table::{find_field_id_by_name, merge_row_ranges, RowRange, Table}; use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::is_vindex_index_type; @@ -219,10 +218,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { }) .collect::>>()?; - let snapshot_manager = SnapshotManager::new( - self.table.file_io().clone(), - self.table.location().to_string(), - ); + let snapshot_manager = self.table.snapshot_manager(); let snapshot = match snapshot_manager.get_latest_snapshot().await? { Some(s) => s, @@ -231,11 +227,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { let index_entries = match snapshot.index_manifest() { Some(index_manifest_name) => { - let manifest_path = format!( - "{}/manifest/{}", - self.table.location().trim_end_matches('/'), - index_manifest_name - ); + let manifest_path = snapshot_manager.manifest_path(index_manifest_name); IndexManifest::read(self.table.file_io(), &manifest_path).await? } None => Vec::new(),