feat(datafusion): support reading branch tables

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()

crates/integrations/datafusion/src/catalog.rs

@@ -364,24 +364,31 @@ 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;
         }
 
         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,38 @@ 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);
+        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(_) => true,
+                    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
+                        }
+                    }
                     Err(paimon::Error::TableNotExist { .. }) => false,
                     Err(e) => {
                         log::error!("failed to check table '{}': {e}", identifier);

crates/integrations/datafusion/src/delete.rs

@@ -114,7 +114,8 @@ 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
             .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()
+            .try_new_commit()
+            .map_err(to_datafusion_error)?;
         commit.commit(messages).await.map_err(to_datafusion_error)?;
     }
 

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)?;
 

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)?,

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;

crates/integrations/datafusion/src/merge_into.rs

@@ -396,7 +396,8 @@ 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
             .map_err(to_datafusion_error)?;
@@ -771,7 +772,8 @@ 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
             .map_err(to_datafusion_error)?;

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.try_new_commit().map_err(to_datafusion_error)?;
 
         if self.overwrite {
             commit

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};
 use paimon::spec::{
     ArrayType as PaimonArrayType, BigIntType, BinaryType, BlobType, BooleanType, CharType,
     DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType,
@@ -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.
@@ -510,10 +511,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) =
+                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 +522,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<dyn TableProvider> = 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 +551,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) =
+                load_table_for_read(&catalog, &identifier).await?;
 
             let millis = Self::parse_timestamp_to_millis(&info.timestamp)?;
 
@@ -555,10 +564,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<dyn TableProvider> = 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)?;
@@ -901,6 +922,7 @@ impl SQLContext {
         operations: &[AlterTableOperation],
         if_exists: bool,
     ) -> DFResult<DataFrame> {
+        Self::ensure_main_branch_write_target(name, "ALTER TABLE")?;
         let identifier = self.resolve_table_name(name)?;
 
         let mut changes = Vec::new();
@@ -1030,6 +1052,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
@@ -1050,6 +1073,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
@@ -1077,6 +1101,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
@@ -1098,6 +1123,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)
@@ -1217,7 +1243,7 @@ impl SQLContext {
         }
 
         let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?;
-        let commit = wb.new_commit();
+        let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
 
         let overwrite_partitions = if static_partitions.is_empty() {
             None
@@ -1242,6 +1268,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,
@@ -1252,7 +1279,7 @@ impl SQLContext {
         };
 
         let wb = table.new_write_builder();
-        let commit = wb.new_commit();
+        let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
 
         if let Some(partitions) = &truncate.partitions {
             if partitions.is_empty() {
@@ -1337,7 +1364,7 @@ impl SQLContext {
         )?;
 
         let wb = table.new_write_builder();
-        let commit = wb.new_commit();
+        let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
         commit
             .truncate_partitions(partition_values)
             .await
@@ -1426,6 +1453,22 @@ 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() {
+            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<Identifier> {
         let (_catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(name)?;