You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by ju...@apache.org on 2026/05/06 13:17:45 UTC
(paimon-rust) branch main updated: feat(datafusion): rename PaimonSqlHandler to SQLContext with multi-catalog support (#293)
This is an automated email from the ASF dual-hosted git repository.
leaves12138 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new a67e8a1 feat(datafusion): rename PaimonSqlHandler to SQLContext with multi-catalog support (#293)
a67e8a1 is described below
commit a67e8a163f3049941c38d55f0692f98b113f4226
Author: Jingsong Lee <ji...@gmail.com>
AuthorDate: Wed May 6 21:17:39 2026 +0800
feat(datafusion): rename PaimonSqlHandler to SQLContext with multi-catalog support (#293)
---
bindings/python/README.md | 14 +-
bindings/python/project-description.md | 26 +-
.../python/python/pypaimon_rust/datafusion.pyi | 13 +-
bindings/python/src/context.rs | 70 ++-
bindings/python/src/error.rs | 4 +
bindings/python/tests/test_datafusion.py | 25 +-
crates/integrations/datafusion/src/catalog.rs | 6 +-
crates/integrations/datafusion/src/delete.rs | 3 +-
crates/integrations/datafusion/src/lib.rs | 6 +-
crates/integrations/datafusion/src/merge_into.rs | 128 ++---
crates/integrations/datafusion/src/procedures.rs | 20 +-
.../src/{sql_handler.rs => sql_context.rs} | 499 ++++++++++++--------
crates/integrations/datafusion/src/update.rs | 99 ++--
.../datafusion/tests/append_merge_into.rs | 522 +++++++++++----------
crates/integrations/datafusion/tests/blob_tests.rs | 159 ++++---
crates/integrations/datafusion/tests/common/mod.rs | 61 ++-
.../datafusion/tests/count_pushdown.rs | 100 ++--
.../datafusion/tests/cross_partition_tables.rs | 98 ++--
.../integrations/datafusion/tests/delete_tests.rs | 240 +++++-----
.../datafusion/tests/dynamic_bucket_tables.rs | 72 +--
.../datafusion/tests/merge_into_tests.rs | 468 +++++++++---------
crates/integrations/datafusion/tests/pk_tables.rs | 287 +++++------
crates/integrations/datafusion/tests/procedures.rs | 96 ++--
.../{sql_handler_tests.rs => sql_context_tests.rs} | 269 +++++++++--
.../integrations/datafusion/tests/update_tests.rs | 179 +++----
docs/src/datafusion.md | 28 +-
26 files changed, 2042 insertions(+), 1450 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 90e184b..58623bb 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -23,7 +23,19 @@ This project builds the Rust-powered core for [PyPaimon](https://paimon.apache.o
## Usage
-For DataFusion queries, use the native `SessionContext` and register a `PaimonCatalog`:
+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
diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md
index 56e3ad2..48e9492 100644
--- a/bindings/python/project-description.md
+++ b/bindings/python/project-description.md
@@ -31,8 +31,30 @@ If you want to use the native Python DataFusion `SessionContext`, install `dataf
## Query Paimon Tables with DataFusion
-`pypaimon-rust` provides a `PaimonCatalog` that can be registered into the native DataFusion `SessionContext`.
-This keeps the standard DataFusion Python API available for regular queries.
+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",
+})
+
+# DDL
+ctx.sql("CREATE SCHEMA paimon.my_db")
+ctx.sql("CREATE TABLE paimon.my_db.users (id INT, name STRING, PRIMARY KEY (id))")
+
+# DML
+ctx.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob')")
+
+# Query tables via SQL (catalog.database.table)
+batches = ctx.sql("SELECT * FROM paimon.my_db.users")
+```
+
+Alternatively, if you want to use the native Python DataFusion `SessionContext`,
+install `datafusion` and register a `PaimonCatalog`:
```python
from datafusion import SessionContext
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 3ca0617..f5a9802 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -15,8 +15,19 @@
# specific language governing permissions and limitations
# under the License.
-from typing import Any, Dict
+from typing import Any, Dict, List
+
+import pyarrow
class PaimonCatalog:
def __init__(self, catalog_options: Dict[str, str]) -> None: ...
def __datafusion_catalog_provider__(self, session: Any) -> object: ...
+
+class SQLContext:
+ def __init__(self) -> None: ...
+ def register_catalog(
+ self, catalog_name: str, catalog_options: Dict[str, str]
+ ) -> None: ...
+ def set_current_catalog(self, catalog_name: str) -> None: ...
+ def set_current_database(self, database_name: str) -> None: ...
+ def sql(self, sql: str) -> List[pyarrow.RecordBatch]: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 2bc1e1d..2df5bd0 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -18,15 +18,16 @@
use std::collections::HashMap;
use std::sync::Arc;
+use arrow::pyarrow::ToPyArrow;
use datafusion::catalog::CatalogProvider;
use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
use paimon::{CatalogFactory, Options};
-use paimon_datafusion::PaimonCatalogProvider;
+use paimon_datafusion::{PaimonCatalogProvider, SQLContext};
use pyo3::prelude::*;
use pyo3::types::PyCapsule;
-use crate::error::to_py_err;
+use crate::error::{df_to_py_err, to_py_err};
use paimon_datafusion::runtime::runtime;
fn build_paimon_catalog_provider(
@@ -86,9 +87,74 @@ impl PaimonCatalog {
}
}
+/// A SQL context that supports registering multiple Paimon catalogs and executing SQL.
+#[pyclass(name = "SQLContext")]
+pub struct PySQLContext {
+ inner: SQLContext,
+}
+
+#[pymethods]
+impl PySQLContext {
+ #[new]
+ fn new() -> Self {
+ Self {
+ inner: SQLContext::new(),
+ }
+ }
+
+ fn register_catalog(
+ &mut self,
+ catalog_name: String,
+ catalog_options: HashMap<String, String>,
+ ) -> PyResult<()> {
+ let rt = runtime();
+ rt.block_on(async {
+ let options = Options::from_map(catalog_options);
+ let catalog = CatalogFactory::create(options).await.map_err(to_py_err)?;
+ self.inner
+ .register_catalog(catalog_name, catalog)
+ .await
+ .map_err(df_to_py_err)
+ })
+ }
+
+ fn set_current_catalog(&mut self, catalog_name: String) -> PyResult<()> {
+ let rt = runtime();
+ rt.block_on(async {
+ self.inner
+ .set_current_catalog(catalog_name)
+ .await
+ .map_err(df_to_py_err)
+ })
+ }
+
+ fn set_current_database(&self, database_name: String) -> PyResult<()> {
+ let rt = runtime();
+ rt.block_on(async {
+ self.inner
+ .set_current_database(&database_name)
+ .await
+ .map_err(df_to_py_err)
+ })
+ }
+
+ fn sql(&self, py: Python<'_>, sql: String) -> PyResult<Vec<Py<PyAny>>> {
+ let rt = runtime();
+ let batches = rt.block_on(async {
+ let df = self.inner.sql(&sql).await.map_err(df_to_py_err)?;
+ df.collect().await.map_err(df_to_py_err)
+ })?;
+ batches
+ .iter()
+ .map(|batch| Ok(batch.to_pyarrow(py)?.unbind()))
+ .collect()
+ }
+}
+
pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
let this = PyModule::new(py, "datafusion")?;
this.add_class::<PaimonCatalog>()?;
+ this.add_class::<PySQLContext>()?;
m.add_submodule(&this)?;
py.import("sys")?
.getattr("modules")?
diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs
index 05dbe2e..992a9f1 100644
--- a/bindings/python/src/error.rs
+++ b/bindings/python/src/error.rs
@@ -21,3 +21,7 @@ use pyo3::PyErr;
pub fn to_py_err(err: paimon::Error) -> PyErr {
PyValueError::new_err(err.to_string())
}
+
+pub fn df_to_py_err(err: datafusion::error::DataFusionError) -> PyErr {
+ PyValueError::new_err(err.to_string())
+}
diff --git a/bindings/python/tests/test_datafusion.py b/bindings/python/tests/test_datafusion.py
index 919ef55..2d8abf8 100644
--- a/bindings/python/tests/test_datafusion.py
+++ b/bindings/python/tests/test_datafusion.py
@@ -16,11 +16,12 @@
# under the License.
import os
+import tempfile
import pyarrow as pa
from datafusion import SessionContext
-from pypaimon_rust.datafusion import PaimonCatalog
+from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
WAREHOUSE = os.environ.get("PAIMON_TEST_WAREHOUSE", "/tmp/paimon-warehouse")
@@ -42,3 +43,25 @@ def test_query_simple_table_via_catalog_provider():
(2, "bob"),
(3, "carol"),
]
+
+
+def test_sql_context_ddl_dml():
+ 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, PRIMARY KEY (id))"
+ )
+
+ ctx.sql("INSERT INTO paimon.test_db.users VALUES (1, 'alice'), (2, 'bob')")
+
+ 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, "alice"), (2, "bob")]
+
+ ctx.sql("DROP TABLE paimon.test_db.users")
+ ctx.sql("DROP SCHEMA paimon.test_db")
diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs
index 7c0bf35..8e58015 100644
--- a/crates/integrations/datafusion/src/catalog.rs
+++ b/crates/integrations/datafusion/src/catalog.rs
@@ -42,7 +42,7 @@ use crate::DynamicOptions;
pub struct PaimonCatalogProvider {
/// Reference to the Paimon catalog.
catalog: Arc<dyn Catalog>,
- /// Session-scoped dynamic options shared with the SQL handler.
+ /// Session-scoped dynamic options shared with the SQL context.
dynamic_options: DynamicOptions,
}
@@ -56,7 +56,7 @@ impl PaimonCatalogProvider {
/// Creates a new [`PaimonCatalogProvider`].
///
/// For standalone use without `SET`/`RESET` support.
- /// When used via [`PaimonSqlHandler`], the handler creates the provider
+ /// When used via [`SQLContext`], the handler creates the provider
/// internally with shared dynamic options.
pub fn new(catalog: Arc<dyn Catalog>) -> Self {
PaimonCatalogProvider {
@@ -178,7 +178,7 @@ pub struct PaimonSchemaProvider {
catalog: Arc<dyn Catalog>,
/// Database name this schema represents.
database: String,
- /// Session-scoped dynamic options shared with the SQL handler.
+ /// Session-scoped dynamic options shared with the SQL context.
dynamic_options: DynamicOptions,
}
diff --git a/crates/integrations/datafusion/src/delete.rs b/crates/integrations/datafusion/src/delete.rs
index e8cd905..cb37c53 100644
--- a/crates/integrations/datafusion/src/delete.rs
+++ b/crates/integrations/datafusion/src/delete.rs
@@ -104,7 +104,8 @@ async fn execute_cow_delete_once(
return ok_result(ctx, 0);
}
- let result = execute_cow_delete_inner(ctx, cow_table_guard.name(), delete, &mut writer).await;
+ let result =
+ execute_cow_delete_inner(ctx, &cow_table_guard.qualified_name(), delete, &mut writer).await;
drop(cow_table_guard);
let total_count = result?;
diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs
index 3a03ba9..47f1bab 100644
--- a/crates/integrations/datafusion/src/lib.rs
+++ b/crates/integrations/datafusion/src/lib.rs
@@ -47,7 +47,7 @@ mod physical_plan;
mod procedures;
mod relation_planner;
pub mod runtime;
-mod sql_handler;
+mod sql_context;
mod system_tables;
mod table;
mod table_function_args;
@@ -59,7 +59,7 @@ use std::sync::{Arc, RwLock};
/// Session-scoped dynamic options set via `SET 'paimon.key' = 'value'`.
///
-/// Shared internally across [`PaimonSqlHandler`] and [`PaimonCatalogProvider`]
+/// Shared internally across [`SQLContext`] and [`PaimonCatalogProvider`]
/// so that SET/RESET mutations are visible to subsequent table scans.
pub(crate) type DynamicOptions = Arc<RwLock<HashMap<String, String>>>;
@@ -69,6 +69,6 @@ pub use error::to_datafusion_error;
pub use full_text_search::{register_full_text_search, FullTextSearchFunction};
pub use physical_plan::PaimonTableScan;
pub use relation_planner::PaimonRelationPlanner;
-pub use sql_handler::PaimonSqlHandler;
+pub use sql_context::SQLContext;
pub use table::PaimonTableProvider;
pub use vector_search::{register_vector_search, VectorSearchFunction};
diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs
index bc852a4..a004aa5 100644
--- a/crates/integrations/datafusion/src/merge_into.rs
+++ b/crates/integrations/datafusion/src/merge_into.rs
@@ -64,6 +64,9 @@ pub(crate) struct CowTableGuard {
table_name: String,
}
+const COW_CATALOG: &str = "datafusion";
+const COW_SCHEMA: &str = "public";
+
impl CowTableGuard {
pub(crate) fn new(ctx: &SessionContext, table_name: String) -> Self {
Self {
@@ -72,14 +75,14 @@ impl CowTableGuard {
}
}
- pub(crate) fn name(&self) -> &str {
- &self.table_name
+ pub(crate) fn qualified_name(&self) -> String {
+ format!("{COW_CATALOG}.{COW_SCHEMA}.{}", self.table_name)
}
}
impl Drop for CowTableGuard {
fn drop(&mut self) {
- let _ = self.ctx.deregister_table(&self.table_name);
+ let _ = self.ctx.deregister_table(self.qualified_name());
}
}
@@ -367,7 +370,7 @@ async fn execute_cow_merge_once(
t_alias,
on_condition: &on_condition,
has_target_data,
- cow_target_name: cow_target_guard.name(),
+ cow_target_name: cow_target_guard.qualified_name(),
update_columns: &update_columns,
};
@@ -401,7 +404,7 @@ struct CowMergeContext<'a> {
t_alias: &'a str,
on_condition: &'a str,
has_target_data: bool,
- cow_target_name: &'a str,
+ cow_target_name: String,
update_columns: &'a [String],
}
@@ -419,7 +422,7 @@ async fn execute_cow_merge_inner(
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_target_name = &merge_ctx.cow_target_name;
let update_columns = merge_ctx.update_columns;
let mut insert_messages = Vec::new();
let mut total_count: u64 = 0;
@@ -833,10 +836,12 @@ async fn build_insert_batches(
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");
- ctx.register_table(&tmp_name, Arc::new(mem_table))?;
+ 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, &tmp_name, table_fields).await;
+ let result =
+ build_insert_batches_inner(ctx, inserts, s_alias, &qualified_tmp, table_fields).await;
result
}
@@ -1252,7 +1257,10 @@ 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(&table_name, Arc::new(mem_table))?;
+ ctx.register_table(
+ format!("{COW_CATALOG}.{COW_SCHEMA}.{table_name}"),
+ Arc::new(mem_table),
+ )?;
}
Ok((has_data, CowTableGuard::new(ctx, table_name)))
@@ -1437,33 +1445,39 @@ mod tests {
use paimon::{CatalogOptions, FileSystemCatalog, Options};
use tempfile::TempDir;
- use crate::{PaimonSqlHandler, PaimonTableProvider};
+ use crate::{PaimonTableProvider, SQLContext};
- async fn setup_handler() -> (TempDir, PaimonSqlHandler, Arc<FileSystemCatalog>) {
+ async fn setup_sql_context() -> (TempDir, SQLContext, Arc<FileSystemCatalog>) {
let temp_dir = TempDir::new().unwrap();
let warehouse = format!("file://{}", temp_dir.path().display());
let mut options = Options::new();
options.set(CatalogOptions::WAREHOUSE, warehouse);
let catalog = Arc::new(FileSystemCatalog::new(options).unwrap());
- let handler =
- PaimonSqlHandler::new(SessionContext::new(), catalog.clone(), "paimon").unwrap();
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
+ let mut sql_context = SQLContext::new();
+ sql_context
+ .register_catalog("paimon", catalog.clone())
+ .await
+ .unwrap();
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
- (temp_dir, handler, catalog)
+ (temp_dir, sql_context, catalog)
}
async fn setup_data_evolution_table(name: &str) -> (TempDir, SessionContext, Table) {
- let (tmp, handler, catalog) = setup_handler().await;
+ let (tmp, sql_context, catalog) = setup_sql_context().await;
- handler
+ sql_context
.sql(&format!(
"CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('row-tracking.enabled' = 'true')"
))
.await
.unwrap();
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)"
))
@@ -1482,9 +1496,10 @@ mod tests {
extra.insert("row-tracking.enabled".to_string(), "true".to_string());
let de_table = table.copy_with_options(extra);
- let ctx = handler.ctx().clone();
+ let ctx = sql_context.ctx().clone();
let provider = PaimonTableProvider::try_new(de_table.clone()).unwrap();
- ctx.register_table("target", Arc::new(provider)).unwrap();
+ ctx.register_table("datafusion.public.target", Arc::new(provider))
+ .unwrap();
(tmp, ctx, de_table)
}
@@ -1504,7 +1519,7 @@ mod tests {
// Create source table with updates
ctx.sql(
- "CREATE TABLE source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
+ "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
)
.await
.unwrap()
@@ -1514,13 +1529,13 @@ mod tests {
// Execute MERGE INTO
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
);
execute_merge_into(&ctx, &merge, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -1563,15 +1578,17 @@ mod tests {
async fn test_merge_into_no_matches() {
let (_tmp, ctx, table) = setup_data_evolution_table("t_merge2").await;
- ctx.sql("CREATE TABLE source (id INT, name VARCHAR) AS VALUES (99, 'nobody')")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+ ctx.sql(
+ "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (99, 'nobody')",
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.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();
@@ -1631,16 +1648,16 @@ mod tests {
// -----------------------------------------------------------------------
async fn setup_append_only_table(name: &str) -> (TempDir, SessionContext, Table) {
- let (tmp, handler, catalog) = setup_handler().await;
+ let (tmp, sql_context, catalog) = setup_sql_context().await;
- handler
+ sql_context
.sql(&format!(
"CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT)"
))
.await
.unwrap();
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)"
))
@@ -1655,9 +1672,10 @@ mod tests {
.await
.unwrap();
- let ctx = handler.ctx().clone();
+ let ctx = sql_context.ctx().clone();
let provider = PaimonTableProvider::try_new(table.clone()).unwrap();
- ctx.register_table("target", Arc::new(provider)).unwrap();
+ ctx.register_table("datafusion.public.target", Arc::new(provider))
+ .unwrap();
(tmp, ctx, table)
}
@@ -1693,7 +1711,7 @@ mod tests {
let (_tmp, ctx, table) = setup_append_only_table("t_cow_upd").await;
ctx.sql(
- "CREATE TABLE source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
+ "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
)
.await
.unwrap()
@@ -1702,13 +1720,13 @@ mod tests {
.unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
);
execute_merge_into(&ctx, &merge, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -1730,7 +1748,7 @@ mod tests {
async fn test_cow_merge_delete_matched_rows() {
let (_tmp, ctx, table) = setup_append_only_table("t_cow_del").await;
- ctx.sql("CREATE TABLE source (id INT) AS VALUES (2)")
+ ctx.sql("CREATE TABLE datafusion.public.source (id INT) AS VALUES (2)")
.await
.unwrap()
.collect()
@@ -1738,13 +1756,13 @@ mod tests {
.unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.source s ON t.id = s.id \
WHEN MATCHED THEN DELETE",
);
execute_merge_into(&ctx, &merge, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -1762,7 +1780,7 @@ mod tests {
async fn test_cow_merge_insert_not_matched() {
let (_tmp, ctx, table) = setup_append_only_table("t_cow_ins").await;
- ctx.sql("CREATE TABLE source (id INT, name VARCHAR, value INT) AS VALUES (4, 'dave', 40), (5, 'eve', 50)")
+ ctx.sql("CREATE TABLE datafusion.public.source (id INT, name VARCHAR, value INT) AS VALUES (4, 'dave', 40), (5, 'eve', 50)")
.await
.unwrap()
.collect()
@@ -1770,13 +1788,13 @@ mod tests {
.unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.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();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -1800,7 +1818,7 @@ mod tests {
async fn test_cow_merge_update_and_insert() {
let (_tmp, ctx, table) = setup_append_only_table("t_cow_upsert").await;
- ctx.sql("CREATE TABLE source (id INT, name VARCHAR, value INT) AS VALUES (2, 'BOB', 200), (4, 'dave', 40)")
+ ctx.sql("CREATE TABLE datafusion.public.source (id INT, name VARCHAR, value INT) AS VALUES (2, 'BOB', 200), (4, 'dave', 40)")
.await
.unwrap()
.collect()
@@ -1808,14 +1826,14 @@ mod tests {
.unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.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();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -1838,15 +1856,17 @@ mod tests {
async fn test_cow_merge_no_matches() {
let (_tmp, ctx, table) = setup_append_only_table("t_cow_nomatch").await;
- ctx.sql("CREATE TABLE source (id INT, name VARCHAR) AS VALUES (99, 'nobody')")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+ ctx.sql(
+ "CREATE TABLE datafusion.public.source (id INT, name VARCHAR) AS VALUES (99, 'nobody')",
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
let merge = parse_merge(
- "MERGE INTO target t USING source s ON t.id = s.id \
+ "MERGE INTO datafusion.public.target t USING datafusion.public.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();
diff --git a/crates/integrations/datafusion/src/procedures.rs b/crates/integrations/datafusion/src/procedures.rs
index 36dcbd3..183788a 100644
--- a/crates/integrations/datafusion/src/procedures.rs
+++ b/crates/integrations/datafusion/src/procedures.rs
@@ -126,11 +126,15 @@ async fn earlier_or_equal_from_all(
pub async fn execute_call(
ctx: &SessionContext,
- catalog: &Arc<dyn Catalog>,
- catalog_name: &str,
+ catalogs: &HashMap<String, Arc<dyn Catalog>>,
+ current_catalog: &str,
func: &Function,
) -> DFResult<DataFrame> {
- let proc_name = extract_procedure_name(&func.name)?;
+ let (explicit_catalog, proc_name) = extract_procedure_name(&func.name)?;
+ let catalog_name = explicit_catalog.as_deref().unwrap_or(current_catalog);
+ let catalog = catalogs
+ .get(catalog_name)
+ .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")))?;
let args = extract_named_args(&func.args)?;
match proc_name.as_str() {
@@ -149,17 +153,19 @@ pub async fn execute_call(
}
}
-fn extract_procedure_name(name: &ObjectName) -> DFResult<String> {
+/// Returns (optional_catalog_name, procedure_name).
+fn extract_procedure_name(name: &ObjectName) -> DFResult<(Option<String>, String)> {
let parts: Vec<String> = name
.0
.iter()
.filter_map(|p| p.as_ident().map(|id| id.value.clone()))
.collect();
match parts.len() {
- 1 => Ok(parts[0].clone()),
- 2 => Ok(parts[1].clone()),
+ 1 => Ok((None, parts[0].clone())),
+ 2 => Ok((None, parts[1].clone())),
+ 3 => Ok((Some(parts[0].clone()), parts[2].clone())),
_ => Err(DataFusionError::Plan(format!(
- "Invalid procedure name: {name}. Expected sys.procedure_name or procedure_name"
+ "Invalid procedure name: {name}. Expected procedure_name, sys.procedure_name, or catalog.sys.procedure_name"
))),
}
}
diff --git a/crates/integrations/datafusion/src/sql_handler.rs b/crates/integrations/datafusion/src/sql_context.rs
similarity index 87%
rename from crates/integrations/datafusion/src/sql_handler.rs
rename to crates/integrations/datafusion/src/sql_context.rs
index f68a210..10974bb 100644
--- a/crates/integrations/datafusion/src/sql_handler.rs
+++ b/crates/integrations/datafusion/src/sql_context.rs
@@ -18,7 +18,7 @@
//! SQL support for Paimon tables.
//!
//! DataFusion does not natively support all SQL statements needed by Paimon.
-//! This module provides [`PaimonSqlHandler`] which intercepts CREATE TABLE,
+//! This module provides [`SQLContext`] which intercepts CREATE TABLE,
//! ALTER TABLE, MERGE INTO, UPDATE and other SQL, translates them to Paimon
//! catalog operations, and delegates everything else (SELECT, CREATE/DROP
//! SCHEMA, DROP TABLE, etc.) to the underlying [`SessionContext`].
@@ -64,54 +64,111 @@ use paimon::spec::{
use crate::error::to_datafusion_error;
use crate::DynamicOptions;
-/// Wraps a [`SessionContext`] and a Paimon [`Catalog`] to handle DDL statements
-/// that DataFusion does not natively support (e.g. ALTER TABLE).
-///
-/// For all other SQL, it delegates to the inner `SessionContext`.
+/// A SQL context that supports registering multiple Paimon catalogs and executing SQL.
///
/// # Example
/// ```ignore
-/// let ctx = SessionContext::new();
-/// let handler = PaimonSqlHandler::new(ctx, catalog, "paimon")?;
-/// let df = handler.sql("ALTER TABLE paimon.db.t ADD COLUMN age INT").await?;
+/// let mut ctx = SQLContext::new();
+/// ctx.register_catalog("paimon", catalog).await?;
+/// ctx.set_current_catalog("paimon").await?;
+/// let df = ctx.sql("ALTER TABLE paimon.db.t ADD COLUMN age INT").await?;
/// ```
-pub struct PaimonSqlHandler {
+pub struct SQLContext {
ctx: SessionContext,
- catalog: Arc<dyn Catalog>,
- /// The catalog name registered in the SessionContext (used to strip the catalog prefix).
- catalog_name: String,
+ catalogs: HashMap<String, Arc<dyn Catalog>>,
/// Session-scoped dynamic options set via `SET 'paimon.key' = 'value'`.
dynamic_options: DynamicOptions,
}
-impl PaimonSqlHandler {
- /// Creates a new handler that registers the Paimon catalog and relation planner
- /// on the given [`SessionContext`].
+impl Default for SQLContext {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SQLContext {
+ /// Creates a new empty SQL context.
+ pub fn new() -> Self {
+ let ctx = SessionContext::new();
+ ctx.register_relation_planner(Arc::new(
+ crate::relation_planner::PaimonRelationPlanner::new(),
+ ))
+ .expect("failed to register relation planner");
+ Self {
+ ctx,
+ catalogs: HashMap::new(),
+ dynamic_options: Default::default(),
+ }
+ }
+
+ /// Registers a Paimon catalog under the given name.
///
- /// Dynamic options for `SET`/`RESET` are managed internally.
- pub fn new(
- ctx: SessionContext,
- catalog: Arc<dyn Catalog>,
+ /// The first registered catalog automatically becomes the current catalog
+ /// for both Paimon-handled SQL and DataFusion-delegated SQL (SELECT, etc.).
+ /// A "default" database is created if it does not already exist (matching
+ /// the behavior of Spark/Flink Paimon catalogs).
+ pub async fn register_catalog(
+ &mut self,
catalog_name: impl Into<String>,
- ) -> DFResult<Self> {
+ catalog: Arc<dyn Catalog>,
+ ) -> DFResult<()> {
let catalog_name = catalog_name.into();
- let dynamic_options: DynamicOptions = Default::default();
- ctx.register_catalog(
+ let is_first = self.catalogs.is_empty();
+
+ catalog
+ .create_database("default", true, Default::default())
+ .await
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+ self.ctx.register_catalog(
&catalog_name,
Arc::new(crate::catalog::PaimonCatalogProvider::with_dynamic_options(
catalog.clone(),
- dynamic_options.clone(),
+ self.dynamic_options.clone(),
)),
);
- ctx.register_relation_planner(Arc::new(
- crate::relation_planner::PaimonRelationPlanner::new(),
- ))?;
- Ok(Self {
- ctx,
- catalog,
- catalog_name,
- dynamic_options,
- })
+ self.catalogs.insert(catalog_name.clone(), catalog);
+ if is_first {
+ self.set_current_catalog(catalog_name).await?;
+ self.set_current_database("default").await?;
+ }
+ Ok(())
+ }
+
+ /// Sets the current catalog for unqualified table references.
+ pub async fn set_current_catalog(&mut self, catalog_name: impl Into<String>) -> DFResult<()> {
+ let catalog_name = catalog_name.into();
+ if !self.catalogs.contains_key(&catalog_name) {
+ return Err(DataFusionError::Plan(format!(
+ "Unknown catalog '{catalog_name}'"
+ )));
+ }
+ if catalog_name.contains('\'') {
+ return Err(DataFusionError::Plan(
+ "Catalog name must not contain single quotes".to_string(),
+ ));
+ }
+ self.ctx
+ .sql(&format!(
+ "SET datafusion.catalog.default_catalog = '{catalog_name}'"
+ ))
+ .await?;
+ Ok(())
+ }
+
+ /// Sets the current database (schema) for unqualified table references.
+ pub async fn set_current_database(&self, database_name: &str) -> DFResult<()> {
+ if database_name.contains('\'') {
+ return Err(DataFusionError::Plan(
+ "Database name must not contain single quotes".to_string(),
+ ));
+ }
+ self.ctx
+ .sql(&format!(
+ "SET datafusion.catalog.default_schema = '{database_name}'"
+ ))
+ .await?;
+ Ok(())
}
/// Returns a reference to the inner [`SessionContext`].
@@ -145,10 +202,16 @@ impl PaimonSqlHandler {
match &statements[0] {
Statement::CreateTable(create_table) => {
- self.handle_create_table(create_table, partition_keys).await
+ 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, _) =
+ self.resolve_catalog_and_table(&alter_table.name)?;
self.handle_alter_table(
+ &catalog,
&alter_table.name,
&alter_table.operations,
alter_table.if_exists,
@@ -200,8 +263,13 @@ impl PaimonSqlHandler {
}
Statement::Truncate(truncate) => self.handle_truncate_table(truncate).await,
Statement::Call(func) => {
- crate::procedures::execute_call(&self.ctx, &self.catalog, &self.catalog_name, func)
- .await
+ crate::procedures::execute_call(
+ &self.ctx,
+ &self.catalogs,
+ &self.current_catalog_name(),
+ func,
+ )
+ .await
}
_ => self.ctx.sql(sql).await,
}
@@ -209,6 +277,7 @@ impl PaimonSqlHandler {
async fn handle_create_table(
&self,
+ catalog: &Arc<dyn Catalog>,
ct: &CreateTable,
partition_keys: Vec<String>,
) -> DFResult<DataFrame> {
@@ -270,7 +339,7 @@ impl PaimonSqlHandler {
let schema = builder.build().map_err(to_datafusion_error)?;
- self.catalog
+ catalog
.create_table(&identifier, schema, ct.if_not_exists)
.await
.map_err(to_datafusion_error)?;
@@ -280,6 +349,7 @@ impl PaimonSqlHandler {
async fn handle_alter_table(
&self,
+ catalog: &Arc<dyn Catalog>,
name: &ObjectName,
operations: &[AlterTableOperation],
if_exists: bool,
@@ -340,6 +410,7 @@ impl PaimonSqlHandler {
} => {
return self
.handle_drop_partitions(
+ catalog,
&identifier,
partitions,
if_exists || *partition_if_exists,
@@ -355,14 +426,14 @@ impl PaimonSqlHandler {
}
if let Some(new_identifier) = rename_to {
- self.catalog
+ catalog
.rename_table(&identifier, &new_identifier, if_exists)
.await
.map_err(to_datafusion_error)?;
}
if !changes.is_empty() {
- self.catalog
+ catalog
.alter_table(&identifier, changes, if_exists)
.await
.map_err(to_datafusion_error)?;
@@ -372,7 +443,6 @@ impl PaimonSqlHandler {
}
async fn handle_merge_into(&self, merge: &Merge) -> DFResult<DataFrame> {
- // Resolve the target table name from the MERGE INTO clause
let table_name = match &merge.table {
TableFactor::Table { name, .. } => name.clone(),
other => {
@@ -381,11 +451,9 @@ impl PaimonSqlHandler {
)))
}
};
- let identifier = self.resolve_table_name(&table_name)?;
+ let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
- // Load the Paimon table from the catalog
- let table = self
- .catalog
+ let table = catalog
.get_table(&identifier)
.await
.map_err(to_datafusion_error)?;
@@ -402,10 +470,9 @@ impl PaimonSqlHandler {
)))
}
};
- let identifier = self.resolve_table_name(&table_name)?;
+ let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
- let table = self
- .catalog
+ let table = catalog
.get_table(&identifier)
.await
.map_err(to_datafusion_error)?;
@@ -429,10 +496,9 @@ impl PaimonSqlHandler {
)))
}
};
- let identifier = self.resolve_table_name(&table_name)?;
+ let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
- let table = self
- .catalog
+ let table = catalog
.get_table(&identifier)
.await
.map_err(to_datafusion_error)?;
@@ -450,9 +516,8 @@ impl PaimonSqlHandler {
)))
}
};
- let identifier = self.resolve_table_name(&table_name)?;
- let table = self
- .catalog
+ let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
+ let table = catalog
.get_table(&identifier)
.await
.map_err(to_datafusion_error)?;
@@ -585,8 +650,8 @@ impl PaimonSqlHandler {
let target = truncate.table_names.first().ok_or_else(|| {
DataFusionError::Plan("TRUNCATE TABLE requires a table name".to_string())
})?;
- let identifier = self.resolve_table_name(&target.name)?;
- let table = match self.catalog.get_table(&identifier).await {
+ let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&target.name)?;
+ let table = match catalog.get_table(&identifier).await {
Ok(t) => t,
Err(e) if truncate.if_exists && is_table_not_exist(&e) => {
return ok_result(&self.ctx);
@@ -621,6 +686,7 @@ impl PaimonSqlHandler {
async fn handle_drop_partitions(
&self,
+ catalog: &Arc<dyn Catalog>,
identifier: &Identifier,
partitions: &[SqlExpr],
if_exists: bool,
@@ -630,7 +696,7 @@ impl PaimonSqlHandler {
"DROP PARTITIONS requires at least one partition specification".to_string(),
));
}
- let table = match self.catalog.get_table(identifier).await {
+ let table = match catalog.get_table(identifier).await {
Ok(t) => t,
Err(e) if if_exists && is_table_not_exist(&e) => {
return ok_result(&self.ctx);
@@ -654,8 +720,30 @@ impl PaimonSqlHandler {
ok_result(&self.ctx)
}
- /// Resolve an ObjectName like `paimon.db.table` or `db.table` to a Paimon Identifier.
- fn resolve_table_name(&self, name: &ObjectName) -> DFResult<Identifier> {
+ /// Returns the name of the current default catalog from DataFusion config.
+ fn current_catalog_name(&self) -> String {
+ self.ctx
+ .state()
+ .config_options()
+ .catalog
+ .default_catalog
+ .clone()
+ }
+
+ fn current_catalog(&self) -> DFResult<Arc<dyn Catalog>> {
+ let name = self.current_catalog_name();
+ self.catalogs.get(&name).cloned().ok_or_else(|| {
+ DataFusionError::Plan(
+ "No catalog registered. Call register_catalog() first.".to_string(),
+ )
+ })
+ }
+
+ /// Resolve an ObjectName like `catalog.db.table` or `db.table` to a catalog and Identifier.
+ fn resolve_catalog_and_table(
+ &self,
+ name: &ObjectName,
+ ) -> DFResult<(Arc<dyn Catalog>, String, Identifier)> {
let parts: Vec<String> = name
.0
.iter()
@@ -663,25 +751,49 @@ impl PaimonSqlHandler {
.collect();
match parts.len() {
3 => {
- // catalog.database.table — strip catalog prefix
- if parts[0] != self.catalog_name {
- return Err(DataFusionError::Plan(format!(
- "Unknown catalog '{}', expected '{}'",
- parts[0], self.catalog_name
- )));
- }
- Ok(Identifier::new(parts[1].clone(), parts[2].clone()))
+ let catalog = self.catalogs.get(&parts[0]).ok_or_else(|| {
+ DataFusionError::Plan(format!("Unknown catalog '{}'", parts[0]))
+ })?;
+ Ok((
+ catalog.clone(),
+ parts[0].clone(),
+ Identifier::new(parts[1].clone(), parts[2].clone()),
+ ))
+ }
+ 2 => {
+ let catalog = self.current_catalog()?;
+ Ok((
+ catalog,
+ self.current_catalog_name(),
+ Identifier::new(parts[0].clone(), parts[1].clone()),
+ ))
+ }
+ 1 => {
+ let catalog = self.current_catalog()?;
+ let default_schema = self
+ .ctx
+ .state()
+ .config_options()
+ .catalog
+ .default_schema
+ .clone();
+ Ok((
+ catalog,
+ self.current_catalog_name(),
+ Identifier::new(default_schema, parts[0].clone()),
+ ))
}
- 2 => Ok(Identifier::new(parts[0].clone(), parts[1].clone())),
- 1 => Err(DataFusionError::Plan(format!(
- "ALTER TABLE requires at least database.table, got: {}",
- parts[0]
- ))),
_ => Err(DataFusionError::Plan(format!(
"Invalid table reference: {name}"
))),
}
}
+
+ /// Resolve an ObjectName to just the Identifier (for backward compat in handle_alter_table).
+ fn resolve_table_name(&self, name: &ObjectName) -> DFResult<Identifier> {
+ let (_catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(name)?;
+ Ok(identifier)
+ }
}
/// Quick check whether the SQL looks like a CREATE TABLE statement.
@@ -1528,8 +1640,10 @@ mod tests {
}
}
- fn make_handler(catalog: Arc<MockCatalog>) -> PaimonSqlHandler {
- PaimonSqlHandler::new(SessionContext::new(), catalog, "paimon").unwrap()
+ async fn make_sql_context(catalog: Arc<MockCatalog>) -> SQLContext {
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog).await.unwrap();
+ ctx
}
fn assert_sql_type_to_paimon(
@@ -1757,10 +1871,10 @@ mod tests {
// ==================== resolve_table_name tests ====================
- #[test]
- fn test_resolve_three_part_name() {
+ #[tokio::test]
+ async fn test_resolve_three_part_name() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
+ let sql_context = make_sql_context(catalog).await;
let dialect = GenericDialect {};
let stmts = Parser::parse_sql(&dialect, "SELECT * FROM paimon.mydb.mytable").unwrap();
if let Statement::Query(q) = &stmts[0] {
@@ -1768,7 +1882,7 @@ mod tests {
if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
&sel.from[0].relation
{
- let id = handler.resolve_table_name(name).unwrap();
+ let id = sql_context.resolve_table_name(name).unwrap();
assert_eq!(id.database(), "mydb");
assert_eq!(id.object(), "mytable");
}
@@ -1776,10 +1890,10 @@ mod tests {
}
}
- #[test]
- fn test_resolve_two_part_name() {
+ #[tokio::test]
+ async fn test_resolve_two_part_name() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
+ let sql_context = make_sql_context(catalog).await;
let dialect = GenericDialect {};
let stmts = Parser::parse_sql(&dialect, "SELECT * FROM mydb.mytable").unwrap();
if let Statement::Query(q) = &stmts[0] {
@@ -1787,7 +1901,7 @@ mod tests {
if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
&sel.from[0].relation
{
- let id = handler.resolve_table_name(name).unwrap();
+ let id = sql_context.resolve_table_name(name).unwrap();
assert_eq!(id.database(), "mydb");
assert_eq!(id.object(), "mytable");
}
@@ -1795,10 +1909,10 @@ mod tests {
}
}
- #[test]
- fn test_resolve_wrong_catalog_name() {
+ #[tokio::test]
+ async fn test_resolve_wrong_catalog_name() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
+ let sql_context = make_sql_context(catalog).await;
let dialect = GenericDialect {};
let stmts = Parser::parse_sql(&dialect, "SELECT * FROM other.mydb.mytable").unwrap();
if let Statement::Query(q) = &stmts[0] {
@@ -1806,17 +1920,17 @@ mod tests {
if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
&sel.from[0].relation
{
- let err = handler.resolve_table_name(name).unwrap_err();
+ let err = sql_context.resolve_table_name(name).unwrap_err();
assert!(err.to_string().contains("Unknown catalog"));
}
}
}
}
- #[test]
- fn test_resolve_single_part_name_error() {
+ #[tokio::test]
+ async fn test_resolve_single_part_name_uses_default_schema() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
+ let sql_context = make_sql_context(catalog).await;
let dialect = GenericDialect {};
let stmts = Parser::parse_sql(&dialect, "SELECT * FROM mytable").unwrap();
if let Statement::Query(q) = &stmts[0] {
@@ -1824,8 +1938,9 @@ mod tests {
if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
&sel.from[0].relation
{
- let err = handler.resolve_table_name(name).unwrap_err();
- assert!(err.to_string().contains("at least database.table"));
+ let id = sql_context.resolve_table_name(name).unwrap();
+ assert_eq!(id.database(), "default");
+ assert_eq!(id.object(), "mytable");
}
}
}
@@ -1855,14 +1970,14 @@ mod tests {
}
}
- // ==================== PaimonSqlHandler::sql integration tests ====================
+ // ==================== SQLContext::sql integration tests ====================
#[tokio::test]
async fn test_create_table_basic() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE TABLE mydb.t1 (id INT NOT NULL, name VARCHAR, PRIMARY KEY (id))")
.await
.unwrap();
@@ -1887,9 +2002,9 @@ mod tests {
#[tokio::test]
async fn test_create_table_if_not_exists() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE TABLE IF NOT EXISTS mydb.t1 (id INT)")
.await
.unwrap();
@@ -1909,9 +2024,9 @@ mod tests {
#[tokio::test]
async fn test_create_table_with_options() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE TABLE mydb.t1 (id INT) WITH ('bucket' = '4', 'file.format' = 'parquet')")
.await
.unwrap();
@@ -1930,9 +2045,9 @@ mod tests {
#[tokio::test]
async fn test_create_table_three_part_name() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.mydb.t1 (id INT)")
.await
.unwrap();
@@ -1949,9 +2064,9 @@ mod tests {
#[tokio::test]
async fn test_create_table_blob_type_preserved() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE TABLE mydb.t1 (id INT, payload BLOB NOT NULL) WITH ('data-evolution.enabled' = 'true')")
.await
.unwrap();
@@ -1973,9 +2088,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_add_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 ADD COLUMN age INT")
.await
.unwrap();
@@ -2002,9 +2117,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_add_blob_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 ADD COLUMN payload BLOB")
.await
.unwrap();
@@ -2029,9 +2144,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_drop_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 DROP COLUMN age")
.await
.unwrap();
@@ -2051,9 +2166,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_rename_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 RENAME COLUMN old_name TO new_name")
.await
.unwrap();
@@ -2075,9 +2190,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_rename_table() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 RENAME TO t2")
.await
.unwrap();
@@ -2097,9 +2212,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_if_exists_add_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE IF EXISTS mydb.t1 ADD COLUMN age INT")
.await
.unwrap();
@@ -2120,9 +2235,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_without_if_exists() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE mydb.t1 ADD COLUMN age INT")
.await
.unwrap();
@@ -2142,9 +2257,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_if_exists_rename() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE IF EXISTS mydb.t1 RENAME TO t2")
.await
.unwrap();
@@ -2168,9 +2283,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_rename_three_part_name() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("ALTER TABLE paimon.mydb.t1 RENAME TO t2")
.await
.unwrap();
@@ -2190,8 +2305,8 @@ mod tests {
#[tokio::test]
async fn test_sql_parse_error() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- let result = handler.sql("NOT VALID SQL !!!").await;
+ let sql_context = make_sql_context(catalog).await;
+ let result = sql_context.sql("NOT VALID SQL !!!").await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("SQL parse error"));
}
@@ -2199,8 +2314,8 @@ mod tests {
#[tokio::test]
async fn test_multiple_statements_error() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- let result = handler.sql("SELECT 1; SELECT 2").await;
+ let sql_context = make_sql_context(catalog).await;
+ let result = sql_context.sql("SELECT 1; SELECT 2").await;
assert!(result.is_err());
assert!(result
.unwrap_err()
@@ -2211,8 +2326,8 @@ mod tests {
#[tokio::test]
async fn test_create_external_table_rejected() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- let result = handler
+ let sql_context = make_sql_context(catalog).await;
+ let result = sql_context
.sql("CREATE EXTERNAL TABLE mydb.t1 (id INT) STORED AS PARQUET")
.await;
assert!(result.is_err());
@@ -2225,9 +2340,9 @@ mod tests {
#[tokio::test]
async fn test_non_ddl_delegates_to_datafusion() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
+ let sql_context = make_sql_context(catalog.clone()).await;
// SELECT should be delegated to DataFusion, not intercepted
- let df = handler.sql("SELECT 1 AS x").await.unwrap();
+ let df = sql_context.sql("SELECT 1 AS x").await.unwrap();
let batches = df.collect().await.unwrap();
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 1);
@@ -2370,8 +2485,8 @@ mod tests {
#[tokio::test]
async fn test_create_table_partition_key_not_in_columns() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- let err = handler
+ let sql_context = make_sql_context(catalog).await;
+ let err = sql_context
.sql("CREATE TABLE mydb.t (id INT, dt STRING) PARTITIONED BY (nonexistent)")
.await
.unwrap_err();
@@ -2381,8 +2496,8 @@ mod tests {
#[tokio::test]
async fn test_create_table_partition_key_matches_column() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog.clone());
- handler
+ let sql_context = make_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE TABLE mydb.t (id INT, dt STRING) PARTITIONED BY (dt)")
.await
.unwrap();
@@ -2400,68 +2515,73 @@ mod tests {
#[tokio::test]
async fn test_set_paimon_option() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- handler
+ let sql_context = make_sql_context(catalog).await;
+ sql_context
.sql("SET 'paimon.scan.version' = '1'")
.await
.unwrap();
- let opts = handler.dynamic_options().read().unwrap();
+ let opts = sql_context.dynamic_options().read().unwrap();
assert_eq!(opts.get("scan.version").unwrap(), "1");
}
#[tokio::test]
async fn test_set_paimon_option_overwrites() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- handler
+ let sql_context = make_sql_context(catalog).await;
+ sql_context
.sql("SET 'paimon.scan.version' = '1'")
.await
.unwrap();
- handler
+ sql_context
.sql("SET 'paimon.scan.version' = '2'")
.await
.unwrap();
- let opts = handler.dynamic_options().read().unwrap();
+ let opts = sql_context.dynamic_options().read().unwrap();
assert_eq!(opts.get("scan.version").unwrap(), "2");
}
#[tokio::test]
async fn test_reset_paimon_option() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- handler
+ let sql_context = make_sql_context(catalog).await;
+ sql_context
.sql("SET 'paimon.scan.version' = '1'")
.await
.unwrap();
- handler.sql("RESET 'paimon.scan.version'").await.unwrap();
- let opts = handler.dynamic_options().read().unwrap();
+ sql_context
+ .sql("RESET 'paimon.scan.version'")
+ .await
+ .unwrap();
+ let opts = sql_context.dynamic_options().read().unwrap();
assert!(opts.get("scan.version").is_none());
}
#[tokio::test]
async fn test_set_non_paimon_option_delegates() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
+ let sql_context = make_sql_context(catalog).await;
// DataFusion handles non-paimon SET; should not error and should not
// appear in dynamic_options.
- let _ = handler.sql("SET datafusion.optimizer.max_passes = 3").await;
- let opts = handler.dynamic_options().read().unwrap();
+ let _ = sql_context
+ .sql("SET datafusion.optimizer.max_passes = 3")
+ .await;
+ let opts = sql_context.dynamic_options().read().unwrap();
assert!(opts.is_empty());
}
#[tokio::test]
async fn test_set_multiple_paimon_options() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- handler
+ let sql_context = make_sql_context(catalog).await;
+ sql_context
.sql("SET 'paimon.scan.version' = '1'")
.await
.unwrap();
- handler
+ sql_context
.sql("SET 'paimon.scan.timestamp-millis' = '1000'")
.await
.unwrap();
- let opts = handler.dynamic_options().read().unwrap();
+ let opts = sql_context.dynamic_options().read().unwrap();
assert_eq!(opts.len(), 2);
assert_eq!(opts.get("scan.version").unwrap(), "1");
assert_eq!(opts.get("scan.timestamp-millis").unwrap(), "1000");
@@ -2470,15 +2590,18 @@ mod tests {
#[tokio::test]
async fn test_reset_nonexistent_paimon_option_is_noop() {
let catalog = Arc::new(MockCatalog::new());
- let handler = make_handler(catalog);
- handler.sql("RESET 'paimon.scan.version'").await.unwrap();
- let opts = handler.dynamic_options().read().unwrap();
+ let sql_context = make_sql_context(catalog).await;
+ sql_context
+ .sql("RESET 'paimon.scan.version'")
+ .await
+ .unwrap();
+ let opts = sql_context.dynamic_options().read().unwrap();
assert!(opts.is_empty());
}
// ==================== TRUNCATE TABLE / DROP PARTITIONS tests ====================
- async fn setup_fs_handler() -> (tempfile::TempDir, PaimonSqlHandler) {
+ async fn setup_fs_sql_context() -> (tempfile::TempDir, SQLContext) {
use paimon::{CatalogOptions, FileSystemCatalog, Options};
let temp_dir = tempfile::TempDir::new().unwrap();
@@ -2487,22 +2610,28 @@ mod tests {
options.set(CatalogOptions::WAREHOUSE, warehouse);
let catalog = Arc::new(FileSystemCatalog::new(options).unwrap());
- let handler =
- PaimonSqlHandler::new(SessionContext::new(), catalog.clone(), "paimon").unwrap();
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
+ let mut sql_context = SQLContext::new();
+ sql_context
+ .register_catalog("paimon", catalog.clone())
+ .await
+ .unwrap();
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
- (temp_dir, handler)
+ (temp_dir, sql_context)
}
#[tokio::test]
async fn test_truncate_table() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.test_db.t1 (id INT, value INT)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t1 VALUES (1, 10), (2, 20)")
.await
.unwrap()
@@ -2510,12 +2639,12 @@ mod tests {
.await
.unwrap();
- handler
+ sql_context
.sql("TRUNCATE TABLE paimon.test_db.t1")
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT * FROM paimon.test_db.t1")
.await
.unwrap()
@@ -2528,13 +2657,13 @@ mod tests {
#[tokio::test]
async fn test_truncate_table_partition() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.test_db.t2 (pt VARCHAR, id INT) PARTITIONED BY (pt)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t2 VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
.await
.unwrap()
@@ -2542,12 +2671,12 @@ mod tests {
.await
.unwrap();
- handler
+ sql_context
.sql("TRUNCATE TABLE paimon.test_db.t2 PARTITION (pt = 'a')")
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT pt, id FROM paimon.test_db.t2 ORDER BY id")
.await
.unwrap()
@@ -2576,13 +2705,13 @@ mod tests {
#[tokio::test]
async fn test_alter_table_drop_partitions() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.test_db.t3 (pt VARCHAR, id INT) PARTITIONED BY (pt)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t3 VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
.await
.unwrap()
@@ -2590,12 +2719,12 @@ mod tests {
.await
.unwrap();
- handler
+ sql_context
.sql("ALTER TABLE paimon.test_db.t3 DROP PARTITION (pt = 'b')")
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT pt, id FROM paimon.test_db.t3 ORDER BY id")
.await
.unwrap()
@@ -2624,13 +2753,13 @@ mod tests {
#[tokio::test]
async fn test_truncate_table_incomplete_partition_spec() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.test_db.t_multi (pt1 VARCHAR, pt2 VARCHAR, id INT) PARTITIONED BY (pt1, pt2)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_multi VALUES ('a', 'x', 1)")
.await
.unwrap()
@@ -2638,7 +2767,7 @@ mod tests {
.await
.unwrap();
- let err = handler
+ let err = sql_context
.sql("TRUNCATE TABLE paimon.test_db.t_multi PARTITION (pt1 = 'a')")
.await
.unwrap_err();
@@ -2650,9 +2779,9 @@ mod tests {
#[tokio::test]
async fn test_truncate_table_if_exists_nonexistent() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("TRUNCATE TABLE IF EXISTS paimon.test_db.nonexistent")
.await
.unwrap();
@@ -2660,9 +2789,9 @@ mod tests {
#[tokio::test]
async fn test_truncate_table_nonexistent_without_if_exists() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- let err = handler
+ let err = sql_context
.sql("TRUNCATE TABLE paimon.test_db.nonexistent")
.await
.unwrap_err();
@@ -2674,9 +2803,9 @@ mod tests {
#[tokio::test]
async fn test_alter_table_if_exists_drop_partition_nonexistent() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("ALTER TABLE IF EXISTS paimon.test_db.nonexistent DROP PARTITION (pt = 'a')")
.await
.unwrap();
@@ -2684,13 +2813,13 @@ mod tests {
#[tokio::test]
async fn test_drop_partition_incomplete_spec() {
- let (_tmp, handler) = setup_fs_handler().await;
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
- handler
+ sql_context
.sql("CREATE TABLE paimon.test_db.t_dp (pt1 VARCHAR, pt2 VARCHAR, id INT) PARTITIONED BY (pt1, pt2)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dp VALUES ('a', 'x', 1)")
.await
.unwrap()
@@ -2698,7 +2827,7 @@ mod tests {
.await
.unwrap();
- let err = handler
+ let err = sql_context
.sql("ALTER TABLE paimon.test_db.t_dp DROP PARTITION (pt1 = 'a')")
.await
.unwrap_err();
diff --git a/crates/integrations/datafusion/src/update.rs b/crates/integrations/datafusion/src/update.rs
index 9bd9445..741462b 100644
--- a/crates/integrations/datafusion/src/update.rs
+++ b/crates/integrations/datafusion/src/update.rs
@@ -204,7 +204,7 @@ async fn execute_cow_update_once(
ctx,
&columns,
&exprs,
- cow_table_guard.name(),
+ &cow_table_guard.qualified_name(),
update,
&mut writer,
)
@@ -335,33 +335,39 @@ mod tests {
use paimon::{CatalogOptions, FileSystemCatalog, Options};
use tempfile::TempDir;
- use crate::{PaimonSqlHandler, PaimonTableProvider};
+ use crate::{PaimonTableProvider, SQLContext};
- async fn setup_handler() -> (TempDir, PaimonSqlHandler, Arc<FileSystemCatalog>) {
+ async fn setup_sql_context() -> (TempDir, SQLContext, Arc<FileSystemCatalog>) {
let temp_dir = TempDir::new().unwrap();
let warehouse = format!("file://{}", temp_dir.path().display());
let mut options = Options::new();
options.set(CatalogOptions::WAREHOUSE, warehouse);
let catalog = Arc::new(FileSystemCatalog::new(options).unwrap());
- let handler =
- PaimonSqlHandler::new(SessionContext::new(), catalog.clone(), "paimon").unwrap();
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
+ let mut sql_context = SQLContext::new();
+ sql_context
+ .register_catalog("paimon", catalog.clone())
+ .await
+ .unwrap();
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
- (temp_dir, handler, catalog)
+ (temp_dir, sql_context, catalog)
}
async fn setup_data_evolution_table(name: &str) -> (TempDir, SessionContext, Table) {
- let (tmp, handler, catalog) = setup_handler().await;
+ let (tmp, sql_context, catalog) = setup_sql_context().await;
- handler
+ sql_context
.sql(&format!(
"CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('row-tracking.enabled' = 'true')"
))
.await
.unwrap();
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)"
))
@@ -380,10 +386,10 @@ mod tests {
extra.insert("row-tracking.enabled".to_string(), "true".to_string());
let de_table = table.copy_with_options(extra);
- // Register the data-evolution table under a simple name so _ROW_ID is visible
- let ctx = handler.ctx().clone();
+ let ctx = sql_context.ctx().clone();
let provider = PaimonTableProvider::try_new(de_table.clone()).unwrap();
- ctx.register_table("target", Arc::new(provider)).unwrap();
+ ctx.register_table("datafusion.public.target", Arc::new(provider))
+ .unwrap();
(tmp, ctx, de_table)
}
@@ -426,11 +432,12 @@ mod tests {
async fn test_update_with_where() {
let (_tmp, ctx, table) = setup_data_evolution_table("t_with_where").await;
- let update = parse_update("UPDATE target SET name = 'ALICE' WHERE id = 1");
+ let update =
+ parse_update("UPDATE datafusion.public.target SET name = 'ALICE' WHERE id = 1");
execute_update(&ctx, &update, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -452,11 +459,11 @@ mod tests {
async fn test_update_without_where() {
let (_tmp, ctx, table) = setup_data_evolution_table("t_without_where").await;
- let update = parse_update("UPDATE target SET value = 99");
+ let update = parse_update("UPDATE datafusion.public.target SET value = 99");
execute_update(&ctx, &update, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -478,11 +485,13 @@ mod tests {
async fn test_update_multiple_columns() {
let (_tmp, ctx, table) = setup_data_evolution_table("t_multi_col").await;
- let update = parse_update("UPDATE target SET name = 'updated', value = 0 WHERE id = 2");
+ let update = parse_update(
+ "UPDATE datafusion.public.target SET name = 'updated', value = 0 WHERE id = 2",
+ );
execute_update(&ctx, &update, table).await.unwrap();
let batches = ctx
- .sql("SELECT id, name, value FROM target ORDER BY id")
+ .sql("SELECT id, name, value FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -504,7 +513,8 @@ mod tests {
async fn test_update_no_matching_rows() {
let (_tmp, ctx, table) = setup_data_evolution_table("t_no_match").await;
- let update = parse_update("UPDATE target SET name = 'nobody' WHERE id = 99");
+ let update =
+ parse_update("UPDATE datafusion.public.target SET name = 'nobody' WHERE id = 99");
let result = execute_update(&ctx, &update, table).await.unwrap();
let batches = result.collect().await.unwrap();
let count = batches[0]
@@ -522,19 +532,20 @@ mod tests {
// Get row IDs before update
let before = ctx
- .sql("SELECT id, \"_ROW_ID\" FROM target ORDER BY id")
+ .sql("SELECT id, \"_ROW_ID\" FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
.await
.unwrap();
- let update = parse_update("UPDATE target SET name = 'ALICE' WHERE id = 1");
+ let update =
+ parse_update("UPDATE datafusion.public.target SET name = 'ALICE' WHERE id = 1");
execute_update(&ctx, &update, table).await.unwrap();
// Get row IDs after update
let after = ctx
- .sql("SELECT id, \"_ROW_ID\" FROM target ORDER BY id")
+ .sql("SELECT id, \"_ROW_ID\" FROM datafusion.public.target ORDER BY id")
.await
.unwrap()
.collect()
@@ -587,17 +598,17 @@ mod tests {
// CoW UPDATE tests (append-only tables)
// -----------------------------------------------------------------------
- async fn setup_append_only_table(name: &str) -> (TempDir, PaimonSqlHandler) {
- let (tmp, handler, _catalog) = setup_handler().await;
+ async fn setup_append_only_table(name: &str) -> (TempDir, SQLContext) {
+ let (tmp, sql_context, _catalog) = setup_sql_context().await;
- handler
+ sql_context
.sql(&format!(
"CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT)"
))
.await
.unwrap();
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)"
))
@@ -607,11 +618,11 @@ mod tests {
.await
.unwrap();
- (tmp, handler)
+ (tmp, sql_context)
}
- async fn query_rows(handler: &PaimonSqlHandler, table: &str) -> Vec<(i32, String, i32)> {
- let batches = handler
+ async fn query_rows(sql_context: &SQLContext, table: &str) -> Vec<(i32, String, i32)> {
+ let batches = sql_context
.sql(&format!("SELECT id, name, value FROM {table} ORDER BY id"))
.await
.unwrap()
@@ -623,14 +634,14 @@ mod tests {
#[tokio::test]
async fn test_cow_update_with_where() {
- let (_tmp, handler) = setup_append_only_table("t_cow_where").await;
+ let (_tmp, sql_context) = setup_append_only_table("t_cow_where").await;
- handler
+ sql_context
.sql("UPDATE paimon.test_db.t_cow_where SET name = 'ALICE' WHERE id = 1")
.await
.unwrap();
- let rows = query_rows(&handler, "paimon.test_db.t_cow_where").await;
+ let rows = query_rows(&sql_context, "paimon.test_db.t_cow_where").await;
assert_eq!(
rows,
vec![
@@ -643,14 +654,14 @@ mod tests {
#[tokio::test]
async fn test_cow_update_without_where() {
- let (_tmp, handler) = setup_append_only_table("t_cow_no_where").await;
+ let (_tmp, sql_context) = setup_append_only_table("t_cow_no_where").await;
- handler
+ sql_context
.sql("UPDATE paimon.test_db.t_cow_no_where SET value = 99")
.await
.unwrap();
- let rows = query_rows(&handler, "paimon.test_db.t_cow_no_where").await;
+ let rows = query_rows(&sql_context, "paimon.test_db.t_cow_no_where").await;
assert_eq!(
rows,
vec![
@@ -663,14 +674,14 @@ mod tests {
#[tokio::test]
async fn test_cow_update_multiple_columns() {
- let (_tmp, handler) = setup_append_only_table("t_cow_multi").await;
+ let (_tmp, sql_context) = setup_append_only_table("t_cow_multi").await;
- handler
+ sql_context
.sql("UPDATE paimon.test_db.t_cow_multi SET name = 'updated', value = 0 WHERE id = 2")
.await
.unwrap();
- let rows = query_rows(&handler, "paimon.test_db.t_cow_multi").await;
+ let rows = query_rows(&sql_context, "paimon.test_db.t_cow_multi").await;
assert_eq!(
rows,
vec![
@@ -683,9 +694,9 @@ mod tests {
#[tokio::test]
async fn test_cow_update_no_matching_rows() {
- let (_tmp, handler) = setup_append_only_table("t_cow_nomatch").await;
+ let (_tmp, sql_context) = setup_append_only_table("t_cow_nomatch").await;
- let result = handler
+ let result = sql_context
.sql("UPDATE paimon.test_db.t_cow_nomatch SET name = 'nobody' WHERE id = 99")
.await
.unwrap()
@@ -703,14 +714,14 @@ mod tests {
#[tokio::test]
async fn test_cow_update_expression() {
- let (_tmp, handler) = setup_append_only_table("t_cow_expr").await;
+ let (_tmp, sql_context) = setup_append_only_table("t_cow_expr").await;
- handler
+ sql_context
.sql("UPDATE paimon.test_db.t_cow_expr SET value = value + 100 WHERE id >= 2")
.await
.unwrap();
- let rows = query_rows(&handler, "paimon.test_db.t_cow_expr").await;
+ let rows = query_rows(&sql_context, "paimon.test_db.t_cow_expr").await;
assert_eq!(
rows,
vec![
diff --git a/crates/integrations/datafusion/tests/append_merge_into.rs b/crates/integrations/datafusion/tests/append_merge_into.rs
index 89b6e30..9e8dc0a 100644
--- a/crates/integrations/datafusion/tests/append_merge_into.rs
+++ b/crates/integrations/datafusion/tests/append_merge_into.rs
@@ -23,25 +23,28 @@
mod common;
use arrow_array::{Array, Int32Array, StringArray};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
use common::{
- collect_int_int_str, collect_int_str, collect_three_ints, create_handler, create_test_env,
+ collect_int_int_str, collect_int_str, collect_three_ints, create_sql_context, create_test_env,
ctx_exec, exec,
};
// ======================= Helpers =======================
-async fn setup(table_ddl: &str) -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup(table_ddl: &str) -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler.sql(table_ddl).await.unwrap();
- (tmp, handler)
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context.sql(table_ddl).await.unwrap();
+ (tmp, sql_context)
}
-async fn query_abc(handler: &PaimonSqlHandler) -> Vec<(i32, i32, String)> {
- let batches = handler
+async fn query_abc(sql_context: &SQLContext) -> Vec<(i32, i32, String)> {
+ let batches = sql_context
.sql("SELECT a, b, c FROM paimon.test_db.target ORDER BY a, b")
.await
.unwrap()
@@ -51,24 +54,24 @@ async fn query_abc(handler: &PaimonSqlHandler) -> Vec<(i32, i32, String)> {
collect_int_int_str(&batches)
}
-async fn setup_abc() -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup_abc() -> (tempfile::TempDir, SQLContext) {
setup("CREATE TABLE paimon.test_db.target (a INT, b INT, c VARCHAR)").await
}
-async fn setup_partitioned() -> (tempfile::TempDir, PaimonSqlHandler) {
- let (tmp, handler) =
+async fn setup_partitioned() -> (tempfile::TempDir, SQLContext) {
+ let (tmp, sql_context) =
setup("CREATE TABLE paimon.test_db.target (a INT, b INT, pt INT) PARTITIONED BY (pt)")
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 1), (2, 20, 1), (3, 30, 2), (4, 40, 2)",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
-async fn query_a_b_pt(handler: &PaimonSqlHandler) -> Vec<(i32, i32, i32)> {
- let batches = handler
+async fn query_a_b_pt(sql_context: &SQLContext) -> Vec<(i32, i32, i32)> {
+ let batches = sql_context
.sql("SELECT a, b, pt FROM paimon.test_db.target ORDER BY pt, a")
.await
.unwrap()
@@ -83,28 +86,28 @@ async fn query_a_b_pt(handler: &PaimonSqlHandler) -> Vec<(i32, i32, i32)> {
/// Paimon MergeInto: only update
#[tokio::test]
async fn test_only_update() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
- WHEN MATCHED THEN UPDATE SET a = source.a, b = source.b, c = source.c",
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.source s ON t.a = s.a \
+ WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c11".into()), (2, 20, "c2".into()),]
);
}
@@ -112,54 +115,54 @@ async fn test_only_update() {
/// Paimon MergeInto: only delete
#[tokio::test]
async fn test_only_delete() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN MATCHED THEN DELETE",
)
.await;
- assert_eq!(query_abc(&handler).await, vec![(2, 20, "c2".into()),]);
+ assert_eq!(query_abc(&sql_context).await, vec![(2, 20, "c2".into()),]);
}
/// Paimon MergeInto: only insert
#[tokio::test]
async fn test_only_insert() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
- WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (a, b, c)",
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.source s ON t.a = s.a \
+ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 10, "c1".into()),
(2, 20, "c2".into()),
@@ -171,29 +174,29 @@ async fn test_only_insert() {
/// Paimon MergeInto: update + insert
#[tokio::test]
async fn test_update_and_insert() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
- WHEN MATCHED THEN UPDATE SET a = source.a, b = source.b, c = source.c \
- WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (a, b, c)",
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 100, "c11".into()),
(2, 20, "c2".into()),
@@ -205,29 +208,29 @@ async fn test_update_and_insert() {
/// Paimon MergeInto: delete + insert
#[tokio::test]
async fn test_delete_and_insert() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN MATCHED THEN DELETE \
- WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (a, b, c)",
+ WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(2, 20, "c2".into()), (3, 300, "c33".into()),]
);
}
@@ -235,27 +238,27 @@ async fn test_delete_and_insert() {
/// Paimon MergeInto: partial insert with null
#[tokio::test]
async fn test_partial_insert_with_null() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
- "MERGE INTO paimon.test_db.target \
- USING source ON paimon.test_db.target.a = source.a \
- WHEN NOT MATCHED THEN INSERT (a) VALUES (a)",
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING datafusion.public.source s ON t.a = s.a \
+ WHEN NOT MATCHED THEN INSERT (a) VALUES (s.a)",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT a, b, c FROM paimon.test_db.target ORDER BY a, b")
.await
.unwrap()
@@ -307,29 +310,29 @@ async fn test_partial_insert_with_null() {
/// Paimon MergeInto: update value from both source and target table
#[tokio::test]
async fn test_update_from_both_source_and_target() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 110, "c11".into()),
(2, 20, "c2".into()),
@@ -341,29 +344,29 @@ async fn test_update_from_both_source_and_target() {
/// Paimon MergeInto: insert/update columns in wrong order
#[tokio::test]
async fn test_columns_in_wrong_order() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 100, "c11".into()),
(2, 20, "c2".into()),
@@ -375,29 +378,29 @@ async fn test_columns_in_wrong_order() {
/// Paimon MergeInto: miss some columns in update
#[tokio::test]
async fn test_partial_update() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 10, "c11".into()),
(2, 20, "c2".into()),
@@ -409,18 +412,18 @@ async fn test_partial_update() {
/// Paimon MergeInto: source is a query (subquery)
#[tokio::test]
async fn test_source_is_subquery() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
- ctx_exec(&handler, "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33'), (4, 400, 'c44')").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(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING (SELECT a, b, c FROM source WHERE a % 2 = 1) AS src \
+ USING (SELECT a, b, c FROM datafusion.public.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)",
@@ -428,7 +431,7 @@ async fn test_source_is_subquery() {
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 100, "c11".into()),
(2, 20, "c2".into()),
@@ -440,48 +443,51 @@ async fn test_source_is_subquery() {
/// Paimon MergeInto: source and target are empty
#[tokio::test]
async fn test_source_and_target_empty() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
// target is empty, source is empty
- ctx_exec(&handler, "CREATE TABLE source (a INT, b INT, c VARCHAR) AS SELECT * FROM (VALUES (1, 1, 'x')) AS t(a, b, c) WHERE 1=0").await;
+ 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(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
.await;
- assert_eq!(query_abc(&handler).await, Vec::<(i32, i32, String)>::new());
+ assert_eq!(
+ query_abc(&sql_context).await,
+ Vec::<(i32, i32, String)>::new()
+ );
}
/// Paimon MergeInto: merge into with alias
#[tokio::test]
async fn test_with_alias() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 100, "c11".into()),
(2, 20, "c2".into()),
@@ -493,28 +499,28 @@ async fn test_with_alias() {
/// Paimon MergeInto: update on source eq target condition (reversed ON clause)
#[tokio::test]
async fn test_reversed_on_condition() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON s.a = t.a \
+ USING datafusion.public.source s ON s.a = t.a \
WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b, c = s.c",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c11".into()), (2, 20, "c2".into()),]
);
}
@@ -522,24 +528,24 @@ async fn test_reversed_on_condition() {
/// Paimon MergeInto: two paimon tables
#[tokio::test]
async fn test_two_paimon_tables() {
- let (_tmp, handler) = setup_abc().await;
- handler
+ let (_tmp, sql_context) = setup_abc().await;
+ sql_context
.sql("CREATE TABLE paimon.test_db.source (a INT, b INT, c VARCHAR)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.source VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
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",
@@ -547,7 +553,7 @@ async fn test_two_paimon_tables() {
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c11".into()), (2, 20, "c2".into()),]
);
}
@@ -555,20 +561,20 @@ async fn test_two_paimon_tables() {
/// Paimon MergeInto: on clause has filter expression
#[tokio::test]
async fn test_on_clause_with_filter() {
- let (_tmp, handler) = setup_abc().await;
- handler
+ let (_tmp, sql_context) = setup_abc().await;
+ sql_context
.sql("CREATE TABLE paimon.test_db.source (a INT, b INT, c VARCHAR)")
.await
.unwrap();
- exec(&handler, "INSERT INTO paimon.test_db.source VALUES (1, 100, 'c11'), (3, 300, 'c11'), (5, 500, 'c55')").await;
+ exec(&sql_context, "INSERT INTO paimon.test_db.source VALUES (1, 100, 'c11'), (3, 300, 'c11'), (5, 500, 'c55')").await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 100, 'cc'), (2, 20, 'cc')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target tgt \
USING (SELECT a, b FROM paimon.test_db.source WHERE c = 'c11') AS src \
ON tgt.a = src.a AND tgt.b = src.b AND tgt.c = 'cc' \
@@ -576,37 +582,37 @@ async fn test_on_clause_with_filter() {
)
.await;
- assert_eq!(query_abc(&handler).await, vec![(2, 20, "cc".into()),]);
+ assert_eq!(query_abc(&sql_context).await, vec![(2, 20, "cc".into()),]);
}
/// Paimon MergeInto: merge into with varchar type
#[tokio::test]
async fn test_with_varchar() {
- let (_tmp, handler) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
- handler
+ let (_tmp, sql_context) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
+ sql_context
.sql("CREATE TABLE paimon.test_db.source (a INT, b VARCHAR)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 'Alice'), (2, 'Bob')",
)
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.source VALUES (1, 'Eve'), (3, 'Cat')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
USING paimon.test_db.source s ON t.a = s.a \
WHEN MATCHED THEN UPDATE SET a = s.a, b = s.b",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT a, b FROM paimon.test_db.target ORDER BY a")
.await
.unwrap()
@@ -622,29 +628,29 @@ async fn test_with_varchar() {
/// Paimon MergeInto: update with coalesce referencing both source and target columns
#[tokio::test]
async fn test_coalesce_source_and_target() {
- let (_tmp, handler) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
+ let (_tmp, sql_context) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 'guid_tgt_1'), (2, 'guid_tgt_2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target AS dest \
- USING source AS src ON dest.a = src.a \
+ USING datafusion.public.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)",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT a, b FROM paimon.test_db.target ORDER BY a")
.await
.unwrap()
@@ -664,29 +670,29 @@ async fn test_coalesce_source_and_target() {
/// Paimon MergeInto: subquery source with coalesce referencing both source and target
#[tokio::test]
async fn test_subquery_source_with_coalesce() {
- let (_tmp, handler) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
+ let (_tmp, sql_context) = setup("CREATE TABLE paimon.test_db.target (a INT, b VARCHAR)").await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 'guid_tgt_1'), (2, 'guid_tgt_2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b VARCHAR) AS VALUES (1, 'guid_src_1'), (3, 'guid_src_3')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target AS dest \
- USING (SELECT * FROM source) AS src ON dest.a = src.a \
+ USING (SELECT * FROM datafusion.public.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)",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT a, b FROM paimon.test_db.target ORDER BY a")
.await
.unwrap()
@@ -706,28 +712,28 @@ async fn test_subquery_source_with_coalesce() {
/// Paimon MergeInto: non pk table insert-only commit kind is APPEND
#[tokio::test]
async fn test_insert_only_is_append_commit() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (2, 2, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 1, 'c1')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 1, 'c1')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 1, "c1".into()), (2, 2, "c2".into()),]
);
}
@@ -735,48 +741,48 @@ async fn test_insert_only_is_append_commit() {
/// Paimon MergeInto: successive merges on append-only table
#[tokio::test]
async fn test_successive_merges() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
// First merge: update b for a=1
ctx_exec(
- &handler,
- "CREATE TABLE src1 (a INT, b INT) AS VALUES (1, 100)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src1 (a INT, b INT) AS VALUES (1, 100)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING src1 s ON t.a = s.a \
+ USING datafusion.public.src1 s ON t.a = s.a \
WHEN MATCHED THEN UPDATE SET b = s.b",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c1".into()), (2, 20, "c2".into()),]
);
// Second merge: update c for a=2
ctx_exec(
- &handler,
- "CREATE TABLE src2 (a INT, c VARCHAR) AS VALUES (2, 'C2_UPDATED')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src2 (a INT, c VARCHAR) AS VALUES (2, 'C2_UPDATED')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING src2 s ON t.a = s.a \
+ USING datafusion.public.src2 s ON t.a = s.a \
WHEN MATCHED THEN UPDATE SET c = s.c",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c1".into()), (2, 20, "C2_UPDATED".into()),]
);
}
@@ -784,28 +790,28 @@ async fn test_successive_merges() {
/// Paimon MergeInto: no match produces no changes
#[tokio::test]
async fn test_no_match_no_change() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (99, 990, 'c99')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (99, 990, 'c99')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN MATCHED THEN UPDATE SET b = s.b, c = s.c",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 10, "c1".into()), (2, 20, "c2".into()),]
);
}
@@ -813,46 +819,53 @@ async fn test_no_match_no_change() {
/// Paimon MergeInto: delete all matched rows
#[tokio::test]
async fn test_delete_all_rows() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
- ctx_exec(&handler, "CREATE TABLE source (a INT) AS VALUES (1), (2)").await;
+ ctx_exec(
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT) AS VALUES (1), (2)",
+ )
+ .await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN MATCHED THEN DELETE",
)
.await;
- assert_eq!(query_abc(&handler).await, Vec::<(i32, i32, String)>::new());
+ assert_eq!(
+ query_abc(&sql_context).await,
+ Vec::<(i32, i32, String)>::new()
+ );
}
/// Paimon MergeInto: multiple inserts from different batches
#[tokio::test]
async fn test_insert_many_rows() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1')",
)
.await;
- ctx_exec(&handler, "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (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 (2, 20, 'c2'), (3, 30, 'c3'), (4, 40, 'c4'), (5, 50, 'c5')").await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.source s ON t.a = s.a \
WHEN NOT MATCHED THEN INSERT (a, b, c) VALUES (s.a, s.b, s.c)",
)
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 10, "c1".into()),
(2, 20, "c2".into()),
@@ -866,22 +879,22 @@ async fn test_insert_many_rows() {
/// Paimon MergeInto: conditional update — rows matching predicate get updated, others get deleted
#[tokio::test]
async fn test_conditional_update() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2'), (3, 30, 'c3')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
@@ -889,7 +902,7 @@ async fn test_conditional_update() {
.await;
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(2, 20, "c2".into()), (3, 300, "c33".into()),]
);
}
@@ -897,22 +910,22 @@ async fn test_conditional_update() {
/// Paimon MergeInto: conditional insert — only insert rows matching predicate
#[tokio::test]
async fn test_conditional_insert() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
@@ -920,7 +933,7 @@ async fn test_conditional_insert() {
// a=1 updated, a=3 NOT inserted (b=300 not < 300)
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![(1, 100, "c11".into()), (2, 20, "c2".into()),]
);
}
@@ -928,22 +941,22 @@ async fn test_conditional_insert() {
/// Paimon MergeInto: conditional delete — only delete rows matching predicate
#[tokio::test]
async fn test_conditional_delete() {
- let (_tmp, handler) = setup_abc().await;
+ let (_tmp, sql_context) = setup_abc().await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.target VALUES (1, 10, 'c1'), (2, 20, 'c2')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, c VARCHAR) AS VALUES (1, 100, 'c11'), (3, 300, 'c33')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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)",
)
@@ -951,7 +964,7 @@ async fn test_conditional_delete() {
// a=1 matched but c='c1' is NOT < 'c1', so not deleted
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(1, 10, "c1".into()),
(2, 20, "c2".into()),
@@ -963,13 +976,13 @@ async fn test_conditional_delete() {
/// Paimon MergeInto: multiple matched clauses with predicates
#[tokio::test]
async fn test_multiple_matched_clauses() {
- let (_tmp, handler) = setup_abc().await;
- exec(&handler, "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(&handler, "CREATE TABLE 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;
+ 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(&handler,
+ exec(&sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a \
+ USING datafusion.public.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 \
@@ -982,7 +995,7 @@ async fn test_multiple_matched_clauses() {
// a=7: not matched, c='c77' not > 'c9' → INSERT *
// a=9: not matched, c='c99' > 'c9' → INSERT with b=990
assert_eq!(
- query_abc(&handler).await,
+ query_abc(&sql_context).await,
vec![
(2, 20, "c2".into()),
(3, 300, "c33".into()),
@@ -998,89 +1011,92 @@ async fn test_multiple_matched_clauses() {
#[tokio::test]
async fn test_partitioned_update_single_partition() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, pt INT) AS VALUES (1, 100, 1)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \
WHEN MATCHED THEN UPDATE SET b = s.b",
)
.await;
assert_eq!(
- query_a_b_pt(&handler).await,
+ query_a_b_pt(&sql_context).await,
vec![(1, 100, 1), (2, 20, 1), (3, 30, 2), (4, 40, 2),]
);
}
#[tokio::test]
async fn test_partitioned_update_multiple_partitions() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (3, 300, 2)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (3, 300, 2)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \
WHEN MATCHED THEN UPDATE SET b = s.b",
)
.await;
assert_eq!(
- query_a_b_pt(&handler).await,
+ query_a_b_pt(&sql_context).await,
vec![(1, 100, 1), (2, 20, 1), (3, 300, 2), (4, 40, 2),]
);
}
#[tokio::test]
async fn test_partitioned_delete() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, pt INT) AS VALUES (1, 1), (3, 2)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, pt INT) AS VALUES (1, 1), (3, 2)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.source s ON t.a = s.a AND t.pt = s.pt \
WHEN MATCHED THEN DELETE",
)
.await;
- assert_eq!(query_a_b_pt(&handler).await, vec![(2, 20, 1), (4, 40, 2),]);
+ assert_eq!(
+ query_a_b_pt(&sql_context).await,
+ vec![(2, 20, 1), (4, 40, 2),]
+ );
}
#[tokio::test]
async fn test_partitioned_insert_new_partition() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, pt INT) AS VALUES (5, 50, 3), (6, 60, 3)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (5, 50, 3), (6, 60, 3)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.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;
assert_eq!(
- query_a_b_pt(&handler).await,
+ query_a_b_pt(&sql_context).await,
vec![
(1, 10, 1),
(2, 20, 1),
@@ -1094,60 +1110,60 @@ async fn test_partitioned_insert_new_partition() {
#[tokio::test]
async fn test_partitioned_update_and_insert() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (5, 50, 2)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source (a INT, b INT, pt INT) AS VALUES (1, 100, 1), (5, 50, 2)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING source s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.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)",
)
.await;
assert_eq!(
- query_a_b_pt(&handler).await,
+ query_a_b_pt(&sql_context).await,
vec![(1, 100, 1), (2, 20, 1), (3, 30, 2), (4, 40, 2), (5, 50, 2),]
);
}
#[tokio::test]
async fn test_partitioned_successive_merges() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
ctx_exec(
- &handler,
- "CREATE TABLE src1 (a INT, b INT, pt INT) AS VALUES (1, 100, 1)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src1 (a INT, b INT, pt INT) AS VALUES (1, 100, 1)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING src1 s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.src1 s ON t.a = s.a AND t.pt = s.pt \
WHEN MATCHED THEN UPDATE SET b = s.b",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE src2 (a INT, b INT, pt INT) AS VALUES (3, 300, 2)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src2 (a INT, b INT, pt INT) AS VALUES (3, 300, 2)",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.target t \
- USING src2 s ON t.a = s.a AND t.pt = s.pt \
+ USING datafusion.public.src2 s ON t.a = s.a AND t.pt = s.pt \
WHEN MATCHED THEN UPDATE SET b = s.b",
)
.await;
assert_eq!(
- query_a_b_pt(&handler).await,
+ query_a_b_pt(&sql_context).await,
vec![(1, 100, 1), (2, 20, 1), (3, 300, 2), (4, 40, 2),]
);
}
diff --git a/crates/integrations/datafusion/tests/blob_tests.rs b/crates/integrations/datafusion/tests/blob_tests.rs
index 3374cb3..286bc6f 100644
--- a/crates/integrations/datafusion/tests/blob_tests.rs
+++ b/crates/integrations/datafusion/tests/blob_tests.rs
@@ -22,9 +22,9 @@
mod common;
use arrow_array::{Array, BinaryArray, Int32Array, RecordBatch, StringArray};
-use common::{create_handler, create_test_env, ctx_exec, exec};
+use common::{create_sql_context, create_test_env, ctx_exec, exec};
use paimon::spec::BlobDescriptor;
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
// ======================= Helpers =======================
@@ -32,12 +32,15 @@ fn to_hex(bytes: &[u8]) -> String {
bytes.iter().map(|b| format!("{b:02X}")).collect()
}
-async fn setup(table_ddl: &str) -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup(table_ddl: &str) -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler.sql(table_ddl).await.unwrap();
- (tmp, handler)
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context.sql(table_ddl).await.unwrap();
+ (tmp, sql_context)
}
fn collect_id_name_picture(batches: &[RecordBatch]) -> Vec<(i32, String, Option<Vec<u8>>)> {
@@ -72,10 +75,10 @@ fn collect_id_name_picture(batches: &[RecordBatch]) -> Vec<(i32, String, Option<
}
async fn query_id_name_picture(
- handler: &PaimonSqlHandler,
+ sql_context: &SQLContext,
sql: &str,
) -> Vec<(i32, String, Option<Vec<u8>>)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
collect_id_name_picture(&batches)
}
@@ -115,16 +118,16 @@ const BLOB_TABLE_DDL: &str = "\
/// Reference: BlobTestBase "Blob: test basic"
#[tokio::test]
async fn test_blob_basic() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Alice', X'48656C6C6F')",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -137,7 +140,7 @@ async fn test_blob_basic() {
/// Reference: BlobTestBase "Blob: test multiple blobs"
#[tokio::test]
async fn test_blob_multiple_columns() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
pic1 BLOB, \
@@ -150,12 +153,12 @@ async fn test_blob_multiple_columns() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, pic1, pic2) VALUES (1, X'414141', X'424242')",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT id, pic1, pic2 FROM paimon.test_db.t")
.await
.unwrap()
@@ -188,10 +191,10 @@ async fn test_blob_multiple_columns() {
/// Blob with NULL values.
#[tokio::test]
async fn test_blob_with_nulls() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'48656C6C6F'), \
(2, 'Bob', NULL), \
@@ -200,7 +203,7 @@ async fn test_blob_with_nulls() {
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -217,21 +220,21 @@ async fn test_blob_with_nulls() {
/// Multiple inserts produce multiple file pairs, all readable.
#[tokio::test]
async fn test_blob_multiple_inserts() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Alice', X'4141')",
)
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (2, 'Bob', X'4242')",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -247,7 +250,7 @@ async fn test_blob_multiple_inserts() {
/// blob-descriptor-field: listed fields are stored inline in parquet (no .blob files).
#[tokio::test]
async fn test_blob_descriptor_field_inline() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -261,13 +264,13 @@ async fn test_blob_descriptor_field_inline() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Alice', X'48656C6C6F')",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -281,10 +284,10 @@ async fn test_blob_descriptor_field_inline() {
/// and the blob data should be preserved.
#[tokio::test]
async fn test_merge_into_updates_non_blob_on_raw_blob_table() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'4141'), \
(2, 'Bob', X'4242')",
@@ -292,21 +295,21 @@ async fn test_merge_into_updates_non_blob_on_raw_blob_table() {
.await;
ctx_exec(
- &handler,
- "CREATE TABLE src (id INT, name VARCHAR) AS VALUES (1, 'Updated')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT, name VARCHAR) AS VALUES (1, 'Updated')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.t t \
- USING src s ON t.id = s.id \
+ USING datafusion.public.src s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -322,24 +325,24 @@ async fn test_merge_into_updates_non_blob_on_raw_blob_table() {
/// Reference: BlobTestBase "Blob: merge-into rejects updating raw-data BLOB column"
#[tokio::test]
async fn test_merge_into_rejects_raw_blob_update() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Alice', X'4141')",
)
.await;
ctx_exec(
- &handler,
- "CREATE TABLE src (id INT, picture BYTEA) AS VALUES (1, X'4242')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT, picture BYTEA) AS VALUES (1, X'4242')",
)
.await;
- let result = handler
+ let result = sql_context
.sql(
"MERGE INTO paimon.test_db.t t \
- USING src s ON t.id = s.id \
+ USING datafusion.public.src s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET picture = s.picture",
)
.await;
@@ -355,7 +358,7 @@ async fn test_merge_into_rejects_raw_blob_update() {
/// Reference: BlobTestBase "Blob: merge-into updates non-blob column on descriptor blob table"
#[tokio::test]
async fn test_merge_into_updates_non_blob_on_descriptor_table() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -369,7 +372,7 @@ async fn test_merge_into_updates_non_blob_on_descriptor_table() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'4141'), \
(2, 'Bob', X'4242')",
@@ -377,20 +380,20 @@ async fn test_merge_into_updates_non_blob_on_descriptor_table() {
.await;
ctx_exec(
- &handler,
- "CREATE TABLE src (id INT, name VARCHAR) AS VALUES (1, 'Updated')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT, name VARCHAR) AS VALUES (1, 'Updated')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.t t \
- USING src s ON t.id = s.id \
+ USING datafusion.public.src s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT id, name FROM paimon.test_db.t ORDER BY id")
.await
.unwrap()
@@ -404,7 +407,7 @@ async fn test_merge_into_updates_non_blob_on_descriptor_table() {
/// Merge-into on a descriptor blob table: updating the blob column should succeed.
#[tokio::test]
async fn test_merge_into_updates_blob_on_descriptor_table() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -418,7 +421,7 @@ async fn test_merge_into_updates_blob_on_descriptor_table() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'4141'), \
(2, 'Bob', X'4242')",
@@ -426,21 +429,21 @@ async fn test_merge_into_updates_blob_on_descriptor_table() {
.await;
ctx_exec(
- &handler,
- "CREATE TABLE src (id INT, picture BYTEA) AS VALUES (1, X'4343')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT, picture BYTEA) AS VALUES (1, X'4343')",
)
.await;
exec(
- &handler,
+ &sql_context,
"MERGE INTO paimon.test_db.t t \
- USING src s ON t.id = s.id \
+ USING datafusion.public.src s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET picture = s.picture",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -456,7 +459,7 @@ async fn test_merge_into_updates_blob_on_descriptor_table() {
/// Blob with partitioned table.
#[tokio::test]
async fn test_blob_with_partition() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
picture BLOB, \
@@ -469,14 +472,14 @@ async fn test_blob_with_partition() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, picture, pt) VALUES \
(1, X'4141', 'a'), \
(2, X'4242', 'b')",
)
.await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT id, picture FROM paimon.test_db.t ORDER BY id")
.await
.unwrap()
@@ -509,7 +512,7 @@ async fn test_blob_with_partition() {
/// resolve it by reading actual data from the referenced URI.
#[tokio::test]
async fn test_blob_resolve_descriptor() {
- let (tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
// Write a source file that the BlobDescriptor will reference.
let source_data = b"ResolvedBlobContent";
@@ -526,10 +529,10 @@ async fn test_blob_resolve_descriptor() {
(1, 'Descriptor', X'{desc_hex}'), \
(2, 'Raw', X'48656C6C6F')"
);
- exec(&handler, &sql).await;
+ exec(&sql_context, &sql).await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -546,7 +549,7 @@ async fn test_blob_resolve_descriptor() {
/// BlobDescriptor with non-zero offset and partial length.
#[tokio::test]
async fn test_blob_resolve_descriptor_with_offset() {
- let (tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
let source_data = b"HEADER_PAYLOAD_TRAILER";
let source_path = tmp.path().join("blob_offset.bin");
@@ -560,10 +563,10 @@ async fn test_blob_resolve_descriptor_with_offset() {
let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Partial', X'{desc_hex}')"
);
- exec(&handler, &sql).await;
+ exec(&sql_context, &sql).await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -574,7 +577,7 @@ async fn test_blob_resolve_descriptor_with_offset() {
/// Blob files roll independently when `blob.target-file-size` is small.
#[tokio::test]
async fn test_blob_rolling() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -589,7 +592,7 @@ async fn test_blob_rolling() {
// Insert multiple rows with blob data large enough to trigger rolling at 50 bytes.
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'A', X'4141414141414141414141414141414141414141414141414141414141414141'), \
(2, 'B', X'4242424242424242424242424242424242424242424242424242424242424242'), \
@@ -598,7 +601,7 @@ async fn test_blob_rolling() {
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -611,7 +614,7 @@ async fn test_blob_rolling() {
/// blob-descriptor-field with multiple inserts: descriptor values are resolved to actual data.
#[tokio::test]
async fn test_blob_descriptor_field_resolve_on_read() {
- let (_tmp, handler) = setup(
+ let (_tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -625,7 +628,7 @@ async fn test_blob_descriptor_field_resolve_on_read() {
.await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'48656C6C6F'), \
(2, 'Bob', NULL), \
@@ -635,14 +638,14 @@ async fn test_blob_descriptor_field_resolve_on_read() {
// Second insert to exercise multi-file merge path.
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(4, 'Dave', X'5061696D6F6E')",
)
.await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -660,7 +663,7 @@ async fn test_blob_descriptor_field_resolve_on_read() {
/// blob-descriptor-field: inserting a serialized BlobDescriptor resolves to actual data on read.
#[tokio::test]
async fn test_blob_descriptor_field_resolve_descriptor_value() {
- let (tmp, handler) = setup(
+ let (tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
@@ -687,10 +690,10 @@ async fn test_blob_descriptor_field_resolve_descriptor_value() {
(1, 'FromDesc', X'{desc_hex}'), \
(2, 'Raw', X'48656C6C6F')"
);
- exec(&handler, &sql).await;
+ exec(&sql_context, &sql).await;
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -707,10 +710,10 @@ async fn test_blob_descriptor_field_resolve_descriptor_value() {
/// bytes instead of the actual blob content.
#[tokio::test]
async fn test_blob_as_descriptor_dynamic_option() {
- let (_tmp, handler) = setup(BLOB_TABLE_DDL).await;
+ let (_tmp, sql_context) = setup(BLOB_TABLE_DDL).await;
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Alice', X'48656C6C6F'), \
(2, 'Bob', X'576F726C64')",
@@ -719,20 +722,20 @@ async fn test_blob_as_descriptor_dynamic_option() {
// Without the option, we get raw blob data.
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
assert_eq!(rows[0].2, Some(b"Hello".to_vec()));
// Enable blob-as-descriptor via dynamic option.
- handler
+ sql_context
.sql("SET 'paimon.blob-as-descriptor' = 'true'")
.await
.unwrap();
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
@@ -749,13 +752,13 @@ async fn test_blob_as_descriptor_dynamic_option() {
assert!(desc.length() > 0);
// RESET should go back to raw data.
- handler
+ sql_context
.sql("RESET 'paimon.blob-as-descriptor'")
.await
.unwrap();
let rows = query_id_name_picture(
- &handler,
+ &sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
diff --git a/crates/integrations/datafusion/tests/common/mod.rs b/crates/integrations/datafusion/tests/common/mod.rs
index 2ab1f08..f330ef3 100644
--- a/crates/integrations/datafusion/tests/common/mod.rs
+++ b/crates/integrations/datafusion/tests/common/mod.rs
@@ -20,9 +20,8 @@
use std::sync::Arc;
use datafusion::arrow::array::{Int32Array, StringArray};
-use datafusion::prelude::SessionContext;
use paimon::{CatalogOptions, FileSystemCatalog, Options};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
use tempfile::TempDir;
use arrow_array::{Array, RecordBatch, UInt64Array};
@@ -36,25 +35,26 @@ pub fn create_test_env() -> (TempDir, Arc<FileSystemCatalog>) {
(temp_dir, Arc::new(catalog))
}
-pub fn create_handler(catalog: Arc<FileSystemCatalog>) -> PaimonSqlHandler {
- let ctx = SessionContext::new();
- PaimonSqlHandler::new(ctx, catalog, "paimon").unwrap()
+pub async fn create_sql_context(catalog: Arc<FileSystemCatalog>) -> SQLContext {
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog).await.unwrap();
+ ctx
}
#[allow(dead_code)]
-pub async fn setup_handler() -> (TempDir, PaimonSqlHandler) {
+pub async fn setup_sql_context() -> (TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- (tmp, handler)
+ (tmp, sql_context)
}
#[allow(dead_code)]
-pub async fn collect_id_name(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32, String)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+pub async fn collect_id_name(sql_context: &SQLContext, sql: &str) -> Vec<(i32, String)> {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
let mut rows = Vec::new();
for batch in &batches {
let ids = batch
@@ -74,8 +74,8 @@ pub async fn collect_id_name(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32,
}
#[allow(dead_code)]
-pub async fn collect_id_value(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32, i32)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+pub async fn collect_id_value(sql_context: &SQLContext, sql: &str) -> Vec<(i32, i32)> {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
let mut rows = Vec::new();
for batch in &batches {
let ids = batch
@@ -95,27 +95,42 @@ pub async fn collect_id_value(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32
}
#[allow(dead_code)]
-pub async fn row_count(handler: &PaimonSqlHandler, sql: &str) -> usize {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+pub async fn row_count(sql_context: &SQLContext, sql: &str) -> usize {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
batches.iter().map(|b| b.num_rows()).sum()
}
/// Execute SQL and collect results, discarding the output.
#[allow(dead_code)]
-pub async fn exec(handler: &PaimonSqlHandler, s: &str) {
- handler.sql(s).await.unwrap().collect().await.unwrap();
+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(handler: &PaimonSqlHandler, s: &str) {
- handler.ctx().sql(s).await.unwrap().collect().await.unwrap();
+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(handler: &PaimonSqlHandler, sql_str: &str) -> u64 {
- let result = handler.sql(sql_str).await.unwrap().collect().await.unwrap();
+pub async fn dml_count(sql_context: &SQLContext, sql_str: &str) -> u64 {
+ let result = sql_context
+ .sql(sql_str)
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
result[0]
.column(0)
.as_any()
@@ -233,7 +248,7 @@ pub fn collect_int_str_int(batches: &[RecordBatch]) -> Vec<(i32, String, i32)> {
/// Query a 3-column (i32, String, i32) table and return sorted rows.
#[allow(dead_code)]
-pub async fn query_int_str_int(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32, String, i32)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+pub async fn query_int_str_int(sql_context: &SQLContext, sql: &str) -> Vec<(i32, String, i32)> {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
collect_int_str_int(&batches)
}
diff --git a/crates/integrations/datafusion/tests/count_pushdown.rs b/crates/integrations/datafusion/tests/count_pushdown.rs
index d5f13c5..fd80061 100644
--- a/crates/integrations/datafusion/tests/count_pushdown.rs
+++ b/crates/integrations/datafusion/tests/count_pushdown.rs
@@ -31,31 +31,31 @@ use std::sync::Arc;
use datafusion::arrow::array::Int64Array;
use datafusion::physical_plan::{displayable, ExecutionPlan};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
-/// Creates a test handler with a table ready for inserts.
-async fn setup_table(schema_sql: &str) -> (tempfile::TempDir, PaimonSqlHandler) {
- let (tmp, handler) = common::setup_handler().await;
- handler
+/// Creates a test sql_context with a table ready for inserts.
+async fn setup_table(schema_sql: &str) -> (tempfile::TempDir, SQLContext) {
+ let (tmp, sql_context) = common::setup_sql_context().await;
+ sql_context
.sql(&format!("CREATE TABLE paimon.test_db.t {schema_sql}"))
.await
.expect("CREATE TABLE should succeed");
- (tmp, handler)
+ (tmp, sql_context)
}
-/// Creates a test handler with a partitioned table ready for inserts.
+/// Creates a test sql_context with a partitioned table ready for inserts.
async fn setup_partitioned_table(
columns: &str,
partition_cols: &str,
-) -> (tempfile::TempDir, PaimonSqlHandler) {
- let (tmp, handler) = common::setup_handler().await;
- handler
+) -> (tempfile::TempDir, SQLContext) {
+ let (tmp, sql_context) = common::setup_sql_context().await;
+ sql_context
.sql(&format!(
"CREATE TABLE paimon.test_db.t ({columns}) PARTITIONED BY ({partition_cols})"
))
.await
.expect("CREATE TABLE should succeed");
- (tmp, handler)
+ (tmp, sql_context)
}
/// Checks if the physical plan contains a PaimonTableScan (meaning data would be scanned).
@@ -67,10 +67,10 @@ fn plan_contains_scan(plan: &Arc<dyn ExecutionPlan>) -> bool {
/// Creates a physical plan for the given SQL and checks if COUNT was pushed down.
/// Returns Ok(plan) if pushdown succeeded (no scan in plan), Err(plan) if it fell back to scanning.
async fn verify_count_pushdown(
- handler: &PaimonSqlHandler,
+ sql_context: &SQLContext,
sql: &str,
) -> Result<Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>> {
- let df = handler.sql(sql).await.expect("Query should succeed");
+ let df = sql_context.sql(sql).await.expect("Query should succeed");
let plan = df
.create_physical_plan()
.await
@@ -84,8 +84,8 @@ async fn verify_count_pushdown(
}
/// Executes a COUNT(*) query and returns the count value.
-async fn run_count_query(handler: &PaimonSqlHandler, sql: &str) -> i64 {
- let batches = handler
+async fn run_count_query(sql_context: &SQLContext, sql: &str) -> i64 {
+ let batches = sql_context
.sql(sql)
.await
.expect("Query should succeed")
@@ -111,9 +111,9 @@ async fn run_count_query(handler: &PaimonSqlHandler, sql: &str) -> i64 {
#[tokio::test]
async fn test_count_star_no_filter_pushes_down() {
- let (_tmp, handler) = setup_table("(id INT, value INT)").await;
+ let (_tmp, sql_context) = setup_table("(id INT, value INT)").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10), (2, 20), (3, 30)")
.await
.unwrap()
@@ -121,11 +121,11 @@ async fn test_count_star_no_filter_pushes_down() {
.await
.unwrap();
- verify_count_pushdown(&handler, "SELECT COUNT(*) FROM paimon.test_db.t")
+ verify_count_pushdown(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t")
.await
.expect("COUNT(*) should push down (no filter)");
- let count = run_count_query(&handler, "SELECT COUNT(*) FROM paimon.test_db.t").await;
+ let count = run_count_query(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t").await;
assert_eq!(count, 3, "COUNT(*) should return 3");
}
@@ -135,9 +135,9 @@ async fn test_count_star_no_filter_pushes_down() {
#[tokio::test]
async fn test_count_star_with_data_filter_does_not_push_down() {
- let (_tmp, handler) = setup_table("(id INT, value INT)").await;
+ let (_tmp, sql_context) = setup_table("(id INT, value INT)").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10), (2, 20), (3, 30)")
.await
.unwrap()
@@ -146,7 +146,7 @@ async fn test_count_star_with_data_filter_does_not_push_down() {
.unwrap();
let plan = verify_count_pushdown(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE id > 1",
)
.await;
@@ -157,7 +157,7 @@ async fn test_count_star_with_data_filter_does_not_push_down() {
);
let count = run_count_query(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE id > 1",
)
.await;
@@ -170,9 +170,9 @@ async fn test_count_star_with_data_filter_does_not_push_down() {
#[tokio::test]
async fn test_count_star_with_non_partition_equality_does_not_push_down() {
- let (_tmp, handler) = setup_table("(id INT, value INT)").await;
+ let (_tmp, sql_context) = setup_table("(id INT, value INT)").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10), (2, 20), (3, 30)")
.await
.unwrap()
@@ -181,7 +181,7 @@ async fn test_count_star_with_non_partition_equality_does_not_push_down() {
.unwrap();
let plan = verify_count_pushdown(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE id = 2",
)
.await;
@@ -192,7 +192,7 @@ async fn test_count_star_with_non_partition_equality_does_not_push_down() {
);
let count = run_count_query(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE id = 2",
)
.await;
@@ -205,9 +205,9 @@ async fn test_count_star_with_non_partition_equality_does_not_push_down() {
#[tokio::test]
async fn test_count_star_single_row_pushes_down() {
- let (_tmp, handler) = setup_table("(id INT, value INT)").await;
+ let (_tmp, sql_context) = setup_table("(id INT, value INT)").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10)")
.await
.unwrap()
@@ -215,11 +215,11 @@ async fn test_count_star_single_row_pushes_down() {
.await
.unwrap();
- verify_count_pushdown(&handler, "SELECT COUNT(*) FROM paimon.test_db.t")
+ verify_count_pushdown(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t")
.await
.expect("COUNT(*) should push down");
- let count = run_count_query(&handler, "SELECT COUNT(*) FROM paimon.test_db.t").await;
+ let count = run_count_query(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t").await;
assert_eq!(count, 1, "COUNT(*) should return 1");
}
@@ -229,13 +229,13 @@ async fn test_count_star_single_row_pushes_down() {
#[tokio::test]
async fn test_count_star_empty_table_pushes_down() {
- let (_tmp, handler) = setup_table("(id INT, value INT)").await;
+ let (_tmp, sql_context) = setup_table("(id INT, value INT)").await;
- verify_count_pushdown(&handler, "SELECT COUNT(*) FROM paimon.test_db.t")
+ verify_count_pushdown(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t")
.await
.expect("COUNT(*) on empty table should push down");
- let count = run_count_query(&handler, "SELECT COUNT(*) FROM paimon.test_db.t").await;
+ let count = run_count_query(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t").await;
assert_eq!(count, 0, "COUNT(*) should return 0");
}
@@ -245,9 +245,9 @@ async fn test_count_star_empty_table_pushes_down() {
#[tokio::test]
async fn test_count_star_with_partition_filter_pushes_down() {
- let (_tmp, handler) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
+ let (_tmp, sql_context) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10, '2024-01-01'), (2, 20, '2024-01-01'), (3, 30, '2024-01-02')")
.await
.unwrap()
@@ -256,14 +256,14 @@ async fn test_count_star_with_partition_filter_pushes_down() {
.unwrap();
verify_count_pushdown(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt = '2024-01-01'",
)
.await
.expect("COUNT(*) with partition filter should push down");
let count = run_count_query(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt = '2024-01-01'",
)
.await;
@@ -276,9 +276,9 @@ async fn test_count_star_with_partition_filter_pushes_down() {
#[tokio::test]
async fn test_count_star_with_mixed_partition_data_filter_does_not_push_down() {
- let (_tmp, handler) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
+ let (_tmp, sql_context) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10, '2024-01-01'), (2, 20, '2024-01-01'), (3, 30, '2024-01-02')")
.await
.unwrap()
@@ -287,7 +287,7 @@ async fn test_count_star_with_mixed_partition_data_filter_does_not_push_down() {
.unwrap();
let plan = verify_count_pushdown(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt = '2024-01-01' AND value > 15",
)
.await;
@@ -298,7 +298,7 @@ async fn test_count_star_with_mixed_partition_data_filter_does_not_push_down() {
);
let count = run_count_query(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt = '2024-01-01' AND value > 15",
)
.await;
@@ -311,9 +311,9 @@ async fn test_count_star_with_mixed_partition_data_filter_does_not_push_down() {
#[tokio::test]
async fn test_count_star_with_partition_in_filter_pushes_down() {
- let (_tmp, handler) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
+ let (_tmp, sql_context) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10, '2024-01-01'), (2, 20, '2024-01-02'), (3, 30, '2024-01-03')")
.await
.unwrap()
@@ -322,14 +322,14 @@ async fn test_count_star_with_partition_in_filter_pushes_down() {
.unwrap();
verify_count_pushdown(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt IN ('2024-01-01', '2024-01-02')",
)
.await
.expect("COUNT(*) with partition IN filter should push down");
let count = run_count_query(
- &handler,
+ &sql_context,
"SELECT COUNT(*) FROM paimon.test_db.t WHERE dt IN ('2024-01-01', '2024-01-02')",
)
.await;
@@ -342,9 +342,9 @@ async fn test_count_star_with_partition_in_filter_pushes_down() {
#[tokio::test]
async fn test_count_star_partitioned_no_filter_pushes_down() {
- let (_tmp, handler) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
+ let (_tmp, sql_context) = setup_partitioned_table("id INT, value INT, dt STRING", "dt").await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 10, '2024-01-01'), (2, 20, '2024-01-02')")
.await
.unwrap()
@@ -352,10 +352,10 @@ async fn test_count_star_partitioned_no_filter_pushes_down() {
.await
.unwrap();
- verify_count_pushdown(&handler, "SELECT COUNT(*) FROM paimon.test_db.t")
+ verify_count_pushdown(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t")
.await
.expect("COUNT(*) without filter should push down");
- let count = run_count_query(&handler, "SELECT COUNT(*) FROM paimon.test_db.t").await;
+ let count = run_count_query(&sql_context, "SELECT COUNT(*) FROM paimon.test_db.t").await;
assert_eq!(count, 2, "COUNT(*) should return 2");
}
diff --git a/crates/integrations/datafusion/tests/cross_partition_tables.rs b/crates/integrations/datafusion/tests/cross_partition_tables.rs
index 47a9323..1318364 100644
--- a/crates/integrations/datafusion/tests/cross_partition_tables.rs
+++ b/crates/integrations/datafusion/tests/cross_partition_tables.rs
@@ -23,7 +23,9 @@
mod common;
-use common::{collect_id_name, collect_id_value, create_handler, create_test_env, setup_handler};
+use common::{
+ collect_id_name, collect_id_value, create_sql_context, create_test_env, setup_sql_context,
+};
use datafusion::arrow::array::Int32Array;
use paimon::catalog::Identifier;
use paimon::Catalog;
@@ -32,10 +34,10 @@ use paimon::Catalog;
/// A record's partition can change — old partition gets a DELETE, new partition gets an ADD.
#[tokio::test]
async fn test_cross_partition_update_basic() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
// PK is only "id", partition is "dt" — PK does NOT include partition field
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_pt (
dt STRING, id INT NOT NULL, value INT,
@@ -46,7 +48,7 @@ async fn test_cross_partition_update_basic() {
.unwrap();
// First commit: id=1,2 in partition "2024-01-01"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_pt VALUES \
('2024-01-01', 1, 10), ('2024-01-01', 2, 20)",
@@ -58,14 +60,14 @@ async fn test_cross_partition_update_basic() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_pt ORDER BY id",
)
.await;
assert_eq!(rows, vec![(1, 10), (2, 20)]);
// Second commit: id=1 moves to partition "2024-01-02"
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_cross_pt VALUES ('2024-01-02', 1, 100)")
.await
.unwrap()
@@ -75,7 +77,7 @@ async fn test_cross_partition_update_basic() {
// Read back — id=1 should have value=100, id=2 unchanged
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_pt ORDER BY id",
)
.await;
@@ -85,9 +87,9 @@ async fn test_cross_partition_update_basic() {
/// Cross-partition: multiple records moving between partitions in one commit.
#[tokio::test]
async fn test_cross_partition_update_multiple_keys() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_multi (
dt STRING, id INT NOT NULL, name STRING,
@@ -98,7 +100,7 @@ async fn test_cross_partition_update_multiple_keys() {
.unwrap();
// First commit: 3 records in partition "a"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_multi VALUES \
('a', 1, 'alice'), ('a', 2, 'bob'), ('a', 3, 'carol')",
@@ -110,7 +112,7 @@ async fn test_cross_partition_update_multiple_keys() {
.unwrap();
// Second commit: id=1 moves to "b", id=2 moves to "c", id=3 stays in "a" with update
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_multi VALUES \
('b', 1, 'alice-v2'), ('c', 2, 'bob-v2'), ('a', 3, 'carol-v2')",
@@ -122,7 +124,7 @@ async fn test_cross_partition_update_multiple_keys() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_cross_multi ORDER BY id",
)
.await;
@@ -140,9 +142,9 @@ async fn test_cross_partition_update_multiple_keys() {
/// Cross-partition: three commits with records bouncing between partitions.
#[tokio::test]
async fn test_cross_partition_update_three_commits() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_3c (
dt STRING, id INT NOT NULL, value INT,
@@ -153,7 +155,7 @@ async fn test_cross_partition_update_three_commits() {
.unwrap();
// Commit 1: id=1 in "a", id=2 in "b"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_3c VALUES \
('a', 1, 10), ('b', 2, 20)",
@@ -165,7 +167,7 @@ async fn test_cross_partition_update_three_commits() {
.unwrap();
// Commit 2: id=1 moves to "b", id=2 moves to "a", add id=3 in "a"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_3c VALUES \
('b', 1, 100), ('a', 2, 200), ('a', 3, 30)",
@@ -177,7 +179,7 @@ async fn test_cross_partition_update_three_commits() {
.unwrap();
// Commit 3: id=1 moves back to "a", id=3 moves to "b"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_3c VALUES \
('a', 1, 1000), ('b', 3, 300)",
@@ -189,7 +191,7 @@ async fn test_cross_partition_update_three_commits() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_3c ORDER BY id",
)
.await;
@@ -200,9 +202,9 @@ async fn test_cross_partition_update_three_commits() {
/// Cross-partition: new records in different partitions (no migration, just new keys).
#[tokio::test]
async fn test_cross_partition_new_keys_no_migration() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_new (
dt STRING, id INT NOT NULL, value INT,
@@ -213,7 +215,7 @@ async fn test_cross_partition_new_keys_no_migration() {
.unwrap();
// All new keys in different partitions — no cross-partition migration
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_new VALUES \
('a', 1, 10), ('b', 2, 20), ('c', 3, 30)",
@@ -225,14 +227,14 @@ async fn test_cross_partition_new_keys_no_migration() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_new ORDER BY id",
)
.await;
assert_eq!(rows, vec![(1, 10), (2, 20), (3, 30)]);
// Add more new keys
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_new VALUES \
('a', 4, 40), ('b', 5, 50)",
@@ -244,7 +246,7 @@ async fn test_cross_partition_new_keys_no_migration() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_new ORDER BY id",
)
.await;
@@ -257,14 +259,14 @@ async fn test_cross_partition_new_keys_no_migration() {
#[tokio::test]
async fn test_cross_partition_delete_file_in_old_partition() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
// PK = (id), partition = dt — cross-partition mode
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_dv (
dt STRING, id INT NOT NULL, value INT,
@@ -275,7 +277,7 @@ async fn test_cross_partition_delete_file_in_old_partition() {
.unwrap();
// Commit 1: id=1,2 in partition "a", id=3 in partition "b"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_dv VALUES \
('a', 1, 10), ('a', 2, 20), ('b', 3, 30)",
@@ -305,7 +307,7 @@ async fn test_cross_partition_delete_file_in_old_partition() {
);
// Commit 2: id=1 moves from partition "a" to partition "c"
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_cross_dv VALUES ('c', 1, 100)")
.await
.unwrap()
@@ -374,14 +376,14 @@ async fn test_cross_partition_delete_file_in_old_partition() {
// Final read: dedup should produce correct results
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_dv ORDER BY id",
)
.await;
assert_eq!(rows, vec![(1, 100), (2, 20), (3, 30)]);
// Commit 3: id=2 also moves from "a" to "b", id=3 stays in "b" with update
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_dv VALUES \
('b', 2, 200), ('b', 3, 300)",
@@ -428,7 +430,7 @@ async fn test_cross_partition_delete_file_in_old_partition() {
// Final read after 3 commits
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_dv ORDER BY id",
)
.await;
@@ -439,14 +441,14 @@ async fn test_cross_partition_delete_file_in_old_partition() {
#[tokio::test]
async fn test_cross_partition_first_row_skip() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
// PK is only "id", partition is "dt", merge-engine = first-row
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_fr (
dt STRING, id INT NOT NULL, value INT,
@@ -458,7 +460,7 @@ async fn test_cross_partition_first_row_skip() {
.unwrap();
// Commit 1: id=1,2 in partition "a"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_fr VALUES \
('a', 1, 10), ('a', 2, 20)",
@@ -471,7 +473,7 @@ async fn test_cross_partition_first_row_skip() {
// Commit 2: try to move id=1 to partition "b" and insert new id=3 in "b"
// FIRST_ROW should discard id=1 (already exists in "a"), but accept id=3
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_fr VALUES \
('b', 1, 100), ('b', 3, 30)",
@@ -543,10 +545,10 @@ async fn test_cross_partition_first_row_skip() {
/// When same id is written under different pt2, old partition should get a DELETE.
#[tokio::test]
async fn test_cross_partition_partial_pk_partition_overlap() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
// PK = (pt1, id), partition = (pt1, pt2) — pt2 is NOT in PK → cross-partition
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_partial (
pt1 STRING, pt2 STRING, id INT NOT NULL, value INT,
@@ -557,7 +559,7 @@ async fn test_cross_partition_partial_pk_partition_overlap() {
.unwrap();
// Commit 1: id=1 in (pt1='a', pt2='x')
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_partial VALUES \
('a', 'x', 1, 10), ('a', 'x', 2, 20)",
@@ -569,14 +571,14 @@ async fn test_cross_partition_partial_pk_partition_overlap() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_partial ORDER BY id",
)
.await;
assert_eq!(rows, vec![(1, 10), (2, 20)]);
// Commit 2: id=1 moves to (pt1='a', pt2='y') — different pt2
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_cross_partial VALUES ('a', 'y', 1, 100)")
.await
.unwrap()
@@ -586,7 +588,7 @@ async fn test_cross_partition_partial_pk_partition_overlap() {
// After dedup, id=1 should have value=100 and only appear once
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_partial ORDER BY id",
)
.await;
@@ -609,9 +611,9 @@ async fn test_cross_partition_partial_pk_partition_overlap() {
/// Both must succeed without schema errors.
#[tokio::test]
async fn test_cross_partition_value_kind_schema_stability() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross_vk (
dt STRING, id INT NOT NULL, value INT,
@@ -622,7 +624,7 @@ async fn test_cross_partition_value_kind_schema_stability() {
.unwrap();
// Commit 1: all new keys, no migration — _VALUE_KIND must still be added
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_vk VALUES \
('a', 1, 10), ('a', 2, 20), ('b', 3, 30)",
@@ -634,7 +636,7 @@ async fn test_cross_partition_value_kind_schema_stability() {
.unwrap();
// Commit 2: id=1 migrates from "a" to "b" — produces DELETE in "a"
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_cross_vk VALUES \
('b', 1, 100), ('a', 2, 200)",
@@ -646,7 +648,7 @@ async fn test_cross_partition_value_kind_schema_stability() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_cross_vk ORDER BY id",
)
.await;
diff --git a/crates/integrations/datafusion/tests/delete_tests.rs b/crates/integrations/datafusion/tests/delete_tests.rs
index ccb522b..17cf61b 100644
--- a/crates/integrations/datafusion/tests/delete_tests.rs
+++ b/crates/integrations/datafusion/tests/delete_tests.rs
@@ -21,55 +21,61 @@
mod common;
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
-use common::{create_handler, create_test_env, dml_count, exec, query_int_str_int};
+use common::{create_sql_context, create_test_env, ctx_exec, dml_count, exec, query_int_str_int};
// ======================= Helpers =======================
-async fn setup() -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup() -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, 'b', 20), (3, 'c', 30), (4, 'd', 40)",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
-async fn setup_partitioned() -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup_partitioned() -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, pt INT) PARTITIONED BY (pt)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t VALUES (1, 'a', 1), (2, 'b', 1), (3, 'c', 2), (4, 'd', 2)",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
-async fn query(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
+async fn query(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
query_int_str_int(
- handler,
+ sql_context,
"SELECT id, name, age FROM paimon.test_db.t ORDER BY id",
)
.await
}
-async fn query_pt(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
+async fn query_pt(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
query_int_str_int(
- handler,
+ sql_context,
"SELECT id, name, pt FROM paimon.test_db.t ORDER BY id",
)
.await
@@ -79,12 +85,16 @@ async fn query_pt(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
#[tokio::test]
async fn test_delete_by_non_partition_column() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE name = 'b'").await;
+ exec(
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE name = 'b'",
+ )
+ .await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(3, "c".into(), 30),
@@ -95,13 +105,13 @@ async fn test_delete_by_non_partition_column() {
#[tokio::test]
async fn test_delete_no_match() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- let cnt = dml_count(&handler, "DELETE FROM paimon.test_db.t WHERE id = 99").await;
+ let cnt = dml_count(&sql_context, "DELETE FROM paimon.test_db.t WHERE id = 99").await;
assert_eq!(cnt, 0);
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(2, "b".into(), 20),
@@ -113,64 +123,68 @@ async fn test_delete_no_match() {
#[tokio::test]
async fn test_delete_full_table() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t").await;
- assert_eq!(query(&handler).await, vec![]);
+ assert_eq!(query(&sql_context).await, vec![]);
}
// ======================= Partitioned table DELETE =======================
#[tokio::test]
async fn test_delete_partitioned_by_name() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE name = 'a'").await;
+ exec(
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE name = 'a'",
+ )
+ .await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![(2, "b".into(), 1), (3, "c".into(), 2), (4, "d".into(), 2),]
);
}
#[tokio::test]
async fn test_delete_partitioned_by_partition_column() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE pt = 1").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t WHERE pt = 1").await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![(3, "c".into(), 2), (4, "d".into(), 2),]
);
}
#[tokio::test]
async fn test_delete_partitioned_in_list() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
exec(
- &handler,
+ &sql_context,
"DELETE FROM paimon.test_db.t WHERE pt IN (1, 2) AND id = 1",
)
.await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![(2, "b".into(), 1), (3, "c".into(), 2), (4, "d".into(), 2),]
);
}
#[tokio::test]
async fn test_delete_partitioned_no_match() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
- let cnt = dml_count(&handler, "DELETE FROM paimon.test_db.t WHERE pt = 99").await;
+ let cnt = dml_count(&sql_context, "DELETE FROM paimon.test_db.t WHERE pt = 99").await;
assert_eq!(cnt, 0);
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![
(1, "a".into(), 1),
(2, "b".into(), 1),
@@ -182,52 +196,56 @@ async fn test_delete_partitioned_no_match() {
#[tokio::test]
async fn test_delete_partitioned_full_table() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
- exec(&handler, "DELETE FROM paimon.test_db.t").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t").await;
- assert_eq!(query_pt(&handler).await, vec![]);
+ assert_eq!(query_pt(&sql_context).await, vec![]);
}
#[tokio::test]
async fn test_delete_partitioned_or_condition() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
exec(
- &handler,
+ &sql_context,
"DELETE FROM paimon.test_db.t WHERE pt = 1 OR id = 3",
)
.await;
- assert_eq!(query_pt(&handler).await, vec![(4, "d".into(), 2)]);
+ assert_eq!(query_pt(&sql_context).await, vec![(4, "d".into(), 2)]);
}
// ======================= IN / NOT IN =======================
#[tokio::test]
async fn test_delete_in_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE id IN (1, 3)").await;
+ exec(
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE id IN (1, 3)",
+ )
+ .await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(2, "b".into(), 20), (4, "d".into(), 40),]
);
}
#[tokio::test]
async fn test_delete_not_in_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"DELETE FROM paimon.test_db.t WHERE id NOT IN (1, 3)",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(1, "a".into(), 10), (3, "c".into(), 30),]
);
}
@@ -236,65 +254,59 @@ async fn test_delete_not_in_condition() {
#[tokio::test]
async fn test_delete_in_subquery() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- handler
- .ctx()
- .sql("CREATE TABLE src (id INT) AS VALUES (1), (3)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+ ctx_exec(
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT) AS VALUES (1), (3)",
+ )
+ .await;
exec(
- &handler,
- "DELETE FROM paimon.test_db.t WHERE id IN (SELECT id FROM src)",
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE id IN (SELECT id FROM datafusion.public.src)",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(2, "b".into(), 20), (4, "d".into(), 40),]
);
}
#[tokio::test]
async fn test_delete_scalar_subquery() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- handler
- .ctx()
- .sql("CREATE TABLE src (id INT) AS VALUES (2)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+ ctx_exec(
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT) AS VALUES (2)",
+ )
+ .await;
exec(
- &handler,
- "DELETE FROM paimon.test_db.t WHERE id >= (SELECT MAX(id) FROM src)",
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE id >= (SELECT MAX(id) FROM datafusion.public.src)",
)
.await;
- assert_eq!(query(&handler).await, vec![(1, "a".into(), 10)]);
+ assert_eq!(query(&sql_context).await, vec![(1, "a".into(), 10)]);
}
// ======================= Range condition =======================
#[tokio::test]
async fn test_delete_range_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"DELETE FROM paimon.test_db.t WHERE id >= 2 AND id <= 3",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(1, "a".into(), 10), (4, "d".into(), 40),]
);
}
@@ -304,13 +316,16 @@ async fn test_delete_range_condition() {
#[tokio::test]
async fn test_delete_should_not_remove_null_rows() {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, NULL, 20), (3, 'c', 30)")
.await
.unwrap()
@@ -318,9 +333,13 @@ async fn test_delete_should_not_remove_null_rows() {
.await
.unwrap();
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE name = 'a'").await;
+ exec(
+ &sql_context,
+ "DELETE FROM paimon.test_db.t WHERE name = 'a'",
+ )
+ .await;
- let rows = query(&handler).await;
+ let rows = query(&sql_context).await;
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].0, 2);
assert_eq!(rows[0].2, 20);
@@ -332,13 +351,13 @@ async fn test_delete_should_not_remove_null_rows() {
#[tokio::test]
async fn test_successive_deletes() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE id = 3").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t WHERE id = 3").await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(2, "b".into(), 20), (4, "d".into(), 40),]
);
}
@@ -347,13 +366,17 @@ async fn test_successive_deletes() {
#[tokio::test]
async fn test_delete_then_insert() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
- exec(&handler, "INSERT INTO paimon.test_db.t VALUES (5, 'e', 50)").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
+ exec(
+ &sql_context,
+ "INSERT INTO paimon.test_db.t VALUES (5, 'e', 50)",
+ )
+ .await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(2, "b".into(), 20),
(3, "c".into(), 30),
@@ -365,17 +388,17 @@ async fn test_delete_then_insert() {
#[tokio::test]
async fn test_delete_all_then_insert() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t").await;
- assert_eq!(query(&handler).await, vec![]);
+ exec(&sql_context, "DELETE FROM paimon.test_db.t").await;
+ assert_eq!(query(&sql_context).await, vec![]);
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t VALUES (10, 'new', 100)",
)
.await;
- assert_eq!(query(&handler).await, vec![(10, "new".into(), 100)]);
+ assert_eq!(query(&sql_context).await, vec![(10, "new".into(), 100)]);
}
// ======================= Empty table =======================
@@ -383,14 +406,17 @@ async fn test_delete_all_then_insert() {
#[tokio::test]
async fn test_delete_empty_table() {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
- let cnt = dml_count(&handler, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
+ let cnt = dml_count(&sql_context, "DELETE FROM paimon.test_db.t WHERE id = 1").await;
assert_eq!(cnt, 0);
drop(tmp);
}
@@ -399,16 +425,16 @@ async fn test_delete_empty_table() {
#[tokio::test]
async fn test_delete_or_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"DELETE FROM paimon.test_db.t WHERE id = 1 OR id = 4",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![(2, "b".into(), 20), (3, "c".into(), 30),]
);
}
@@ -417,9 +443,9 @@ async fn test_delete_or_condition() {
#[tokio::test]
async fn test_delete_multiple_rows_from_single_commit() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "DELETE FROM paimon.test_db.t WHERE id >= 2").await;
+ exec(&sql_context, "DELETE FROM paimon.test_db.t WHERE id >= 2").await;
- assert_eq!(query(&handler).await, vec![(1, "a".into(), 10)]);
+ assert_eq!(query(&sql_context).await, vec![(1, "a".into(), 10)]);
}
diff --git a/crates/integrations/datafusion/tests/dynamic_bucket_tables.rs b/crates/integrations/datafusion/tests/dynamic_bucket_tables.rs
index 6ab7195..59804b6 100644
--- a/crates/integrations/datafusion/tests/dynamic_bucket_tables.rs
+++ b/crates/integrations/datafusion/tests/dynamic_bucket_tables.rs
@@ -19,7 +19,9 @@
mod common;
-use common::{collect_id_name, collect_id_value, create_handler, create_test_env, setup_handler};
+use common::{
+ collect_id_name, collect_id_value, create_sql_context, create_test_env, setup_sql_context,
+};
use datafusion::arrow::array::{Int32Array, StringArray};
use paimon::catalog::Identifier;
use paimon::spec::{IndexManifest, IndexManifestEntry};
@@ -28,9 +30,9 @@ use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options, SnapshotManage
/// PK table with bucket=-1 (dynamic bucket) should write and read correctly.
#[tokio::test]
async fn test_pk_dynamic_bucket() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn (
id INT NOT NULL, name STRING,
@@ -41,7 +43,7 @@ async fn test_pk_dynamic_bucket() {
.unwrap();
// First insert — no 'bucket' option, defaults to -1 (dynamic bucket)
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn VALUES (1, 'alice'), (2, 'bob')")
.await
.unwrap()
@@ -50,14 +52,14 @@ async fn test_pk_dynamic_bucket() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_dyn ORDER BY id",
)
.await;
assert_eq!(rows, vec![(1, "alice".to_string()), (2, "bob".to_string())]);
// Second insert with overlapping key — dedup should keep latest
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn VALUES (2, 'bobby'), (3, 'carol')")
.await
.unwrap()
@@ -66,7 +68,7 @@ async fn test_pk_dynamic_bucket() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_dyn ORDER BY id",
)
.await;
@@ -83,9 +85,9 @@ async fn test_pk_dynamic_bucket() {
/// Dynamic bucket with partitioned table.
#[tokio::test]
async fn test_pk_dynamic_bucket_partitioned() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn_part (
dt STRING, id INT NOT NULL, value INT,
@@ -95,7 +97,7 @@ async fn test_pk_dynamic_bucket_partitioned() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_dyn_part VALUES \
('2024-01-01', 1, 10), ('2024-01-01', 2, 20), \
@@ -108,7 +110,7 @@ async fn test_pk_dynamic_bucket_partitioned() {
.unwrap();
// Update within each partition
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_dyn_part VALUES \
('2024-01-01', 1, 11), ('2024-01-02', 2, 222)",
@@ -119,7 +121,7 @@ async fn test_pk_dynamic_bucket_partitioned() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, value FROM paimon.test_db.t_dyn_part ORDER BY dt, id")
.await
.unwrap()
@@ -161,9 +163,9 @@ async fn test_pk_dynamic_bucket_partitioned() {
/// Dynamic bucket with three commits — verifies sequence number tracking.
#[tokio::test]
async fn test_pk_dynamic_bucket_three_commits() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn3 (
id INT NOT NULL, value INT,
@@ -173,7 +175,7 @@ async fn test_pk_dynamic_bucket_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn3 VALUES (1, 10), (2, 20)")
.await
.unwrap()
@@ -181,7 +183,7 @@ async fn test_pk_dynamic_bucket_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn3 VALUES (2, 200), (3, 30)")
.await
.unwrap()
@@ -189,7 +191,7 @@ async fn test_pk_dynamic_bucket_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn3 VALUES (1, 100), (3, 300), (4, 40)")
.await
.unwrap()
@@ -198,7 +200,7 @@ async fn test_pk_dynamic_bucket_three_commits() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_dyn3 ORDER BY id",
)
.await;
@@ -259,13 +261,13 @@ async fn collect_all_hashes(table: &paimon::Table) -> Vec<i32> {
#[tokio::test]
async fn test_pk_dynamic_bucket_insert_overwrite() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn_ow (
id INT NOT NULL, name STRING,
@@ -276,7 +278,7 @@ async fn test_pk_dynamic_bucket_insert_overwrite() {
.unwrap();
// Commit 1: insert 3 rows
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn_ow VALUES (1, 'alice'), (2, 'bob'), (3, 'carol')")
.await
.unwrap()
@@ -292,7 +294,7 @@ async fn test_pk_dynamic_bucket_insert_overwrite() {
assert_eq!(hashes_before.len(), 3, "Should have 3 hash entries");
// INSERT OVERWRITE with only 2 rows — old index entries must be cleared
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_dyn_ow VALUES (10, 'new_a'), (20, 'new_b')")
.await
.unwrap()
@@ -302,7 +304,7 @@ async fn test_pk_dynamic_bucket_insert_overwrite() {
// Verify data
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_dyn_ow ORDER BY id",
)
.await;
@@ -338,13 +340,13 @@ async fn test_pk_dynamic_bucket_insert_overwrite() {
#[tokio::test]
async fn test_pk_dynamic_bucket_partitioned_insert_overwrite() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn_part_ow (
dt STRING, id INT NOT NULL, value INT,
@@ -355,7 +357,7 @@ async fn test_pk_dynamic_bucket_partitioned_insert_overwrite() {
.unwrap();
// Commit 1: two partitions
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_dyn_part_ow VALUES \
('a', 1, 10), ('a', 2, 20), \
@@ -379,7 +381,7 @@ async fn test_pk_dynamic_bucket_partitioned_insert_overwrite() {
);
// INSERT OVERWRITE partition 'a' with only 1 row
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_dyn_part_ow VALUES ('a', 5, 50)")
.await
.unwrap()
@@ -389,7 +391,7 @@ async fn test_pk_dynamic_bucket_partitioned_insert_overwrite() {
// Verify data: partition 'a' replaced, partition 'b' untouched
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_dyn_part_ow ORDER BY id",
)
.await;
@@ -488,9 +490,9 @@ async fn test_read_spark_dynamic_bucket_and_compare_index() {
);
// --- Write the same data into a new dynamic bucket table ---
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn_cmp (
id INT NOT NULL, name STRING,
@@ -501,7 +503,7 @@ async fn test_read_spark_dynamic_bucket_and_compare_index() {
.unwrap();
// Replicate the same two commits as provision.py
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dyn_cmp VALUES (1, 'alice'), (2, 'bob'), (3, 'carol')")
.await
.unwrap()
@@ -509,7 +511,7 @@ async fn test_read_spark_dynamic_bucket_and_compare_index() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_dyn_cmp VALUES \
(2, 'bob-v2'), (3, 'carol-v2'), (4, 'dave')",
@@ -522,7 +524,7 @@ async fn test_read_spark_dynamic_bucket_and_compare_index() {
// Verify written data matches
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_dyn_cmp ORDER BY id",
)
.await;
diff --git a/crates/integrations/datafusion/tests/merge_into_tests.rs b/crates/integrations/datafusion/tests/merge_into_tests.rs
index c74b5c8..55ad852 100644
--- a/crates/integrations/datafusion/tests/merge_into_tests.rs
+++ b/crates/integrations/datafusion/tests/merge_into_tests.rs
@@ -24,11 +24,10 @@
use std::sync::Arc;
use arrow_array::{Int32Array, Int64Array, StringArray};
-use datafusion::prelude::SessionContext;
use paimon::catalog::Identifier;
use paimon::table::SnapshotManager;
use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
use tempfile::TempDir;
// ======================= Helpers =======================
@@ -42,17 +41,18 @@ fn create_test_env() -> (TempDir, Arc<FileSystemCatalog>) {
(temp_dir, Arc::new(catalog))
}
-fn create_handler(catalog: Arc<FileSystemCatalog>) -> PaimonSqlHandler {
- let ctx = SessionContext::new();
- PaimonSqlHandler::new(ctx, catalog, "paimon").unwrap()
+async fn create_sql_context(catalog: Arc<FileSystemCatalog>) -> SQLContext {
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog).await.unwrap();
+ ctx
}
-async fn setup_data_evolution_table(handler: &PaimonSqlHandler) {
- handler
+async fn setup_data_evolution_table(sql_context: &SQLContext) {
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.target (\
id INT NOT NULL, name STRING, value INT\
@@ -64,15 +64,15 @@ async fn setup_data_evolution_table(handler: &PaimonSqlHandler) {
.expect("CREATE TABLE failed");
}
-async fn enable_data_evolution(handler: &PaimonSqlHandler) {
- handler
+async fn enable_data_evolution(sql_context: &SQLContext) {
+ sql_context
.sql("ALTER TABLE paimon.test_db.target SET TBLPROPERTIES('data-evolution.enabled' = 'true')")
.await
.expect("ALTER TABLE failed");
}
-async fn collect_rows_3col(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32, String, i32)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+async fn collect_rows_3col(sql_context: &SQLContext, sql: &str) -> Vec<(i32, String, i32)> {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
let mut rows = Vec::new();
for batch in &batches {
let ids = batch
@@ -97,8 +97,8 @@ async fn collect_rows_3col(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i32, S
rows
}
-async fn collect_row_ids(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i64, i32)> {
- let batches = handler.sql(sql).await.unwrap().collect().await.unwrap();
+async fn collect_row_ids(sql_context: &SQLContext, sql: &str) -> Vec<(i64, i32)> {
+ let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap();
let mut rows = Vec::new();
for batch in &batches {
let row_ids = batch
@@ -118,8 +118,8 @@ async fn collect_row_ids(handler: &PaimonSqlHandler, sql: &str) -> Vec<(i64, i32
rows
}
-async fn assert_merge_error(handler: &PaimonSqlHandler, sql: &str, expected_substring: &str) {
- let result = handler.sql(sql).await;
+async fn assert_merge_error(sql_context: &SQLContext, sql: &str, expected_substring: &str) {
+ let result = sql_context.sql(sql).await;
assert!(
result.is_err(),
"Expected error containing '{expected_substring}', but got Ok"
@@ -131,15 +131,9 @@ async fn assert_merge_error(handler: &PaimonSqlHandler, sql: &str, expected_subs
);
}
-async fn register_source(handler: &PaimonSqlHandler, sql: &str) {
- handler
- .ctx()
- .sql(sql)
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+async fn register_source(sql_context: &SQLContext, sql: &str) {
+ let ctx = sql_context.ctx();
+ ctx.sql(sql).await.unwrap().collect().await.unwrap();
}
// ======================= Functional Tests =======================
@@ -147,16 +141,16 @@ async fn register_source(handler: &PaimonSqlHandler, sql: &str) {
#[tokio::test]
async fn test_row_id_values_after_insert() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
- let batches = handler
+ let batches = sql_context
.sql("SELECT \"_ROW_ID\" FROM paimon.test_db.target ORDER BY \"_ROW_ID\"")
.await
.unwrap()
@@ -195,32 +189,32 @@ async fn test_row_id_values_after_insert() {
#[tokio::test]
async fn test_row_id_stability_after_merge_into() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Capture _ROW_ID -> id mapping before merge
let before = collect_row_ids(
- &handler,
+ &sql_context,
"SELECT \"_ROW_ID\", id FROM paimon.test_db.target ORDER BY id",
)
.await;
// Register source and execute MERGE INTO
register_source(
- &handler,
- "CREATE TABLE source1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.source1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (3, 'CHARLIE')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING source1 s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.source1 s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -231,7 +225,7 @@ async fn test_row_id_stability_after_merge_into() {
// Capture _ROW_ID -> id mapping after merge
let after = collect_row_ids(
- &handler,
+ &sql_context,
"SELECT \"_ROW_ID\", id FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -244,7 +238,7 @@ async fn test_row_id_stability_after_merge_into() {
// Verify data correctness
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -261,24 +255,24 @@ async fn test_row_id_stability_after_merge_into() {
#[tokio::test]
async fn test_multiple_merge_into_different_columns() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// First MERGE: update name for id=1
register_source(
- &handler,
- "CREATE TABLE src_name (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_name (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_name s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_name s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -289,13 +283,13 @@ async fn test_multiple_merge_into_different_columns() {
// Second MERGE: update value for id=2
register_source(
- &handler,
- "CREATE TABLE src_value (id INT, value INT) AS VALUES (2, 200)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_value (id INT, value INT) AS VALUES (2, 200)",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_value s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_value s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET value = s.value",
)
.await
@@ -305,7 +299,7 @@ async fn test_multiple_merge_into_different_columns() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -318,25 +312,25 @@ async fn test_multiple_merge_into_different_columns() {
#[tokio::test]
async fn test_merge_into_with_non_paimon_source() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Source is a plain DataFusion in-memory table, not Paimon
register_source(
- &handler,
- "CREATE TABLE df_source (id INT, name VARCHAR) AS VALUES (2, 'BOB_UPDATED')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.df_source (id INT, name VARCHAR) AS VALUES (2, 'BOB_UPDATED')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING df_source s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.df_source s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -346,7 +340,7 @@ async fn test_merge_into_with_non_paimon_source() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -362,18 +356,18 @@ async fn test_merge_into_with_non_paimon_source() {
#[tokio::test]
async fn test_merge_into_join_on_row_id() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Get _ROW_ID for id=2
let row_id_map = collect_row_ids(
- &handler,
+ &sql_context,
"SELECT \"_ROW_ID\", id FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -381,16 +375,16 @@ async fn test_merge_into_join_on_row_id() {
// Create source with that row_id
register_source(
- &handler,
+ &sql_context,
&format!(
- "CREATE TABLE rid_source (row_id BIGINT, name VARCHAR) AS VALUES ({row_id_of_2}, 'BOB')"
+ "CREATE TABLE datafusion.public.rid_source (row_id BIGINT, name VARCHAR) AS VALUES ({row_id_of_2}, 'BOB')"
),
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING rid_source s ON t.\"_ROW_ID\" = s.row_id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.rid_source s ON t.\"_ROW_ID\" = s.row_id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -400,7 +394,7 @@ async fn test_merge_into_join_on_row_id() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -417,17 +411,17 @@ async fn test_merge_into_join_on_row_id() {
#[tokio::test]
async fn test_self_merge() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Self-merge: target used as both target and source
- handler
+ sql_context
.sql(
"MERGE INTO paimon.test_db.target t USING paimon.test_db.target s \
ON t.\"_ROW_ID\" = s.\"_ROW_ID\" \
@@ -440,7 +434,7 @@ async fn test_self_merge() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -453,14 +447,14 @@ async fn test_self_merge() {
#[tokio::test]
async fn test_row_count_after_merge() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Snapshot 1: 3 rows inserted
let table = catalog
@@ -473,13 +467,13 @@ async fn test_row_count_after_merge() {
// MERGE INTO: update 1 row
register_source(
- &handler,
- "CREATE TABLE src_count (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_count (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_count s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_count s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -504,24 +498,24 @@ async fn test_row_count_after_merge() {
#[tokio::test]
async fn test_merge_into_update_and_insert() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
register_source(
- &handler,
- "CREATE TABLE src_ui (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_ui (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_ui s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.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)",
)
@@ -532,7 +526,7 @@ async fn test_merge_into_update_and_insert() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -549,25 +543,25 @@ async fn test_merge_into_update_and_insert() {
#[tokio::test]
async fn test_merge_into_insert_only() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
register_source(
- &handler,
- "CREATE TABLE src_io (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_io (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
)
.await;
// Only INSERT, no MATCHED clause — matched row id=2 should be untouched
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_io s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_io s ON t.id = s.id \
WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)",
)
.await
@@ -577,7 +571,7 @@ async fn test_merge_into_insert_only() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -594,10 +588,10 @@ async fn test_merge_into_insert_only() {
#[tokio::test]
async fn test_merge_into_insert_all_columns() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (2, 'bob', 20)")
.await
.unwrap()
@@ -605,18 +599,18 @@ async fn test_merge_into_insert_all_columns() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Source schema matches target: (id, name, value)
register_source(
- &handler,
- "CREATE TABLE src_star (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 10), (2, 'BOB', 22)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_star (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 10), (2, 'BOB', 22)",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_star s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.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)",
)
@@ -627,7 +621,7 @@ async fn test_merge_into_insert_all_columns() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -643,10 +637,10 @@ async fn test_merge_into_insert_all_columns() {
#[tokio::test]
async fn test_merge_into_insert_partial_columns() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (2, 'bob', 20)")
.await
.unwrap()
@@ -654,18 +648,18 @@ async fn test_merge_into_insert_partial_columns() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
register_source(
- &handler,
- "CREATE TABLE src_partial (id INT, name VARCHAR) AS VALUES (1, 'alice'), (2, 'BOB')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_partial (id INT, name VARCHAR) AS VALUES (1, 'alice'), (2, 'BOB')",
)
.await;
// INSERT only id and name, value should be NULL (but our schema has INT, so this tests partial insert)
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_partial s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_partial s ON t.id = s.id \
WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, 0)",
)
.await
@@ -675,7 +669,7 @@ async fn test_merge_into_insert_partial_columns() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -691,10 +685,10 @@ async fn test_merge_into_insert_partial_columns() {
#[tokio::test]
async fn test_merge_into_insert_with_predicate() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10)")
.await
.unwrap()
@@ -702,19 +696,19 @@ async fn test_merge_into_insert_with_predicate() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Source has 3 rows, only id=1 matches target
register_source(
- &handler,
- "CREATE TABLE src_pred (id INT, name VARCHAR, value INT) AS VALUES (1, 'ALICE', 11), (2, 'bob', 20), (3, 'charlie', 30)",
+ &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)",
)
.await;
// Only insert when value > 25
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_pred s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.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
@@ -724,7 +718,7 @@ async fn test_merge_into_insert_with_predicate() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -741,32 +735,32 @@ async fn test_merge_into_insert_with_predicate() {
#[tokio::test]
async fn test_merge_into_row_id_for_inserted_rows() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (2, 'bob', 20), (3, 'charlie', 30)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// Before merge: _ROW_ID should be 0, 1
let before = collect_row_ids(
- &handler,
+ &sql_context,
"SELECT \"_ROW_ID\", id FROM paimon.test_db.target ORDER BY id",
)
.await;
assert_eq!(before, vec![(0, 2), (1, 3)]);
register_source(
- &handler,
- "CREATE TABLE src_rid (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_rid (id INT, name VARCHAR, value INT) AS VALUES (1, 'alice', 11), (2, 'BOB', 22)",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_rid s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.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)",
)
@@ -778,7 +772,7 @@ async fn test_merge_into_row_id_for_inserted_rows() {
// After merge: existing rows keep their _ROW_ID, new row gets next available
let after = collect_row_ids(
- &handler,
+ &sql_context,
"SELECT \"_ROW_ID\", id FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -805,10 +799,10 @@ async fn test_merge_into_row_id_for_inserted_rows() {
#[tokio::test]
async fn test_rejects_when_matched_delete() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10)")
.await
.unwrap()
@@ -816,13 +810,17 @@ async fn test_rejects_when_matched_delete() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
- register_source(&handler, "CREATE TABLE src_del (id INT) AS VALUES (1)").await;
+ register_source(
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_del (id INT) AS VALUES (1)",
+ )
+ .await;
assert_merge_error(
- &handler,
- "MERGE INTO paimon.test_db.target t USING src_del s ON t.id = s.id \
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_del s ON t.id = s.id \
WHEN MATCHED THEN DELETE",
"WHEN MATCHED THEN DELETE is not supported",
)
@@ -832,10 +830,10 @@ async fn test_rejects_when_matched_delete() {
#[tokio::test]
async fn test_rejects_multiple_when_matched() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10)")
.await
.unwrap()
@@ -843,17 +841,17 @@ async fn test_rejects_multiple_when_matched() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
register_source(
- &handler,
- "CREATE TABLE src_multi (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_multi (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
assert_merge_error(
- &handler,
- "MERGE INTO paimon.test_db.target t USING src_multi s ON t.id = s.id \
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.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 <predicate> is not yet supported",
@@ -864,10 +862,13 @@ async fn test_rejects_multiple_when_matched() {
#[tokio::test]
async fn test_rejects_partition_column_in_set() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.part_target (\
pt STRING, id INT NOT NULL, name STRING\
@@ -878,7 +879,7 @@ async fn test_rejects_partition_column_in_set() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.part_target (pt, id, name) VALUES ('a', 1, 'alice')")
.await
.unwrap()
@@ -886,17 +887,17 @@ async fn test_rejects_partition_column_in_set() {
.await
.unwrap();
- handler.sql("ALTER TABLE paimon.test_db.part_target SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
+ sql_context.sql("ALTER TABLE paimon.test_db.part_target SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
register_source(
- &handler,
- "CREATE TABLE src_pt (id INT, pt VARCHAR) AS VALUES (1, 'b')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_pt (id INT, pt VARCHAR) AS VALUES (1, 'b')",
)
.await;
assert_merge_error(
- &handler,
- "MERGE INTO paimon.test_db.part_target t USING src_pt s ON t.id = s.id \
+ &sql_context,
+ "MERGE INTO paimon.test_db.part_target t USING datafusion.public.src_pt s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET pt = s.pt",
"Cannot update partition column",
)
@@ -906,10 +907,13 @@ async fn test_rejects_partition_column_in_set() {
#[tokio::test]
async fn test_rejects_table_without_row_tracking() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.no_tracking (\
id INT NOT NULL, name STRING\
@@ -918,7 +922,7 @@ async fn test_rejects_table_without_row_tracking() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.no_tracking (id, name) VALUES (1, 'alice')")
.await
.unwrap()
@@ -926,17 +930,17 @@ async fn test_rejects_table_without_row_tracking() {
.await
.unwrap();
- handler.sql("ALTER TABLE paimon.test_db.no_tracking SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
+ sql_context.sql("ALTER TABLE paimon.test_db.no_tracking SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
register_source(
- &handler,
- "CREATE TABLE src_nrt (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_nrt (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
assert_merge_error(
- &handler,
- "MERGE INTO paimon.test_db.no_tracking t USING src_nrt s ON t.id = s.id \
+ &sql_context,
+ "MERGE INTO paimon.test_db.no_tracking t USING datafusion.public.src_nrt s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
"row-tracking.enabled",
)
@@ -948,24 +952,24 @@ async fn test_successive_merges_read_file_group() {
// Verifies that a second MERGE INTO correctly reads columns from the file group
// (base file + partial-column files created by the first merge), not just a single file.
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// First MERGE: update 'name' column → creates a partial-column file for 'name'
register_source(
- &handler,
- "CREATE TABLE src_m1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (2, 'BOB')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_m1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE'), (2, 'BOB')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_m1 s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_m1 s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -976,7 +980,7 @@ async fn test_successive_merges_read_file_group() {
// Verify first merge result
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -988,13 +992,13 @@ async fn test_successive_merges_read_file_group() {
// Second MERGE: update 'name' again → must read the merged 'name' from file group
// (base file has original 'name', partial file has updated 'name' from first merge)
register_source(
- &handler,
- "CREATE TABLE src_m2 (id INT, name VARCHAR) AS VALUES (1, 'Alice_v2')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_m2 (id INT, name VARCHAR) AS VALUES (1, 'Alice_v2')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_m2 s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_m2 s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -1004,7 +1008,7 @@ async fn test_successive_merges_read_file_group() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -1020,24 +1024,24 @@ async fn test_successive_merges_different_columns_read_file_group() {
// The second merge must correctly read 'value' from the file group
// even though a partial-column file for 'name' now exists.
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await.unwrap().collect().await.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
// First MERGE: update 'name'
register_source(
- &handler,
- "CREATE TABLE src_dc1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_dc1 (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_dc1 s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_dc1 s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name",
)
.await
@@ -1048,13 +1052,13 @@ async fn test_successive_merges_different_columns_read_file_group() {
// Second MERGE: update 'value' — reads from file group (base + name-partial)
register_source(
- &handler,
- "CREATE TABLE src_dc2 (id INT, value INT) AS VALUES (1, 100), (2, 200)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_dc2 (id INT, value INT) AS VALUES (1, 100), (2, 200)",
)
.await;
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_dc2 s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_dc2 s ON t.id = s.id \
WHEN MATCHED THEN UPDATE SET value = s.value",
)
.await
@@ -1064,7 +1068,7 @@ async fn test_successive_merges_different_columns_read_file_group() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -1081,10 +1085,10 @@ async fn test_merge_insert_reordered_columns() {
// Table schema: (id INT, name STRING, value INT)
// INSERT specifies: (value, name, id) — reversed order
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- setup_data_evolution_table(&handler).await;
+ let sql_context = create_sql_context(catalog.clone()).await;
+ setup_data_evolution_table(&sql_context).await;
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.target (id, name, value) VALUES (1, 'alice', 10)")
.await
.unwrap()
@@ -1092,18 +1096,18 @@ async fn test_merge_insert_reordered_columns() {
.await
.unwrap();
- enable_data_evolution(&handler).await;
+ enable_data_evolution(&sql_context).await;
register_source(
- &handler,
- "CREATE TABLE src_reorder (id INT, name VARCHAR, value INT) AS VALUES (2, 'bob', 20), (1, 'ALICE', 11)",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_reorder (id INT, name VARCHAR, value INT) AS VALUES (2, 'bob', 20), (1, 'ALICE', 11)",
)
.await;
// INSERT columns in reversed order: (value, name, id)
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.target t USING src_reorder s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.target t USING datafusion.public.src_reorder s ON t.id = s.id \
WHEN NOT MATCHED THEN INSERT (value, name, id) VALUES (s.value, s.name, s.id)",
)
.await
@@ -1113,7 +1117,7 @@ async fn test_merge_insert_reordered_columns() {
.unwrap();
let rows = collect_rows_3col(
- &handler,
+ &sql_context,
"SELECT id, name, value FROM paimon.test_db.target ORDER BY id",
)
.await;
@@ -1131,10 +1135,13 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() {
// Verifies column reordering on a partitioned table where mis-mapping
// would cause data to land in the wrong partition.
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.part_tbl (\
dt STRING, id INT NOT NULL, name STRING\
@@ -1145,7 +1152,7 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.part_tbl (dt, id, name) VALUES ('2024-01-01', 1, 'alice')")
.await
.unwrap()
@@ -1153,18 +1160,18 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() {
.await
.unwrap();
- handler.sql("ALTER TABLE paimon.test_db.part_tbl SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
+ sql_context.sql("ALTER TABLE paimon.test_db.part_tbl SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
register_source(
- &handler,
- "CREATE TABLE src_pt_reorder (id INT, name VARCHAR, dt VARCHAR) AS VALUES (2, 'bob', '2024-02-01'), (1, 'ALICE', '2024-01-01')",
+ &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')",
)
.await;
// INSERT with columns in different order than table schema: (name, id, dt) vs table (dt, id, name)
- handler
+ sql_context
.sql(
- "MERGE INTO paimon.test_db.part_tbl t USING src_pt_reorder s ON t.id = s.id \
+ "MERGE INTO paimon.test_db.part_tbl t USING datafusion.public.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
@@ -1173,7 +1180,7 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.part_tbl ORDER BY id")
.await
.unwrap()
@@ -1219,10 +1226,13 @@ async fn test_merge_insert_reordered_columns_on_partitioned_table() {
#[tokio::test]
async fn test_rejects_table_with_primary_keys() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.pk_target (\
id INT NOT NULL, name STRING, PRIMARY KEY (id)\
@@ -1234,16 +1244,16 @@ async fn test_rejects_table_with_primary_keys() {
.unwrap();
register_source(
- &handler,
- "CREATE TABLE src_pk (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
+ &sql_context,
+ "CREATE TABLE datafusion.public.src_pk (id INT, name VARCHAR) AS VALUES (1, 'ALICE')",
)
.await;
- handler.sql("ALTER TABLE paimon.test_db.pk_target SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
+ sql_context.sql("ALTER TABLE paimon.test_db.pk_target SET TBLPROPERTIES('data-evolution.enabled' = 'true')").await.unwrap();
assert_merge_error(
- &handler,
- "MERGE INTO paimon.test_db.pk_target t USING src_pk s ON t.id = s.id \
+ &sql_context,
+ "MERGE INTO paimon.test_db.pk_target t USING datafusion.public.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/pk_tables.rs b/crates/integrations/datafusion/tests/pk_tables.rs
index e5d3d92..fe09f98 100644
--- a/crates/integrations/datafusion/tests/pk_tables.rs
+++ b/crates/integrations/datafusion/tests/pk_tables.rs
@@ -28,7 +28,8 @@
mod common;
use common::{
- collect_id_name, collect_id_value, create_handler, create_test_env, row_count, setup_handler,
+ collect_id_name, collect_id_value, create_sql_context, create_test_env, row_count,
+ setup_sql_context,
};
use datafusion::arrow::array::{Int32Array, StringArray};
use paimon::catalog::Identifier;
@@ -39,9 +40,9 @@ use paimon::Catalog;
/// Basic: CREATE TABLE with PK, INSERT, SELECT — verifies round-trip.
#[tokio::test]
async fn test_pk_basic_write_read() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t1 (
id INT NOT NULL, name STRING,
@@ -51,7 +52,7 @@ async fn test_pk_basic_write_read() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t1 VALUES (1, 'alice'), (2, 'bob'), (3, 'carol')")
.await
.unwrap()
@@ -60,7 +61,7 @@ async fn test_pk_basic_write_read() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t1 ORDER BY id",
)
.await;
@@ -80,9 +81,9 @@ async fn test_pk_basic_write_read() {
/// Duplicate keys in a single INSERT — last value wins (Deduplicate engine).
#[tokio::test]
async fn test_pk_dedup_within_single_commit() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dedup (
id INT NOT NULL, value INT,
@@ -92,7 +93,7 @@ async fn test_pk_dedup_within_single_commit() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dedup VALUES (1, 10), (2, 20), (1, 100), (2, 200)")
.await
.unwrap()
@@ -101,7 +102,7 @@ async fn test_pk_dedup_within_single_commit() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_dedup ORDER BY id",
)
.await;
@@ -115,9 +116,9 @@ async fn test_pk_dedup_within_single_commit() {
/// Two commits with overlapping keys — second commit's values win.
#[tokio::test]
async fn test_pk_dedup_across_commits() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_cross (
id INT NOT NULL, name STRING,
@@ -128,7 +129,7 @@ async fn test_pk_dedup_across_commits() {
.unwrap();
// First commit
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_cross VALUES (1, 'alice'), (2, 'bob'), (3, 'carol')")
.await
.unwrap()
@@ -137,7 +138,7 @@ async fn test_pk_dedup_across_commits() {
.unwrap();
// Second commit: update id=1,3, add id=4
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_cross VALUES (1, 'alice-v2'), (3, 'carol-v2'), (4, 'dave')")
.await
.unwrap()
@@ -146,7 +147,7 @@ async fn test_pk_dedup_across_commits() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_cross ORDER BY id",
)
.await;
@@ -167,9 +168,9 @@ async fn test_pk_dedup_across_commits() {
/// Three successive commits — verifies sequence number tracking across commits.
#[tokio::test]
async fn test_pk_three_commits() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_three (
id INT NOT NULL, value INT,
@@ -179,7 +180,7 @@ async fn test_pk_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_three VALUES (1, 10), (2, 20)")
.await
.unwrap()
@@ -187,7 +188,7 @@ async fn test_pk_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_three VALUES (2, 200), (3, 30)")
.await
.unwrap()
@@ -195,7 +196,7 @@ async fn test_pk_three_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_three VALUES (1, 100), (3, 300), (4, 40)")
.await
.unwrap()
@@ -204,7 +205,7 @@ async fn test_pk_three_commits() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_three ORDER BY id",
)
.await;
@@ -217,9 +218,9 @@ async fn test_pk_three_commits() {
/// Partitioned PK table: dedup happens per-partition independently.
#[tokio::test]
async fn test_pk_partitioned_write_read() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_part (
dt STRING, id INT NOT NULL, name STRING,
@@ -230,7 +231,7 @@ async fn test_pk_partitioned_write_read() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -243,7 +244,7 @@ async fn test_pk_partitioned_write_read() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_part ORDER BY id, name",
)
.await;
@@ -262,9 +263,9 @@ async fn test_pk_partitioned_write_read() {
/// Partitioned PK table: dedup across commits within same partition.
#[tokio::test]
async fn test_pk_partitioned_dedup_across_commits() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_part_dedup (
dt STRING, id INT NOT NULL, name STRING,
@@ -275,7 +276,7 @@ async fn test_pk_partitioned_dedup_across_commits() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part_dedup VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -288,7 +289,7 @@ async fn test_pk_partitioned_dedup_across_commits() {
.unwrap();
// Update within partition 2024-01-01
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part_dedup VALUES \
('2024-01-01', 1, 'alice-v2'), ('2024-01-02', 2, 'dave')",
@@ -299,7 +300,7 @@ async fn test_pk_partitioned_dedup_across_commits() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_part_dedup ORDER BY dt, id")
.await
.unwrap()
@@ -344,9 +345,9 @@ async fn test_pk_partitioned_dedup_across_commits() {
/// Partition filter on PK table — only matching partition returned.
#[tokio::test]
async fn test_pk_partitioned_filter() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_part_filter (
dt STRING, id INT NOT NULL, name STRING,
@@ -357,7 +358,7 @@ async fn test_pk_partitioned_filter() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part_filter VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -370,7 +371,7 @@ async fn test_pk_partitioned_filter() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_part_filter WHERE dt = '2024-01-01' ORDER BY id",
)
.await;
@@ -383,9 +384,9 @@ async fn test_pk_partitioned_filter() {
/// Multiple buckets: rows are distributed by PK hash, dedup still works.
#[tokio::test]
async fn test_pk_multi_bucket() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_mbucket (
id INT NOT NULL, value INT,
@@ -395,7 +396,7 @@ async fn test_pk_multi_bucket() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_mbucket VALUES \
(1, 10), (2, 20), (3, 30), (4, 40), (5, 50), (6, 60), (7, 70), (8, 80)",
@@ -407,7 +408,7 @@ async fn test_pk_multi_bucket() {
.unwrap();
// Update some keys
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_mbucket VALUES (2, 200), (5, 500), (8, 800)")
.await
.unwrap()
@@ -416,7 +417,7 @@ async fn test_pk_multi_bucket() {
.unwrap();
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_mbucket ORDER BY id",
)
.await;
@@ -441,9 +442,9 @@ async fn test_pk_multi_bucket() {
/// SELECT only a subset of columns from a PK table.
#[tokio::test]
async fn test_pk_column_projection() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_proj (
id INT NOT NULL, name STRING, value INT,
@@ -453,7 +454,7 @@ async fn test_pk_column_projection() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_proj VALUES (1, 'alice', 10), (2, 'bob', 20)")
.await
.unwrap()
@@ -462,7 +463,7 @@ async fn test_pk_column_projection() {
.unwrap();
// Update id=1
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_proj VALUES (1, 'alice-v2', 100)")
.await
.unwrap()
@@ -471,7 +472,7 @@ async fn test_pk_column_projection() {
.unwrap();
// Project only name
- let batches = handler
+ let batches = sql_context
.sql("SELECT name FROM paimon.test_db.t_proj ORDER BY name")
.await
.unwrap()
@@ -495,7 +496,7 @@ async fn test_pk_column_projection() {
// Project only value
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_proj ORDER BY id",
)
.await;
@@ -507,9 +508,9 @@ async fn test_pk_column_projection() {
/// sequence.field: dedup uses the specified field instead of system sequence number.
#[tokio::test]
async fn test_pk_sequence_field() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_seqf (
id INT NOT NULL, version INT, name STRING,
@@ -520,7 +521,7 @@ async fn test_pk_sequence_field() {
.unwrap();
// First commit: version=2 for id=1
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_seqf VALUES (1, 2, 'alice-v2'), (2, 1, 'bob-v1')")
.await
.unwrap()
@@ -529,7 +530,7 @@ async fn test_pk_sequence_field() {
.unwrap();
// Second commit: version=1 for id=1 (older), version=2 for id=2 (newer)
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_seqf VALUES (1, 1, 'alice-v1'), (2, 2, 'bob-v2')")
.await
.unwrap()
@@ -538,7 +539,7 @@ async fn test_pk_sequence_field() {
.unwrap();
let rows = collect_id_name(
- &handler,
+ &sql_context,
"SELECT id, name FROM paimon.test_db.t_seqf ORDER BY id",
)
.await;
@@ -557,9 +558,9 @@ async fn test_pk_sequence_field() {
/// INSERT OVERWRITE on a partitioned PK table replaces the partition.
#[tokio::test]
async fn test_pk_insert_overwrite_partitioned() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_overwrite (
dt STRING, id INT NOT NULL, name STRING,
@@ -570,7 +571,7 @@ async fn test_pk_insert_overwrite_partitioned() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_overwrite VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -583,7 +584,7 @@ async fn test_pk_insert_overwrite_partitioned() {
.unwrap();
// Overwrite partition 2024-01-01
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_overwrite VALUES ('2024-01-01', 10, 'new_alice')")
.await
.unwrap()
@@ -591,7 +592,7 @@ async fn test_pk_insert_overwrite_partitioned() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_overwrite ORDER BY dt, id")
.await
.unwrap()
@@ -634,9 +635,9 @@ async fn test_pk_insert_overwrite_partitioned() {
/// INSERT OVERWRITE with explicit PARTITION clause (Hive-style static partition).
#[tokio::test]
async fn test_pk_insert_overwrite_with_partition_clause() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_ow_part (
dt STRING, id INT NOT NULL, name STRING,
@@ -648,7 +649,7 @@ async fn test_pk_insert_overwrite_with_partition_clause() {
.unwrap();
// Insert data into two partitions
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_ow_part VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -662,7 +663,7 @@ async fn test_pk_insert_overwrite_with_partition_clause() {
// Overwrite partition dt='2024-01-01' using Hive-style PARTITION clause.
// The SELECT only provides non-partition columns (id, name).
- handler
+ sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_ow_part PARTITION (dt = '2024-01-01') \
VALUES (10, 'new_alice'), (20, 'new_bob')",
@@ -673,7 +674,7 @@ async fn test_pk_insert_overwrite_with_partition_clause() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_ow_part ORDER BY dt, id")
.await
.unwrap()
@@ -719,9 +720,9 @@ async fn test_pk_insert_overwrite_with_partition_clause() {
/// Only specifies dt, region comes from the source query (dynamic partition).
#[tokio::test]
async fn test_pk_insert_overwrite_partial_partition_clause() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_multi_part (
dt STRING, region STRING, id INT NOT NULL, name STRING,
@@ -733,7 +734,7 @@ async fn test_pk_insert_overwrite_partial_partition_clause() {
.unwrap();
// Insert data into multiple partitions
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_multi_part VALUES \
('2024-01-01', 'us', 1, 'alice'), \
@@ -748,7 +749,7 @@ async fn test_pk_insert_overwrite_partial_partition_clause() {
// Overwrite only dt='2024-01-01', region comes from VALUES (dynamic).
// This should overwrite all sub-partitions under dt='2024-01-01' that appear in the data.
- handler
+ sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_multi_part PARTITION (dt = '2024-01-01') \
VALUES ('us', 10, 'new_alice')",
@@ -759,7 +760,7 @@ async fn test_pk_insert_overwrite_partial_partition_clause() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, region, id, name FROM paimon.test_db.t_multi_part ORDER BY dt, region, id")
.await
.unwrap()
@@ -820,9 +821,9 @@ async fn test_pk_insert_overwrite_partial_partition_clause() {
/// INSERT OVERWRITE with PARTITION clause and empty source truncates the partition.
#[tokio::test]
async fn test_pk_insert_overwrite_partition_truncate() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_trunc (
dt STRING, id INT NOT NULL, name STRING,
@@ -833,7 +834,7 @@ async fn test_pk_insert_overwrite_partition_truncate() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_trunc VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -846,7 +847,7 @@ async fn test_pk_insert_overwrite_partition_truncate() {
.unwrap();
// Overwrite dt='2024-01-01' with empty source — should truncate that partition
- handler
+ sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_trunc PARTITION (dt = '2024-01-01') \
SELECT id, name FROM paimon.test_db.t_trunc WHERE false",
@@ -857,7 +858,7 @@ async fn test_pk_insert_overwrite_partition_truncate() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_trunc ORDER BY dt, id")
.await
.unwrap()
@@ -898,9 +899,9 @@ async fn test_pk_insert_overwrite_partition_truncate() {
/// PARTITION clause with a non-partition column should fail.
#[tokio::test]
async fn test_pk_insert_overwrite_partition_non_partition_column_error() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_err (
dt STRING, id INT NOT NULL, name STRING,
@@ -911,7 +912,7 @@ async fn test_pk_insert_overwrite_partition_non_partition_column_error() {
.await
.unwrap();
- let result = handler
+ let result = sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_err PARTITION (name = 'alice') \
VALUES (1)",
@@ -930,9 +931,9 @@ async fn test_pk_insert_overwrite_partition_non_partition_column_error() {
/// not drop all partitions.
#[tokio::test]
async fn test_pk_insert_overwrite_dynamic_partition_preserves_other_partitions() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dyn (
dt STRING, id INT NOT NULL, name STRING,
@@ -943,7 +944,7 @@ async fn test_pk_insert_overwrite_dynamic_partition_preserves_other_partitions()
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_dyn VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-02', 2, 'bob')",
@@ -956,7 +957,7 @@ async fn test_pk_insert_overwrite_dynamic_partition_preserves_other_partitions()
// Dynamic partition overwrite: PARTITION (dt) with no static value.
// Should only overwrite partitions present in the source data.
- handler
+ sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_dyn PARTITION (dt) \
VALUES ('2024-01-01', 10, 'new_alice')",
@@ -967,7 +968,7 @@ async fn test_pk_insert_overwrite_dynamic_partition_preserves_other_partitions()
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_dyn ORDER BY dt, id")
.await
.unwrap()
@@ -1011,9 +1012,9 @@ async fn test_pk_insert_overwrite_dynamic_partition_preserves_other_partitions()
/// Source query with wrong column count should fail even when the result is empty.
#[tokio::test]
async fn test_pk_insert_overwrite_empty_source_wrong_columns_error() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_empty_err (
dt STRING, id INT NOT NULL, name STRING,
@@ -1024,7 +1025,7 @@ async fn test_pk_insert_overwrite_empty_source_wrong_columns_error() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_empty_err VALUES \
('2024-01-01', 1, 'alice')",
@@ -1036,7 +1037,7 @@ async fn test_pk_insert_overwrite_empty_source_wrong_columns_error() {
.unwrap();
// Source only produces `id` but target expects `id, name` — should fail
- let result = handler
+ let result = sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_empty_err PARTITION (dt = '2024-01-01') \
SELECT id FROM paimon.test_db.t_empty_err WHERE false",
@@ -1054,9 +1055,9 @@ async fn test_pk_insert_overwrite_empty_source_wrong_columns_error() {
/// Explicit target column list after PARTITION should reorder source columns to match schema.
#[tokio::test]
async fn test_pk_insert_overwrite_with_after_columns_reorder() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_reorder (
dt STRING, id INT NOT NULL, name STRING,
@@ -1068,7 +1069,7 @@ async fn test_pk_insert_overwrite_with_after_columns_reorder() {
.unwrap();
// Insert with columns in reversed order: (name, id) instead of schema order (id, name)
- handler
+ sql_context
.sql(
"INSERT OVERWRITE paimon.test_db.t_reorder (name, id) PARTITION (dt = '2024-01-01') \
VALUES ('alice', 1), ('bob', 2)",
@@ -1079,7 +1080,7 @@ async fn test_pk_insert_overwrite_with_after_columns_reorder() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, name FROM paimon.test_db.t_reorder ORDER BY id")
.await
.unwrap()
@@ -1125,9 +1126,9 @@ async fn test_pk_insert_overwrite_with_after_columns_reorder() {
/// Composite PK with multiple columns.
#[tokio::test]
async fn test_pk_composite_key() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_composite (
region STRING NOT NULL, id INT NOT NULL, value INT,
@@ -1137,7 +1138,7 @@ async fn test_pk_composite_key() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_composite VALUES \
('us', 1, 10), ('eu', 1, 20), ('us', 2, 30)",
@@ -1149,7 +1150,7 @@ async fn test_pk_composite_key() {
.unwrap();
// Update (us, 1) — (eu, 1) should be untouched
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_composite VALUES ('us', 1, 100)")
.await
.unwrap()
@@ -1157,7 +1158,7 @@ async fn test_pk_composite_key() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT region, id, value FROM paimon.test_db.t_composite ORDER BY region, id")
.await
.unwrap()
@@ -1199,9 +1200,9 @@ async fn test_pk_composite_key() {
/// Reading an empty PK table returns zero rows.
#[tokio::test]
async fn test_pk_empty_table_read() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_empty (
id INT NOT NULL, name STRING,
@@ -1211,7 +1212,7 @@ async fn test_pk_empty_table_read() {
.await
.unwrap();
- let count = row_count(&handler, "SELECT id, name FROM paimon.test_db.t_empty").await;
+ let count = row_count(&sql_context, "SELECT id, name FROM paimon.test_db.t_empty").await;
assert_eq!(count, 0);
}
@@ -1220,9 +1221,9 @@ async fn test_pk_empty_table_read() {
/// Many rows with overlapping keys in a single commit.
#[tokio::test]
async fn test_pk_large_batch_dedup() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_large (
id INT NOT NULL, value INT,
@@ -1240,7 +1241,7 @@ async fn test_pk_large_batch_dedup() {
values2.push(format!("({i}, {})", i * 10)); // id=i, value=i*10
}
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.t_large VALUES {}",
values1.join(", ")
@@ -1251,7 +1252,7 @@ async fn test_pk_large_batch_dedup() {
.await
.unwrap();
- handler
+ sql_context
.sql(&format!(
"INSERT INTO paimon.test_db.t_large VALUES {}",
values2.join(", ")
@@ -1262,12 +1263,12 @@ async fn test_pk_large_batch_dedup() {
.await
.unwrap();
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.t_large").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t_large").await;
assert_eq!(count, 100, "Dedup should keep exactly 100 unique keys");
// Spot-check a few values
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_large WHERE id IN (1, 50, 100) ORDER BY id",
)
.await;
@@ -1279,9 +1280,9 @@ async fn test_pk_large_batch_dedup() {
/// Partitioned PK table with multiple buckets.
#[tokio::test]
async fn test_pk_partitioned_multi_bucket() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_part_mb (
dt STRING, id INT NOT NULL, value INT,
@@ -1292,7 +1293,7 @@ async fn test_pk_partitioned_multi_bucket() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part_mb VALUES \
('2024-01-01', 1, 10), ('2024-01-01', 2, 20), \
@@ -1306,7 +1307,7 @@ async fn test_pk_partitioned_multi_bucket() {
.unwrap();
// Update across partitions
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_part_mb VALUES \
('2024-01-01', 2, 222), ('2024-01-02', 1, 111)",
@@ -1317,7 +1318,7 @@ async fn test_pk_partitioned_multi_bucket() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT dt, id, value FROM paimon.test_db.t_part_mb ORDER BY dt, id")
.await
.unwrap()
@@ -1362,9 +1363,9 @@ async fn test_pk_partitioned_multi_bucket() {
/// PK table with changelog-producer=input should be rejected.
#[tokio::test]
async fn test_pk_reject_changelog_producer_input() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_changelog (
id INT NOT NULL, name STRING,
@@ -1374,7 +1375,7 @@ async fn test_pk_reject_changelog_producer_input() {
.await
.unwrap();
- let result = handler
+ let result = sql_context
.sql("INSERT INTO paimon.test_db.t_changelog VALUES (1, 'alice')")
.await;
@@ -1393,9 +1394,9 @@ async fn test_pk_reject_changelog_producer_input() {
/// PK table with STRING primary key.
#[tokio::test]
async fn test_pk_string_key() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_strpk (
code STRING NOT NULL, name STRING,
@@ -1405,7 +1406,7 @@ async fn test_pk_string_key() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_strpk VALUES \
('A001', 'alice'), ('B002', 'bob'), ('C003', 'carol')",
@@ -1417,7 +1418,7 @@ async fn test_pk_string_key() {
.unwrap();
// Update A001
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_strpk VALUES ('A001', 'alice-v2')")
.await
.unwrap()
@@ -1425,7 +1426,7 @@ async fn test_pk_string_key() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT code, name FROM paimon.test_db.t_strpk ORDER BY code")
.await
.unwrap()
@@ -1463,9 +1464,9 @@ async fn test_pk_string_key() {
/// PK table with many value columns — verifies all columns survive dedup.
#[tokio::test]
async fn test_pk_multiple_value_columns() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_multi_val (
id INT NOT NULL, col_a INT, col_b STRING, col_c INT,
@@ -1475,7 +1476,7 @@ async fn test_pk_multiple_value_columns() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_multi_val VALUES (1, 10, 'x', 100), (2, 20, 'y', 200)")
.await
.unwrap()
@@ -1483,7 +1484,7 @@ async fn test_pk_multiple_value_columns() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_multi_val VALUES (1, 11, 'xx', 111)")
.await
.unwrap()
@@ -1491,7 +1492,7 @@ async fn test_pk_multiple_value_columns() {
.await
.unwrap();
- let batches = handler
+ let batches = sql_context
.sql("SELECT id, col_a, col_b, col_c FROM paimon.test_db.t_multi_val ORDER BY id")
.await
.unwrap()
@@ -1547,13 +1548,13 @@ async fn test_pk_multiple_value_columns() {
#[tokio::test]
async fn test_pk_first_row_insert_overwrite() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_fr_ow (
dt STRING, id INT NOT NULL, name STRING,
@@ -1565,7 +1566,7 @@ async fn test_pk_first_row_insert_overwrite() {
.unwrap();
// First commit: two partitions, creates level-0 files
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_fr_ow VALUES \
('2024-01-01', 1, 'alice'), ('2024-01-01', 2, 'bob'), \
@@ -1596,7 +1597,7 @@ async fn test_pk_first_row_insert_overwrite() {
);
// INSERT OVERWRITE partition 2024-01-01 — must delete old level-0 file
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_fr_ow VALUES ('2024-01-01', 10, 'new_alice')")
.await
.unwrap()
@@ -1622,7 +1623,7 @@ async fn test_pk_first_row_insert_overwrite() {
);
// Second overwrite on the same partition — no stale files should accumulate
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_fr_ow VALUES ('2024-01-01', 20, 'newer_alice')")
.await
.unwrap()
@@ -1654,13 +1655,13 @@ async fn test_pk_first_row_insert_overwrite() {
#[tokio::test]
async fn test_postpone_write_invisible_to_select() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_postpone (
id INT NOT NULL, value INT,
@@ -1671,7 +1672,7 @@ async fn test_postpone_write_invisible_to_select() {
.unwrap();
// Write data
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_postpone VALUES (1, 10), (2, 20), (3, 30)")
.await
.unwrap()
@@ -1695,7 +1696,7 @@ async fn test_postpone_write_invisible_to_select() {
assert_eq!(file_count, 1, "scan_all_files should find 1 postpone file");
// Normal SELECT should return 0 rows (postpone files are invisible)
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.t_postpone").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t_postpone").await;
assert_eq!(count, 0, "SELECT should return 0 rows for postpone table");
}
@@ -1703,13 +1704,13 @@ async fn test_postpone_write_invisible_to_select() {
#[tokio::test]
async fn test_postpone_insert_overwrite() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
- handler
+ let sql_context = create_sql_context(catalog.clone()).await;
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA failed");
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_postpone_ow (
id INT NOT NULL, value INT,
@@ -1720,7 +1721,7 @@ async fn test_postpone_insert_overwrite() {
.unwrap();
// First commit
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_postpone_ow VALUES (1, 10), (2, 20)")
.await
.unwrap()
@@ -1743,7 +1744,7 @@ async fn test_postpone_insert_overwrite() {
assert_eq!(file_count, 1, "After INSERT: 1 postpone file");
// INSERT OVERWRITE should replace old file
- handler
+ sql_context
.sql("INSERT OVERWRITE paimon.test_db.t_postpone_ow VALUES (3, 30)")
.await
.unwrap()
@@ -1777,9 +1778,9 @@ async fn test_postpone_insert_overwrite() {
/// `trimmed_primary_keys()`, causing bucket pruning to target the wrong bucket.
#[tokio::test]
async fn test_pk_partitioned_fixed_bucket_predicate_query() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_bk_pred (
pt STRING, id INT NOT NULL, value INT,
@@ -1790,7 +1791,7 @@ async fn test_pk_partitioned_fixed_bucket_predicate_query() {
.await
.unwrap();
- handler
+ sql_context
.sql(
"INSERT INTO paimon.test_db.t_bk_pred VALUES \
('a', 1, 10), ('a', 2, 20), ('b', 3, 30), ('b', 4, 40)",
@@ -1803,14 +1804,14 @@ async fn test_pk_partitioned_fixed_bucket_predicate_query() {
// Query with both partition and PK columns in predicate
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_bk_pred WHERE pt = 'a' AND id = 1",
)
.await;
assert_eq!(rows, vec![(1, 10)], "Predicate query must find the row");
let rows = collect_id_value(
- &handler,
+ &sql_context,
"SELECT id, value FROM paimon.test_db.t_bk_pred WHERE pt = 'b' AND id = 4",
)
.await;
@@ -1826,9 +1827,9 @@ async fn test_pk_partitioned_fixed_bucket_predicate_query() {
/// on compaction to produce higher-level files).
#[tokio::test]
async fn test_pk_dv_deduplicate_read_no_error() {
- let (_tmp, handler) = setup_handler().await;
+ let (_tmp, sql_context) = setup_sql_context().await;
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.test_db.t_dv_dedup (
id INT NOT NULL, value INT,
@@ -1838,7 +1839,7 @@ async fn test_pk_dv_deduplicate_read_no_error() {
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dv_dedup VALUES (1, 10), (2, 20)")
.await
.unwrap()
@@ -1847,7 +1848,7 @@ async fn test_pk_dv_deduplicate_read_no_error() {
.unwrap();
// Second commit with overlapping key — creates level-0 files
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t_dv_dedup VALUES (2, 200), (3, 30)")
.await
.unwrap()
@@ -1860,7 +1861,7 @@ async fn test_pk_dv_deduplicate_read_no_error() {
// and get skipped — count may be 0, but the read must succeed without error.
// Before the fix, this would hard-fail with "KeyValueFileReader does not
// support deletion vectors".
- let result = handler
+ let result = sql_context
.sql("SELECT * FROM paimon.test_db.t_dv_dedup")
.await
.unwrap()
diff --git a/crates/integrations/datafusion/tests/procedures.rs b/crates/integrations/datafusion/tests/procedures.rs
index 532978a..c1c670c 100644
--- a/crates/integrations/datafusion/tests/procedures.rs
+++ b/crates/integrations/datafusion/tests/procedures.rs
@@ -17,46 +17,50 @@
mod common;
-use common::{exec, row_count, setup_handler};
+use common::{exec, row_count, setup_sql_context};
-async fn setup_table_with_snapshots() -> (tempfile::TempDir, paimon_datafusion::PaimonSqlHandler) {
- let (tmp, handler) = setup_handler().await;
+async fn setup_table_with_snapshots() -> (tempfile::TempDir, paimon_datafusion::SQLContext) {
+ let (tmp, sql_context) = setup_sql_context().await;
exec(
- &handler,
+ &sql_context,
"CREATE TABLE paimon.test_db.t1 (id INT, name VARCHAR(100), PRIMARY KEY (id))",
)
.await;
// Insert data to create snapshot 1
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t1 VALUES (1, 'alice')",
)
.await;
// Insert data to create snapshot 2
- exec(&handler, "INSERT INTO paimon.test_db.t1 VALUES (2, 'bob')").await;
+ exec(
+ &sql_context,
+ "INSERT INTO paimon.test_db.t1 VALUES (2, 'bob')",
+ )
+ .await;
// Insert data to create snapshot 3
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t1 VALUES (3, 'charlie')",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
#[tokio::test]
async fn test_create_tag() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// Create tag from latest snapshot
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1')",
)
.await;
// Verify tag exists via $tags system table
let count = row_count(
- &handler,
+ &sql_context,
"SELECT * FROM paimon.test_db.`t1$tags` WHERE tag_name = 'v1'",
)
.await;
@@ -65,16 +69,16 @@ async fn test_create_tag() {
#[tokio::test]
async fn test_create_tag_with_snapshot_id() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1', snapshot_id => '1')",
)
.await;
let count = row_count(
- &handler,
+ &sql_context,
"SELECT * FROM paimon.test_db.`t1$tags` WHERE tag_name = 'v1'",
)
.await;
@@ -83,15 +87,15 @@ async fn test_create_tag_with_snapshot_id() {
#[tokio::test]
async fn test_create_tag_already_exists() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1')",
)
.await;
- let result = handler
+ let result = sql_context
.sql("CALL sys.create_tag(table => 'test_db.t1', tag => 'v1')")
.await;
assert!(result.is_err());
@@ -100,21 +104,21 @@ async fn test_create_tag_already_exists() {
#[tokio::test]
async fn test_delete_tag() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1')",
)
.await;
exec(
- &handler,
+ &sql_context,
"CALL sys.delete_tag(table => 'test_db.t1', tag => 'v1')",
)
.await;
let count = row_count(
- &handler,
+ &sql_context,
"SELECT * FROM paimon.test_db.`t1$tags` WHERE tag_name = 'v1'",
)
.await;
@@ -123,73 +127,73 @@ async fn test_delete_tag() {
#[tokio::test]
async fn test_delete_multiple_tags() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1')",
)
.await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v2', snapshot_id => '1')",
)
.await;
exec(
- &handler,
+ &sql_context,
"CALL sys.delete_tag(table => 'test_db.t1', tag => 'v1,v2')",
)
.await;
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.`t1$tags`").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.`t1$tags`").await;
assert_eq!(count, 0);
}
#[tokio::test]
async fn test_rollback_to_snapshot() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// We have 3 snapshots. Rollback to snapshot 1.
exec(
- &handler,
+ &sql_context,
"CALL sys.rollback_to(table => 'test_db.t1', snapshot_id => '1')",
)
.await;
// After rollback, only snapshot 1 data should be visible
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.t1").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t1").await;
assert_eq!(count, 1);
}
#[tokio::test]
async fn test_rollback_to_tag() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// Create tag on snapshot 1
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v1', snapshot_id => '1')",
)
.await;
// Rollback to tag
exec(
- &handler,
+ &sql_context,
"CALL sys.rollback_to(table => 'test_db.t1', tag => 'v1')",
)
.await;
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.t1").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t1").await;
assert_eq!(count, 1);
}
#[tokio::test]
async fn test_rollback_to_timestamp() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// Get the timestamp of snapshot 1 from $snapshots system table
- let batches = handler
+ let batches = sql_context
.sql("SELECT snapshot_id, commit_time FROM paimon.test_db.`t1$snapshots` ORDER BY snapshot_id")
.await
.unwrap()
@@ -207,23 +211,23 @@ async fn test_rollback_to_timestamp() {
.value(0);
exec(
- &handler,
+ &sql_context,
&format!(
"CALL sys.rollback_to_timestamp(table => 'test_db.t1', timestamp => '{snap1_time}')"
),
)
.await;
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.t1").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t1").await;
assert_eq!(count, 1);
}
#[tokio::test]
async fn test_create_tag_from_timestamp() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// Get the timestamp of snapshot 2
- let batches = handler
+ let batches = sql_context
.sql("SELECT snapshot_id, commit_time FROM paimon.test_db.`t1$snapshots` ORDER BY snapshot_id")
.await
.unwrap()
@@ -240,7 +244,7 @@ async fn test_create_tag_from_timestamp() {
.value(1);
exec(
- &handler,
+ &sql_context,
&format!(
"CALL sys.create_tag_from_timestamp(table => 'test_db.t1', tag => 'ts_tag', timestamp => '{snap2_time}')"
),
@@ -248,7 +252,7 @@ async fn test_create_tag_from_timestamp() {
.await;
let count = row_count(
- &handler,
+ &sql_context,
"SELECT * FROM paimon.test_db.`t1$tags` WHERE tag_name = 'ts_tag'",
)
.await;
@@ -257,27 +261,27 @@ async fn test_create_tag_from_timestamp() {
#[tokio::test]
async fn test_rollback_cleans_newer_tags() {
- let (_tmp, handler) = setup_table_with_snapshots().await;
+ let (_tmp, sql_context) = setup_table_with_snapshots().await;
// Create tags on snapshot 2 and 3
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v2', snapshot_id => '2')",
)
.await;
exec(
- &handler,
+ &sql_context,
"CALL sys.create_tag(table => 'test_db.t1', tag => 'v3', snapshot_id => '3')",
)
.await;
// Rollback to snapshot 1 — tags v2 and v3 should be cleaned
exec(
- &handler,
+ &sql_context,
"CALL sys.rollback_to(table => 'test_db.t1', snapshot_id => '1')",
)
.await;
- let count = row_count(&handler, "SELECT * FROM paimon.test_db.`t1$tags`").await;
+ let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.`t1$tags`").await;
assert_eq!(count, 0);
}
diff --git a/crates/integrations/datafusion/tests/sql_handler_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs
similarity index 65%
rename from crates/integrations/datafusion/tests/sql_handler_tests.rs
rename to crates/integrations/datafusion/tests/sql_context_tests.rs
index cd78b20..7d27422 100644
--- a/crates/integrations/datafusion/tests/sql_handler_tests.rs
+++ b/crates/integrations/datafusion/tests/sql_context_tests.rs
@@ -15,16 +15,15 @@
// specific language governing permissions and limitations
// under the License.
-//! SQL handler integration tests for paimon-datafusion.
+//! SQL context integration tests for paimon-datafusion.
use std::sync::Arc;
use datafusion::catalog::CatalogProvider;
-use datafusion::prelude::SessionContext;
use paimon::catalog::Identifier;
use paimon::spec::{ArrayType, BlobType, DataType, IntType, MapType, VarCharType};
use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
-use paimon_datafusion::{PaimonCatalogProvider, PaimonSqlHandler};
+use paimon_datafusion::{PaimonCatalogProvider, SQLContext};
use tempfile::TempDir;
fn create_test_env() -> (TempDir, Arc<FileSystemCatalog>) {
@@ -36,9 +35,10 @@ fn create_test_env() -> (TempDir, Arc<FileSystemCatalog>) {
(temp_dir, Arc::new(catalog))
}
-fn create_handler(catalog: Arc<FileSystemCatalog>) -> PaimonSqlHandler {
- let ctx = SessionContext::new();
- PaimonSqlHandler::new(ctx, catalog, "paimon").unwrap()
+async fn create_sql_context(catalog: Arc<FileSystemCatalog>) -> SQLContext {
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog).await.unwrap();
+ ctx
}
// ======================= CREATE / DROP SCHEMA =======================
@@ -46,9 +46,9 @@ fn create_handler(catalog: Arc<FileSystemCatalog>) -> PaimonSqlHandler {
#[tokio::test]
async fn test_create_schema() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
- handler
+ sql_context
.sql("CREATE SCHEMA paimon.test_db")
.await
.expect("CREATE SCHEMA should succeed");
@@ -63,14 +63,14 @@ async fn test_create_schema() {
#[tokio::test]
async fn test_drop_schema() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("drop_me", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql("DROP SCHEMA paimon.drop_me CASCADE")
.await
.expect("DROP SCHEMA should succeed");
@@ -106,14 +106,14 @@ async fn test_schema_names_via_catalog_provider() {
#[tokio::test]
async fn test_create_table() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.mydb.users (
id INT NOT NULL,
@@ -144,14 +144,14 @@ async fn test_create_table() {
#[tokio::test]
async fn test_create_table_with_blob_type() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.mydb.assets (
id INT NOT NULL,
@@ -178,14 +178,14 @@ async fn test_create_table_with_blob_type() {
#[tokio::test]
async fn test_create_table_with_partition() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.mydb.events (
id INT NOT NULL,
@@ -215,14 +215,14 @@ async fn test_create_table_with_partition() {
#[tokio::test]
async fn test_create_table_partitioned_by_rejects_typed_columns() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- let err = handler
+ let err = sql_context
.sql(
"CREATE TABLE paimon.mydb.events (
id INT NOT NULL,
@@ -242,7 +242,7 @@ async fn test_create_table_partitioned_by_rejects_typed_columns() {
#[tokio::test]
async fn test_create_table_if_not_exists() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
@@ -254,10 +254,13 @@ async fn test_create_table_if_not_exists() {
)";
// First create should succeed
- handler.sql(sql).await.expect("First CREATE should succeed");
+ sql_context
+ .sql(sql)
+ .await
+ .expect("First CREATE should succeed");
// Second create with IF NOT EXISTS should also succeed
- handler
+ sql_context
.sql(sql)
.await
.expect("Second CREATE with IF NOT EXISTS should succeed");
@@ -266,14 +269,14 @@ async fn test_create_table_if_not_exists() {
#[tokio::test]
async fn test_create_external_table_rejected() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- let result = handler
+ let result = sql_context
.sql(
"CREATE EXTERNAL TABLE paimon.mydb.bad (
id INT NOT NULL
@@ -295,14 +298,14 @@ async fn test_create_external_table_rejected() {
#[tokio::test]
async fn test_create_table_with_array_and_map() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.mydb.complex_types (
id INT NOT NULL,
@@ -349,14 +352,14 @@ async fn test_create_table_with_array_and_map() {
#[tokio::test]
async fn test_create_table_with_row_type() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
.await
.unwrap();
- handler
+ sql_context
.sql(
"CREATE TABLE paimon.mydb.row_table (
id INT NOT NULL,
@@ -393,7 +396,7 @@ async fn test_create_table_with_row_type() {
#[tokio::test]
async fn test_drop_table() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
@@ -419,7 +422,7 @@ async fn test_drop_table() {
.unwrap()
.contains(&"to_drop".to_string()));
- handler
+ sql_context
.sql("DROP TABLE paimon.mydb.to_drop")
.await
.expect("DROP TABLE should succeed");
@@ -439,7 +442,7 @@ async fn test_drop_table() {
#[tokio::test]
async fn test_alter_table_add_column() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
@@ -463,7 +466,7 @@ async fn test_alter_table_add_column() {
.unwrap();
// ALTER TABLE is not yet implemented in FileSystemCatalog, so we expect an error
- let result = handler
+ let result = sql_context
.sql("ALTER TABLE paimon.mydb.alter_test ADD COLUMN age INT")
.await;
@@ -484,7 +487,7 @@ async fn test_alter_table_add_column() {
#[tokio::test]
async fn test_alter_table_rename() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
@@ -503,7 +506,7 @@ async fn test_alter_table_rename() {
.await
.unwrap();
- handler
+ sql_context
.sql("ALTER TABLE mydb.old_name RENAME TO new_name")
.await
.expect("ALTER TABLE RENAME should succeed");
@@ -520,9 +523,9 @@ async fn test_alter_table_rename() {
}
#[tokio::test]
-async fn test_ddl_handler_delegates_select() {
+async fn test_ddl_context_delegates_select() {
let (_tmp, catalog) = create_test_env();
- let handler = create_handler(catalog.clone());
+ let sql_context = create_sql_context(catalog.clone()).await;
catalog
.create_database("mydb", false, Default::default())
@@ -542,7 +545,7 @@ async fn test_ddl_handler_delegates_select() {
.unwrap();
// SELECT should be delegated to DataFusion
- let df = handler
+ let df = sql_context
.sql("SELECT * FROM paimon.mydb.t1")
.await
.expect("SELECT should be delegated to DataFusion");
@@ -552,3 +555,197 @@ async fn test_ddl_handler_delegates_select() {
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 0, "Empty table should return 0 rows");
}
+
+// ======================= MULTI-CATALOG =======================
+
+#[tokio::test]
+async fn test_multi_catalog_register_and_query() {
+ let (_tmp1, catalog1) = create_test_env();
+ let (_tmp2, catalog2) = create_test_env();
+
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("cat1", catalog1).await.unwrap();
+ ctx.register_catalog("cat2", catalog2).await.unwrap();
+
+ ctx.sql("CREATE SCHEMA cat1.db1").await.unwrap();
+ ctx.sql("CREATE SCHEMA cat2.db2").await.unwrap();
+
+ ctx.sql("CREATE TABLE cat1.db1.t (id INT NOT NULL, name STRING, PRIMARY KEY (id))")
+ .await
+ .unwrap();
+ ctx.sql("CREATE TABLE cat2.db2.t (id INT NOT NULL, value STRING, PRIMARY KEY (id))")
+ .await
+ .unwrap();
+
+ ctx.sql("INSERT INTO cat1.db1.t VALUES (1, 'alice')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ ctx.sql("INSERT INTO cat2.db2.t VALUES (2, 'hello')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+
+ let batches = ctx
+ .sql("SELECT id, name FROM cat1.db1.t")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
+
+ let batches = ctx
+ .sql("SELECT id, value FROM cat2.db2.t")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
+}
+
+#[tokio::test]
+async fn test_set_current_catalog() {
+ let (_tmp1, catalog1) = create_test_env();
+ let (_tmp2, catalog2) = create_test_env();
+
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("cat1", catalog1).await.unwrap();
+ ctx.register_catalog("cat2", catalog2).await.unwrap();
+
+ ctx.sql("CREATE SCHEMA cat1.mydb").await.unwrap();
+ ctx.sql("CREATE TABLE cat1.mydb.t (id INT NOT NULL, PRIMARY KEY (id))")
+ .await
+ .unwrap();
+
+ ctx.set_current_catalog("cat1").await.unwrap();
+ ctx.set_current_database("mydb").await.unwrap();
+
+ // Unqualified query should resolve against cat1.mydb
+ let df = ctx.sql("SELECT * FROM t").await;
+ assert!(
+ df.is_ok(),
+ "Unqualified table should resolve via current catalog/database"
+ );
+
+ // Switching to unknown catalog should fail
+ let err = ctx.set_current_catalog("nonexistent").await;
+ assert!(err.is_err());
+}
+
+#[tokio::test]
+async fn test_set_default_catalog_via_datafusion_config() {
+ let (_tmp1, catalog1) = create_test_env();
+ let (_tmp2, catalog2) = create_test_env();
+
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("cat1", catalog1).await.unwrap();
+ ctx.register_catalog("cat2", catalog2).await.unwrap();
+
+ // Create a table in cat2
+ ctx.sql("CREATE SCHEMA cat2.mydb").await.unwrap();
+ ctx.sql("CREATE TABLE cat2.mydb.t (id INT NOT NULL, name VARCHAR, PRIMARY KEY (id))")
+ .await
+ .unwrap();
+ ctx.sql("INSERT INTO cat2.mydb.t VALUES (1, 'hello')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+
+ // Switch default catalog via raw DataFusion SET instead of set_current_catalog()
+ ctx.sql("SET datafusion.catalog.default_catalog = 'cat2'")
+ .await
+ .unwrap();
+ ctx.set_current_database("mydb").await.unwrap();
+
+ // Unqualified query should now resolve against cat2.mydb
+ let batches = ctx
+ .sql("SELECT id, name FROM t")
+ .await
+ .expect("Unqualified table should resolve via DataFusion default_catalog config")
+ .collect()
+ .await
+ .unwrap();
+ let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(total_rows, 1);
+
+ // DDL on unqualified name should also go to cat2.
+ // CREATE TABLE in mydb should succeed because cat2.mydb exists.
+ ctx.sql("CREATE TABLE mydb.t2 (id INT NOT NULL, PRIMARY KEY (id))")
+ .await
+ .expect("CREATE TABLE should resolve against cat2 after SET default_catalog");
+
+ // Verify the table was created in cat2 by querying with fully qualified name
+ let df = ctx.sql("SELECT * FROM cat2.mydb.t2").await;
+ assert!(
+ df.is_ok(),
+ "Table t2 should exist in cat2.mydb after unqualified CREATE TABLE"
+ );
+}
+
+#[tokio::test]
+async fn test_first_registered_catalog_is_default() {
+ let (_tmp, catalog) = create_test_env();
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog).await.unwrap();
+
+ ctx.sql("CREATE SCHEMA paimon.mydb").await.unwrap();
+ ctx.sql("CREATE TABLE paimon.mydb.t (id INT NOT NULL, PRIMARY KEY (id))")
+ .await
+ .unwrap();
+
+ ctx.set_current_database("mydb").await.unwrap();
+
+ // Should resolve to paimon.mydb.t without calling set_current_catalog
+ let df = ctx.sql("SELECT * FROM t").await;
+ assert!(
+ df.is_ok(),
+ "First registered catalog should be the default for unqualified queries"
+ );
+}
+
+#[tokio::test]
+async fn test_one_part_table_name_uses_current_database() {
+ let (_tmp, catalog) = create_test_env();
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog.clone())
+ .await
+ .unwrap();
+
+ catalog
+ .create_database("mydb", false, Default::default())
+ .await
+ .unwrap();
+ ctx.set_current_database("mydb").await.unwrap();
+
+ // 1-part name: "users" should resolve to paimon.mydb.users
+ ctx.sql(
+ "CREATE TABLE users (
+ id INT NOT NULL,
+ name STRING,
+ PRIMARY KEY (id)
+ )",
+ )
+ .await
+ .expect("CREATE TABLE with 1-part name should succeed");
+
+ let tables = catalog.list_tables("mydb").await.unwrap();
+ assert!(
+ tables.contains(&"users".to_string()),
+ "Table should be created in the current database"
+ );
+
+ // SELECT with 1-part name should also work
+ let df = ctx.sql("SELECT * FROM users").await;
+ assert!(
+ df.is_ok(),
+ "SELECT with 1-part name should resolve correctly"
+ );
+}
diff --git a/crates/integrations/datafusion/tests/update_tests.rs b/crates/integrations/datafusion/tests/update_tests.rs
index 6fa4f47..839bc26 100644
--- a/crates/integrations/datafusion/tests/update_tests.rs
+++ b/crates/integrations/datafusion/tests/update_tests.rs
@@ -21,55 +21,61 @@
mod common;
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
-use common::{create_handler, create_test_env, dml_count, exec, query_int_str_int};
+use common::{create_sql_context, create_test_env, ctx_exec, dml_count, exec, query_int_str_int};
// ======================= Helpers =======================
-async fn setup() -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup() -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, 'b', 20), (3, 'c', 30), (4, 'd', 40)",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
-async fn setup_partitioned() -> (tempfile::TempDir, PaimonSqlHandler) {
+async fn setup_partitioned() -> (tempfile::TempDir, SQLContext) {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, pt INT) PARTITIONED BY (pt)")
.await
.unwrap();
exec(
- &handler,
+ &sql_context,
"INSERT INTO paimon.test_db.t VALUES (1, 'a', 1), (2, 'b', 1), (3, 'c', 2), (4, 'd', 2)",
)
.await;
- (tmp, handler)
+ (tmp, sql_context)
}
-async fn query(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
+async fn query(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
query_int_str_int(
- handler,
+ sql_context,
"SELECT id, name, age FROM paimon.test_db.t ORDER BY id",
)
.await
}
-async fn query_pt(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
+async fn query_pt(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
query_int_str_int(
- handler,
+ sql_context,
"SELECT id, name, pt FROM paimon.test_db.t ORDER BY id",
)
.await
@@ -79,16 +85,16 @@ async fn query_pt(handler: &PaimonSqlHandler) -> Vec<(i32, String, i32)> {
#[tokio::test]
async fn test_update_with_where() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'a_new' WHERE id = 1",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a_new".into(), 10),
(2, "b".into(), 20),
@@ -100,16 +106,16 @@ async fn test_update_with_where() {
#[tokio::test]
async fn test_update_with_expression() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET age = age + 1 WHERE id = 1",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 11),
(2, "b".into(), 20),
@@ -121,17 +127,17 @@ async fn test_update_with_expression() {
#[tokio::test]
async fn test_update_no_match_empty_commit() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
let cnt = dml_count(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE id = 99",
)
.await;
assert_eq!(cnt, 0);
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(2, "b".into(), 20),
@@ -143,12 +149,12 @@ async fn test_update_no_match_empty_commit() {
#[tokio::test]
async fn test_update_without_where() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "UPDATE paimon.test_db.t SET age = 0").await;
+ exec(&sql_context, "UPDATE paimon.test_db.t SET age = 0").await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 0),
(2, "b".into(), 0),
@@ -162,16 +168,16 @@ async fn test_update_without_where() {
#[tokio::test]
async fn test_update_partitioned_by_non_partition_column() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE id = 1",
)
.await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![
(1, "x".into(), 1),
(2, "b".into(), 1),
@@ -183,16 +189,16 @@ async fn test_update_partitioned_by_non_partition_column() {
#[tokio::test]
async fn test_update_partitioned_by_partition_column() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE pt = 1",
)
.await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![
(1, "x".into(), 1),
(2, "x".into(), 1),
@@ -204,16 +210,16 @@ async fn test_update_partitioned_by_partition_column() {
#[tokio::test]
async fn test_update_partitioned_by_both_columns() {
- let (_tmp, handler) = setup_partitioned().await;
+ let (_tmp, sql_context) = setup_partitioned().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE pt = 2 AND id = 3",
)
.await;
assert_eq!(
- query_pt(&handler).await,
+ query_pt(&sql_context).await,
vec![
(1, "a".into(), 1),
(2, "b".into(), 1),
@@ -227,16 +233,16 @@ async fn test_update_partitioned_by_both_columns() {
#[tokio::test]
async fn test_update_in_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE id IN (1, 3)",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "x".into(), 10),
(2, "b".into(), 20),
@@ -248,16 +254,16 @@ async fn test_update_in_condition() {
#[tokio::test]
async fn test_update_not_in_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE id NOT IN (1, 3)",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(2, "x".into(), 20),
@@ -271,16 +277,16 @@ async fn test_update_not_in_condition() {
#[tokio::test]
async fn test_update_cross_column_reference() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET age = id * 100 WHERE id <= 2",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 100),
(2, "b".into(), 200),
@@ -292,16 +298,16 @@ async fn test_update_cross_column_reference() {
#[tokio::test]
async fn test_update_multiple_columns_with_expressions() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'updated', age = age + id WHERE id >= 3",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(2, "b".into(), 20),
@@ -315,17 +321,21 @@ async fn test_update_multiple_columns_with_expressions() {
#[tokio::test]
async fn test_successive_updates() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- exec(&handler, "UPDATE paimon.test_db.t SET age = 0 WHERE id = 1").await;
exec(
- &handler,
+ &sql_context,
+ "UPDATE paimon.test_db.t SET age = 0 WHERE id = 1",
+ )
+ .await;
+ exec(
+ &sql_context,
"UPDATE paimon.test_db.t SET age = 99 WHERE id = 1",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 99),
(2, "b".into(), 20),
@@ -339,16 +349,16 @@ async fn test_successive_updates() {
#[tokio::test]
async fn test_update_or_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET name = 'x' WHERE id = 1 OR id = 4",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "x".into(), 10),
(2, "b".into(), 20),
@@ -362,16 +372,16 @@ async fn test_update_or_condition() {
#[tokio::test]
async fn test_update_range_condition() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET age = 0 WHERE id >= 2 AND id <= 3",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "a".into(), 10),
(2, "b".into(), 0),
@@ -385,25 +395,22 @@ async fn test_update_range_condition() {
#[tokio::test]
async fn test_update_in_subquery() {
- let (_tmp, handler) = setup().await;
+ let (_tmp, sql_context) = setup().await;
- handler
- .ctx()
- .sql("CREATE TABLE src (id INT) AS VALUES (1), (3)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
+ ctx_exec(
+ &sql_context,
+ "CREATE TABLE datafusion.public.src (id INT) AS VALUES (1), (3)",
+ )
+ .await;
exec(
- &handler,
- "UPDATE paimon.test_db.t SET name = 'sub' WHERE id IN (SELECT id FROM src)",
+ &sql_context,
+ "UPDATE paimon.test_db.t SET name = 'sub' WHERE id IN (SELECT id FROM datafusion.public.src)",
)
.await;
assert_eq!(
- query(&handler).await,
+ query(&sql_context).await,
vec![
(1, "sub".into(), 10),
(2, "b".into(), 20),
@@ -418,13 +425,16 @@ async fn test_update_in_subquery() {
#[tokio::test]
async fn test_update_with_null_values() {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
- handler
+ sql_context
.sql("INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, NULL, 20), (3, 'c', NULL)")
.await
.unwrap()
@@ -433,12 +443,12 @@ async fn test_update_with_null_values() {
.unwrap();
exec(
- &handler,
+ &sql_context,
"UPDATE paimon.test_db.t SET age = 0 WHERE name = 'a'",
)
.await;
- let rows = query(&handler).await;
+ let rows = query(&sql_context).await;
assert_eq!(rows[0], (1, "a".into(), 0));
assert_eq!(rows[1].0, 2);
assert_eq!(rows[1].2, 20);
@@ -451,14 +461,17 @@ async fn test_update_with_null_values() {
#[tokio::test]
async fn test_update_empty_table() {
let (tmp, catalog) = create_test_env();
- let handler = create_handler(catalog);
- handler.sql("CREATE SCHEMA paimon.test_db").await.unwrap();
- handler
+ let sql_context = create_sql_context(catalog).await;
+ sql_context
+ .sql("CREATE SCHEMA paimon.test_db")
+ .await
+ .unwrap();
+ sql_context
.sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
.await
.unwrap();
- let cnt = dml_count(&handler, "UPDATE paimon.test_db.t SET name = 'x'").await;
+ let cnt = dml_count(&sql_context, "UPDATE paimon.test_db.t SET name = 'x'").await;
assert_eq!(cnt, 0);
drop(tmp);
}
diff --git a/docs/src/datafusion.md b/docs/src/datafusion.md
index 5664b7c..38675c1 100644
--- a/docs/src/datafusion.md
+++ b/docs/src/datafusion.md
@@ -37,24 +37,23 @@ Register an entire Paimon catalog so all databases and tables are accessible via
```rust
use std::sync::Arc;
-use datafusion::prelude::SessionContext;
use paimon::{CatalogOptions, FileSystemCatalog, Options};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
async fn example() -> Result<(), Box<dyn std::error::Error>> {
let mut options = Options::new();
options.set(CatalogOptions::WAREHOUSE, "file:///tmp/paimon-warehouse");
let catalog = Arc::new(FileSystemCatalog::new(options)?);
- let ctx = SessionContext::new();
- let handler = PaimonSqlHandler::new(ctx, catalog, "paimon")?;
- let df = handler.sql("SELECT * FROM paimon.default.my_table").await?;
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog)?;
+ let df = ctx.sql("SELECT * FROM paimon.default.my_table").await?;
df.show().await?;
Ok(())
}
```
-`PaimonSqlHandler::new` automatically registers the Paimon catalog provider and relation planner on the session context. It also manages session-scoped dynamic options internally for `SET`/`RESET` support.
+`SQLContext::new` creates a session context with the Paimon relation planner pre-registered. Use `register_catalog` to add one or more Paimon catalogs. It also manages session-scoped dynamic options internally for `SET`/`RESET` support.
## Data Types
@@ -706,9 +705,8 @@ Set via `WITH ('key' = 'value')` at table creation time, or dynamically via `SET
```rust
use std::sync::Arc;
-use datafusion::prelude::SessionContext;
use paimon::{CatalogOptions, FileSystemCatalog, Options};
-use paimon_datafusion::PaimonSqlHandler;
+use paimon_datafusion::SQLContext;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -717,13 +715,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
options.set(CatalogOptions::WAREHOUSE, "file:///tmp/paimon-warehouse");
let catalog = Arc::new(FileSystemCatalog::new(options)?);
- // Create handler (registers catalog provider and relation planner automatically)
- let ctx = SessionContext::new();
- let handler = PaimonSqlHandler::new(ctx, catalog, "paimon")?;
+ // Create SQL context and register catalog
+ let mut ctx = SQLContext::new();
+ ctx.register_catalog("paimon", catalog)?;
// Create database and table
- handler.sql("CREATE SCHEMA paimon.my_db").await?;
- handler.sql(
+ ctx.sql("CREATE SCHEMA paimon.my_db").await?;
+ ctx.sql(
"CREATE TABLE paimon.my_db.users (
id INT NOT NULL,
name STRING,
@@ -732,11 +730,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
).await?;
// Insert data
- handler.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob')")
+ ctx.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob')")
.await?.collect().await?;
// Query
- let df = handler.sql("SELECT * FROM paimon.my_db.users ORDER BY id").await?;
+ let df = ctx.sql("SELECT * FROM paimon.my_db.users ORDER BY id").await?;
df.show().await?;
Ok(())