feat(datafusion): support reading branch tables

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/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,
@@ -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<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 +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<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)?;
@@ -649,6 +669,37 @@ impl SQLContext {
         }
     }
 
+    async fn load_table_for_read(
+        catalog: &Arc<dyn Catalog>,
+        identifier: &Identifier,
+    ) -> DFResult<(paimon::table::Table, Identifier, Option<String>)> {
+        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<dyn Catalog>,
@@ -901,6 +952,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 +1082,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 +1103,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 +1131,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 +1153,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 +1273,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 +1298,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 +1309,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 +1394,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 +1483,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)?;

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,17 +160,22 @@ impl TableProvider for ManifestsTable {
 
 async fn collect_manifests(table: &Table) -> paimon::Result<Vec<ManifestFileMeta>> {
     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? {
-        Some(s) => s,
-        None => return Ok(Vec::new()),
+    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 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 {