From eb419460aba639c08f5feda6855c8e3a8b5e1f29 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 6 May 2026 21:55:06 +0800 Subject: [PATCH 1/2] feat(datafusion): add catalog-level temporary table support --- .gitignore | 1 + bindings/python/README.md | 28 - bindings/python/project-description.md | 19 + .../python/pypaimon_rust/datafusion.pyi | 1 + bindings/python/src/context.rs | 12 +- bindings/python/tests/test_datafusion.py | 99 ++ crates/integrations/datafusion/Cargo.toml | 1 + crates/integrations/datafusion/src/catalog.rs | 180 +++- crates/integrations/datafusion/src/delete.rs | 20 +- .../integrations/datafusion/src/merge_into.rs | 341 ++++--- .../datafusion/src/sql_context.rs | 888 +++++++++++++++++- crates/integrations/datafusion/src/update.rs | 110 +-- .../datafusion/tests/append_merge_into.rs | 186 ++-- .../datafusion/tests/blob_tests.rs | 26 +- .../datafusion/tests/common/mod.rs | 15 - .../datafusion/tests/delete_tests.rs | 14 +- .../datafusion/tests/merge_into_tests.rs | 95 +- .../datafusion/tests/read_tables.rs | 223 ++--- .../datafusion/tests/sql_context_tests.rs | 504 ++++++++++ .../datafusion/tests/system_tables.rs | 16 +- .../datafusion/tests/update_tests.rs | 8 +- docs/mkdocs.yml | 2 +- docs/src/{datafusion.md => sql.md} | 159 +++- 23 files changed, 2325 insertions(+), 623 deletions(-) rename docs/src/{datafusion.md => sql.md} (80%) diff --git a/.gitignore b/.gitignore index 25b80322..64007203 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ .vscode **/.DS_Store dist/* +.qoder diff --git a/bindings/python/README.md b/bindings/python/README.md index 58623bb4..71e160c2 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -21,34 +21,6 @@ This project builds the Rust-powered core for [PyPaimon](https://paimon.apache.org/docs/master/pypaimon/overview/) while also providing DataFusion integration for querying Paimon tables. -## Usage - -The recommended way to query Paimon tables is through `SQLContext`, which supports -multi-catalog registration, DDL, DML, and all Paimon-specific SQL extensions: - -```python -from pypaimon_rust.datafusion import SQLContext - -ctx = SQLContext() -ctx.register_catalog("paimon", {"warehouse": "/path/to/warehouse"}) - -batches = ctx.sql("SELECT * FROM paimon.default.my_table") -``` - -Alternatively, you can register a `PaimonCatalog` into DataFusion's native `SessionContext`: - -```python -from datafusion import SessionContext -from pypaimon_rust.datafusion import PaimonCatalog - -catalog = PaimonCatalog({"warehouse": "/path/to/warehouse"}) -ctx = SessionContext() -ctx.register_catalog_provider("paimon", catalog) - -df = ctx.sql("SELECT * FROM paimon.default.my_table") -df.show() -``` - ## Setup Install [uv](https://docs.astral.sh/uv/getting-started/installation/): diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md index 48e94928..0cab3532 100644 --- a/bindings/python/project-description.md +++ b/bindings/python/project-description.md @@ -53,6 +53,25 @@ ctx.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob')") batches = ctx.sql("SELECT * FROM paimon.my_db.users") ``` +### Temporary Tables + +You can register temporary in-memory tables programmatically. Names support the same resolution rules as SQL: bare names use the current catalog and database, partially qualified names use the current catalog, and fully qualified names specify catalog.database.table. + +Register a single PyArrow RecordBatch as a temporary table: + +```python +import pyarrow as pa + +batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"]) + +ctx.register_batch("paimon.default.my_temp", batch) + +batches = ctx.sql("SELECT * FROM paimon.default.my_temp") + +# Drop it via SQL when no longer needed +ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp") +``` + Alternatively, if you want to use the native Python DataFusion `SessionContext`, install `datafusion` and register a `PaimonCatalog`: diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index f5a98023..4d0e973a 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -30,4 +30,5 @@ class SQLContext: ) -> None: ... def set_current_catalog(self, catalog_name: str) -> None: ... def set_current_database(self, database_name: str) -> None: ... + def register_batch(self, name: str, batch: pyarrow.RecordBatch) -> None: ... def sql(self, sql: str) -> List[pyarrow.RecordBatch]: ... diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 2df5bd0e..e1050d38 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -18,7 +18,7 @@ use std::collections::HashMap; use std::sync::Arc; -use arrow::pyarrow::ToPyArrow; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; use datafusion::catalog::CatalogProvider; use datafusion_ffi::catalog_provider::FFI_CatalogProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; @@ -138,6 +138,16 @@ impl PySQLContext { }) } + fn register_batch(&self, name: String, batch: Bound<'_, PyAny>) -> PyResult<()> { + let batch = datafusion::arrow::record_batch::RecordBatch::from_pyarrow_bound(&batch)?; + let schema = batch.schema(); + let mem_table = datafusion::datasource::MemTable::try_new(schema, vec![vec![batch]]) + .map_err(df_to_py_err)?; + self.inner + .register_temp_table(&name, Arc::new(mem_table)) + .map_err(df_to_py_err) + } + fn sql(&self, py: Python<'_>, sql: String) -> PyResult>> { let rt = runtime(); let batches = rt.block_on(async { diff --git a/bindings/python/tests/test_datafusion.py b/bindings/python/tests/test_datafusion.py index 2d8abf80..95623855 100644 --- a/bindings/python/tests/test_datafusion.py +++ b/bindings/python/tests/test_datafusion.py @@ -65,3 +65,102 @@ def test_sql_context_ddl_dml(): ctx.sql("DROP TABLE paimon.test_db.users") ctx.sql("DROP SCHEMA paimon.test_db") + + +def test_register_batch_fully_qualified(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + + batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"]) + ctx.register_batch("paimon.default.my_temp", batch) + + batches = ctx.sql("SELECT id, name FROM paimon.default.my_temp") + table = pa.Table.from_batches(batches) + rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist())) + assert rows == [(1, "alice"), (2, "bob")] + + ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp") + + +def test_register_batch_bare_name(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + + batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"]) + # Bare name uses current catalog and current database + ctx.register_batch("my_temp", batch) + + batches = ctx.sql("SELECT id, name FROM paimon.default.my_temp") + table = pa.Table.from_batches(batches) + rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist())) + assert rows == [(1, "alice"), (2, "bob")] + + ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp") + + +def test_temp_table_shadows_paimon_table(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + + ctx.sql("CREATE SCHEMA paimon.test_db") + ctx.sql("CREATE TABLE paimon.test_db.users (id INT, name STRING)") + ctx.sql("INSERT INTO paimon.test_db.users VALUES (1, 'real')") + + batch = pa.record_batch([[2], ["temp"]], names=["id", "name"]) + ctx.register_batch("paimon.test_db.users", batch) + + # Temp table should shadow the real Paimon table + batches = ctx.sql("SELECT id, name FROM paimon.test_db.users") + table = pa.Table.from_batches(batches) + rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist())) + assert rows == [(2, "temp")] + + ctx.sql("DROP TEMPORARY TABLE paimon.test_db.users") + + # After dropping, the real table is visible again + batches = ctx.sql("SELECT id, name FROM paimon.test_db.users") + table = pa.Table.from_batches(batches) + rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist())) + assert rows == [(1, "real")] + + ctx.sql("DROP TABLE paimon.test_db.users") + ctx.sql("DROP SCHEMA paimon.test_db") + + +def test_drop_temp_table_if_exists(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + + batch = pa.record_batch([[1]], names=["id"]) + ctx.register_batch("paimon.default.my_temp", batch) + + ctx.sql("DROP TEMPORARY TABLE IF EXISTS paimon.default.my_temp") + + # Should be able to drop again without error + ctx.sql("DROP TEMPORARY TABLE IF EXISTS paimon.default.my_temp") + + +def test_multi_catalog_temp_table(): + with tempfile.TemporaryDirectory() as wh1, tempfile.TemporaryDirectory() as wh2: + ctx = SQLContext() + ctx.register_catalog("cat1", {"warehouse": wh1}) + ctx.register_catalog("cat2", {"warehouse": wh2}) + + batch1 = pa.record_batch([[1]], names=["id"]) + batch2 = pa.record_batch([[2]], names=["id"]) + + ctx.register_batch("cat1.default.t1", batch1) + ctx.register_batch("cat2.default.t2", batch2) + + result1 = ctx.sql("SELECT id FROM cat1.default.t1") + assert pa.Table.from_batches(result1)["id"].to_pylist() == [1] + + result2 = ctx.sql("SELECT id FROM cat2.default.t2") + assert pa.Table.from_batches(result2)["id"].to_pylist() == [2] + + ctx.sql("DROP TEMPORARY TABLE cat1.default.t1") + ctx.sql("DROP TEMPORARY TABLE cat2.default.t2") diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index 832326e1..b9c2b268 100644 --- a/crates/integrations/datafusion/Cargo.toml +++ b/crates/integrations/datafusion/Cargo.toml @@ -41,6 +41,7 @@ futures = "0.3" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { workspace = true, features = ["rt", "time", "fs"] } +uuid = { version = "1", features = ["v4"] } [dev-dependencies] arrow-array = { workspace = true } diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 8e58015a..58e6a2e1 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -21,9 +21,11 @@ use std::any::Any; use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; +use std::sync::RwLock; use async_trait::async_trait; -use datafusion::catalog::{CatalogProvider, SchemaProvider}; +use datafusion::catalog::{CatalogProvider, MemorySchemaProvider, SchemaProvider}; +use datafusion::common::plan_datafusion_err; use datafusion::datasource::TableProvider; use datafusion::error::Result as DFResult; use paimon::catalog::{Catalog, Identifier}; @@ -44,6 +46,14 @@ pub struct PaimonCatalogProvider { catalog: Arc, /// Session-scoped dynamic options shared with the SQL context. dynamic_options: DynamicOptions, + /// Temporary in-memory tables and views stored in MemorySchemaProvider per database. + /// + /// Uses `RwLock` with poison recovery (`unwrap_or_else(|e| e.into_inner())`) throughout. + /// This is a deliberate choice: since temp tables are session-scoped and non-critical, + /// it is preferable to continue with potentially stale data after a panic rather than + /// propagate the panic to all subsequent operations. The worst case is a temp table + /// becoming invisible or stale, which is recoverable by re-registering it. + temp_tables: Arc>>>, } impl Debug for PaimonCatalogProvider { @@ -62,6 +72,7 @@ impl PaimonCatalogProvider { PaimonCatalogProvider { catalog, dynamic_options: Default::default(), + temp_tables: Arc::new(RwLock::new(HashMap::new())), } } @@ -72,6 +83,7 @@ impl PaimonCatalogProvider { PaimonCatalogProvider { catalog, dynamic_options, + temp_tables: Arc::new(RwLock::new(HashMap::new())), } } } @@ -85,13 +97,10 @@ impl CatalogProvider for PaimonCatalogProvider { let catalog = Arc::clone(&self.catalog); block_on_with_runtime( async move { - match catalog.list_databases().await { - Ok(names) => names, - Err(e) => { - log::error!("failed to list databases: {e}"); - vec![] - } - } + catalog.list_databases().await.unwrap_or_else(|e| { + log::error!("failed to list databases: {e}"); + vec![] + }) }, "paimon catalog access thread panicked", ) @@ -101,6 +110,12 @@ impl CatalogProvider for PaimonCatalogProvider { let catalog = Arc::clone(&self.catalog); let dynamic_options = Arc::clone(&self.dynamic_options); let name = name.to_string(); + + let temp_provider = { + let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); + databases.get(&name).cloned() + }; + block_on_with_runtime( async move { match catalog.get_database(&name).await { @@ -108,8 +123,20 @@ impl CatalogProvider for PaimonCatalogProvider { Arc::clone(&catalog), name, dynamic_options, + temp_provider, )) as Arc), - Err(paimon::Error::DatabaseNotExist { .. }) => None, + Err(paimon::Error::DatabaseNotExist { .. }) => { + if temp_provider.is_some() { + Some(Arc::new(PaimonSchemaProvider::new( + Arc::clone(&catalog), + name, + dynamic_options, + temp_provider, + )) as Arc) + } else { + None + } + } Err(e) => { log::error!("failed to get database '{}': {e}", name); None @@ -138,6 +165,7 @@ impl CatalogProvider for PaimonCatalogProvider { Arc::clone(&catalog), name, dynamic_options, + None, )) as Arc)) }, "paimon catalog access thread panicked", @@ -162,6 +190,7 @@ impl CatalogProvider for PaimonCatalogProvider { Arc::clone(&catalog), name, dynamic_options, + None, )) as Arc)) }, "paimon catalog access thread panicked", @@ -169,6 +198,94 @@ impl CatalogProvider for PaimonCatalogProvider { } } +impl PaimonCatalogProvider { + /// Creates or returns an existing temporary in-memory database for temp tables/views. + fn get_or_create_temp_database(&self, name: &str) -> Arc { + let mut databases = self.temp_tables.write().unwrap_or_else(|e| e.into_inner()); + databases + .entry(name.to_string()) + .or_insert_with(|| Arc::new(MemorySchemaProvider::new())) + .clone() + } + + /// Registers a temporary table or view in the specified database. + /// Creates the database if it does not exist. + /// + /// Returns an error if a temp table with the same name already exists in + /// the same database. Logs a warning if the name shadows a real Paimon table. + pub fn register_temp_table( + &self, + database: &str, + table_name: &str, + table: Arc, + ) -> DFResult<()> { + // Check if a temp table with this name already exists + { + let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); + if let Some(mem_db) = databases.get(database) { + if mem_db.table_exist(table_name) { + return Err(plan_datafusion_err!( + "Temporary table '{database}.{table_name}' already exists" + )); + } + } + } + + // Warn if this shadows a real Paimon table + let catalog = Arc::clone(&self.catalog); + let db = database.to_string(); + let tbl = table_name.to_string(); + let identifier = Identifier::new(db, tbl); + if let Ok(true) = block_on_with_runtime( + async move { + match catalog.get_table(&identifier).await { + Ok(_) => Ok::(true), + Err(paimon::Error::TableNotExist { .. }) => Ok(false), + Err(_) => Ok(false), + } + }, + "paimon catalog access thread panicked", + ) { + log::warn!( + "Temporary table '{database}.{table_name}' shadows an existing Paimon table" + ); + } + + let mem_database = self.get_or_create_temp_database(database); + mem_database.register_table(table_name.to_string(), table)?; + Ok(()) + } + + /// Deregisters a temporary table or view from the specified database. + pub fn deregister_temp_table( + &self, + database: &str, + table_name: &str, + ) -> DFResult>> { + let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); + let mem_database = databases + .get(database) + .ok_or_else(|| plan_datafusion_err!("Unknown temp database '{database}'"))?; + mem_database.deregister_table(table_name) + } + + /// Returns whether a temp table database exists with the given name. + pub fn has_temp_table_database(&self, name: &str) -> bool { + self.temp_tables + .read() + .unwrap_or_else(|e| e.into_inner()) + .contains_key(name) + } + + /// Returns whether a temp table with the given name exists in the specified database. + pub fn temp_table_exist(&self, database: &str, table_name: &str) -> bool { + let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); + databases + .get(database) + .is_some_and(|db| db.table_exist(table_name)) + } +} + /// Represents a [`SchemaProvider`] for the Paimon [`Catalog`], managing /// access to table providers within a specific database. /// @@ -180,12 +297,15 @@ pub struct PaimonSchemaProvider { database: String, /// Session-scoped dynamic options shared with the SQL context. dynamic_options: DynamicOptions, + /// Optional temporary in-memory provider for temp tables and views. + temp_provider: Option>, } impl Debug for PaimonSchemaProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PaimonSchemaProvider") .field("database", &self.database) + .field("has_temp_provider", &self.temp_provider.is_some()) .finish() } } @@ -196,11 +316,13 @@ impl PaimonSchemaProvider { catalog: Arc, database: String, dynamic_options: DynamicOptions, + temp_provider: Option>, ) -> Self { PaimonSchemaProvider { catalog, database, dynamic_options, + temp_provider, } } } @@ -214,21 +336,39 @@ impl SchemaProvider for PaimonSchemaProvider { fn table_names(&self) -> Vec { let catalog = Arc::clone(&self.catalog); let database = self.database.clone(); - block_on_with_runtime( - async move { - match catalog.list_tables(&database).await { - Ok(names) => names, - Err(e) => { - log::error!("failed to list tables in '{}': {e}", database); - vec![] + let mut names = block_on_with_runtime( + { + let db = database.clone(); + async move { + match catalog.list_tables(&db).await { + Ok(names) => names, + Err(e) => { + log::error!("failed to list tables in '{}': {e}", db); + vec![] + } } } }, "paimon catalog access thread panicked", - ) + ); + + if let Some(temp) = &self.temp_provider { + names.extend(temp.table_names()); + } + + let mut seen = std::collections::HashSet::new(); + names.retain(|name| seen.insert(name.clone())); + + names } async fn table(&self, name: &str) -> DFResult>> { + if let Some(temp) = &self.temp_provider { + if let Some(table) = temp.table(name).await? { + return Ok(Some(table)); + } + } + let (base, system_name) = system_tables::split_object_name(name); if let Some(system_name) = system_name { return await_with_runtime(system_tables::load( @@ -263,6 +403,12 @@ impl SchemaProvider for PaimonSchemaProvider { } fn table_exist(&self, name: &str) -> bool { + if let Some(temp) = &self.temp_provider { + if temp.table_exist(name) { + return true; + } + } + let (base, system_name) = system_tables::split_object_name(name); if let Some(system_name) = system_name { if !system_tables::is_registered(system_name) { diff --git a/crates/integrations/datafusion/src/delete.rs b/crates/integrations/datafusion/src/delete.rs index cb37c532..c890cf69 100644 --- a/crates/integrations/datafusion/src/delete.rs +++ b/crates/integrations/datafusion/src/delete.rs @@ -27,17 +27,19 @@ use paimon::spec::CoreOptions; use paimon::table::{CopyOnWriteMergeWriter, Table}; use crate::error::to_datafusion_error; +use crate::merge_into::TempTableTracker; use crate::merge_into::{ build_partition_set_from_where, extract_tracking_columns, is_delete_conflict, ok_result, register_cow_target_table, retry_on_conflict, }; +use crate::sql_context::SQLContext; /// Execute a DELETE statement on a Paimon table. /// /// `table_ref` is the SQL table reference string (e.g. `"paimon.test_db.t"`), /// already extracted by the caller for catalog resolution. pub(crate) async fn execute_delete( - ctx: &SessionContext, + ctx: &SQLContext, delete: &Delete, table: Table, table_ref: &str, @@ -73,7 +75,7 @@ pub(crate) async fn execute_delete( /// Execute DELETE on an append-only table with retry on delete conflict. async fn execute_cow_delete( - ctx: &SessionContext, + ctx: &SQLContext, delete: &Delete, table: &Table, table_ref: &str, @@ -86,7 +88,7 @@ async fn execute_cow_delete( /// Single attempt of CoW DELETE execution. async fn execute_cow_delete_once( - ctx: &SessionContext, + ctx: &SQLContext, delete: &Delete, table: &Table, table_ref: &str, @@ -99,14 +101,16 @@ async fn execute_cow_delete_once( .await .map_err(to_datafusion_error)?; - let (has_data, cow_table_guard) = register_cow_target_table(ctx, table, &writer).await?; + let mut temp_tracker = TempTableTracker::new(ctx); + let (has_data, cow_table_name) = + register_cow_target_table(ctx, table, &writer, &mut temp_tracker).await?; if !has_data { - return ok_result(ctx, 0); + return ok_result(ctx.ctx(), 0); } + let cow_target_qualified = cow_table_name; let result = - execute_cow_delete_inner(ctx, &cow_table_guard.qualified_name(), delete, &mut writer).await; - drop(cow_table_guard); + execute_cow_delete_inner(ctx.ctx(), &cow_target_qualified, delete, &mut writer).await; let total_count = result?; let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; @@ -115,7 +119,7 @@ async fn execute_cow_delete_once( commit.commit(messages).await.map_err(to_datafusion_error)?; } - ok_result(ctx, total_count) + ok_result(ctx.ctx(), total_count) } async fn execute_cow_delete_inner( diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index a004aa56..cd78ddec 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -29,7 +29,6 @@ use std::sync::Arc; use datafusion::arrow::array::{Array, Int32Array, RecordBatch, UInt32Array, UInt64Array}; use datafusion::arrow::compute; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; -use datafusion::datasource::MemTable; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::sqlparser::ast::{ @@ -41,6 +40,7 @@ use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions}; use paimon::table::{CopyOnWriteMergeWriter, DataEvolutionWriter, DataSplitBuilder, Table}; use crate::error::to_datafusion_error; +use crate::sql_context::SQLContext; /// Maximum number of retries when DML conflicts with concurrent compaction. const DML_MAX_RETRIES: u32 = 5; @@ -57,32 +57,33 @@ fn next_cow_table_name(prefix: &str) -> String { format!("{prefix}_{id}") } -/// RAII guard that deregisters a MemTable from the SessionContext on drop. -/// Prevents leaks when the future is cancelled between register and deregister. -pub(crate) struct CowTableGuard { - ctx: SessionContext, - table_name: String, +/// Tracks registered temporary table names and auto-cleans them up on drop. +/// +/// This RAII guard ensures temp tables are always deregistered, even if the +/// enclosing function panics or returns early with an error. +pub(crate) struct TempTableTracker<'a> { + tables: Vec, + ctx: &'a SQLContext, } -const COW_CATALOG: &str = "datafusion"; -const COW_SCHEMA: &str = "public"; - -impl CowTableGuard { - pub(crate) fn new(ctx: &SessionContext, table_name: String) -> Self { +impl<'a> TempTableTracker<'a> { + pub(crate) fn new(ctx: &'a SQLContext) -> Self { Self { - ctx: ctx.clone(), - table_name, + tables: Vec::new(), + ctx, } } - pub(crate) fn qualified_name(&self) -> String { - format!("{COW_CATALOG}.{COW_SCHEMA}.{}", self.table_name) + pub(crate) fn register(&mut self, table_name: &str) { + self.tables.push(table_name.to_string()); } } -impl Drop for CowTableGuard { +impl Drop for TempTableTracker<'_> { fn drop(&mut self) { - let _ = self.ctx.deregister_table(self.qualified_name()); + for table in &self.tables { + let _ = self.ctx.deregister_temp_table(table); + } } } @@ -121,7 +122,7 @@ where /// - Data evolution tables → partial-column writes via `DataEvolutionWriter` /// - Append-only tables (no PK) → copy-on-write file rewriting via `CopyOnWriteMergeWriter` pub(crate) async fn execute_merge_into( - ctx: &SessionContext, + ctx: &SQLContext, merge: &Merge, table: Table, ) -> DFResult { @@ -161,7 +162,7 @@ pub(crate) fn is_delete_conflict(err: &DataFusionError) -> bool { /// Execute MERGE INTO on a data evolution table with retry on row ID conflict. async fn execute_data_evolution_merge( - ctx: &SessionContext, + ctx: &SQLContext, merge: &Merge, table: Table, ) -> DFResult { @@ -295,11 +296,7 @@ fn extract_cow_merge_clauses(merge: &Merge) -> DFResult { } /// Execute MERGE INTO on an append-only table with retry on delete conflict. -async fn execute_cow_merge( - ctx: &SessionContext, - merge: &Merge, - table: Table, -) -> DFResult { +async fn execute_cow_merge(ctx: &SQLContext, merge: &Merge, table: Table) -> DFResult { retry_on_conflict("CoW MERGE INTO", is_delete_conflict, || { execute_cow_merge_once(ctx, merge, &table) }) @@ -308,7 +305,7 @@ async fn execute_cow_merge( /// Execute a single attempt of CoW MERGE INTO. async fn execute_cow_merge_once( - ctx: &SessionContext, + ctx: &SQLContext, merge: &Merge, table: &Table, ) -> DFResult { @@ -361,8 +358,9 @@ async fn execute_cow_merge_once( } // Read each target file individually, attach __paimon_file_idx and __paimon_row_offset - let (has_target_data, cow_target_guard) = - register_cow_target_table(ctx, table, &writer).await?; + let mut temp_tracker = TempTableTracker::new(ctx); + let (has_target_data, cow_table_name) = + register_cow_target_table(ctx, table, &writer, &mut temp_tracker).await?; let merge_ctx = CowMergeContext { source_ref: &source_ref, @@ -370,13 +368,19 @@ async fn execute_cow_merge_once( t_alias, on_condition: &on_condition, has_target_data, - cow_target_name: cow_target_guard.qualified_name(), + cow_table_name, update_columns: &update_columns, }; - let result = execute_cow_merge_inner(ctx, &clauses, &mut writer, table, &merge_ctx).await; - - drop(cow_target_guard); + let result = execute_cow_merge_inner( + ctx, + &clauses, + &mut writer, + table, + &merge_ctx, + &mut temp_tracker, + ) + .await; let (insert_messages, total_count) = result?; @@ -394,7 +398,7 @@ async fn execute_cow_merge_once( .map_err(to_datafusion_error)?; } - ok_result(ctx, total_count) + ok_result(ctx.ctx(), total_count) } /// Context for CoW merge inner execution — groups join-related parameters. @@ -404,25 +408,27 @@ struct CowMergeContext<'a> { t_alias: &'a str, on_condition: &'a str, has_target_data: bool, - cow_target_name: String, + cow_table_name: String, update_columns: &'a [String], } /// Inner function that populates the CoW writer with matched operations and handles INSERT. /// Returns (insert_commit_messages, total_affected_count). async fn execute_cow_merge_inner( - ctx: &SessionContext, + ctx: &SQLContext, clauses: &CowMergeClauses, writer: &mut CopyOnWriteMergeWriter, table: &Table, merge_ctx: &CowMergeContext<'_>, + temp_tracker: &mut TempTableTracker<'_>, ) -> DFResult<(Vec, u64)> { let source_ref = merge_ctx.source_ref; let s_alias = merge_ctx.s_alias; let t_alias = merge_ctx.t_alias; let on_condition = merge_ctx.on_condition; let has_target_data = merge_ctx.has_target_data; - let cow_target_name = &merge_ctx.cow_target_name; + let cow_table_name = &merge_ctx.cow_table_name; + let cow_target_name = cow_table_name.clone(); let update_columns = merge_ctx.update_columns; let mut insert_messages = Vec::new(); let mut total_count: u64 = 0; @@ -481,7 +487,7 @@ async fn execute_cow_merge_inner( INNER JOIN {cow_target_name} AS {t_alias} ON {on_condition}{where_clause}" ); - let join_result = ctx.sql(&join_sql).await?.collect().await?; + let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; for batch in &join_result { if batch.num_rows() == 0 { @@ -530,7 +536,7 @@ async fn execute_cow_merge_inner( INNER JOIN {cow_target_name} AS {t_alias} ON {on_condition}{where_clause}" ); - let join_result = ctx.sql(&join_sql).await?.collect().await?; + let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; for batch in &join_result { if batch.num_rows() == 0 { @@ -574,7 +580,7 @@ async fn execute_cow_merge_inner( format!("SELECT * FROM {source_ref} AS {s_alias}") }; - let not_matched_batches = ctx.sql(&insert_sql).await?.collect().await?; + let not_matched_batches = ctx.ctx().sql(&insert_sql).await?.collect().await?; if !not_matched_batches.is_empty() { let insert_batches = build_insert_batches( @@ -584,6 +590,7 @@ async fn execute_cow_merge_inner( s_alias, &[], &table_fields, + temp_tracker, ) .await?; @@ -617,7 +624,7 @@ async fn execute_cow_merge_inner( // --------------------------------------------------------------------------- async fn execute_merge_into_once( - ctx: &SessionContext, + ctx: &SQLContext, merge: &Merge, table: &Table, ) -> DFResult { @@ -666,7 +673,7 @@ async fn execute_merge_into_once( LEFT JOIN {target_ref} AS {t_alias} ON {on_condition}" ); - let join_result = ctx.sql(&join_sql).await?.collect().await?; + let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; // 3. Split by _ROW_ID null/not-null let mut all_messages = Vec::new(); @@ -709,6 +716,7 @@ async fn execute_merge_into_once( .iter() .map(|f| f.name().to_string()) .collect(); + let mut temp_tracker = TempTableTracker::new(ctx); let insert_batches = build_insert_batches( ctx, ¬_matched_batches, @@ -716,6 +724,7 @@ async fn execute_merge_into_once( s_alias, &injected_columns, &table_fields, + &mut temp_tracker, ) .await?; let insert_count: usize = insert_batches.iter().map(|b| b.num_rows()).sum(); @@ -748,7 +757,7 @@ async fn execute_merge_into_once( .map_err(to_datafusion_error)?; } - ok_result(ctx, total_count) + ok_result(ctx.ctx(), total_count) } /// Split join result into matched (_ROW_ID not null) and not-matched (_ROW_ID null) batches. @@ -818,12 +827,13 @@ pub(crate) fn project_update_columns( /// Build insert batches from not-matched rows, applying INSERT clause projections and predicates. async fn build_insert_batches( - ctx: &SessionContext, + ctx: &SQLContext, not_matched_batches: &[RecordBatch], inserts: &[MergeInsertClause], s_alias: &str, injected_columns: &[String], table_fields: &[String], + temp_tracker: &mut TempTableTracker<'_>, ) -> DFResult> { if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) { return Ok(Vec::new()); @@ -834,21 +844,20 @@ async fn build_insert_batches( // Register as temp table for SQL-based projection/filtering let first_schema = source_batches[0].schema(); - let mem_table = MemTable::try_new(first_schema, vec![source_batches])?; let tmp_name = next_cow_table_name("__merge_not_matched"); - let qualified_tmp = format!("{COW_CATALOG}.{COW_SCHEMA}.{tmp_name}"); - ctx.register_table(&qualified_tmp, Arc::new(mem_table))?; - let _guard = CowTableGuard::new(ctx, tmp_name.clone()); - let result = - build_insert_batches_inner(ctx, inserts, s_alias, &qualified_tmp, table_fields).await; + let mem_table = datafusion::datasource::MemTable::try_new(first_schema, vec![source_batches])?; + ctx.register_temp_table(&tmp_name, Arc::new(mem_table))?; + temp_tracker.register(&tmp_name); + + let result = build_insert_batches_inner(ctx, inserts, s_alias, &tmp_name, table_fields).await; result } /// Execute INSERT clause queries against the registered temp table. async fn build_insert_batches_inner( - ctx: &SessionContext, + ctx: &SQLContext, inserts: &[MergeInsertClause], s_alias: &str, tmp_name: &str, @@ -878,7 +887,7 @@ async fn build_insert_batches_inner( let select_clause = insert_select_clause(ins, table_fields); let sql = format!("SELECT {select_clause} FROM {tmp_name} AS {s_alias}{where_clause}"); - let batches = ctx.sql(&sql).await?.collect().await?; + let batches = ctx.ctx().sql(&sql).await?.collect().await?; all_batches.extend(batches); } @@ -1144,20 +1153,22 @@ pub(crate) fn extract_tracking_columns( /// Read all files from a table via the CoW writer's file index, attach `__paimon_file_idx` /// and `__paimon_row_offset` tracking columns, and register the result as a MemTable. /// -/// Returns `(has_data, guard)`. The guard deregisters the table on drop. +/// Returns `(has_data, cow_table_name)`. The caller is responsible for deregistering +/// via `TempTableTracker`. /// /// Note: all matching partition files are loaded into memory at once. For partitions /// with many large files this may cause significant memory pressure. A future /// optimization could stream or batch-process files instead of materializing everything. pub(crate) async fn register_cow_target_table( - ctx: &SessionContext, + ctx: &SQLContext, table: &Table, writer: &CopyOnWriteMergeWriter, -) -> DFResult<(bool, CowTableGuard)> { + temp_tracker: &mut TempTableTracker<'_>, +) -> DFResult<(bool, String)> { let file_index = writer.file_index(); if file_index.is_empty() { let table_name = next_cow_table_name("__cow_target"); - return Ok((false, CowTableGuard::new(ctx, table_name))); + return Ok((false, table_name)); } // Read all files in parallel @@ -1256,14 +1267,12 @@ pub(crate) async fn register_cow_target_table( if has_data { let s = schema.unwrap(); - let mem_table = MemTable::try_new(s, vec![all_batches])?; - ctx.register_table( - format!("{COW_CATALOG}.{COW_SCHEMA}.{table_name}"), - Arc::new(mem_table), - )?; + let mem_table = datafusion::datasource::MemTable::try_new(s, vec![all_batches])?; + ctx.register_temp_table(&table_name, Arc::new(mem_table))?; + temp_tracker.register(&table_name); } - Ok((has_data, CowTableGuard::new(ctx, table_name))) + Ok((has_data, table_name)) } /// Build a partition set from Arrow batches containing partition column values. @@ -1308,7 +1317,7 @@ pub(crate) fn build_partition_set_from_batches( /// /// Returns `None` for non-partitioned tables. pub(crate) async fn build_partition_set_from_where( - ctx: &SessionContext, + ctx: &SQLContext, table: &Table, table_ref: &str, where_clause: Option<&str>, @@ -1328,7 +1337,7 @@ pub(crate) async fn build_partition_set_from_where( None => String::new(), }; let sql = format!("SELECT DISTINCT {cols} FROM {table_ref}{where_part}"); - let batches = ctx.sql(&sql).await?.collect().await?; + let batches = ctx.ctx().sql(&sql).await?.collect().await?; build_partition_set_from_batches(table, &batches) } @@ -1338,7 +1347,7 @@ pub(crate) async fn build_partition_set_from_where( /// Returns `None` for non-partitioned tables or when the source lacks matching /// partition key columns (falls back to full-partition scan). async fn build_source_partition_set( - ctx: &SessionContext, + ctx: &SQLContext, table: &Table, source_ref: &str, s_alias: &str, @@ -1354,7 +1363,7 @@ async fn build_source_partition_set( .collect::>() .join(", "); let sql = format!("SELECT DISTINCT {cols} FROM {source_ref} AS {s_alias}"); - match ctx.sql(&sql).await { + match ctx.ctx().sql(&sql).await { Ok(df) => { let batches = df.collect().await?; build_partition_set_from_batches(table, &batches) @@ -1436,7 +1445,6 @@ pub(crate) fn ok_result(ctx: &SessionContext, count: u64) -> DFResult #[cfg(test)] mod tests { use super::*; - use datafusion::prelude::SessionContext; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; use paimon::catalog::{Catalog, Identifier}; @@ -1445,7 +1453,7 @@ mod tests { use paimon::{CatalogOptions, FileSystemCatalog, Options}; use tempfile::TempDir; - use crate::{PaimonTableProvider, SQLContext}; + use crate::SQLContext; async fn setup_sql_context() -> (TempDir, SQLContext, Arc) { let temp_dir = TempDir::new().unwrap(); @@ -1467,12 +1475,12 @@ mod tests { (temp_dir, sql_context, catalog) } - async fn setup_data_evolution_table(name: &str) -> (TempDir, SessionContext, Table) { + async fn setup_data_evolution_table(name: &str) -> (TempDir, SQLContext, Table) { let (tmp, sql_context, catalog) = setup_sql_context().await; sql_context .sql(&format!( - "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('row-tracking.enabled' = 'true')" + "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')" )) .await .unwrap(); @@ -1496,12 +1504,7 @@ mod tests { extra.insert("row-tracking.enabled".to_string(), "true".to_string()); let de_table = table.copy_with_options(extra); - let ctx = sql_context.ctx().clone(); - let provider = PaimonTableProvider::try_new(de_table.clone()).unwrap(); - ctx.register_table("datafusion.public.target", Arc::new(provider)) - .unwrap(); - - (tmp, ctx, de_table) + (tmp, sql_context, de_table) } fn parse_merge(sql: &str) -> Merge { @@ -1515,27 +1518,36 @@ mod tests { #[tokio::test] async fn test_merge_into_updates_matched_rows() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_merge").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_merge").await; // Create source table with updates - ctx.sql( - "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')", - ) - .await - .unwrap() - .collect() - .await - .unwrap(); + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (1, 'ALICE'), (3, 'CHARLIE')") + .await + .unwrap() + .collect() + .await + .unwrap(); // Execute MERGE INTO let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_merge t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ); - execute_merge_into(&ctx, &merge, table).await.unwrap(); + execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_merge ORDER BY id") .await .unwrap() .collect() @@ -1576,22 +1588,31 @@ mod tests { #[tokio::test] async fn test_merge_into_no_matches() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_merge2").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_merge2").await; - ctx.sql( - "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (99, 'nobody')", - ) - .await - .unwrap() - .collect() - .await - .unwrap(); + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (99, 'nobody')") + .await + .unwrap() + .collect() + .await + .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_merge2 t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ); - let result = execute_merge_into(&ctx, &merge, table).await.unwrap(); + let result = execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); let batches = result.collect().await.unwrap(); let count = batches[0] .column(0) @@ -1630,12 +1651,12 @@ mod tests { None, ); - let ctx = SessionContext::new(); + let sql_context = SQLContext::new(); let merge = parse_merge( "MERGE INTO t USING s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET id = s.id", ); - let result = execute_merge_into(&ctx, &merge, table).await; + let result = execute_merge_into(&sql_context, &merge, table).await; assert!(result.is_err()); assert!(result .unwrap_err() @@ -1647,7 +1668,7 @@ mod tests { // CoW MERGE INTO tests (append-only tables) // ----------------------------------------------------------------------- - async fn setup_append_only_table(name: &str) -> (TempDir, SessionContext, Table) { + async fn setup_append_only_table(name: &str) -> (TempDir, SQLContext, Table) { let (tmp, sql_context, catalog) = setup_sql_context().await; sql_context @@ -1672,12 +1693,7 @@ mod tests { .await .unwrap(); - let ctx = sql_context.ctx().clone(); - let provider = PaimonTableProvider::try_new(table.clone()).unwrap(); - ctx.register_table("datafusion.public.target", Arc::new(provider)) - .unwrap(); - - (tmp, ctx, table) + (tmp, sql_context, table) } fn collect_rows(batches: &[RecordBatch]) -> Vec<(i32, String, i32)> { @@ -1708,25 +1724,30 @@ mod tests { #[tokio::test] async fn test_cow_merge_update_matched_rows() { - let (_tmp, ctx, table) = setup_append_only_table("t_cow_upd").await; + let (_tmp, sql_context, table) = setup_append_only_table("t_cow_upd").await; - ctx.sql( - "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')", - ) - .await - .unwrap() - .collect() - .await - .unwrap(); + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source (id, name) VALUES (1, 'ALICE'), (3, 'CHARLIE')") + .await + .unwrap() + .collect() + .await + .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_cow_upd t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ); - execute_merge_into(&ctx, &merge, table).await.unwrap(); + execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upd ORDER BY id") .await .unwrap() .collect() @@ -1746,9 +1767,14 @@ mod tests { #[tokio::test] async fn test_cow_merge_delete_matched_rows() { - let (_tmp, ctx, table) = setup_append_only_table("t_cow_del").await; + let (_tmp, sql_context, table) = setup_append_only_table("t_cow_del").await; - ctx.sql("CREATE TABLE datafusion.public.source (id INT) AS VALUES (2)") + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source (id) VALUES (2)") .await .unwrap() .collect() @@ -1756,13 +1782,15 @@ mod tests { .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_cow_del t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN DELETE", ); - execute_merge_into(&ctx, &merge, table).await.unwrap(); + execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_del ORDER BY id") .await .unwrap() .collect() @@ -1778,9 +1806,14 @@ mod tests { #[tokio::test] async fn test_cow_merge_insert_not_matched() { - let (_tmp, ctx, table) = setup_append_only_table("t_cow_ins").await; + let (_tmp, sql_context, table) = setup_append_only_table("t_cow_ins").await; - ctx.sql("CREATE TABLE datafusion.public.source (id INT, name VARCHAR, value INT) AS VALUES (4, 'dave', 40), (5, 'eve', 50)") + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (4, 'dave', 40), (5, 'eve', 50)") .await .unwrap() .collect() @@ -1788,13 +1821,15 @@ mod tests { .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_cow_ins t USING paimon.test_db.source s ON t.id = s.id \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ); - execute_merge_into(&ctx, &merge, table).await.unwrap(); + execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins ORDER BY id") .await .unwrap() .collect() @@ -1816,9 +1851,14 @@ mod tests { #[tokio::test] async fn test_cow_merge_update_and_insert() { - let (_tmp, ctx, table) = setup_append_only_table("t_cow_upsert").await; + let (_tmp, sql_context, table) = setup_append_only_table("t_cow_upsert").await; - ctx.sql("CREATE TABLE datafusion.public.source (id INT, name VARCHAR, value INT) AS VALUES (2, 'BOB', 200), (4, 'dave', 40)") + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (2, 'BOB', 200), (4, 'dave', 40)") .await .unwrap() .collect() @@ -1826,14 +1866,16 @@ mod tests { .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_cow_upsert t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name, value = s.value \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ); - execute_merge_into(&ctx, &merge, table).await.unwrap(); + execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upsert ORDER BY id") .await .unwrap() .collect() @@ -1854,22 +1896,23 @@ mod tests { #[tokio::test] async fn test_cow_merge_no_matches() { - let (_tmp, ctx, table) = setup_append_only_table("t_cow_nomatch").await; + let (_tmp, sql_context, table) = setup_append_only_table("t_cow_nomatch").await; - ctx.sql( - "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (99, 'nobody')", - ) - .await - .unwrap() - .collect() - .await - .unwrap(); + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); let merge = parse_merge( - "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.t_cow_nomatch t USING paimon.test_db.source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ); - let result = execute_merge_into(&ctx, &merge, table).await.unwrap(); + let result = execute_merge_into(&sql_context, &merge, table) + .await + .unwrap(); let batches = result.collect().await.unwrap(); let count = batches[0] .column(0) diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 10974bbd..3c82ba6a 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -43,12 +43,15 @@ use datafusion::arrow::array::{ use datafusion::arrow::compute::cast; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::TableReference; +use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::sqlparser::ast::{ - AlterTableOperation, ColumnDef, CreateTable, CreateTableOptions, Delete, Expr as SqlExpr, - FromTable, Insert, Merge, ObjectName, RenameTableNameKind, Reset, ResetStatement, Set, - SqlOption, Statement, TableFactor, TableObject, Truncate, Update, Value as SqlValue, + AlterTableOperation, ColumnDef, CreateTable, CreateTableOptions, CreateView, Delete, + Expr as SqlExpr, FromTable, Insert, Merge, ObjectName, ObjectType, RenameTableNameKind, Reset, + ResetStatement, Set, SqlOption, Statement, TableFactor, TableObject, Truncate, Update, + Value as SqlValue, }; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; @@ -156,7 +159,7 @@ impl SQLContext { Ok(()) } - /// Sets the current database (schema) for unqualified table references. + /// Sets the current database for unqualified table references. pub async fn set_current_database(&self, database_name: &str) -> DFResult<()> { if database_name.contains('\'') { return Err(DataFusionError::Plan( @@ -176,6 +179,104 @@ impl SQLContext { &self.ctx } + /// Registers a temporary in-memory table or view. + /// + /// The `name` parameter accepts flexible table references, similar to DataFusion: + /// - `"my_table"` — uses the current catalog and current database + /// - `"database.my_table"` — uses the current catalog with the specified database + /// - `"catalog.database.my_table"` — fully qualified + /// + /// The table exists only for the lifetime of this SQLContext instance. + pub fn register_temp_table( + &self, + name: impl Into, + table: Arc, + ) -> DFResult<()> { + let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?; + let catalog_provider = self + .ctx + .catalog(&catalog) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?; + + let paimon_provider = catalog_provider + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog")) + })?; + + paimon_provider.register_temp_table(&database, &table_name, table) + } + + /// Deregisters a temporary table or view. + /// + /// Accepts the same flexible name format as `register_temp_table`. + pub fn deregister_temp_table( + &self, + name: impl Into, + ) -> DFResult>> { + let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?; + let catalog_provider = self + .ctx + .catalog(&catalog) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?; + + let paimon_provider = catalog_provider + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog")) + })?; + + paimon_provider.deregister_temp_table(&database, &table_name) + } + + /// Returns whether a temporary table or view with the given name already exists. + /// + /// Accepts the same flexible name format as `register_temp_table`. + pub fn temp_table_exist(&self, name: impl Into) -> DFResult { + let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?; + let catalog_provider = self + .ctx + .catalog(&catalog) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?; + + let paimon_provider = catalog_provider + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog")) + })?; + + Ok(paimon_provider.temp_table_exist(&database, &table_name)) + } + + /// Resolve a TableReference into (catalog, database, table_name). + fn resolve_temp_table_name(&self, name: TableReference) -> DFResult<(String, String, String)> { + match name { + TableReference::Bare { table } => { + let catalog = self.current_catalog_name(); + let database = self + .ctx + .state() + .config_options() + .catalog + .default_schema + .clone(); + Ok((catalog, database, table.to_string())) + } + TableReference::Partial { schema, table } => { + let catalog = self.current_catalog_name(); + Ok((catalog, schema.to_string(), table.to_string())) + } + TableReference::Full { + catalog, + schema, + table, + } => Ok((catalog.to_string(), schema.to_string(), table.to_string())), + } + } + #[cfg(test)] pub(crate) fn dynamic_options(&self) -> &DynamicOptions { &self.dynamic_options @@ -190,8 +291,16 @@ impl SQLContext { } else { (sql.to_string(), vec![]) }; - let dialect = GenericDialect {}; - let statements = Parser::parse_sql(&dialect, &rewritten_sql) + let sql_lower = rewritten_sql.to_lowercase(); + let has_time_travel = + sql_lower.contains("version as of") || sql_lower.contains("timestamp as of"); + + if has_time_travel { + // Time-travel queries are not DDL; skip our own parsing and handle directly. + return self.handle_time_travel_query(&rewritten_sql).await; + } + + let statements = Parser::parse_sql(&GenericDialect {}, &rewritten_sql) .map_err(|e| DataFusionError::Plan(format!("SQL parse error: {e}")))?; if statements.len() != 1 { @@ -202,10 +311,14 @@ impl SQLContext { match &statements[0] { Statement::CreateTable(create_table) => { - let (catalog, _catalog_name, _) = - self.resolve_catalog_and_table(&create_table.name)?; - self.handle_create_table(&catalog, create_table, partition_keys) - .await + if create_table.temporary { + self.handle_create_temp_table(create_table).await + } else { + let (catalog, _catalog_name, _) = + self.resolve_catalog_and_table(&create_table.name)?; + self.handle_create_table(&catalog, create_table, partition_keys) + .await + } } Statement::AlterTable(alter_table) => { let (catalog, _catalog_name, _) = @@ -262,6 +375,23 @@ impl SQLContext { self.ctx.sql(sql).await } Statement::Truncate(truncate) => self.handle_truncate_table(truncate).await, + Statement::CreateView(create_view) => self.handle_create_view(create_view).await, + Statement::Drop { + object_type, + if_exists, + names, + temporary, + .. + } if matches!(*object_type, ObjectType::Table | ObjectType::View) => { + if *temporary { + self.handle_drop_temp_table(names, *if_exists) + } else if *object_type == ObjectType::Table { + let (catalog, _catalog_name, _) = self.resolve_catalog_and_table(&names[0])?; + self.handle_drop_table(&catalog, names, *if_exists).await + } else { + self.ctx.sql(sql).await + } + } Statement::Call(func) => { crate::procedures::execute_call( &self.ctx, @@ -275,6 +405,130 @@ impl SQLContext { } } + /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). + /// + /// DataFusion's default SQL parser does not support these clauses, so we: + /// 1. Extract the table name and version/timestamp value via regex + /// 2. Strip the time-travel clause from the SQL + /// 3. Create a `PaimonTableProvider` with the appropriate scan options + /// 4. Register it, execute the stripped SQL, then deregister + async fn handle_time_travel_query(&self, sql: &str) -> DFResult { + use crate::table::PaimonTableProvider; + use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION}; + + let (table_name, options, clause_range) = if let Some(info) = extract_version_as_of(sql) { + let options = HashMap::from([(SCAN_VERSION_OPTION.to_string(), info.version)]); + (info.table_name, options, info.clause_range) + } else if let Some(info) = extract_timestamp_as_of(sql) { + let millis = Self::parse_timestamp_to_millis(&info.timestamp)?; + let options = + HashMap::from([(SCAN_TIMESTAMP_MILLIS_OPTION.to_string(), millis.to_string())]); + (info.table_name, options, info.clause_range) + } else { + return Err(DataFusionError::Plan( + "Failed to parse time-travel clause in SQL".to_string(), + )); + }; + + // Resolve the table from our catalog and create a provider with scan options + let table_ref: datafusion::common::TableReference = table_name.as_str().into(); + 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 table_with_options = paimon_table.copy_with_options(options); + let provider = Arc::new(PaimonTableProvider::try_new(table_with_options)?); + + // Use a UUID-based temp table name to avoid conflicts with existing tables. + let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); + + // Replace the original table name + time-travel clause with just the UUID name + let rewritten_sql = format!( + "{}{}{}", + &sql[..clause_range.0], + uuid_name, + &sql[clause_range.1..] + ); + + // Register the provider under the UUID temp table name + self.register_temp_table(uuid_name.as_str(), provider)?; + + // Execute the rewritten SQL + let result = self.ctx.sql(&rewritten_sql).await; + + // Clean up the temp table + let _ = self.deregister_temp_table(uuid_name.as_str()); + + result + } + + /// Parse a timestamp string to milliseconds since epoch (using local timezone). + fn parse_timestamp_to_millis(ts: &str) -> DFResult { + use chrono::{Local, NaiveDateTime, TimeZone}; + + let naive = NaiveDateTime::parse_from_str(ts, "%Y-%m-%d %H:%M:%S").map_err(|e| { + DataFusionError::Plan(format!( + "Cannot parse time travel timestamp '{ts}': {e}. Expected format: YYYY-MM-DD HH:MM:SS" + )) + })?; + let local = Local.from_local_datetime(&naive).single().ok_or_else(|| { + DataFusionError::Plan(format!("Ambiguous or invalid local time: '{ts}'")) + })?; + Ok(local.timestamp_millis()) + } + + /// Resolve a TableReference to (catalog, catalog_name, Identifier). + fn resolve_table_name_from_ref( + &self, + table_ref: &datafusion::common::TableReference, + ) -> DFResult<(Arc, String, Identifier)> { + match table_ref { + datafusion::common::TableReference::Full { + catalog, + schema, + table, + } => { + let catalog_arc = self + .catalogs + .get(catalog.as_ref()) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?; + Ok(( + catalog_arc.clone(), + catalog.to_string(), + Identifier::new(schema.as_ref(), table.as_ref()), + )) + } + datafusion::common::TableReference::Partial { schema, table } => { + let catalog = self.current_catalog()?; + let catalog_name = self.current_catalog_name(); + Ok(( + catalog, + catalog_name, + Identifier::new(schema.as_ref(), table.as_ref()), + )) + } + datafusion::common::TableReference::Bare { table } => { + let catalog = self.current_catalog()?; + let catalog_name = self.current_catalog_name(); + let default_schema = self + .ctx + .state() + .config_options() + .catalog + .default_schema + .clone(); + Ok(( + catalog, + catalog_name, + Identifier::new(default_schema, table.as_ref()), + )) + } + } + } + async fn handle_create_table( &self, catalog: &Arc, @@ -347,6 +601,150 @@ impl SQLContext { ok_result(&self.ctx) } + async fn handle_create_temp_table(&self, ct: &CreateTable) -> DFResult { + let table_ref: TableReference = ct.name.to_string().as_str().into(); + + if ct.if_not_exists && self.temp_table_exist(table_ref.clone())? { + return ok_result(&self.ctx); + } + + // Build the schema from column definitions if provided + let declared_schema = if !ct.columns.is_empty() { + let fields: Vec = ct + .columns + .iter() + .map(|col| { + let paimon_type = + sql_data_type_to_paimon_type(&col.data_type, column_def_nullable(col))?; + let arrow_type = paimon::arrow::paimon_type_to_arrow(&paimon_type) + .map_err(to_datafusion_error)?; + Ok(Field::new( + &col.name.value, + arrow_type, + column_def_nullable(col), + )) + }) + .collect::>>()?; + Some(Arc::new(Schema::new(fields))) + } else { + None + }; + + if let Some(query) = &ct.query { + // CREATE TEMPORARY TABLE ... AS SELECT ... + let query_sql = query.to_string(); + let df = self.ctx.sql(&query_sql).await?; + let schema = df.schema().inner().clone(); + let batches = df.collect().await?; + + // If column types are specified, cast each column to the declared type + let batches = if ct.columns.is_empty() { + batches + } else { + let target_fields: Vec<(String, ArrowDataType)> = ct + .columns + .iter() + .map(|col| { + let paimon_type = + sql_data_type_to_paimon_type(&col.data_type, column_def_nullable(col))?; + let arrow_type = paimon::arrow::paimon_type_to_arrow(&paimon_type) + .map_err(to_datafusion_error)?; + Ok((col.name.value.clone(), arrow_type)) + }) + .collect::>>()?; + + let select_col_count = schema.fields().len(); + let declared_col_count = target_fields.len(); + if select_col_count < declared_col_count { + return Err(DataFusionError::Plan(format!( + "CREATE TEMPORARY TABLE AS SELECT: declared {declared_col_count} column(s) \ + but SELECT query returns only {select_col_count} column(s)" + ))); + } + + batches + .into_iter() + .map(|batch| { + let columns = batch + .columns() + .iter() + .enumerate() + .map(|(i, col)| { + if i < target_fields.len() { + let target_dt = &target_fields[i].1; + if *col.data_type() != *target_dt { + cast(col, target_dt) + .map_err(|e| DataFusionError::External(e.into())) + } else { + Ok(col.clone()) + } + } else { + Ok(col.clone()) + } + }) + .collect::>>()?; + let new_fields = target_fields + .iter() + .zip(schema.fields().iter()) + .map(|((name, dt), _)| Field::new(name, dt.clone(), true)) + .chain( + schema + .fields() + .iter() + .skip(target_fields.len()) + .map(|f| f.as_ref().clone()), + ) + .collect::>(); + let new_schema = Schema::new(new_fields); + RecordBatch::try_new(Arc::new(new_schema), columns) + .map_err(|e| DataFusionError::External(e.into())) + }) + .collect::>>()? + }; + + let schema = batches.first().map(|b| b.schema()).unwrap_or(schema); + let mem_table = MemTable::try_new(schema, vec![batches])?; + self.register_temp_table(table_ref, Arc::new(mem_table))?; + } else if let Some(schema) = declared_schema { + // CREATE TEMPORARY TABLE (col1 TYPE, col2 TYPE, ...) — no data, just the schema + let mem_table = MemTable::try_new(schema, vec![vec![]])?; + self.register_temp_table(table_ref, Arc::new(mem_table))?; + } else { + return Err(DataFusionError::Plan( + "CREATE TEMPORARY TABLE requires column definitions or AS SELECT".to_string(), + )); + } + + ok_result(&self.ctx) + } + + fn handle_drop_temp_table(&self, names: &[ObjectName], if_exists: bool) -> DFResult { + for name in names { + let table_ref: TableReference = name.to_string().as_str().into(); + if if_exists && !self.temp_table_exist(table_ref.clone())? { + continue; + } + self.deregister_temp_table(table_ref)?; + } + ok_result(&self.ctx) + } + + async fn handle_drop_table( + &self, + catalog: &Arc, + names: &[ObjectName], + if_exists: bool, + ) -> DFResult { + for name in names { + let identifier = self.resolve_table_name(name)?; + catalog + .drop_table(&identifier, if_exists) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + } + ok_result(&self.ctx) + } + async fn handle_alter_table( &self, catalog: &Arc, @@ -458,7 +856,7 @@ impl SQLContext { .await .map_err(to_datafusion_error)?; - crate::merge_into::execute_merge_into(&self.ctx, merge, table).await + crate::merge_into::execute_merge_into(self, merge, table).await } async fn handle_update(&self, update: &Update) -> DFResult { @@ -477,7 +875,7 @@ impl SQLContext { .await .map_err(to_datafusion_error)?; - crate::update::execute_update(&self.ctx, update, table).await + crate::update::execute_update(self, update, table).await } async fn handle_delete(&self, delete: &Delete) -> DFResult { @@ -504,7 +902,7 @@ impl SQLContext { .map_err(to_datafusion_error)?; let table_ref = table_name.to_string(); - crate::delete::execute_delete(&self.ctx, delete, table, &table_ref).await + crate::delete::execute_delete(self, delete, table, &table_ref).await } async fn handle_insert_overwrite_partition(&self, insert: &Insert) -> DFResult { @@ -684,6 +1082,40 @@ impl SQLContext { ok_result(&self.ctx) } + async fn handle_create_view(&self, create_view: &CreateView) -> DFResult { + if create_view.materialized { + return Err(DataFusionError::Plan( + "CREATE MATERIALIZED VIEW is not supported".to_string(), + )); + } + + let view_name = create_view.name.to_string(); + let table_ref: TableReference = view_name.as_str().into(); + let (catalog, database, name) = self.resolve_temp_table_name(table_ref)?; + + // Use DataFusion's SQL planner to convert the sqlparser Query into a LogicalPlan + let query_sql = create_view.query.to_string(); + let df = self.ctx.sql(&query_sql).await?; + let logical_plan = df.logical_plan().clone(); + + if create_view.temporary { + if create_view.if_not_exists + && self.temp_table_exist(format!("{catalog}.{database}.{name}"))? + { + return ok_result(&self.ctx); + } + // Create a ViewTable and register it as a temp table + let view_table = datafusion::datasource::ViewTable::new(logical_plan, Some(query_sql)); + self.register_temp_table(format!("{catalog}.{database}.{name}"), Arc::new(view_table))?; + ok_result(&self.ctx) + } else { + Err(DataFusionError::Plan( + "CREATE VIEW (non-temporary) is not supported. Use CREATE TEMPORARY VIEW instead." + .to_string(), + )) + } + } + async fn handle_drop_partitions( &self, catalog: &Arc, @@ -721,7 +1153,7 @@ impl SQLContext { } /// Returns the name of the current default catalog from DataFusion config. - fn current_catalog_name(&self) -> String { + pub(crate) fn current_catalog_name(&self) -> String { self.ctx .state() .config_options() @@ -827,7 +1259,7 @@ fn looks_like_create_table(sql: &str) -> bool { } break; } - // Match "CREATE" then whitespace then "TABLE" (all ASCII, byte-safe) + // Match "CREATE" then whitespace then optional "TEMPORARY"/"TEMP" then "TABLE" (all ASCII, byte-safe) if i + 6 > len || !bytes[i..i + 6].eq_ignore_ascii_case(b"CREATE") { return false; } @@ -838,6 +1270,18 @@ fn looks_like_create_table(sql: &str) -> bool { while i < len && bytes[i].is_ascii_whitespace() { i += 1; } + // Skip optional TEMPORARY or TEMP keyword + if i + 9 <= len && bytes[i..i + 9].eq_ignore_ascii_case(b"TEMPORARY") { + i += 9; + while i < len && bytes[i].is_ascii_whitespace() { + i += 1; + } + } else if i + 4 <= len && bytes[i..i + 4].eq_ignore_ascii_case(b"TEMP") { + i += 4; + while i < len && bytes[i].is_ascii_whitespace() { + i += 1; + } + } i + 5 <= len && bytes[i..i + 5].eq_ignore_ascii_case(b"TABLE") } @@ -1498,6 +1942,115 @@ fn datum_to_constant_array( } } +struct VersionAsOfInfo { + table_name: String, + version: String, + /// Byte range (start, end) covering "table_name VERSION AS OF n" + clause_range: (usize, usize), +} + +struct TimestampAsOfInfo { + table_name: String, + timestamp: String, + /// Byte range (start, end) covering "table_name TIMESTAMP AS OF 'ts'" + clause_range: (usize, usize), +} + +/// Extract `VERSION AS OF ` or `VERSION AS OF ''` from a SQL string. +/// +/// Looks for the pattern (case-insensitive): +/// - `... VERSION AS OF ` — numeric snapshot ID +/// - `... VERSION AS OF ''` — tag name (quoted string) +/// +/// Returns the table name, version/tag value, and byte range of the full clause. +fn extract_version_as_of(sql: &str) -> Option { + let lower = sql.to_lowercase(); + let keyword = "version as of "; + let kw_start = lower.find(keyword)?; + let val_start = kw_start + keyword.len(); + + let remaining = &sql[val_start..]; + + // Parse either a quoted tag name or a numeric snapshot ID + let version = if let Some(after_quote) = remaining.strip_prefix('\'') { + // Tag name: VERSION AS OF 'tagname' + let close_quote = after_quote.find('\'')?; + after_quote[..close_quote].to_string() + } else { + // Numeric snapshot ID: VERSION AS OF 1 + let v: String = remaining + .chars() + .take_while(|c| c.is_ascii_digit()) + .collect(); + if v.is_empty() { + return None; + } + v + }; + + let is_quoted = remaining.starts_with('\''); + let val_end = if is_quoted { + val_start + version.len() + 2 // 2 quotes + } else { + val_start + version.len() + }; + + // Walk backwards from kw_start to find the table name boundary + let table_end = sql[..kw_start].trim_end_matches(' ').len(); + let table_start = sql[..table_end] + .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') + .map(|i| i + 1) + .unwrap_or(0); + let table_name = sql[table_start..table_end].to_string(); + if table_name.is_empty() { + return None; + } + + Some(VersionAsOfInfo { + table_name, + version, + clause_range: (table_start, val_end), + }) +} + +/// Extract `TIMESTAMP AS OF ''` from a SQL string. +/// +/// Looks for the pattern `... TIMESTAMP AS OF ''` (case-insensitive), +/// returns the table name, timestamp string, and byte range of the full clause +/// (from table name start to closing quote end). +fn extract_timestamp_as_of(sql: &str) -> Option { + let lower = sql.to_lowercase(); + let keyword = "timestamp as of "; + let kw_start = lower.find(keyword)?; + let val_start = kw_start + keyword.len(); + + // Read the quoted timestamp string + let remaining = &sql[val_start..]; + if !remaining.starts_with('\'') { + return None; + } + let close_quote = remaining[1..].find('\'')?; + let timestamp = remaining[1..close_quote + 1].to_string(); + let val_end = val_start + close_quote + 2; // skip both quotes + + // Walk backwards to find the table name boundary + let table_end = sql[..kw_start].trim_end_matches(' ').len(); + let table_start = sql[..table_end] + .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') + .map(|i| i + 1) + .unwrap_or(0); + let table_name = sql[table_start..table_end].to_string(); + if table_name.is_empty() { + return None; + } + + Some(TimestampAsOfInfo { + table_name, + timestamp, + clause_range: (table_start, val_end), + }) +} + /// Return an empty DataFrame with a single "result" column containing "OK". fn ok_result(ctx: &SessionContext) -> DFResult { let schema = Arc::new(Schema::new(vec![Field::new( @@ -1576,7 +2129,9 @@ mod tests { Ok(()) } async fn get_database(&self, _name: &str) -> paimon::Result { - unimplemented!() + Err(paimon::Error::DatabaseNotExist { + database: _name.to_string(), + }) } async fn drop_database( &self, @@ -1587,7 +2142,9 @@ mod tests { Ok(()) } async fn get_table(&self, _identifier: &Identifier) -> paimon::Result { - unimplemented!() + Err(paimon::Error::TableNotExist { + full_name: _identifier.to_string(), + }) } async fn list_tables(&self, _database_name: &str) -> paimon::Result> { Ok(vec![]) @@ -2836,4 +3393,301 @@ mod tests { "Expected incomplete partition spec error, got: {err}" ); } + + #[tokio::test] + async fn test_create_temp_table_if_not_exists() { + let catalog = Arc::new(MockCatalog::new()); + let sql_context = make_sql_context(catalog).await; + + // First creation succeeds + sql_context + .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)") + .await + .unwrap(); + + // Second creation without IF NOT EXISTS should fail + let err = sql_context + .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)") + .await + .unwrap_err(); + assert!( + err.to_string().contains("already exists"), + "Expected already-exists error, got: {err}" + ); + + // With IF NOT EXISTS, it should succeed silently + sql_context + .sql("CREATE TEMPORARY TABLE IF NOT EXISTS mydb.t1 (id INT)") + .await + .unwrap(); + } + + #[tokio::test] + async fn test_create_temp_table_if_not_exists_as_select() { + let catalog = Arc::new(MockCatalog::new()); + let sql_context = make_sql_context(catalog).await; + + // Create temp table with AS SELECT + sql_context + .sql("CREATE TEMPORARY TABLE mydb.t2 AS SELECT 1 AS id") + .await + .unwrap(); + + // IF NOT EXISTS should skip when the table already exists + sql_context + .sql("CREATE TEMPORARY TABLE IF NOT EXISTS mydb.t2 AS SELECT 2 AS id") + .await + .unwrap(); + + // Verify the original data is still there (not overwritten) + let df = sql_context.sql("SELECT * FROM mydb.t2").await.unwrap(); + let batches = df.collect().await.unwrap(); + let val = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(val.value(0), 1); + } + + #[tokio::test] + async fn test_create_temp_view_if_not_exists() { + let catalog = Arc::new(MockCatalog::new()); + let sql_context = make_sql_context(catalog).await; + + // First creation succeeds + sql_context + .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 1 AS id") + .await + .unwrap(); + + // Second creation without IF NOT EXISTS should fail + let err = sql_context + .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 2 AS id") + .await + .unwrap_err(); + assert!( + err.to_string().contains("already exists"), + "Expected already-exists error, got: {err}" + ); + + // With IF NOT EXISTS, it should succeed silently + sql_context + .sql("CREATE TEMPORARY VIEW IF NOT EXISTS mydb.v1 AS SELECT 3 AS id") + .await + .unwrap(); + + // Verify the original view is still intact + let df = sql_context.sql("SELECT * FROM mydb.v1").await.unwrap(); + let batches = df.collect().await.unwrap(); + let val = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(val.value(0), 1); + } + + #[tokio::test] + async fn test_drop_temp_table_if_exists() { + let catalog = Arc::new(MockCatalog::new()); + let sql_context = make_sql_context(catalog).await; + + // Dropping a nonexistent temp table without IF EXISTS should error + let err = sql_context + .sql("DROP TEMPORARY TABLE mydb.nonexistent") + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("doesn't exist") + || msg.contains("does not exist") + || msg.contains("Unknown temp database"), + "Expected table-not-exist error, got: {msg}" + ); + + // Dropping with IF EXISTS should succeed silently + sql_context + .sql("DROP TEMPORARY TABLE IF EXISTS mydb.nonexistent") + .await + .unwrap(); + + // Create, then drop with IF EXISTS should actually drop it + sql_context + .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)") + .await + .unwrap(); + + sql_context + .sql("DROP TEMPORARY TABLE IF EXISTS mydb.t1") + .await + .unwrap(); + + // Verify the table is gone + assert!( + !sql_context.temp_table_exist("mydb.t1").unwrap(), + "Expected temp table to be gone after DROP" + ); + } + + #[tokio::test] + async fn test_drop_temp_view_if_exists() { + let catalog = Arc::new(MockCatalog::new()); + let sql_context = make_sql_context(catalog).await; + + // Dropping a nonexistent temp view without IF EXISTS should error + let err = sql_context + .sql("DROP TEMPORARY VIEW mydb.nonexistent") + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("doesn't exist") + || msg.contains("does not exist") + || msg.contains("Unknown temp database"), + "Expected view-not-exist error, got: {msg}" + ); + + // Dropping with IF EXISTS should succeed silently + sql_context + .sql("DROP TEMPORARY VIEW IF EXISTS mydb.nonexistent") + .await + .unwrap(); + + // Create a temp view, then drop with IF EXISTS + sql_context + .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 1 AS id") + .await + .unwrap(); + + sql_context + .sql("DROP TEMPORARY VIEW IF EXISTS mydb.v1") + .await + .unwrap(); + + // Verify the view is gone + assert!( + !sql_context.temp_table_exist("mydb.v1").unwrap(), + "Expected temp view to be gone after DROP" + ); + } + + #[test] + fn test_extract_version_as_of() { + let sql = "SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 1"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "1"); + assert_eq!(info.table_name, "paimon.default.time_travel_table"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT id, name FROM __uuid"); + } + + #[test] + fn test_extract_version_as_of_multi_digit() { + let sql = "SELECT * FROM mydb.t VERSION AS OF 42"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "42"); + assert_eq!(info.table_name, "mydb.t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT * FROM __uuid"); + } + + #[test] + fn test_extract_version_as_of_case_insensitive() { + let sql = "SELECT * FROM t version as of 5"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "5"); + assert_eq!(info.table_name, "t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT * FROM __uuid"); + } + + #[test] + fn test_extract_version_as_of_not_present() { + let sql = "SELECT * FROM t"; + assert!(extract_version_as_of(sql).is_none()); + } + + #[test] + fn test_extract_version_as_of_tag() { + let sql = "SELECT id, name FROM paimon.default.t VERSION AS OF 'snapshot1'"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "snapshot1"); + assert_eq!(info.table_name, "paimon.default.t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT id, name FROM __uuid"); + } + + #[test] + fn test_extract_version_as_of_tag_case_insensitive() { + let sql = "SELECT * FROM t version as of 'my_tag'"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "my_tag"); + assert_eq!(info.table_name, "t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT * FROM __uuid"); + } + + #[test] + fn test_extract_version_as_of_numeric_still_works() { + // Ensure numeric snapshot ID still works alongside tag support + let sql = "SELECT * FROM t VERSION AS OF 123"; + let info = extract_version_as_of(sql).unwrap(); + assert_eq!(info.version, "123"); + assert_eq!(info.table_name, "t"); + } + + #[test] + fn test_extract_timestamp_as_of() { + let sql = "SELECT * FROM paimon.default.t TIMESTAMP AS OF '2024-01-15 10:30:00'"; + let info = extract_timestamp_as_of(sql).unwrap(); + assert_eq!(info.timestamp, "2024-01-15 10:30:00"); + assert_eq!(info.table_name, "paimon.default.t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT * FROM __uuid"); + } + + #[test] + fn test_extract_timestamp_as_of_case_insensitive() { + let sql = "SELECT * FROM t timestamp as of '2024-06-01 00:00:00'"; + let info = extract_timestamp_as_of(sql).unwrap(); + assert_eq!(info.timestamp, "2024-06-01 00:00:00"); + assert_eq!(info.table_name, "t"); + let rewritten = format!( + "{}__uuid{}", + &sql[..info.clause_range.0], + &sql[info.clause_range.1..] + ); + assert_eq!(rewritten, "SELECT * FROM __uuid"); + } + + #[test] + fn test_extract_timestamp_as_of_not_present() { + let sql = "SELECT * FROM t"; + assert!(extract_timestamp_as_of(sql).is_none()); + } } diff --git a/crates/integrations/datafusion/src/update.rs b/crates/integrations/datafusion/src/update.rs index 741462b4..8803dcc6 100644 --- a/crates/integrations/datafusion/src/update.rs +++ b/crates/integrations/datafusion/src/update.rs @@ -37,16 +37,13 @@ use crate::error::to_datafusion_error; use crate::merge_into::{ build_partition_set_from_where, extract_tracking_columns, is_delete_conflict, is_row_id_conflict, ok_result, project_update_columns, quote_identifier, - register_cow_target_table, retry_on_conflict, + register_cow_target_table, retry_on_conflict, TempTableTracker, }; +use crate::sql_context::SQLContext; /// Execute an UPDATE statement on a Paimon table. -/// -/// Dispatches to the appropriate execution path based on table type: -/// - Data evolution tables → partial-column writes via `DataEvolutionWriter` -/// - Append-only tables (no PK) → copy-on-write file rewriting via `CopyOnWriteMergeWriter` pub(crate) async fn execute_update( - ctx: &SessionContext, + ctx: &SQLContext, update: &Update, table: Table, ) -> DFResult { @@ -77,7 +74,7 @@ pub(crate) async fn execute_update( /// Execute UPDATE on a data evolution table with retry on row ID conflict. async fn execute_data_evolution_update( - ctx: &SessionContext, + ctx: &SQLContext, update: &Update, table: Table, ) -> DFResult { @@ -89,7 +86,7 @@ async fn execute_data_evolution_update( /// Single attempt of UPDATE execution. async fn execute_update_once( - ctx: &SessionContext, + ctx: &SQLContext, update: &Update, table: &Table, ) -> DFResult { @@ -137,12 +134,12 @@ async fn execute_update_once( }; let query_sql = format!("SELECT {select_clause} FROM {table_ref}{where_clause}"); - let batches = ctx.sql(&query_sql).await?.collect().await?; + let batches = ctx.ctx().sql(&query_sql).await?.collect().await?; // 4. Project update columns (rename __upd_X → X) let total_count: u64 = batches.iter().map(|b| b.num_rows() as u64).sum(); if total_count == 0 { - return ok_result(ctx, 0); + return ok_result(ctx.ctx(), 0); } let update_batches = project_update_columns(&batches, &columns)?; @@ -159,7 +156,7 @@ async fn execute_update_once( commit.commit(messages).await.map_err(to_datafusion_error)?; } - ok_result(ctx, total_count) + ok_result(ctx.ctx(), total_count) } // --------------------------------------------------------------------------- @@ -168,7 +165,7 @@ async fn execute_update_once( /// Execute UPDATE on an append-only table with retry on delete conflict. async fn execute_cow_update( - ctx: &SessionContext, + ctx: &SQLContext, update: &Update, table: &Table, ) -> DFResult { @@ -180,7 +177,7 @@ async fn execute_cow_update( /// Single attempt of CoW UPDATE execution. async fn execute_cow_update_once( - ctx: &SessionContext, + ctx: &SQLContext, update: &Update, table: &Table, ) -> DFResult { @@ -195,21 +192,23 @@ async fn execute_cow_update_once( .await .map_err(to_datafusion_error)?; - let (has_data, cow_table_guard) = register_cow_target_table(ctx, table, &writer).await?; + let mut temp_tracker = TempTableTracker::new(ctx); + let (has_data, cow_table_name) = + register_cow_target_table(ctx, table, &writer, &mut temp_tracker).await?; if !has_data { - return ok_result(ctx, 0); + return ok_result(ctx.ctx(), 0); } + let cow_target_name = cow_table_name; let result = execute_cow_update_inner( - ctx, + ctx.ctx(), &columns, &exprs, - &cow_table_guard.qualified_name(), + &cow_target_name, update, &mut writer, ) .await; - drop(cow_table_guard); let total_count = result?; let messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; @@ -218,7 +217,7 @@ async fn execute_cow_update_once( commit.commit(messages).await.map_err(to_datafusion_error)?; } - ok_result(ctx, total_count) + ok_result(ctx.ctx(), total_count) } async fn execute_cow_update_inner( @@ -326,7 +325,6 @@ mod tests { use std::sync::Arc; use datafusion::arrow::array::{Int32Array, StringArray, UInt64Array}; - use datafusion::prelude::SessionContext; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; use paimon::catalog::{Catalog, Identifier}; @@ -335,8 +333,6 @@ mod tests { use paimon::{CatalogOptions, FileSystemCatalog, Options}; use tempfile::TempDir; - use crate::{PaimonTableProvider, SQLContext}; - async fn setup_sql_context() -> (TempDir, SQLContext, Arc) { let temp_dir = TempDir::new().unwrap(); let warehouse = format!("file://{}", temp_dir.path().display()); @@ -357,12 +353,12 @@ mod tests { (temp_dir, sql_context, catalog) } - async fn setup_data_evolution_table(name: &str) -> (TempDir, SessionContext, Table) { + async fn setup_data_evolution_table(name: &str) -> (TempDir, SQLContext, Table) { let (tmp, sql_context, catalog) = setup_sql_context().await; sql_context .sql(&format!( - "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('row-tracking.enabled' = 'true')" + "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true')" )) .await .unwrap(); @@ -381,17 +377,8 @@ mod tests { .get_table(&Identifier::new("test_db", name)) .await .unwrap(); - let mut extra = std::collections::HashMap::new(); - extra.insert("data-evolution.enabled".to_string(), "true".to_string()); - extra.insert("row-tracking.enabled".to_string(), "true".to_string()); - let de_table = table.copy_with_options(extra); - - let ctx = sql_context.ctx().clone(); - let provider = PaimonTableProvider::try_new(de_table.clone()).unwrap(); - ctx.register_table("datafusion.public.target", Arc::new(provider)) - .unwrap(); - (tmp, ctx, de_table) + (tmp, sql_context, table) } fn parse_update(sql: &str) -> Update { @@ -430,14 +417,14 @@ mod tests { #[tokio::test] async fn test_update_with_where() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_with_where").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_with_where").await; let update = - parse_update("UPDATE datafusion.public.target SET name = 'ALICE' WHERE id = 1"); - execute_update(&ctx, &update, table).await.unwrap(); + parse_update("UPDATE paimon.test_db.t_with_where SET name = 'ALICE' WHERE id = 1"); + execute_update(&sql_context, &update, table).await.unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_with_where ORDER BY id") .await .unwrap() .collect() @@ -457,13 +444,13 @@ mod tests { #[tokio::test] async fn test_update_without_where() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_without_where").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_without_where").await; - let update = parse_update("UPDATE datafusion.public.target SET value = 99"); - execute_update(&ctx, &update, table).await.unwrap(); + let update = parse_update("UPDATE paimon.test_db.t_without_where SET value = 99"); + execute_update(&sql_context, &update, table).await.unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_without_where ORDER BY id") .await .unwrap() .collect() @@ -483,15 +470,15 @@ mod tests { #[tokio::test] async fn test_update_multiple_columns() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_multi_col").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_multi_col").await; let update = parse_update( - "UPDATE datafusion.public.target SET name = 'updated', value = 0 WHERE id = 2", + "UPDATE paimon.test_db.t_multi_col SET name = 'updated', value = 0 WHERE id = 2", ); - execute_update(&ctx, &update, table).await.unwrap(); + execute_update(&sql_context, &update, table).await.unwrap(); - let batches = ctx - .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id") + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_multi_col ORDER BY id") .await .unwrap() .collect() @@ -511,11 +498,11 @@ mod tests { #[tokio::test] async fn test_update_no_matching_rows() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_no_match").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_no_match").await; let update = - parse_update("UPDATE datafusion.public.target SET name = 'nobody' WHERE id = 99"); - let result = execute_update(&ctx, &update, table).await.unwrap(); + parse_update("UPDATE paimon.test_db.t_no_match SET name = 'nobody' WHERE id = 99"); + let result = execute_update(&sql_context, &update, table).await.unwrap(); let batches = result.collect().await.unwrap(); let count = batches[0] .column(0) @@ -528,24 +515,23 @@ mod tests { #[tokio::test] async fn test_update_row_id_stability() { - let (_tmp, ctx, table) = setup_data_evolution_table("t_row_id").await; + let (_tmp, sql_context, table) = setup_data_evolution_table("t_row_id").await; // Get row IDs before update - let before = ctx - .sql("SELECT id, \"_ROW_ID\" FROM datafusion.public.target ORDER BY id") + let before = sql_context + .sql("SELECT id, \"_ROW_ID\" FROM paimon.test_db.t_row_id ORDER BY id") .await .unwrap() .collect() .await .unwrap(); - let update = - parse_update("UPDATE datafusion.public.target SET name = 'ALICE' WHERE id = 1"); - execute_update(&ctx, &update, table).await.unwrap(); + let update = parse_update("UPDATE paimon.test_db.t_row_id SET name = 'ALICE' WHERE id = 1"); + execute_update(&sql_context, &update, table).await.unwrap(); // Get row IDs after update - let after = ctx - .sql("SELECT id, \"_ROW_ID\" FROM datafusion.public.target ORDER BY id") + let after = sql_context + .sql("SELECT id, \"_ROW_ID\" FROM paimon.test_db.t_row_id ORDER BY id") .await .unwrap() .collect() @@ -584,9 +570,9 @@ mod tests { None, ); - let ctx = SessionContext::new(); + let sql_context = SQLContext::new(); let update = parse_update("UPDATE t SET id = 1"); - let result = execute_update(&ctx, &update, table).await; + let result = execute_update(&sql_context, &update, table).await; assert!(result.is_err()); assert!(result .unwrap_err() diff --git a/crates/integrations/datafusion/tests/append_merge_into.rs b/crates/integrations/datafusion/tests/append_merge_into.rs index 9e8dc0a7..c14936c6 100644 --- a/crates/integrations/datafusion/tests/append_merge_into.rs +++ b/crates/integrations/datafusion/tests/append_merge_into.rs @@ -27,7 +27,7 @@ use paimon_datafusion::SQLContext; use common::{ collect_int_int_str, collect_int_str, collect_three_ints, create_sql_context, create_test_env, - ctx_exec, exec, + exec, }; // ======================= Helpers ======================= @@ -92,16 +92,16 @@ async fn test_only_update() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c", ) .await; @@ -121,16 +121,16 @@ async fn test_only_delete() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN DELETE", ) .await; @@ -147,16 +147,16 @@ async fn test_only_insert() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) .await; @@ -180,16 +180,16 @@ async fn test_update_and_insert() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -214,16 +214,16 @@ async fn test_delete_and_insert() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN DELETE \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -244,16 +244,16 @@ async fn test_partial_insert_with_null() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN NOT MATCHED THEN INSERT (a) VALUES (s.a)", ) .await; @@ -316,16 +316,16 @@ async fn test_update_from_both_source_and_target() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET b = t.b * 11, c = s.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b * 2, s.c)", ) @@ -350,16 +350,16 @@ async fn test_columns_in_wrong_order() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET c = s.c, b = s.b \ WHEN NOT MATCHED THEN INSERT (b, c, a) VALUES (b, c, a)", ) @@ -384,16 +384,16 @@ async fn test_partial_update() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET c = s.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -418,12 +418,12 @@ async fn test_source_is_subquery() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec(&sql_context, "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33'), (4, 400, 'c44')").await; + exec(&sql_context, "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33'), (4, 400, 'c44')) AS t(a, b, c)").await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING (SELECT a, b, c FROM datafusion.public.source WHERE a % 2 = 1) AS src \ + USING (SELECT a, b, c FROM paimon.test_db.source WHERE a % 2 = 1) AS src \ ON t.a = src.a \ WHEN MATCHED THEN UPDATE SET b = src.b, c = src.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (src.a, src.b, src.c)", @@ -445,12 +445,12 @@ async fn test_source_is_subquery() { async fn test_source_and_target_empty() { let (_tmp, sql_context) = setup_abc().await; // target is empty, source is empty - ctx_exec(&sql_context, "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS SELECT * FROM (VALUES (1, 1, 'x')) AS t(a, b, c) WHERE 1=0").await; + exec(&sql_context, "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 1, 'x')) AS t(a, b, c) WHERE 1=0").await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -471,16 +471,16 @@ async fn test_with_alias() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET b = s.b, c = s.c \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -505,16 +505,16 @@ async fn test_reversed_on_condition() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON s.a = t.a \ + USING paimon.test_db.source s ON s.a = t.a \ WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c", ) .await; @@ -634,16 +634,16 @@ async fn test_coalesce_source_and_target() { "INSERT INTO paimon.test_db.target VALUES (1, 'guid_tgt_1'), (2, 'guid_tgt_2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 'guid_src_1'), (3, 'guid_src_3')) AS t(a, b)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target AS dest \ - USING datafusion.public.source AS src ON dest.a = src.a \ + USING paimon.test_db.source AS src ON dest.a = src.a \ WHEN MATCHED AND (nullif(cast(src.b as STRING), '') IS NOT NULL) THEN \ UPDATE SET b = COALESCE(nullif(cast(src.b as STRING), ''), dest.b) \ WHEN NOT MATCHED THEN INSERT (a, b) VALUES (src.a, src.b)", @@ -676,16 +676,16 @@ async fn test_subquery_source_with_coalesce() { "INSERT INTO paimon.test_db.target VALUES (1, 'guid_tgt_1'), (2, 'guid_tgt_2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 'guid_src_1'), (3, 'guid_src_3')) AS t(a, b)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target AS dest \ - USING (SELECT * FROM datafusion.public.source) AS src ON dest.a = src.a \ + USING (SELECT * FROM paimon.test_db.source) AS src ON dest.a = src.a \ WHEN MATCHED AND (nullif(cast(src.b as STRING), '') IS NOT NULL) THEN \ UPDATE SET b = COALESCE(nullif(cast(src.b as STRING), ''), dest.b) \ WHEN NOT MATCHED THEN INSERT (a, b) VALUES (src.a, src.b)", @@ -718,16 +718,16 @@ async fn test_insert_only_is_append_commit() { "INSERT INTO paimon.test_db.target VALUES (2, 2, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 1, 'c1')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 1, 'c1')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) .await; @@ -749,15 +749,15 @@ async fn test_successive_merges() { .await; // First merge: update b for a=1 - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src1 (a INT, b INT) AS VALUES (1, 100)", + "CREATE TEMPORARY TABLE paimon.test_db.src1 AS SELECT * FROM (VALUES (1, 100)) AS t(a, b)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.src1 s ON t.a = s.a \ + USING paimon.test_db.src1 s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET b = s.b", ) .await; @@ -768,15 +768,15 @@ async fn test_successive_merges() { ); // Second merge: update c for a=2 - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src2 (a INT, c VARCHAR) AS VALUES (2, 'C2_UPDATED')", + "CREATE TEMPORARY TABLE paimon.test_db.src2 AS SELECT * FROM (VALUES (2, 'C2_UPDATED')) AS t(a, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.src2 s ON t.a = s.a \ + USING paimon.test_db.src2 s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET c = s.c", ) .await; @@ -796,16 +796,16 @@ async fn test_no_match_no_change() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (99, 990, 'c99')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (99, 990, 'c99')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET b = s.b, c = s.c", ) .await; @@ -825,16 +825,16 @@ async fn test_delete_all_rows() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT) AS VALUES (1), (2)", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1), (2)) AS t(a)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN DELETE", ) .await; @@ -854,12 +854,12 @@ async fn test_insert_many_rows() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1')", ) .await; - ctx_exec(&sql_context, "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (2, 20, 'c2'), (3, 30, 'c3'), (4, 40, 'c4'), (5, 50, 'c5')").await; + exec(&sql_context, "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (2, 20, 'c2'), (3, 30, 'c3'), (4, 40, 'c4'), (5, 50, 'c5')) AS t(a, b, c)").await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) .await; @@ -885,16 +885,16 @@ async fn test_conditional_update() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2'), (3, 30, 'c3')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED AND s.b > 200 THEN UPDATE SET b = s.b, c = s.c \ WHEN MATCHED THEN DELETE \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", @@ -916,16 +916,16 @@ async fn test_conditional_insert() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED THEN UPDATE SET b = s.b, c = s.c \ WHEN NOT MATCHED AND s.b < 300 THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -947,16 +947,16 @@ async fn test_conditional_delete() { "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')", + "CREATE TEMPORARY TABLE paimon.test_db.source AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33')) AS t(a, b, c)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED AND t.c < 'c1' THEN DELETE \ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)", ) @@ -978,11 +978,11 @@ async fn test_conditional_delete() { async fn test_multiple_matched_clauses() { let (_tmp, sql_context) = setup_abc().await; exec(&sql_context, "INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2'), (3, 30, 'c3'), (4, 40, 'c4'), (5, 50, 'c5')").await; - ctx_exec(&sql_context, "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33'), (5, 500, 'c55'), (7, 700, 'c77'), (9, 900, 'c99')").await; + exec(&sql_context, "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, c VARCHAR) AS SELECT * FROM (VALUES (1, 100, 'c11'), (3, 300, 'c33'), (5, 500, 'c55'), (7, 700, 'c77'), (9, 900, 'c99')) AS t(a, b, c)").await; exec(&sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a \ + USING paimon.test_db.source s ON t.a = s.a \ WHEN MATCHED AND t.a = 5 THEN UPDATE SET b = s.b + t.b \ WHEN MATCHED AND s.c > 'c2' THEN UPDATE SET a = s.a, b = s.b, c = s.c \ WHEN MATCHED THEN DELETE \ @@ -1012,16 +1012,16 @@ async fn test_multiple_matched_clauses() { #[tokio::test] async fn test_partitioned_update_single_partition() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1)", + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (1, 100, 1)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.source s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN UPDATE SET b = s.b", ) .await; @@ -1035,16 +1035,16 @@ async fn test_partitioned_update_single_partition() { #[tokio::test] async fn test_partitioned_update_multiple_partitions() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (3, 300, 2)", + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (1, 100, 1), (3, 300, 2)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.source s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN UPDATE SET b = s.b", ) .await; @@ -1058,16 +1058,16 @@ async fn test_partitioned_update_multiple_partitions() { #[tokio::test] async fn test_partitioned_delete() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, pt INT) AS VALUES (1, 1), (3, 2)", + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, pt INT) AS SELECT * FROM (VALUES (1, 1), (3, 2)) AS t(a, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.source s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN DELETE", ) .await; @@ -1081,16 +1081,16 @@ async fn test_partitioned_delete() { #[tokio::test] async fn test_partitioned_insert_new_partition() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (5, 50, 3), (6, 60, 3)", + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (5, 50, 3), (6, 60, 3)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.source s ON t.a = s.a AND t.pt = s.pt \ WHEN NOT MATCHED THEN INSERT (a, b, pt) VALUES (s.a, s.b, s.pt)", ) .await; @@ -1111,16 +1111,16 @@ async fn test_partitioned_insert_new_partition() { #[tokio::test] async fn test_partitioned_update_and_insert() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (5, 50, 2)", + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (1, 100, 1), (5, 50, 2)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.source s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN UPDATE SET b = s.b \ WHEN NOT MATCHED THEN INSERT (a, b, pt) VALUES (s.a, s.b, s.pt)", ) @@ -1136,28 +1136,28 @@ async fn test_partitioned_update_and_insert() { async fn test_partitioned_successive_merges() { let (_tmp, sql_context) = setup_partitioned().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src1 (a INT, b INT, pt INT) AS VALUES (1, 100, 1)", + "CREATE TEMPORARY TABLE paimon.test_db.src1 (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (1, 100, 1)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.src1 s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.src1 s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN UPDATE SET b = s.b", ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src2 (a INT, b INT, pt INT) AS VALUES (3, 300, 2)", + "CREATE TEMPORARY TABLE paimon.test_db.src2 (a INT, b INT, pt INT) AS SELECT * FROM (VALUES (3, 300, 2)) AS t(a, b, pt)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.target t \ - USING datafusion.public.src2 s ON t.a = s.a AND t.pt = s.pt \ + USING paimon.test_db.src2 s ON t.a = s.a AND t.pt = s.pt \ WHEN MATCHED THEN UPDATE SET b = s.b", ) .await; diff --git a/crates/integrations/datafusion/tests/blob_tests.rs b/crates/integrations/datafusion/tests/blob_tests.rs index 286bc6fc..3c3d3c6d 100644 --- a/crates/integrations/datafusion/tests/blob_tests.rs +++ b/crates/integrations/datafusion/tests/blob_tests.rs @@ -22,7 +22,7 @@ mod common; use arrow_array::{Array, BinaryArray, Int32Array, RecordBatch, StringArray}; -use common::{create_sql_context, create_test_env, ctx_exec, exec}; +use common::{create_sql_context, create_test_env, exec}; use paimon::spec::BlobDescriptor; use paimon_datafusion::SQLContext; @@ -294,16 +294,16 @@ async fn test_merge_into_updates_non_blob_on_raw_blob_table() { ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT, name VARCHAR) AS VALUES (1, 'Updated')", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1, 'Updated')) AS t(id, name)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.t t \ - USING datafusion.public.src s ON t.id = s.id \ + USING paimon.test_db.src s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await; @@ -333,16 +333,16 @@ async fn test_merge_into_rejects_raw_blob_update() { ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT, picture BYTEA) AS VALUES (1, X'4242')", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1, X'4242')) AS t(id, picture)", ) .await; let result = sql_context .sql( "MERGE INTO paimon.test_db.t t \ - USING datafusion.public.src s ON t.id = s.id \ + USING paimon.test_db.src s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET picture = s.picture", ) .await; @@ -379,16 +379,16 @@ async fn test_merge_into_updates_non_blob_on_descriptor_table() { ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT, name VARCHAR) AS VALUES (1, 'Updated')", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1, 'Updated')) AS t(id, name)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.t t \ - USING datafusion.public.src s ON t.id = s.id \ + USING paimon.test_db.src s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await; @@ -428,16 +428,16 @@ async fn test_merge_into_updates_blob_on_descriptor_table() { ) .await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT, picture BYTEA) AS VALUES (1, X'4343')", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1, X'4343')) AS t(id, picture)", ) .await; exec( &sql_context, "MERGE INTO paimon.test_db.t t \ - USING datafusion.public.src s ON t.id = s.id \ + USING paimon.test_db.src s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET picture = s.picture", ) .await; diff --git a/crates/integrations/datafusion/tests/common/mod.rs b/crates/integrations/datafusion/tests/common/mod.rs index f330ef3e..d32734c3 100644 --- a/crates/integrations/datafusion/tests/common/mod.rs +++ b/crates/integrations/datafusion/tests/common/mod.rs @@ -106,21 +106,6 @@ pub async fn exec(sql_context: &SQLContext, s: &str) { sql_context.sql(s).await.unwrap().collect().await.unwrap(); } -/// Execute SQL on the raw DataFusion context (for non-Paimon source tables). -/// Temporarily switches to datafusion.public so CREATE TABLE lands in the -/// built-in catalog, then restores the Paimon defaults. -#[allow(dead_code)] -pub async fn ctx_exec(sql_context: &SQLContext, s: &str) { - sql_context - .ctx() - .sql(s) - .await - .unwrap() - .collect() - .await - .unwrap(); -} - /// Extract the count from a DML result (returns a single UInt64 column). #[allow(dead_code)] pub async fn dml_count(sql_context: &SQLContext, sql_str: &str) -> u64 { diff --git a/crates/integrations/datafusion/tests/delete_tests.rs b/crates/integrations/datafusion/tests/delete_tests.rs index 17cf61b4..0ed3205c 100644 --- a/crates/integrations/datafusion/tests/delete_tests.rs +++ b/crates/integrations/datafusion/tests/delete_tests.rs @@ -23,7 +23,7 @@ mod common; use paimon_datafusion::SQLContext; -use common::{create_sql_context, create_test_env, ctx_exec, dml_count, exec, query_int_str_int}; +use common::{create_sql_context, create_test_env, dml_count, exec, query_int_str_int}; // ======================= Helpers ======================= @@ -256,15 +256,15 @@ async fn test_delete_not_in_condition() { async fn test_delete_in_subquery() { let (_tmp, sql_context) = setup().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT) AS VALUES (1), (3)", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1), (3)) AS t(id)", ) .await; exec( &sql_context, - "DELETE FROM paimon.test_db.t WHERE id IN (SELECT id FROM datafusion.public.src)", + "DELETE FROM paimon.test_db.t WHERE id IN (SELECT id FROM paimon.test_db.src)", ) .await; @@ -278,15 +278,15 @@ async fn test_delete_in_subquery() { async fn test_delete_scalar_subquery() { let (_tmp, sql_context) = setup().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT) AS VALUES (2)", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (2)) AS t(id)", ) .await; exec( &sql_context, - "DELETE FROM paimon.test_db.t WHERE id >= (SELECT MAX(id) FROM datafusion.public.src)", + "DELETE FROM paimon.test_db.t WHERE id >= (SELECT MAX(id) FROM paimon.test_db.src)", ) .await; diff --git a/crates/integrations/datafusion/tests/merge_into_tests.rs b/crates/integrations/datafusion/tests/merge_into_tests.rs index 55ad8526..97d355bb 100644 --- a/crates/integrations/datafusion/tests/merge_into_tests.rs +++ b/crates/integrations/datafusion/tests/merge_into_tests.rs @@ -132,8 +132,7 @@ async fn assert_merge_error(sql_context: &SQLContext, sql: &str, expected_substr } async fn register_source(sql_context: &SQLContext, sql: &str) { - let ctx = sql_context.ctx(); - ctx.sql(sql).await.unwrap().collect().await.unwrap(); + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); } // ======================= Functional Tests ======================= @@ -208,13 +207,13 @@ async fn test_row_id_stability_after_merge_into() { // Register source and execute MERGE INTO register_source( &sql_context, - "CREATE TABLE datafusion.public.source1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')", + "CREATE TEMPORARY TABLE paimon.test_db.source1 AS SELECT * FROM (VALUES (1, 'ALICE'), (3, 'CHARLIE')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.source1 s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.source1 s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -267,12 +266,12 @@ async fn test_multiple_merge_into_different_columns() { // First MERGE: update name for id=1 register_source( &sql_context, - "CREATE TABLE datafusion.public.src_name (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_name AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_name s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_name s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -284,12 +283,12 @@ async fn test_multiple_merge_into_different_columns() { // Second MERGE: update value for id=2 register_source( &sql_context, - "CREATE TABLE datafusion.public.src_value (id INT, value INT) AS VALUES (2, 200)", + "CREATE TEMPORARY TABLE paimon.test_db.src_value AS SELECT * FROM (VALUES (2, 200)) AS t(id, value)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_value s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_value s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET value = s.value", ) .await @@ -324,13 +323,13 @@ async fn test_merge_into_with_non_paimon_source() { // Source is a plain DataFusion in-memory table, not Paimon register_source( &sql_context, - "CREATE TABLE datafusion.public.df_source (id INT, name VARCHAR) AS VALUES (2, 'BOB_UPDATED')", + "CREATE TEMPORARY TABLE paimon.test_db.df_source AS SELECT * FROM (VALUES (2, 'BOB_UPDATED')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.df_source s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.df_source s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -377,14 +376,14 @@ async fn test_merge_into_join_on_row_id() { register_source( &sql_context, &format!( - "CREATE TABLE datafusion.public.rid_source (row_id BIGINT, name VARCHAR) AS VALUES ({row_id_of_2}, 'BOB')" + "CREATE TEMPORARY TABLE paimon.test_db.rid_source AS SELECT * FROM (VALUES ({row_id_of_2}, 'BOB')) AS t(row_id, name)" ), ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.rid_source s ON t.\"_ROW_ID\" = s.row_id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.rid_source s ON t.\"_ROW_ID\" = s.row_id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -468,12 +467,12 @@ async fn test_row_count_after_merge() { // MERGE INTO: update 1 row register_source( &sql_context, - "CREATE TABLE datafusion.public.src_count (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_count AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_count s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_count s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -509,13 +508,13 @@ async fn test_merge_into_update_and_insert() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_ui (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)", + "CREATE TEMPORARY TABLE paimon.test_db.src_ui AS SELECT * FROM (VALUES (1, 'alice', 11), (2, 'BOB', 22)) AS t(id, name, value)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_ui s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_ui s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ) @@ -554,14 +553,14 @@ async fn test_merge_into_insert_only() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_io (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)", + "CREATE TEMPORARY TABLE paimon.test_db.src_io AS SELECT * FROM (VALUES (1, 'alice', 11), (2, 'BOB', 22)) AS t(id, name, value)", ) .await; // Only INSERT, no MATCHED clause — matched row id=2 should be untouched sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_io s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_io s ON t.id = s.id \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ) .await @@ -604,13 +603,13 @@ async fn test_merge_into_insert_all_columns() { // Source schema matches target: (id, name, value) register_source( &sql_context, - "CREATE TABLE datafusion.public.src_star (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 10), (2, 'BOB', 22)", + "CREATE TEMPORARY TABLE paimon.test_db.src_star AS SELECT * FROM (VALUES (1, 'alice', 10), (2, 'BOB', 22)) AS t(id, name, value)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_star s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_star s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ) @@ -652,14 +651,14 @@ async fn test_merge_into_insert_partial_columns() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_partial (id INT, name VARCHAR) AS VALUES (1, 'alice'), (2, 'BOB')", + "CREATE TEMPORARY TABLE paimon.test_db.src_partial AS SELECT * FROM (VALUES (1, 'alice'), (2, 'BOB')) AS t(id, name)", ) .await; // INSERT only id and name, value should be NULL (but our schema has INT, so this tests partial insert) sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_partial s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_partial s ON t.id = s.id \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, 0)", ) .await @@ -701,14 +700,14 @@ async fn test_merge_into_insert_with_predicate() { // Source has 3 rows, only id=1 matches target register_source( &sql_context, - "CREATE TABLE datafusion.public.src_pred (id INT, name VARCHAR, value INT) AS VALUES (1, 'ALICE', 11), (2, 'bob', 20), (3, 'charlie', 30)", + "CREATE TEMPORARY TABLE paimon.test_db.src_pred AS SELECT * FROM (VALUES (1, 'ALICE', 11), (2, 'bob', 20), (3, 'charlie', 30)) AS t(id, name, value)", ) .await; // Only insert when value > 25 sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_pred s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_pred s ON t.id = s.id \ WHEN NOT MATCHED AND s.value > 25 THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ) .await @@ -754,13 +753,13 @@ async fn test_merge_into_row_id_for_inserted_rows() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_rid (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)", + "CREATE TEMPORARY TABLE paimon.test_db.src_rid AS SELECT * FROM (VALUES (1, 'alice', 11), (2, 'BOB', 22)) AS t(id, name, value)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_rid s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_rid s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name \ WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", ) @@ -814,13 +813,13 @@ async fn test_rejects_when_matched_delete() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_del (id INT) AS VALUES (1)", + "CREATE TEMPORARY TABLE paimon.test_db.src_del AS SELECT * FROM (VALUES (1)) AS t(id)", ) .await; assert_merge_error( &sql_context, - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_del s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_del s ON t.id = s.id \ WHEN MATCHED THEN DELETE", "WHEN MATCHED THEN DELETE is not supported", ) @@ -845,13 +844,13 @@ async fn test_rejects_multiple_when_matched() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_multi (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_multi AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; assert_merge_error( &sql_context, - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_multi s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_multi s ON t.id = s.id \ WHEN MATCHED AND t.id = 1 THEN UPDATE SET name = s.name \ WHEN MATCHED THEN UPDATE SET name = 'default'", "WHEN MATCHED AND is not yet supported", @@ -891,13 +890,13 @@ async fn test_rejects_partition_column_in_set() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_pt (id INT, pt VARCHAR) AS VALUES (1, 'b')", + "CREATE TEMPORARY TABLE paimon.test_db.src_pt AS SELECT * FROM (VALUES (1, 'b')) AS t(id, pt)", ) .await; assert_merge_error( &sql_context, - "MERGE INTO paimon.test_db.part_target t USING datafusion.public.src_pt s ON t.id = s.id \ + "MERGE INTO paimon.test_db.part_target t USING paimon.test_db.src_pt s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET pt = s.pt", "Cannot update partition column", ) @@ -934,13 +933,13 @@ async fn test_rejects_table_without_row_tracking() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_nrt (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_nrt AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; assert_merge_error( &sql_context, - "MERGE INTO paimon.test_db.no_tracking t USING datafusion.public.src_nrt s ON t.id = s.id \ + "MERGE INTO paimon.test_db.no_tracking t USING paimon.test_db.src_nrt s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", "row-tracking.enabled", ) @@ -964,12 +963,12 @@ async fn test_successive_merges_read_file_group() { // First MERGE: update 'name' column → creates a partial-column file for 'name' register_source( &sql_context, - "CREATE TABLE datafusion.public.src_m1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (2, 'BOB')", + "CREATE TEMPORARY TABLE paimon.test_db.src_m1 AS SELECT * FROM (VALUES (1, 'ALICE'), (2, 'BOB')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_m1 s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_m1 s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -993,12 +992,12 @@ async fn test_successive_merges_read_file_group() { // (base file has original 'name', partial file has updated 'name' from first merge) register_source( &sql_context, - "CREATE TABLE datafusion.public.src_m2 (id INT, name VARCHAR) AS VALUES (1, 'Alice_v2')", + "CREATE TEMPORARY TABLE paimon.test_db.src_m2 AS SELECT * FROM (VALUES (1, 'Alice_v2')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_m2 s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_m2 s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -1036,12 +1035,12 @@ async fn test_successive_merges_different_columns_read_file_group() { // First MERGE: update 'name' register_source( &sql_context, - "CREATE TABLE datafusion.public.src_dc1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_dc1 AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_dc1 s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_dc1 s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", ) .await @@ -1053,12 +1052,12 @@ async fn test_successive_merges_different_columns_read_file_group() { // Second MERGE: update 'value' — reads from file group (base + name-partial) register_source( &sql_context, - "CREATE TABLE datafusion.public.src_dc2 (id INT, value INT) AS VALUES (1, 100), (2, 200)", + "CREATE TEMPORARY TABLE paimon.test_db.src_dc2 AS SELECT * FROM (VALUES (1, 100), (2, 200)) AS t(id, value)", ) .await; sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_dc2 s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_dc2 s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET value = s.value", ) .await @@ -1100,14 +1099,14 @@ async fn test_merge_insert_reordered_columns() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_reorder (id INT, name VARCHAR, value INT) AS VALUES (2, 'bob', 20), (1, 'ALICE', 11)", + "CREATE TEMPORARY TABLE paimon.test_db.src_reorder AS SELECT * FROM (VALUES (2, 'bob', 20), (1, 'ALICE', 11)) AS t(id, name, value)", ) .await; // INSERT columns in reversed order: (value, name, id) sql_context .sql( - "MERGE INTO paimon.test_db.target t USING datafusion.public.src_reorder s ON t.id = s.id \ + "MERGE INTO paimon.test_db.target t USING paimon.test_db.src_reorder s ON t.id = s.id \ WHEN NOT MATCHED THEN INSERT (value, name, id) VALUES (s.value, s.name, s.id)", ) .await @@ -1164,14 +1163,14 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_pt_reorder (id INT, name VARCHAR, dt VARCHAR) AS VALUES (2, 'bob', '2024-02-01'), (1, 'ALICE', '2024-01-01')", + "CREATE TEMPORARY TABLE paimon.test_db.src_pt_reorder AS SELECT * FROM (VALUES (2, 'bob', '2024-02-01'), (1, 'ALICE', '2024-01-01')) AS t(id, name, dt)", ) .await; // INSERT with columns in different order than table schema: (name, id, dt) vs table (dt, id, name) sql_context .sql( - "MERGE INTO paimon.test_db.part_tbl t USING datafusion.public.src_pt_reorder s ON t.id = s.id \ + "MERGE INTO paimon.test_db.part_tbl t USING paimon.test_db.src_pt_reorder s ON t.id = s.id \ WHEN NOT MATCHED THEN INSERT (name, id, dt) VALUES (s.name, s.id, s.dt)", ) .await @@ -1245,7 +1244,7 @@ async fn test_rejects_table_with_primary_keys() { register_source( &sql_context, - "CREATE TABLE datafusion.public.src_pk (id INT, name VARCHAR) AS VALUES (1, 'ALICE')", + "CREATE TEMPORARY TABLE paimon.test_db.src_pk AS SELECT * FROM (VALUES (1, 'ALICE')) AS t(id, name)", ) .await; @@ -1253,7 +1252,7 @@ async fn test_rejects_table_with_primary_keys() { assert_merge_error( &sql_context, - "MERGE INTO paimon.test_db.pk_target t USING datafusion.public.src_pk s ON t.id = s.id \ + "MERGE INTO paimon.test_db.pk_target t USING paimon.test_db.src_pk s ON t.id = s.id \ WHEN MATCHED THEN UPDATE SET name = s.name", "does not support primary keys", ) diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index 9318a9a0..396e9514 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -23,10 +23,10 @@ use datafusion::catalog::CatalogProvider; use datafusion::datasource::TableProvider; use datafusion::logical_expr::{col, lit, TableProviderFilterPushDown}; use datafusion::physical_plan::{displayable, ExecutionPlan}; -use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::prelude::SessionConfig; use paimon::catalog::Identifier; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; -use paimon_datafusion::{PaimonCatalogProvider, PaimonRelationPlanner, PaimonTableProvider}; +use paimon_datafusion::{PaimonCatalogProvider, PaimonTableProvider, SQLContext}; fn get_test_warehouse() -> String { std::env::var("PAIMON_TEST_WAREHOUSE").unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string()) @@ -39,12 +39,13 @@ fn create_catalog() -> FileSystemCatalog { FileSystemCatalog::new(options).expect("Failed to create catalog") } -async fn create_context(table_name: &str) -> SessionContext { - let provider = create_provider(table_name).await; - let ctx = SessionContext::new(); - ctx.register_table(table_name, Arc::new(provider)) - .expect("Failed to register table"); - +async fn create_context() -> SQLContext { + let catalog = create_catalog(); + let catalog: Arc = Arc::new(catalog); + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog) + .await + .expect("Failed to register catalog"); ctx } @@ -75,7 +76,8 @@ async fn create_provider_with_options( } async fn read_rows(table_name: &str) -> Vec<(i32, String)> { - let batches = collect_query(table_name, &format!("SELECT id, name FROM {table_name}")) + let sql = format!("SELECT id, name FROM paimon.default.{table_name}"); + let batches = collect_query(&sql) .await .expect("Failed to collect query result"); @@ -90,19 +92,14 @@ async fn read_rows(table_name: &str) -> Vec<(i32, String)> { } async fn collect_query( - table_name: &str, sql: &str, ) -> datafusion::error::Result> { - let ctx = create_context(table_name).await; - + let ctx = create_context().await; ctx.sql(sql).await?.collect().await } -async fn create_physical_plan( - table_name: &str, - sql: &str, -) -> datafusion::error::Result> { - let ctx = create_context(table_name).await; +async fn create_physical_plan(sql: &str) -> datafusion::error::Result> { + let ctx = create_context().await; ctx.sql(sql).await?.create_physical_plan().await } @@ -176,7 +173,7 @@ async fn test_read_primary_key_table_via_datafusion() { #[tokio::test] async fn test_projection_via_datafusion() { - let batches = collect_query("simple_log_table", "SELECT id FROM simple_log_table") + let batches = collect_query("SELECT id FROM paimon.default.simple_log_table") .await .expect("Subset projection should succeed"); @@ -242,7 +239,7 @@ async fn test_scan_partition_count_respects_session_config() { // With generous target_partitions, the plan should expose more than one partition. let config = SessionConfig::new().with_target_partitions(8); - let ctx = SessionContext::new_with_config(config); + let ctx = datafusion::prelude::SessionContext::new_with_config(config); let state = ctx.state(); let plan = provider .scan(&state, None, &[], None) @@ -257,7 +254,7 @@ async fn test_scan_partition_count_respects_session_config() { // With target_partitions=1, all splits must be coalesced into a single partition let config_single = SessionConfig::new().with_target_partitions(1); - let ctx_single = SessionContext::new_with_config(config_single); + let ctx_single = datafusion::prelude::SessionContext::new_with_config(config_single); let state_single = ctx_single.state(); let plan_single = provider .scan(&state_single, None, &[], None) @@ -277,8 +274,7 @@ async fn test_scan_partition_count_respects_session_config() { #[tokio::test] async fn test_partition_filter_query_via_datafusion() { let batches = collect_query( - "partitioned_log_table", - "SELECT id, name FROM partitioned_log_table WHERE dt = '2024-01-01'", + "SELECT id, name FROM paimon.default.partitioned_log_table WHERE dt = '2024-01-01'", ) .await .expect("Partition filter query should succeed"); @@ -294,8 +290,7 @@ async fn test_partition_filter_query_via_datafusion() { #[tokio::test] async fn test_multi_partition_filter_query_via_datafusion() { let batches = collect_query( - "multi_partitioned_log_table", - "SELECT id, name FROM multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr = 10", + "SELECT id, name FROM paimon.default.multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr = 10", ) .await .expect("Multi-partition filter query should succeed"); @@ -311,8 +306,7 @@ async fn test_multi_partition_filter_query_via_datafusion() { #[tokio::test] async fn test_mixed_and_filter_keeps_residual_datafusion_filter() { let batches = collect_query( - "partitioned_log_table", - "SELECT id, name FROM partitioned_log_table WHERE dt = '2024-01-01' AND id > 1", + "SELECT id, name FROM paimon.default.partitioned_log_table WHERE dt = '2024-01-01' AND id > 1", ) .await .expect("Mixed filter query should succeed"); @@ -324,8 +318,8 @@ async fn test_mixed_and_filter_keeps_residual_datafusion_filter() { #[tokio::test] async fn test_partially_translated_filter_keeps_partition_pruning_and_correctness() { - let sql = "SELECT id, name FROM multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr + 1 > 20 LIMIT 1"; - let plan = create_physical_plan("multi_partitioned_log_table", sql) + let sql = "SELECT id, name FROM paimon.default.multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr + 1 > 20 LIMIT 1"; + let plan = create_physical_plan(sql) .await .expect("Physical plan creation should succeed"); let plan_text = format_physical_plan(&plan); @@ -346,7 +340,7 @@ async fn test_partially_translated_filter_keeps_partition_pruning_and_correctnes "Partially translated filters should not revive the removed fetch contract, plan:\n{plan_text}" ); - let batches = collect_query("multi_partitioned_log_table", sql) + let batches = collect_query(sql) .await .expect("Partially translated filter + LIMIT query should succeed"); let rows = extract_id_name_rows(&batches); @@ -360,12 +354,9 @@ async fn test_partially_translated_filter_keeps_partition_pruning_and_correctnes #[tokio::test] async fn test_limit_pushdown_on_data_evolution_table_returns_merged_rows() { - let batches = collect_query( - "data_evolution_table", - "SELECT id, name FROM data_evolution_table LIMIT 3", - ) - .await - .expect("Limit query on data evolution table should succeed"); + let batches = collect_query("SELECT id, name FROM paimon.default.data_evolution_table LIMIT 3") + .await + .expect("Limit query on data evolution table should succeed"); let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!( @@ -389,8 +380,8 @@ async fn test_limit_pushdown_on_data_evolution_table_returns_merged_rows() { #[tokio::test] async fn test_limit_pushdown_marks_safe_scan_limit_hint_and_keeps_correctness() { - let sql = "SELECT id, name FROM simple_log_table LIMIT 2"; - let plan = create_physical_plan("simple_log_table", sql) + let sql = "SELECT id, name FROM paimon.default.simple_log_table LIMIT 2"; + let plan = create_physical_plan(sql) .await .expect("Physical plan creation should succeed"); let plan_text = format_physical_plan(&plan); @@ -403,7 +394,7 @@ async fn test_limit_pushdown_marks_safe_scan_limit_hint_and_keeps_correctness() "Safe LIMIT query should push a scan limit hint into PaimonTableScan, plan:\n{plan_text}" ); - let batches = collect_query("simple_log_table", sql) + let batches = collect_query(sql) .await .expect("LIMIT query should succeed"); let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); @@ -412,8 +403,8 @@ async fn test_limit_pushdown_marks_safe_scan_limit_hint_and_keeps_correctness() #[tokio::test] async fn test_offset_limit_pushdown_keeps_correctness_without_fetch_contract() { - let sql = "SELECT id, name FROM partitioned_log_table OFFSET 1 LIMIT 1"; - let plan = create_physical_plan("partitioned_log_table", sql) + let sql = "SELECT id, name FROM paimon.default.partitioned_log_table OFFSET 1 LIMIT 1"; + let plan = create_physical_plan(sql) .await .expect("Physical plan creation should succeed"); let plan_text = format_physical_plan(&plan); @@ -428,7 +419,7 @@ async fn test_offset_limit_pushdown_keeps_correctness_without_fetch_contract() { "OFFSET + LIMIT should not rely on the removed DataFusion fetch contract in PaimonTableScan, plan:\n{plan_text}" ); - let batches = collect_query("partitioned_log_table", sql) + let batches = collect_query(sql) .await .expect("OFFSET + LIMIT query should succeed"); @@ -441,8 +432,8 @@ async fn test_offset_limit_pushdown_keeps_correctness_without_fetch_contract() { #[tokio::test] async fn test_inexact_filter_limit_keeps_correctness_without_fetch_contract() { - let sql = "SELECT id, name FROM partitioned_log_table WHERE id > 1 LIMIT 1"; - let plan = create_physical_plan("partitioned_log_table", sql) + let sql = "SELECT id, name FROM paimon.default.partitioned_log_table WHERE id > 1 LIMIT 1"; + let plan = create_physical_plan(sql) .await .expect("Physical plan creation should succeed"); let plan_text = format_physical_plan(&plan); @@ -457,7 +448,7 @@ async fn test_inexact_filter_limit_keeps_correctness_without_fetch_contract() { "Inexact filter queries should not revive the removed fetch contract, plan:\n{plan_text}" ); - let batches = collect_query("partitioned_log_table", sql) + let batches = collect_query(sql) .await .expect("Inexact filter + LIMIT query should succeed"); let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); @@ -469,8 +460,8 @@ async fn test_inexact_filter_limit_keeps_correctness_without_fetch_contract() { #[tokio::test] async fn test_residual_filter_limit_keeps_connector_limit_and_correctness() { - let sql = "SELECT id, name FROM simple_log_table WHERE id + 1 > 3 LIMIT 1"; - let plan = create_physical_plan("simple_log_table", sql) + let sql = "SELECT id, name FROM paimon.default.simple_log_table WHERE id + 1 > 3 LIMIT 1"; + let plan = create_physical_plan(sql) .await .expect("Physical plan creation should succeed"); let plan_text = format_physical_plan(&plan); @@ -491,7 +482,7 @@ async fn test_residual_filter_limit_keeps_connector_limit_and_correctness() { "Residual filter queries should not push a scan limit hint when residual filters stay above the scan, plan:\n{plan_text}" ); - let batches = collect_query("simple_log_table", sql) + let batches = collect_query(sql) .await .expect("Residual filter + LIMIT query should succeed"); let rows = extract_id_name_rows(&batches); @@ -507,10 +498,11 @@ async fn test_residual_filter_limit_keeps_connector_limit_and_correctness() { #[tokio::test] async fn test_query_via_catalog_provider() { let catalog = create_catalog(); - let provider = PaimonCatalogProvider::new(Arc::new(catalog)); - - let ctx = SessionContext::new(); - ctx.register_catalog("paimon", Arc::new(provider)); + let catalog: Arc = Arc::new(catalog); + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog) + .await + .expect("Failed to register catalog"); let df = ctx .sql("SELECT id, name FROM paimon.default.simple_log_table") @@ -535,18 +527,14 @@ async fn test_missing_database_returns_no_schema() { // ======================= Time Travel Tests ======================= -/// Helper: create a SessionContext with catalog + relation planner for time travel. -/// Uses Databricks dialect to enable `VERSION AS OF` and `TIMESTAMP AS OF` syntax. -async fn create_time_travel_context() -> SessionContext { +/// Helper: create a SQLContext with catalog + relation planner for time travel. +async fn create_time_travel_context() -> SQLContext { let catalog = create_catalog(); - let config = SessionConfig::new().set_str("datafusion.sql_parser.dialect", "Databricks"); - let ctx = SessionContext::new_with_config(config); - ctx.register_catalog( - "paimon", - Arc::new(PaimonCatalogProvider::new(Arc::new(catalog))), - ); - ctx.register_relation_planner(Arc::new(PaimonRelationPlanner::new())) - .expect("Failed to register relation planner"); + let catalog: Arc = Arc::new(catalog); + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog) + .await + .expect("Failed to register catalog"); ctx } @@ -596,21 +584,11 @@ async fn test_time_travel_by_snapshot_id() { #[tokio::test] async fn test_time_travel_by_tag_name() { - // Tag-based time travel uses `scan.version` option directly since - // `VERSION AS OF` in SQL only accepts numeric values. - let provider = create_provider_with_options( - "time_travel_table", - HashMap::from([("scan.version".to_string(), "snapshot1".to_string())]), - ) - .await; - - let ctx = SessionContext::new(); - ctx.register_table("time_travel_table", Arc::new(provider)) - .expect("Failed to register table"); + let ctx = create_time_travel_context().await; // Tag 'snapshot1' points to snapshot 1: should contain only (alice, bob) let batches = ctx - .sql("SELECT id, name FROM time_travel_table") + .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 'snapshot1'") .await .expect("tag time travel query should parse") .collect() @@ -626,18 +604,8 @@ async fn test_time_travel_by_tag_name() { ); // Tag 'snapshot2' points to snapshot 2: should contain all rows - let provider2 = create_provider_with_options( - "time_travel_table", - HashMap::from([("scan.version".to_string(), "snapshot2".to_string())]), - ) - .await; - - let ctx2 = SessionContext::new(); - ctx2.register_table("time_travel_table", Arc::new(provider2)) - .expect("Failed to register table"); - - let batches = ctx2 - .sql("SELECT id, name FROM time_travel_table") + let batches = ctx + .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 'snapshot2'") .await .expect("tag time travel query should parse") .collect() @@ -660,23 +628,25 @@ async fn test_time_travel_by_tag_name() { #[tokio::test] async fn test_time_travel_conflicting_selectors_fail() { + // When both scan.version and scan.timestamp-millis are set on the same + // provider, Paimon rejects the combination at scan time. let provider = create_provider_with_options( "time_travel_table", - HashMap::from([("scan.timestamp-millis".to_string(), "1234".to_string())]), + HashMap::from([ + ("scan.version".to_string(), "1".to_string()), + ("scan.timestamp-millis".to_string(), "1234".to_string()), + ]), ) .await; - let config = SessionConfig::new().set_str("datafusion.sql_parser.dialect", "Databricks"); - let ctx = SessionContext::new_with_config(config); - ctx.register_table("time_travel_table", Arc::new(provider)) - .expect("Failed to register table"); - ctx.register_relation_planner(Arc::new(PaimonRelationPlanner::new())) - .expect("Failed to register relation planner"); + let ctx = create_context().await; + ctx.register_temp_table("paimon.default.time_travel_table", Arc::new(provider)) + .expect("Failed to register temp table"); let err = ctx - .sql("SELECT id, name FROM time_travel_table VERSION AS OF 2") + .sql("SELECT id, name FROM paimon.default.time_travel_table") .await - .expect("time travel query should parse") + .expect("query should parse") .collect() .await .expect_err("conflicting time-travel selectors should fail"); @@ -700,12 +670,12 @@ async fn test_time_travel_invalid_version_fails() { ) .await; - let ctx = SessionContext::new(); - ctx.register_table("time_travel_table", Arc::new(provider)) - .expect("Failed to register table"); + let ctx = create_context().await; + ctx.register_temp_table("paimon.default.time_travel_table", Arc::new(provider)) + .expect("Failed to register temp table"); let err = ctx - .sql("SELECT id, name FROM time_travel_table") + .sql("SELECT id, name FROM paimon.default.time_travel_table") .await .expect("query should parse") .collect() @@ -724,12 +694,10 @@ async fn test_time_travel_invalid_version_fails() { /// Without the fix, `active_file_indices` would be empty and rows would be silently lost. #[tokio::test] async fn test_data_evolution_drop_column_null_fill() { - let batches = collect_query( - "data_evolution_drop_column", - "SELECT id, name, extra FROM data_evolution_drop_column", - ) - .await - .expect("data_evolution_drop_column query should succeed"); + let batches = + collect_query("SELECT id, name, extra FROM paimon.default.data_evolution_drop_column") + .await + .expect("data_evolution_drop_column query should succeed"); let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!( @@ -779,8 +747,7 @@ async fn test_data_evolution_drop_column_null_fill() { #[tokio::test] async fn test_read_complex_type_table_via_datafusion() { let batches = collect_query( - "complex_type_table", - "SELECT id, int_array, string_map, row_field FROM complex_type_table ORDER BY id", + "SELECT id, int_array, string_map, row_field FROM paimon.default.complex_type_table ORDER BY id", ) .await .expect("Complex type query should succeed"); @@ -854,10 +821,10 @@ async fn test_read_complex_type_table_via_datafusion() { async fn test_select_row_id_from_data_evolution_table() { use datafusion::arrow::array::Int64Array; - let ctx = create_context("data_evolution_table").await; + let ctx = create_context().await; let batches = ctx - .sql(r#"SELECT "_ROW_ID", id, name FROM data_evolution_table"#) + .sql(r#"SELECT "_ROW_ID", id, name FROM paimon.default.data_evolution_table"#) .await .expect("SQL should parse") .collect() @@ -890,10 +857,10 @@ async fn test_select_row_id_from_data_evolution_table() { async fn test_filter_row_id_from_data_evolution_table() { use datafusion::arrow::array::Int64Array; - let ctx = create_context("data_evolution_table").await; + let ctx = create_context().await; let all_batches = ctx - .sql(r#"SELECT "_ROW_ID" FROM data_evolution_table"#) + .sql(r#"SELECT "_ROW_ID" FROM paimon.default.data_evolution_table"#) .await .expect("SQL") .collect() @@ -902,7 +869,7 @@ async fn test_filter_row_id_from_data_evolution_table() { let all_count: usize = all_batches.iter().map(|b| b.num_rows()).sum(); let filtered_batches = ctx - .sql(r#"SELECT "_ROW_ID", id FROM data_evolution_table WHERE "_ROW_ID" = 0"#) + .sql(r#"SELECT "_ROW_ID", id FROM paimon.default.data_evolution_table WHERE "_ROW_ID" = 0"#) .await .expect("SQL") .collect() @@ -931,9 +898,8 @@ mod fulltext_tests { use std::sync::Arc; use datafusion::arrow::array::{Int32Array, StringArray}; - use datafusion::prelude::SessionContext; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; - use paimon_datafusion::{register_full_text_search, PaimonCatalogProvider}; + use paimon_datafusion::{register_full_text_search, SQLContext}; /// Extract the bundled tar.gz into a temp dir and return (tempdir, warehouse_path). fn extract_test_warehouse() -> (tempfile::TempDir, String) { @@ -953,19 +919,18 @@ mod fulltext_tests { (tmp, warehouse) } - async fn create_fulltext_context() -> (SessionContext, tempfile::TempDir) { + async fn create_fulltext_context() -> (SQLContext, 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 ctx = SessionContext::new(); - ctx.register_catalog( - "paimon", - Arc::new(PaimonCatalogProvider::new(Arc::clone(&catalog))), - ); - register_full_text_search(&ctx, catalog, "default"); + 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) } @@ -1053,9 +1018,8 @@ mod vector_search_tests { use std::sync::Arc; use datafusion::arrow::array::Int32Array; - use datafusion::prelude::SessionContext; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; - use paimon_datafusion::{register_vector_search, PaimonCatalogProvider}; + use paimon_datafusion::{register_vector_search, SQLContext}; fn extract_test_warehouse() -> (tempfile::TempDir, String) { let archive_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) @@ -1074,19 +1038,18 @@ mod vector_search_tests { (tmp, warehouse) } - async fn create_vector_search_context() -> (SessionContext, tempfile::TempDir) { + async fn create_vector_search_context() -> (SQLContext, 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 ctx = SessionContext::new(); - ctx.register_catalog( - "paimon", - Arc::new(PaimonCatalogProvider::new(Arc::clone(&catalog))), - ); - register_vector_search(&ctx, catalog, "default"); + 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) } diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 7d274222..5a866737 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use datafusion::catalog::CatalogProvider; +use datafusion::datasource::MemTable; use paimon::catalog::Identifier; use paimon::spec::{ArrayType, BlobType, DataType, IntType, MapType, VarCharType}; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; @@ -749,3 +750,506 @@ async fn test_one_part_table_name_uses_current_database() { "SELECT with 1-part name should resolve correctly" ); } + +// ======================= TEMP TABLE ======================= + +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField}; + +#[tokio::test] +async fn test_register_temp_table_fully_qualified() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + // Fully qualified: catalog.database.table + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("paimon.my_db.my_temp", Arc::new(mem_table)) + .unwrap(); + + // Query the temp table via SQL + let batches = ctx + .sql("SELECT * FROM paimon.my_db.my_temp") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); +} + +#[tokio::test] +async fn test_register_temp_table_database_qualified() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec![ + Some("alice"), + Some("bob"), + Some("charlie"), + Some("dave"), + ])), + ], + ) + .unwrap(); + + // Database-qualified: database.table (uses current catalog) + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("my_db.users", Arc::new(mem_table)) + .unwrap(); + + let batches = ctx + .sql("SELECT id, name FROM paimon.my_db.users WHERE id > 2") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 2); +} + +#[tokio::test] +async fn test_register_temp_table_bare() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + // Create a database and set it as current database + ctx.sql("CREATE DATABASE paimon.my_db").await.unwrap(); + ctx.set_current_database("my_db").await.unwrap(); + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + // Bare: just table name (uses current catalog + current database) + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("my_temp", Arc::new(mem_table)) + .unwrap(); + + // Query via paimon.my_db.my_temp + let batches = ctx + .sql("SELECT * FROM paimon.my_db.my_temp") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); +} + +#[tokio::test] +async fn test_register_temp_table_unknown_catalog() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + + let mem_table = MemTable::try_new(schema, vec![vec![]]).unwrap(); + let result = ctx.register_temp_table("nonexistent.my_db.t", Arc::new(mem_table)); + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("Unknown catalog")); +} + +#[tokio::test] +async fn test_deregister_temp_table() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + + let mem_table = MemTable::try_new(schema.clone(), vec![vec![batch]]).unwrap(); + ctx.register_temp_table("paimon.my_db.my_temp", Arc::new(mem_table)) + .unwrap(); + + // Deregister with flexible name + ctx.deregister_temp_table("paimon.my_db.my_temp").unwrap(); + + // Query should fail + let result = ctx.sql("SELECT * FROM paimon.my_db.my_temp").await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_multiple_temp_tables_in_same_database() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema1 = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch1 = RecordBatch::try_new( + schema1.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .unwrap(); + + let schema2 = Arc::new(Schema::new(vec![ArrowField::new( + "value", + ArrowDataType::Int32, + false, + )])); + let batch2 = RecordBatch::try_new( + schema2.clone(), + vec![Arc::new(Int32Array::from(vec![10, 20, 30]))], + ) + .unwrap(); + + let mem_table = MemTable::try_new(schema1, vec![vec![batch1]]).unwrap(); + ctx.register_temp_table("my_db.t1", Arc::new(mem_table)) + .unwrap(); + let mem_table = MemTable::try_new(schema2, vec![vec![batch2]]).unwrap(); + ctx.register_temp_table("my_db.t2", Arc::new(mem_table)) + .unwrap(); + + // Both should be queryable + let rows1 = ctx + .sql("SELECT * FROM paimon.my_db.t1") + .await + .unwrap() + .collect() + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(); + assert_eq!(rows1, 2); + + let rows2 = ctx + .sql("SELECT * FROM paimon.my_db.t2") + .await + .unwrap() + .collect() + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(); + assert_eq!(rows2, 3); +} + +use datafusion::arrow::array::StringArray; +use datafusion::arrow::datatypes::Schema; +use datafusion::arrow::record_batch::RecordBatch; + +#[tokio::test] +async fn test_create_temporary_table_as_select() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + // Create a temporary table via SQL + ctx.sql("CREATE TEMPORARY TABLE paimon.my_db.source AS SELECT * FROM (VALUES (1, 'alice'), (2, 'bob')) AS t(id, name)") + .await + .unwrap(); + + // Query the temporary table + let batches = ctx + .sql("SELECT * FROM paimon.my_db.source ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); +} + +#[tokio::test] +async fn test_drop_temporary_table() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + // Create a temporary table + ctx.sql("CREATE TEMPORARY TABLE paimon.my_db.source AS SELECT * FROM (VALUES (1, 'alice'), (2, 'bob')) AS t(id, name)") + .await + .unwrap(); + + // Verify it exists + let batches = ctx + .sql("SELECT * FROM paimon.my_db.source ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + + // Drop it + ctx.sql("DROP TEMPORARY TABLE paimon.my_db.source") + .await + .unwrap(); + + // Verify it no longer exists + let result = ctx.sql("SELECT * FROM paimon.my_db.source").await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_drop_temporary_table_if_exists() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + // DROP TEMPORARY TABLE on non-existent table with IF EXISTS should succeed + ctx.sql("DROP TEMPORARY TABLE IF EXISTS paimon.my_db.nonexistent") + .await + .unwrap(); + + // Without IF EXISTS, it should fail + let result = ctx + .sql("DROP TEMPORARY TABLE paimon.my_db.nonexistent") + .await; + assert!(result.is_err()); +} + +// ======================= TEMP VIEW ======================= + +#[tokio::test] +async fn test_create_temporary_view_fully_qualified() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])), + ], + ) + .unwrap(); + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("paimon.my_db.users", Arc::new(mem_table)) + .unwrap(); + + ctx.sql("CREATE TEMPORARY VIEW paimon.my_db.my_view AS SELECT * FROM paimon.my_db.users WHERE id > 0") + .await + .unwrap(); + + let batches = ctx + .sql("SELECT * FROM paimon.my_db.my_view") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 2); +} + +#[tokio::test] +async fn test_create_temporary_view_database_qualified() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "value", + ArrowDataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![10, 20, 30]))], + ) + .unwrap(); + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("paimon.my_db.data", Arc::new(mem_table)) + .unwrap(); + + ctx.sql("CREATE TEMPORARY VIEW my_db.summary AS SELECT value FROM paimon.my_db.data WHERE value > 5") + .await + .unwrap(); + + let batches = ctx + .sql("SELECT value FROM paimon.my_db.summary WHERE value > 15") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 2); +} + +#[tokio::test] +async fn test_create_temporary_view_bare() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + ctx.sql("CREATE DATABASE paimon.my_db").await.unwrap(); + ctx.set_current_database("my_db").await.unwrap(); + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![100, 200]))], + ) + .unwrap(); + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("my_db.source", Arc::new(mem_table)) + .unwrap(); + + ctx.sql("CREATE TEMPORARY VIEW my_view AS SELECT id FROM paimon.my_db.source") + .await + .unwrap(); + + let batches = ctx + .sql("SELECT * FROM paimon.my_db.my_view") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 2); +} + +#[tokio::test] +async fn test_drop_temporary_view() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + let mem_table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_temp_table("paimon.my_db.source", Arc::new(mem_table)) + .unwrap(); + + ctx.sql("CREATE TEMPORARY VIEW paimon.my_db.my_view AS SELECT * FROM paimon.my_db.source") + .await + .unwrap(); + + // Drop via SQL + ctx.sql("DROP TEMPORARY VIEW paimon.my_db.my_view") + .await + .unwrap(); + + let result = ctx.sql("SELECT * FROM paimon.my_db.my_view").await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_multiple_temporary_views_in_same_database() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog.clone()).await; + + let schema1 = Arc::new(Schema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int32, + false, + )])); + let batch1 = RecordBatch::try_new( + schema1.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .unwrap(); + let mem_table = MemTable::try_new(schema1, vec![vec![batch1]]).unwrap(); + ctx.register_temp_table("paimon.my_db.t1", Arc::new(mem_table)) + .unwrap(); + + let schema2 = Arc::new(Schema::new(vec![ArrowField::new( + "name", + ArrowDataType::Utf8, + true, + )])); + let batch2 = RecordBatch::try_new( + schema2.clone(), + vec![Arc::new(StringArray::from(vec![ + Some("x"), + Some("y"), + Some("z"), + ]))], + ) + .unwrap(); + let mem_table = MemTable::try_new(schema2, vec![vec![batch2]]).unwrap(); + ctx.register_temp_table("paimon.my_db.t2", Arc::new(mem_table)) + .unwrap(); + + ctx.sql("CREATE TEMPORARY VIEW my_db.v1 AS SELECT id FROM paimon.my_db.t1") + .await + .unwrap(); + ctx.sql("CREATE TEMPORARY VIEW my_db.v2 AS SELECT name FROM paimon.my_db.t2") + .await + .unwrap(); + + let rows1 = ctx + .sql("SELECT * FROM paimon.my_db.v1") + .await + .unwrap() + .collect() + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(); + assert_eq!(rows1, 2); + + let rows2 = ctx + .sql("SELECT * FROM paimon.my_db.v2") + .await + .unwrap() + .collect() + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(); + assert_eq!(rows2, 3); +} diff --git a/crates/integrations/datafusion/tests/system_tables.rs b/crates/integrations/datafusion/tests/system_tables.rs index 0d0afa40..d39e14fa 100644 --- a/crates/integrations/datafusion/tests/system_tables.rs +++ b/crates/integrations/datafusion/tests/system_tables.rs @@ -22,10 +22,9 @@ use std::sync::Arc; use datafusion::arrow::array::{Array, Int64Array, StringArray}; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::prelude::SessionContext; use paimon::catalog::Identifier; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; -use paimon_datafusion::PaimonCatalogProvider; +use paimon_datafusion::SQLContext; const FIXTURE_TABLE: &str = "test_tantivy_fulltext"; @@ -46,22 +45,21 @@ fn extract_test_warehouse() -> (tempfile::TempDir, String) { (tmp, warehouse) } -async fn create_context() -> (SessionContext, Arc, tempfile::TempDir) { +async fn create_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 ctx = SessionContext::new(); - ctx.register_catalog( - "paimon", - Arc::new(PaimonCatalogProvider::new(Arc::clone(&catalog))), - ); + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", catalog.clone()) + .await + .expect("Failed to register catalog"); (ctx, catalog, tmp) } -async fn run_sql(ctx: &SessionContext, sql: &str) -> Vec { +async fn run_sql(ctx: &SQLContext, sql: &str) -> Vec { ctx.sql(sql) .await .unwrap_or_else(|e| panic!("Failed to plan `{sql}`: {e}")) diff --git a/crates/integrations/datafusion/tests/update_tests.rs b/crates/integrations/datafusion/tests/update_tests.rs index 839bc26a..676e41ca 100644 --- a/crates/integrations/datafusion/tests/update_tests.rs +++ b/crates/integrations/datafusion/tests/update_tests.rs @@ -23,7 +23,7 @@ mod common; use paimon_datafusion::SQLContext; -use common::{create_sql_context, create_test_env, ctx_exec, dml_count, exec, query_int_str_int}; +use common::{create_sql_context, create_test_env, dml_count, exec, query_int_str_int}; // ======================= Helpers ======================= @@ -397,15 +397,15 @@ async fn test_update_range_condition() { async fn test_update_in_subquery() { let (_tmp, sql_context) = setup().await; - ctx_exec( + exec( &sql_context, - "CREATE TABLE datafusion.public.src (id INT) AS VALUES (1), (3)", + "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1), (3)) AS t(id)", ) .await; exec( &sql_context, - "UPDATE paimon.test_db.t SET name = 'sub' WHERE id IN (SELECT id FROM datafusion.public.src)", + "UPDATE paimon.test_db.t SET name = 'sub' WHERE id IN (SELECT id FROM paimon.test_db.src)", ) .await; diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index a23c05a9..e4cca707 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -48,7 +48,7 @@ theme: nav: - Home: index.md - Getting Started: getting-started.md - - DataFusion Integration: datafusion.md + - SQL Integration: sql.md - Go Integration: go-binding.md - Architecture: architecture.md - Contributing: contributing.md diff --git a/docs/src/datafusion.md b/docs/src/sql.md similarity index 80% rename from docs/src/datafusion.md rename to docs/src/sql.md index 38675c15..5fb03773 100644 --- a/docs/src/datafusion.md +++ b/docs/src/sql.md @@ -17,7 +17,7 @@ specific language governing permissions and limitations under the License. --> -# DataFusion Integration +# SQL Integration [Apache DataFusion](https://datafusion.apache.org/) is a fast, extensible query engine for building data-centric systems in Rust. The `paimon-datafusion` crate provides a full SQL integration that lets you create, query, and modify Paimon tables. @@ -152,6 +152,52 @@ CREATE TABLE paimon.my_db.complex_types ( ```sql DROP TABLE paimon.my_db.users; +DROP TABLE IF EXISTS paimon.my_db.users; +``` + +### CREATE TEMPORARY TABLE + +Create an in-memory temporary table from a query result. Temporary tables exist only for the lifetime of the `SQLContext` instance and are automatically cleaned up when the context is dropped. + +```sql +-- Without column types (types inferred from the query) +CREATE TEMPORARY TABLE paimon.my_db.source AS SELECT * FROM (VALUES (1, 'alice'), (2, 'bob')) AS t(id, name); + +-- With explicit column types (recommended when integer precision matters) +CREATE TEMPORARY TABLE paimon.my_db.source (id INT, name STRING) AS SELECT * FROM (VALUES (1, 'alice'), (2, 'bob')) AS t(id, name); +``` + +`IF NOT EXISTS` is supported — if the table already exists, the statement is silently ignored: + +```sql +CREATE TEMPORARY TABLE IF NOT EXISTS paimon.my_db.source AS SELECT 1; +``` + +> **Note:** When using `VALUES` without explicit column types, DataFusion infers integer literals as `Int64`. If the temporary table will be used as a source in `MERGE INTO` against a Paimon table with `Int32` columns, specify the column types explicitly to avoid type mismatch errors. + +### CREATE TEMPORARY VIEW + +Create a temporary view from a query: + +```sql +CREATE TEMPORARY VIEW paimon.my_db.active_users AS SELECT * FROM paimon.my_db.users WHERE id > 0; +``` + +`IF NOT EXISTS` is supported: + +```sql +CREATE TEMPORARY VIEW IF NOT EXISTS paimon.my_db.active_users AS SELECT * FROM paimon.my_db.users WHERE id > 0; +``` + +### DROP TEMPORARY TABLE / DROP TEMPORARY VIEW + +Remove a temporary table or view: + +```sql +DROP TEMPORARY TABLE paimon.my_db.source; +DROP TEMPORARY TABLE IF EXISTS paimon.my_db.source; +DROP TEMPORARY VIEW paimon.my_db.active_users; +DROP TEMPORARY VIEW IF EXISTS paimon.my_db.active_users; ``` ### ALTER TABLE @@ -503,12 +549,10 @@ SELECT * FROM paimon.default.my_table VERSION AS OF 1; ### By Tag Name -`VERSION AS OF` in SQL only accepts numeric values. Tag-based time travel is done via the `scan.version` table option: +Use a quoted tag name with `VERSION AS OF`: -```rust -let table = table.copy_with_options(HashMap::from([ - ("scan.version".to_string(), "my_tag".to_string()), -])); +```sql +SELECT * FROM paimon.default.my_table VERSION AS OF 'my_tag'; ``` Resolution order: first checks if a tag with that name exists, then tries to parse it as a snapshot ID. @@ -523,21 +567,6 @@ SELECT * FROM paimon.default.my_table TIMESTAMP AS OF '2024-01-01 00:00:00'; This finds the latest snapshot whose commit time is less than or equal to the given timestamp. The timestamp is interpreted in the local timezone. -### Enabling Time Travel Syntax - -DataFusion requires the Databricks SQL dialect to parse `VERSION AS OF` and `TIMESTAMP AS OF`: - -```rust -use datafusion::prelude::{SessionConfig, SessionContext}; - -let config = SessionConfig::new() - .set_str("datafusion.sql_parser.dialect", "Databricks"); -let ctx = SessionContext::new_with_config(config); - -ctx.register_catalog("paimon", Arc::new(PaimonCatalogProvider::new(catalog))); -ctx.register_relation_planner(Arc::new(PaimonRelationPlanner::new()))?; -``` - ## Dynamic Options (SET / RESET) Use `SET` to configure session-scoped Paimon dynamic options that apply to subsequent table loads: @@ -560,6 +589,94 @@ SELECT * FROM paimon.my_db.assets; RESET 'paimon.blob-as-descriptor'; ``` +## Temporary Tables + +You can register in-memory temporary tables under any catalog. Temporary tables exist only for the lifetime of the `SQLContext` instance and are automatically cleaned up when the context is dropped. + +The table name accepts flexible references, similar to DataFusion: +- `"my_table"` — uses the current catalog and current database +- `"database.my_table"` — uses the current catalog with the specified database +- `"catalog.database.my_table"` — fully qualified + +### register_mem_table + +The easiest way to register a temporary table from a schema and record batches: + +```rust +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; + +let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("name", ArrowDataType::Utf8, true), +])); +let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["alice", "bob", "carol"])), + ], +)?; + +// Fully qualified +ctx.register_mem_table("paimon.my_db.users", schema.clone(), vec![batch.clone()])?; +let df = ctx.sql("SELECT * FROM paimon.my_db.users WHERE id > 1").await?; +df.show().await?; + +// Database-qualified (uses current catalog) +ctx.register_mem_table("my_db.other_table", schema.clone(), vec![batch.clone()])?; + +// Bare table name (uses current catalog + current database) +ctx.register_mem_table("quick_lookup", schema, vec![batch])?; +``` + +### register_temp_table + +For advanced use cases, `register_temp_table` accepts any `Arc` (including `ViewTable`, custom providers, etc.): + +```rust +use datafusion::datasource::ViewTable; + +// Register a view as a temporary table +let view_table = ViewTable::new(logical_plan, Some(query_sql)); +ctx.register_temp_table("paimon.my_db.my_view", Arc::new(view_table))?; +``` + +### CREATE TEMPORARY TABLE + +You can also create temporary tables directly from SQL. See the [DDL section](#create-temporary-table) for details. + +```sql +CREATE TEMPORARY TABLE paimon.my_db.source (id INT, name STRING) AS SELECT * FROM (VALUES (1, 'alice'), (2, 'bob')) AS t(id, name); +``` + +### CREATE TEMPORARY VIEW + +Create a temporary view directly from SQL. See the [DDL section](#create-temporary-view) for details. + +```sql +CREATE TEMPORARY VIEW paimon.my_db.active_users AS SELECT * FROM paimon.my_db.users WHERE id > 0; +``` + +### Deregister + +Use `deregister_temp_table` to remove a temporary table or view programmatically, or use the `DROP TEMPORARY TABLE` / `DROP TEMPORARY VIEW` SQL statements (see the [DDL section](#drop-temporary-table--drop-temporary-view)): + +```rust +ctx.deregister_temp_table("paimon.my_db.users")?; +``` + +Multiple temporary tables can share the same database — the database is created automatically on first use: + +```rust +ctx.register_mem_table("my_db.table_a", schema_a, vec![batch_a])?; +ctx.register_mem_table("my_db.table_b", schema_b, vec![batch_b])?; + +// Join two temp tables +let df = ctx.sql("SELECT * FROM paimon.my_db.table_a JOIN paimon.my_db.table_b ON a.id = b.id").await?; +``` + ## System Tables Access table metadata via the `$` syntax. From bbdb70b2b3f5e3d3316e444839e0a7976053d23a Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 8 May 2026 15:06:59 +0800 Subject: [PATCH 2/2] Fix comments --- bindings/python/README.md | 29 + bindings/python/tests/test_datafusion.py | 13 + crates/integrations/datafusion/src/catalog.rs | 38 +- .../datafusion/src/sql_context.rs | 583 +++++++++++++----- docs/src/sql.md | 35 +- 5 files changed, 507 insertions(+), 191 deletions(-) diff --git a/bindings/python/README.md b/bindings/python/README.md index 71e160c2..1c908bce 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -21,6 +21,35 @@ This project builds the Rust-powered core for [PyPaimon](https://paimon.apache.org/docs/master/pypaimon/overview/) while also providing DataFusion integration for querying Paimon tables. +## Usage + +```python +import pyarrow as pa +from pypaimon_rust.datafusion import SQLContext + +# Create a SQL context and register a Paimon catalog +ctx = SQLContext() +ctx.register_catalog("paimon", {"warehouse": "/tmp/paimon-warehouse"}) + +# Create a table and insert data +ctx.sql("CREATE SCHEMA paimon.my_db") +ctx.sql("CREATE TABLE paimon.my_db.users (id INT, name STRING, PRIMARY KEY (id))") +ctx.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob')") + +# Query data +batches = ctx.sql("SELECT id, name FROM paimon.my_db.users ORDER BY id") + +# Register a temporary table from a PyArrow RecordBatch +batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"]) +ctx.register_batch("paimon.default.my_temp", batch) +batches = ctx.sql("SELECT * FROM paimon.default.my_temp") + +# Drop it via SQL when no longer needed +ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp") +``` + +For the full SQL reference, see the [SQL Integration docs](https://paimon.apache.org/docs/master/sql/). + ## Setup Install [uv](https://docs.astral.sh/uv/getting-started/installation/): diff --git a/bindings/python/tests/test_datafusion.py b/bindings/python/tests/test_datafusion.py index 95623855..5e4e5e99 100644 --- a/bindings/python/tests/test_datafusion.py +++ b/bindings/python/tests/test_datafusion.py @@ -164,3 +164,16 @@ def test_multi_catalog_temp_table(): ctx.sql("DROP TEMPORARY TABLE cat1.default.t1") ctx.sql("DROP TEMPORARY TABLE cat2.default.t2") + + +def test_register_batch_invalid_catalog(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + + batch = pa.record_batch([[1]], names=["id"]) + try: + ctx.register_batch("unknown_catalog.default.my_temp", batch) + assert False, "Expected an error for unknown catalog" + except Exception as e: + assert "unknown_catalog" in str(e).lower() or "not a paimon" in str(e).lower() or "unknown" in str(e).lower() diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 58e6a2e1..a93201aa 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -199,15 +199,6 @@ impl CatalogProvider for PaimonCatalogProvider { } impl PaimonCatalogProvider { - /// Creates or returns an existing temporary in-memory database for temp tables/views. - fn get_or_create_temp_database(&self, name: &str) -> Arc { - let mut databases = self.temp_tables.write().unwrap_or_else(|e| e.into_inner()); - databases - .entry(name.to_string()) - .or_insert_with(|| Arc::new(MemorySchemaProvider::new())) - .clone() - } - /// Registers a temporary table or view in the specified database. /// Creates the database if it does not exist. /// @@ -219,19 +210,7 @@ impl PaimonCatalogProvider { table_name: &str, table: Arc, ) -> DFResult<()> { - // Check if a temp table with this name already exists - { - let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); - if let Some(mem_db) = databases.get(database) { - if mem_db.table_exist(table_name) { - return Err(plan_datafusion_err!( - "Temporary table '{database}.{table_name}' already exists" - )); - } - } - } - - // Warn if this shadows a real Paimon table + // Warn if this shadows a real Paimon table (outside the lock — not critical) let catalog = Arc::clone(&self.catalog); let db = database.to_string(); let tbl = table_name.to_string(); @@ -251,8 +230,19 @@ impl PaimonCatalogProvider { ); } - let mem_database = self.get_or_create_temp_database(database); - mem_database.register_table(table_name.to_string(), table)?; + // Atomically check-then-register under a single write lock to avoid TOCTOU + let mut databases = self.temp_tables.write().unwrap_or_else(|e| e.into_inner()); + let mem_database = databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(MemorySchemaProvider::new())); + + // register_table returns Ok(Some(old_table)) if the name already existed + let old = mem_database.register_table(table_name.to_string(), table)?; + if old.is_some() { + return Err(plan_datafusion_err!( + "Temporary table '{database}.{table_name}' already exists" + )); + } Ok(()) } diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 3c82ba6a..87f5ccb6 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -291,11 +291,7 @@ impl SQLContext { } else { (sql.to_string(), vec![]) }; - let sql_lower = rewritten_sql.to_lowercase(); - let has_time_travel = - sql_lower.contains("version as of") || sql_lower.contains("timestamp as of"); - - if has_time_travel { + if contains_time_travel_keyword(&rewritten_sql) { // Time-travel queries are not DDL; skip our own parsing and handle directly. return self.handle_time_travel_query(&rewritten_sql).await; } @@ -375,7 +371,21 @@ impl SQLContext { self.ctx.sql(sql).await } Statement::Truncate(truncate) => self.handle_truncate_table(truncate).await, - Statement::CreateView(create_view) => self.handle_create_view(create_view).await, + Statement::CreateView(create_view) => { + if create_view.temporary { + // Temporary views are always handled by us (Paimon catalog temp storage) + self.handle_create_view(create_view).await + } else { + // Non-temporary views: only intercept if the target catalog is Paimon + let view_name = create_view.name.to_string(); + let table_ref: TableReference = view_name.as_str().into(); + if self.is_paimon_catalog_ref(&table_ref) { + self.handle_create_view(create_view).await + } else { + self.ctx.sql(sql).await + } + } + } Statement::Drop { object_type, if_exists, @@ -386,8 +396,15 @@ impl SQLContext { if *temporary { self.handle_drop_temp_table(names, *if_exists) } else if *object_type == ObjectType::Table { - let (catalog, _catalog_name, _) = self.resolve_catalog_and_table(&names[0])?; - self.handle_drop_table(&catalog, names, *if_exists).await + // Only intercept DROP TABLE for Paimon catalogs; fall through for others + let table_ref: TableReference = names[0].to_string().as_str().into(); + if self.is_paimon_catalog_ref(&table_ref) { + let (catalog, _catalog_name, _) = + self.resolve_catalog_and_table(&names[0])?; + self.handle_drop_table(&catalog, names, *if_exists).await + } else { + self.ctx.sql(sql).await + } } else { self.ctx.sql(sql).await } @@ -408,61 +425,96 @@ impl SQLContext { /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). /// /// DataFusion's default SQL parser does not support these clauses, so we: - /// 1. Extract the table name and version/timestamp value via regex - /// 2. Strip the time-travel clause from the SQL - /// 3. Create a `PaimonTableProvider` with the appropriate scan options - /// 4. Register it, execute the stripped SQL, then deregister + /// 1. Extract all table name + version/timestamp pairs (skipping string literals and comments) + /// 2. Strip the time-travel clauses from the SQL + /// 3. For each table, create a `PaimonTableProvider` with the appropriate scan options + /// (merged with session-scoped dynamic options) + /// 4. Register them as UUID-named temp tables, execute the rewritten SQL, then deregister async fn handle_time_travel_query(&self, sql: &str) -> DFResult { use crate::table::PaimonTableProvider; use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION}; - let (table_name, options, clause_range) = if let Some(info) = extract_version_as_of(sql) { - let options = HashMap::from([(SCAN_VERSION_OPTION.to_string(), info.version)]); - (info.table_name, options, info.clause_range) - } else if let Some(info) = extract_timestamp_as_of(sql) { - let millis = Self::parse_timestamp_to_millis(&info.timestamp)?; - let options = - HashMap::from([(SCAN_TIMESTAMP_MILLIS_OPTION.to_string(), millis.to_string())]); - (info.table_name, options, info.clause_range) - } else { + let mut tracker = crate::merge_into::TempTableTracker::new(self); + + let version_clauses = extract_all_version_as_of(sql); + let timestamp_clauses = extract_all_timestamp_as_of(sql); + + if version_clauses.is_empty() && timestamp_clauses.is_empty() { return Err(DataFusionError::Plan( "Failed to parse time-travel clause in SQL".to_string(), )); - }; + } - // Resolve the table from our catalog and create a provider with scan options - let table_ref: datafusion::common::TableReference = table_name.as_str().into(); - let (catalog, _catalog_name, identifier) = self.resolve_table_name_from_ref(&table_ref)?; + // Collect all replacements: (clause_range, uuid_name) + let mut replacements: Vec<((usize, usize), String)> = Vec::new(); - let paimon_table = catalog - .get_table(&identifier) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + // Process all VERSION AS OF clauses + for info in &version_clauses { + let table_ref: datafusion::common::TableReference = info.table_name.as_str().into(); + let (catalog, _catalog_name, identifier) = + self.resolve_table_name_from_ref(&table_ref)?; - let table_with_options = paimon_table.copy_with_options(options); - let provider = Arc::new(PaimonTableProvider::try_new(table_with_options)?); + let paimon_table = catalog + .get_table(&identifier) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; - // Use a UUID-based temp table name to avoid conflicts with existing tables. - let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); + // Merge dynamic options with time-travel options + let mut options = self.dynamic_options.read().unwrap().clone(); + options.insert(SCAN_VERSION_OPTION.to_string(), info.version.clone()); - // Replace the original table name + time-travel clause with just the UUID name - let rewritten_sql = format!( - "{}{}{}", - &sql[..clause_range.0], - uuid_name, - &sql[clause_range.1..] - ); + let table_with_options = paimon_table.copy_with_options(options); + let provider = Arc::new(PaimonTableProvider::try_new(table_with_options)?); - // Register the provider under the UUID temp table name - self.register_temp_table(uuid_name.as_str(), provider)?; + let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); + self.register_temp_table(uuid_name.as_str(), provider)?; + tracker.register(&uuid_name); + replacements.push((info.clause_range, uuid_name)); + } + + // Process all TIMESTAMP AS OF clauses + for info in ×tamp_clauses { + let table_ref: datafusion::common::TableReference = info.table_name.as_str().into(); + 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 millis = Self::parse_timestamp_to_millis(&info.timestamp)?; - // Execute the rewritten SQL - let result = self.ctx.sql(&rewritten_sql).await; + // Merge dynamic options with time-travel options + let mut options = self.dynamic_options.read().unwrap().clone(); + options.insert(SCAN_TIMESTAMP_MILLIS_OPTION.to_string(), millis.to_string()); - // Clean up the temp table - let _ = self.deregister_temp_table(uuid_name.as_str()); + let table_with_options = paimon_table.copy_with_options(options); + let provider = Arc::new(PaimonTableProvider::try_new(table_with_options)?); - result + let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); + self.register_temp_table(uuid_name.as_str(), provider)?; + tracker.register(&uuid_name); + replacements.push((info.clause_range, uuid_name)); + } + + // Sort replacements by position (descending) so that replacements + // from right to left don't shift indices of earlier ones + replacements.sort_by(|a, b| b.0 .0.cmp(&a.0 .0)); + + // Build the rewritten SQL by replacing each clause from right to left + let mut rewritten_sql = sql.to_string(); + for ((start, end), uuid_name) in &replacements { + rewritten_sql = format!( + "{}{}{}", + &rewritten_sql[..*start], + uuid_name, + &rewritten_sql[*end..] + ); + } + + // Execute the rewritten SQL; tracker auto-deregisters on drop + self.ctx.sql(&rewritten_sql).await } /// Parse a timestamp string to milliseconds since epoch (using local timezone). @@ -1171,6 +1223,17 @@ impl SQLContext { }) } + /// Check whether a TableReference targets a registered Paimon catalog. + fn is_paimon_catalog_ref(&self, table_ref: &TableReference) -> bool { + let catalog_name = match table_ref { + TableReference::Full { catalog, .. } => catalog.to_string(), + TableReference::Partial { .. } | TableReference::Bare { .. } => { + self.current_catalog_name() + } + }; + self.catalogs.contains_key(&catalog_name) + } + /// Resolve an ObjectName like `catalog.db.table` or `db.table` to a catalog and Identifier. fn resolve_catalog_and_table( &self, @@ -1282,6 +1345,10 @@ fn looks_like_create_table(sql: &str) -> bool { i += 1; } } + // After optional TEMPORARY/TEMP, reject CREATE TEMPORARY VIEW / CREATE TEMP VIEW + if i + 4 <= len && bytes[i..i + 4].eq_ignore_ascii_case(b"VIEW") { + return false; + } i + 5 <= len && bytes[i..i + 5].eq_ignore_ascii_case(b"TABLE") } @@ -1956,99 +2023,262 @@ struct TimestampAsOfInfo { clause_range: (usize, usize), } -/// Extract `VERSION AS OF ` or `VERSION AS OF ''` from a SQL string. -/// -/// Looks for the pattern (case-insensitive): -/// - `... VERSION AS OF ` — numeric snapshot ID -/// - `... VERSION AS OF ''` — tag name (quoted string) -/// -/// Returns the table name, version/tag value, and byte range of the full clause. -fn extract_version_as_of(sql: &str) -> Option { +/// Check whether a SQL string contains a time-travel keyword (`VERSION AS OF` or +/// `TIMESTAMP AS OF`) **outside** of single-quoted string literals, `--` line +/// comments, and `/* */` block comments. +fn contains_time_travel_keyword(sql: &str) -> bool { let lower = sql.to_lowercase(); - let keyword = "version as of "; - let kw_start = lower.find(keyword)?; - let val_start = kw_start + keyword.len(); + let bytes = lower.as_bytes(); + let len = bytes.len(); + let mut i = 0; + while i < len { + match bytes[i] { + b'\'' => { + // Skip string literal + i += 1; + while i < len { + if bytes[i] == b'\'' { + i += 1; + if i < len && bytes[i] == b'\'' { + i += 1; // escaped quote + } else { + break; + } + } else { + i += 1; + } + } + } + b'-' if i + 1 < len && bytes[i + 1] == b'-' => { + // Skip line comment + i += 2; + while i < len && bytes[i] != b'\n' { + i += 1; + } + } + b'/' if i + 1 < len && bytes[i + 1] == b'*' => { + // Skip block comment + i += 2; + while i + 1 < len { + if bytes[i] == b'*' && bytes[i + 1] == b'/' { + i += 2; + break; + } + i += 1; + } + } + _ => { + // Check for keywords + if i + 14 <= len && bytes[i..i + 14].eq_ignore_ascii_case(b"version as of ") { + return true; + } + if i + 16 <= len && bytes[i..i + 16].eq_ignore_ascii_case(b"timestamp as of ") { + return true; + } + i += 1; + } + } + } + false +} - let remaining = &sql[val_start..]; +/// Extract **all** `VERSION AS OF ` or `VERSION AS OF ''` clauses from a +/// SQL string, skipping string literals and comments. +fn extract_all_version_as_of(sql: &str) -> Vec { + let lower = sql.to_lowercase(); + let bytes = lower.as_bytes(); + let len = bytes.len(); + let sql_bytes = sql.as_bytes(); + let mut i = 0; + let mut results = Vec::new(); - // Parse either a quoted tag name or a numeric snapshot ID - let version = if let Some(after_quote) = remaining.strip_prefix('\'') { - // Tag name: VERSION AS OF 'tagname' - let close_quote = after_quote.find('\'')?; - after_quote[..close_quote].to_string() - } else { - // Numeric snapshot ID: VERSION AS OF 1 - let v: String = remaining - .chars() - .take_while(|c| c.is_ascii_digit()) - .collect(); - if v.is_empty() { - return None; - } - v - }; + while i < len { + match bytes[i] { + b'\'' => { + // Skip string literal + i += 1; + while i < len { + if sql_bytes[i] == b'\'' { + i += 1; + if i < len && sql_bytes[i] == b'\'' { + i += 1; // escaped quote + } else { + break; + } + } else { + i += 1; + } + } + } + b'-' if i + 1 < len && bytes[i + 1] == b'-' => { + // Skip line comment + i += 2; + while i < len && bytes[i] != b'\n' { + i += 1; + } + } + b'/' if i + 1 < len && bytes[i + 1] == b'*' => { + // Skip block comment + i += 2; + while i + 1 < len { + if bytes[i] == b'*' && bytes[i + 1] == b'/' { + i += 2; + break; + } + i += 1; + } + } + _ => { + if i + 14 <= len && bytes[i..i + 14].eq_ignore_ascii_case(b"version as of ") { + let kw_start = i; + let val_start = i + 14; + let remaining = &sql[val_start..]; + + // Parse either a quoted tag name or a numeric snapshot ID + let version = if let Some(after_quote) = remaining.strip_prefix('\'') { + // Tag name: VERSION AS OF 'tagname' + if let Some(close_quote) = after_quote.find('\'') { + after_quote[..close_quote].to_string() + } else { + i += 1; + continue; + } + } else { + // Numeric snapshot ID: VERSION AS OF 1 + let v: String = remaining + .chars() + .take_while(|c| c.is_ascii_digit()) + .collect(); + if v.is_empty() { + i += 1; + continue; + } + v + }; - let is_quoted = remaining.starts_with('\''); - let val_end = if is_quoted { - val_start + version.len() + 2 // 2 quotes - } else { - val_start + version.len() - }; + let is_quoted = remaining.starts_with('\''); + let val_end = if is_quoted { + val_start + version.len() + 2 // 2 quotes + } else { + val_start + version.len() + }; - // Walk backwards from kw_start to find the table name boundary - let table_end = sql[..kw_start].trim_end_matches(' ').len(); - let table_start = sql[..table_end] - .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') - .map(|i| i + 1) - .unwrap_or(0); - let table_name = sql[table_start..table_end].to_string(); - if table_name.is_empty() { - return None; - } - - Some(VersionAsOfInfo { - table_name, - version, - clause_range: (table_start, val_end), - }) + // Walk backwards from kw_start to find the table name boundary + let table_end = sql[..kw_start].trim_end_matches(' ').len(); + let table_start = sql[..table_end] + .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') + .map(|idx| idx + 1) + .unwrap_or(0); + let table_name = sql[table_start..table_end].to_string(); + + if !table_name.is_empty() { + results.push(VersionAsOfInfo { + table_name, + version, + clause_range: (table_start, val_end), + }); + } + + i = val_end; + } else { + i += 1; + } + } + } + } + + results } -/// Extract `TIMESTAMP AS OF ''` from a SQL string. -/// -/// Looks for the pattern `... TIMESTAMP AS OF ''` (case-insensitive), -/// returns the table name, timestamp string, and byte range of the full clause -/// (from table name start to closing quote end). -fn extract_timestamp_as_of(sql: &str) -> Option { +/// Extract **all** `TIMESTAMP AS OF ''` clauses from a SQL string, skipping +/// string literals and comments. +fn extract_all_timestamp_as_of(sql: &str) -> Vec { let lower = sql.to_lowercase(); - let keyword = "timestamp as of "; - let kw_start = lower.find(keyword)?; - let val_start = kw_start + keyword.len(); - - // Read the quoted timestamp string - let remaining = &sql[val_start..]; - if !remaining.starts_with('\'') { - return None; - } - let close_quote = remaining[1..].find('\'')?; - let timestamp = remaining[1..close_quote + 1].to_string(); - let val_end = val_start + close_quote + 2; // skip both quotes - - // Walk backwards to find the table name boundary - let table_end = sql[..kw_start].trim_end_matches(' ').len(); - let table_start = sql[..table_end] - .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') - .map(|i| i + 1) - .unwrap_or(0); - let table_name = sql[table_start..table_end].to_string(); - if table_name.is_empty() { - return None; - } - - Some(TimestampAsOfInfo { - table_name, - timestamp, - clause_range: (table_start, val_end), - }) + let bytes = lower.as_bytes(); + let len = bytes.len(); + let sql_bytes = sql.as_bytes(); + let mut i = 0; + let mut results = Vec::new(); + + while i < len { + match bytes[i] { + b'\'' => { + // Skip string literal + i += 1; + while i < len { + if sql_bytes[i] == b'\'' { + i += 1; + if i < len && sql_bytes[i] == b'\'' { + i += 1; // escaped quote + } else { + break; + } + } else { + i += 1; + } + } + } + b'-' if i + 1 < len && bytes[i + 1] == b'-' => { + // Skip line comment + i += 2; + while i < len && bytes[i] != b'\n' { + i += 1; + } + } + b'/' if i + 1 < len && bytes[i + 1] == b'*' => { + // Skip block comment + i += 2; + while i + 1 < len { + if bytes[i] == b'*' && bytes[i + 1] == b'/' { + i += 2; + break; + } + i += 1; + } + } + _ => { + if i + 16 <= len && bytes[i..i + 16].eq_ignore_ascii_case(b"timestamp as of ") { + let kw_start = i; + let val_start = i + 16; + let remaining = &sql[val_start..]; + + // Read the quoted timestamp string + if !remaining.starts_with('\'') { + i += 1; + continue; + } + if let Some(close_quote) = remaining[1..].find('\'') { + let timestamp = remaining[1..close_quote + 1].to_string(); + let val_end = val_start + close_quote + 2; // skip both quotes + + // Walk backwards to find the table name boundary + let table_end = sql[..kw_start].trim_end_matches(' ').len(); + let table_start = sql[..table_end] + .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(') + .map(|idx| idx + 1) + .unwrap_or(0); + let table_name = sql[table_start..table_end].to_string(); + + if !table_name.is_empty() { + results.push(TimestampAsOfInfo { + table_name, + timestamp, + clause_range: (table_start, val_end), + }); + } + + i = val_end; + } else { + i += 1; + } + } else { + i += 1; + } + } + } + } + + results } /// Return an empty DataFrame with a single "result" column containing "OK". @@ -3575,7 +3805,9 @@ mod tests { #[test] fn test_extract_version_as_of() { let sql = "SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 1"; - let info = extract_version_as_of(sql).unwrap(); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.version, "1"); assert_eq!(info.table_name, "paimon.default.time_travel_table"); let rewritten = format!( @@ -3589,7 +3821,9 @@ mod tests { #[test] fn test_extract_version_as_of_multi_digit() { let sql = "SELECT * FROM mydb.t VERSION AS OF 42"; - let info = extract_version_as_of(sql).unwrap(); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.version, "42"); assert_eq!(info.table_name, "mydb.t"); let rewritten = format!( @@ -3603,7 +3837,9 @@ mod tests { #[test] fn test_extract_version_as_of_case_insensitive() { let sql = "SELECT * FROM t version as of 5"; - let info = extract_version_as_of(sql).unwrap(); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.version, "5"); assert_eq!(info.table_name, "t"); let rewritten = format!( @@ -3617,13 +3853,15 @@ mod tests { #[test] fn test_extract_version_as_of_not_present() { let sql = "SELECT * FROM t"; - assert!(extract_version_as_of(sql).is_none()); + assert!(extract_all_version_as_of(sql).is_empty()); } #[test] fn test_extract_version_as_of_tag() { let sql = "SELECT id, name FROM paimon.default.t VERSION AS OF 'snapshot1'"; - let info = extract_version_as_of(sql).unwrap(); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.version, "snapshot1"); assert_eq!(info.table_name, "paimon.default.t"); let rewritten = format!( @@ -3637,7 +3875,9 @@ mod tests { #[test] fn test_extract_version_as_of_tag_case_insensitive() { let sql = "SELECT * FROM t version as of 'my_tag'"; - let info = extract_version_as_of(sql).unwrap(); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.version, "my_tag"); assert_eq!(info.table_name, "t"); let rewritten = format!( @@ -3650,17 +3890,68 @@ mod tests { #[test] fn test_extract_version_as_of_numeric_still_works() { - // Ensure numeric snapshot ID still works alongside tag support let sql = "SELECT * FROM t VERSION AS OF 123"; - let info = extract_version_as_of(sql).unwrap(); - assert_eq!(info.version, "123"); - assert_eq!(info.table_name, "t"); + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 1); + assert_eq!(infos[0].version, "123"); + assert_eq!(infos[0].table_name, "t"); + } + + #[test] + fn test_extract_version_as_of_multiple() { + // JOIN two time-travel tables + let sql = "SELECT * FROM t1 VERSION AS OF 1 JOIN t2 VERSION AS OF 2 ON t1.id = t2.id"; + let infos = extract_all_version_as_of(sql); + assert_eq!(infos.len(), 2); + assert_eq!(infos[0].version, "1"); + assert_eq!(infos[0].table_name, "t1"); + assert_eq!(infos[1].version, "2"); + assert_eq!(infos[1].table_name, "t2"); + } + + #[test] + fn test_extract_version_as_of_skips_string_literal() { + let sql = "SELECT * FROM t WHERE note = 'version as of 1'"; + let infos = extract_all_version_as_of(sql); + assert!(infos.is_empty()); + } + + #[test] + fn test_extract_version_as_of_skips_comment() { + let sql = "SELECT * FROM t -- version as of 1\n WHERE id > 0"; + let infos = extract_all_version_as_of(sql); + assert!(infos.is_empty()); + } + + #[test] + fn test_contains_time_travel_keyword() { + assert!(contains_time_travel_keyword( + "SELECT * FROM t VERSION AS OF 1" + )); + assert!(contains_time_travel_keyword( + "SELECT * FROM t TIMESTAMP AS OF '2024-01-01 00:00:00'" + )); + // Inside string literal — should NOT match + assert!(!contains_time_travel_keyword( + "SELECT * FROM t WHERE note = 'version as of 1'" + )); + // Inside comment — should NOT match + assert!(!contains_time_travel_keyword( + "SELECT * FROM t -- version as of 1" + )); + assert!(!contains_time_travel_keyword( + "SELECT * FROM t /* timestamp as of now */ WHERE id > 0" + )); + // No keyword at all + assert!(!contains_time_travel_keyword("SELECT * FROM t")); } #[test] fn test_extract_timestamp_as_of() { let sql = "SELECT * FROM paimon.default.t TIMESTAMP AS OF '2024-01-15 10:30:00'"; - let info = extract_timestamp_as_of(sql).unwrap(); + let infos = extract_all_timestamp_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.timestamp, "2024-01-15 10:30:00"); assert_eq!(info.table_name, "paimon.default.t"); let rewritten = format!( @@ -3674,7 +3965,9 @@ mod tests { #[test] fn test_extract_timestamp_as_of_case_insensitive() { let sql = "SELECT * FROM t timestamp as of '2024-06-01 00:00:00'"; - let info = extract_timestamp_as_of(sql).unwrap(); + let infos = extract_all_timestamp_as_of(sql); + assert_eq!(infos.len(), 1); + let info = &infos[0]; assert_eq!(info.timestamp, "2024-06-01 00:00:00"); assert_eq!(info.table_name, "t"); let rewritten = format!( @@ -3688,6 +3981,6 @@ mod tests { #[test] fn test_extract_timestamp_as_of_not_present() { let sql = "SELECT * FROM t"; - assert!(extract_timestamp_as_of(sql).is_none()); + assert!(extract_all_timestamp_as_of(sql).is_empty()); } } diff --git a/docs/src/sql.md b/docs/src/sql.md index 5fb03773..e8d45188 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -598,14 +598,15 @@ The table name accepts flexible references, similar to DataFusion: - `"database.my_table"` — uses the current catalog with the specified database - `"catalog.database.my_table"` — fully qualified -### register_mem_table +### register_temp_table -The easiest way to register a temporary table from a schema and record batches: +Register any `Arc` as a temporary table (including `MemTable`, `ViewTable`, custom providers, etc.): ```rust use datafusion::arrow::array::Int32Array; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::MemTable; let schema = Arc::new(Schema::new(vec![ Field::new("id", ArrowDataType::Int32, false), @@ -619,28 +620,16 @@ let batch = RecordBatch::try_new( ], )?; -// Fully qualified -ctx.register_mem_table("paimon.my_db.users", schema.clone(), vec![batch.clone()])?; +// Register a MemTable as a temp table +let mem_table = Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]])?); +ctx.register_temp_table("paimon.my_db.users", mem_table)?; let df = ctx.sql("SELECT * FROM paimon.my_db.users WHERE id > 1").await?; df.show().await?; -// Database-qualified (uses current catalog) -ctx.register_mem_table("my_db.other_table", schema.clone(), vec![batch.clone()])?; - -// Bare table name (uses current catalog + current database) -ctx.register_mem_table("quick_lookup", schema, vec![batch])?; -``` - -### register_temp_table - -For advanced use cases, `register_temp_table` accepts any `Arc` (including `ViewTable`, custom providers, etc.): - -```rust +// Register a ViewTable as a temp table use datafusion::datasource::ViewTable; - -// Register a view as a temporary table -let view_table = ViewTable::new(logical_plan, Some(query_sql)); -ctx.register_temp_table("paimon.my_db.my_view", Arc::new(view_table))?; +let view_table = Arc::new(ViewTable::new(logical_plan, Some(query_sql))); +ctx.register_temp_table("paimon.my_db.my_view", view_table)?; ``` ### CREATE TEMPORARY TABLE @@ -670,8 +659,10 @@ ctx.deregister_temp_table("paimon.my_db.users")?; Multiple temporary tables can share the same database — the database is created automatically on first use: ```rust -ctx.register_mem_table("my_db.table_a", schema_a, vec![batch_a])?; -ctx.register_mem_table("my_db.table_b", schema_b, vec![batch_b])?; +let mem_a = Arc::new(MemTable::try_new(schema_a, vec![vec![batch_a]])?); +let mem_b = Arc::new(MemTable::try_new(schema_b, vec![vec![batch_b]])?); +ctx.register_temp_table("my_db.table_a", mem_a)?; +ctx.register_temp_table("my_db.table_b", mem_b)?; // Join two temp tables let df = ctx.sql("SELECT * FROM paimon.my_db.table_a JOIN paimon.my_db.table_b ON a.id = b.id").await?;