diff --git a/crates/integrations/datafusion/README.md b/crates/integrations/datafusion/README.md index b95837264..1b6523e22 100644 --- a/crates/integrations/datafusion/README.md +++ b/crates/integrations/datafusion/README.md @@ -24,6 +24,18 @@ This crate contains the integration of [Apache DataFusion](https://datafusion.apache.org/) and [Apache Paimon](https://paimon.apache.org/). +## External engine snapshots + +Implement `TableEngineResolver::resolve_table_snapshot` to support +`VERSION AS OF ` on external data tables (integer IDs only). +Providers must use the snapshot's data and schema. Missing snapshots are errors; +`Ok(None)` means the table is missing. Resolved providers remain read-only. + +`SQLContext` needs no dialect setting and resolves snapshots independently of the +current provider. Raw `SessionContext` needs `register_catalog_table_engine` and a +dialect supporting `VERSION AS OF`; it preloads the current provider and propagates +loading errors. + ## Reader memory Paimon Parquet scans use the memory pool from the executing DataFusion `TaskContext`. diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 9ba17d958..4232ae2b0 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -49,9 +49,7 @@ pub(crate) type SessionStateProvider = Arc Option + Se /// providers, so registrations stay visible to schemas obtained earlier. type TableEngines = Arc>>>; -/// What an engine is asked to resolve. Non-exhaustive so later releases can -/// carry more of the request — a snapshot selector, say — without breaking -/// existing resolvers. +/// Identifies an engine-owned table. #[derive(Debug, Clone)] #[non_exhaustive] pub struct EngineTableRequest { @@ -71,10 +69,7 @@ impl EngineTableRequest { } } -/// Resolves tables owned by another engine (see -/// [`PaimonCatalogProvider::register_table_engine`]). `Ok(None)` means not -/// found; errors propagate, so an engine failure never looks like a missing -/// table. +/// Resolves external tables. `Ok(None)` means missing; engine errors propagate. #[async_trait] pub trait TableEngineResolver: Debug + Send + Sync { /// Resolve a request to the engine's table provider. @@ -83,6 +78,21 @@ pub trait TableEngineResolver: Debug + Send + Sync { request: &EngineTableRequest, ) -> DFResult>>; + /// Resolve an exact snapshot and its schema. + /// Missing snapshots are errors; `Ok(None)` means the table is missing. + async fn resolve_table_snapshot( + &self, + request: &EngineTableRequest, + _snapshot_id: i64, + ) -> DFResult>> { + Err(plan_datafusion_err!( + "time travel is not supported for routed '{}' tables ('{}.{}')", + request.declared, + request.database, + request.table + )) + } + /// Resolves `$`; returns `Ok(None)` if missing or unsupported. /// The request names the base table; `metadata_table` is ASCII-lowercase. /// Session scan options are checked before this call. @@ -95,14 +105,82 @@ pub trait TableEngineResolver: Debug + Send + Sync { } } -/// Read-only wrapper around an engine-resolved provider: reads delegate, -/// DML is rejected even when the engine's own provider is writable. -#[derive(Debug)] +/// Delegates reads to the engine and rejects DML. +#[derive(Debug, Clone)] pub(crate) struct ReadOnlyTableProvider { inner: Arc, pub(crate) declared: PaimonTableType, pub(crate) table_name: String, metadata_table: bool, + request: EngineTableRequest, + resolver: Arc, + dynamic_options: DynamicOptions, +} + +impl ReadOnlyTableProvider { + pub(crate) fn unsupported_time_travel(&self) -> datafusion::error::DataFusionError { + plan_datafusion_err!( + "time travel is not supported for routed '{}' tables ('{}'); \ + only VERSION AS OF with an integer snapshot id is supported on data tables", + self.declared, + self.table_name + ) + } + + pub(crate) async fn at_snapshot(&self, snapshot_id: i64) -> DFResult> { + if self.metadata_table { + return Err(self.unsupported_time_travel()); + } + resolve_engine_snapshot( + self.request.clone(), + Arc::clone(&self.resolver), + Arc::clone(&self.dynamic_options), + snapshot_id, + ) + .await + } +} + +pub(crate) fn parse_engine_snapshot_id(version: &str) -> DFResult { + version.parse().map_err(|_| { + plan_datafusion_err!( + "external table time travel requires an integer snapshot id in VERSION AS OF, got '{version}'" + ) + }) +} + +async fn resolve_engine_snapshot( + request: EngineTableRequest, + resolver: Arc, + dynamic_options: DynamicOptions, + snapshot_id: i64, +) -> DFResult> { + await_with_runtime(async move { + let table_name = format!("{}.{}", request.database, request.table); + let options = dynamic_options + .read() + .unwrap_or_else(|e| e.into_inner()) + .clone(); + paimon::spec::CoreOptions::new(&options) + .ensure_engine_can_serve(&table_name) + .map_err(to_datafusion_error)?; + let inner = resolver + .resolve_table_snapshot(&request, snapshot_id) + .await? + .ok_or_else(|| { + plan_datafusion_err!("table '{table_name}' was not found at snapshot {snapshot_id}") + })?; + Ok(Arc::new(ReadOnlyTableProvider { + inner, + declared: request.declared, + table_name, + metadata_table: false, + request, + resolver, + dynamic_options, + }) as Arc) + }) + .await } #[async_trait] @@ -170,20 +248,16 @@ impl TableProvider for ReadOnlyTableProvider { } } -/// Metadata-only provider for an external table whose engine is not registered. -/// -/// DataFusion's information schema asks every catalog table for its schema. -/// Returning this provider keeps those metadata queries available without -/// weakening the fail-closed behavior for actual table reads or writes. +/// Exposes an unregistered engine's table schema while rejecting scans and writes. #[derive(Debug)] -struct UnavailableEngineTableProvider { +pub(crate) struct UnavailableEngineTableProvider { schema: datafusion::arrow::datatypes::SchemaRef, declared: PaimonTableType, table_name: String, } impl UnavailableEngineTableProvider { - fn unavailable_error(&self) -> datafusion::error::DataFusionError { + pub(crate) fn unavailable_error(&self) -> datafusion::error::DataFusionError { plan_datafusion_err!( "no table engine is registered for '{}' tables ('{}')", self.declared, @@ -349,6 +423,34 @@ impl PaimonCatalogProvider { Arc::clone(&self.table_engines) } + pub(crate) async fn table_at_snapshot( + &self, + request: EngineTableRequest, + snapshot_id: i64, + ) -> DFResult> { + let resolver = self + .table_engines + .read() + .unwrap_or_else(|e| e.into_inner()) + .get(&request.declared) + .cloned() + .ok_or_else(|| { + plan_datafusion_err!( + "no table engine is registered for '{}' tables ('{}.{}')", + request.declared, + request.database, + request.table + ) + })?; + resolve_engine_snapshot( + request, + resolver, + Arc::clone(&self.dynamic_options), + snapshot_id, + ) + .await + } + fn paimon_schema(&self, name: &str) -> Option> { let catalog = Arc::clone(&self.catalog); let table_engines = self.table_engines(); @@ -812,6 +914,9 @@ impl SchemaProvider for PaimonSchemaProvider { declared, table_name, metadata_table, + request, + resolver: Arc::clone(resolver), + dynamic_options, }) as Arc })) } diff --git a/crates/integrations/datafusion/src/relation_planner.rs b/crates/integrations/datafusion/src/relation_planner.rs index d9cf04925..2078d48ec 100644 --- a/crates/integrations/datafusion/src/relation_planner.rs +++ b/crates/integrations/datafusion/src/relation_planner.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Custom [`RelationPlanner`] for Paimon time travel via `VERSION AS OF` and `TIMESTAMP AS OF`. +//! Time-travel planning for Paimon tables and registered engines. use std::collections::HashMap; use std::fmt::Debug; @@ -24,6 +24,7 @@ use std::sync::Arc; use datafusion::catalog::default_table_source::{provider_as_source, source_as_provider}; use datafusion::common::plan_datafusion_err; use datafusion::common::TableReference; +use datafusion::datasource::TableProvider; use datafusion::error::Result as DFResult; use datafusion::logical_expr::builder::LogicalPlanBuilder; use datafusion::logical_expr::planner::{ @@ -32,15 +33,12 @@ use datafusion::logical_expr::planner::{ use datafusion::sql::sqlparser::ast::{self, TableFactor, TableVersion}; use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION}; -use crate::catalog::ReadOnlyTableProvider; +use crate::catalog::{ + parse_engine_snapshot_id, ReadOnlyTableProvider, UnavailableEngineTableProvider, +}; use crate::table::{ObjectTableProvider, PaimonTableProvider}; -/// A [`RelationPlanner`] that intercepts `VERSION AS OF` and `TIMESTAMP AS OF` -/// clauses on Paimon tables and resolves them to time travel options. -/// -/// - `VERSION AS OF ` → sets `scan.version` option on the table. -/// At scan time, the version is resolved: tag name (if exists) → snapshot id → error. -/// - `TIMESTAMP AS OF ` → parsed as a timestamp, sets `scan.timestamp-millis`. +/// Resolves Paimon time travel and registered engine snapshot IDs. #[derive(Debug)] pub struct PaimonRelationPlanner; @@ -62,7 +60,6 @@ impl RelationPlanner for PaimonRelationPlanner { relation: TableFactor, context: &mut dyn RelationPlannerContext, ) -> DFResult { - // Only handle Table factors with a version clause. let TableFactor::Table { ref name, ref version, @@ -76,27 +73,33 @@ impl RelationPlanner for PaimonRelationPlanner { return Ok(RelationPlanning::Original(Box::new(relation))); }; - // Resolve the table reference. let table_ref = object_name_to_table_reference(name, context)?; let source = context .context_provider() .get_table_source(table_ref.clone())?; let provider = source_as_provider(&source)?; - // Every version clause, not just the two resolved below: falling - // through drops the clause silently. + if let Some(unavailable) = provider.downcast_ref::() { + return Err(unavailable.unavailable_error()); + } + if let Some(routed) = provider.downcast_ref::() { - return Err(plan_datafusion_err!( - "time travel is not supported for routed '{}' tables ('{}')", - routed.declared, - routed.table_name - )); + let TableVersion::VersionAsOf(expr) = version else { + return Err(routed.unsupported_time_travel()); + }; + let options = resolve_version_as_of(expr)?; + let snapshot_id = parse_engine_snapshot_id(&options[SCAN_VERSION_OPTION])?; + let routed = routed.clone(); + let new_provider = crate::runtime::block_on_with_runtime( + async move { routed.at_snapshot(snapshot_id).await }, + "engine snapshot resolution thread panicked", + )?; + return plan_scan(relation, table_ref, new_provider); } let Some(paimon_provider) = provider.downcast_ref::() else { - // Ours, but with no snapshot to rewrite onto: refuse, because - // handing them on reaches the default planner, which drops the - // clause. Any other provider is another engine's, clause and all. + // Reject unsupported Paimon providers before the default planner drops the clause. + // Leave foreign providers to their own planner. if provider.is::() || crate::system_tables::is_system_table_provider(provider.as_ref()) { @@ -109,12 +112,10 @@ impl RelationPlanner for PaimonRelationPlanner { let extra_options = match version { TableVersion::VersionAsOf(expr) => resolve_version_as_of(expr)?, - // Same timestamp expression as `TIMESTAMP AS OF`. TableVersion::TimestampAsOf(expr) | TableVersion::ForSystemTimeAsOf(expr) => { resolve_timestamp_as_of(expr)? } - // `AT(...)` and `CHANGES(...)` name neither a snapshot nor a - // timestamp; dropping the clause would answer with current rows. + // Never let unsupported version clauses fall through to current data. _ => { return Err(plan_datafusion_err!( "this time-travel syntax is not supported for Paimon tables; \ @@ -123,9 +124,7 @@ impl RelationPlanner for PaimonRelationPlanner { } }; - // Resolving time travel may switch the table to the snapshot's schema, - // which requires async IO; this planner hook is synchronous, so bridge - // through the shared runtime like other sync DataFusion callbacks. + // Snapshot schema resolution needs async I/O inside this synchronous hook. let table = paimon_provider.table().clone(); let new_table = crate::runtime::block_on_with_runtime( async move { table.copy_with_time_travel(extra_options).await }, @@ -133,20 +132,24 @@ impl RelationPlanner for PaimonRelationPlanner { ) .map_err(crate::to_datafusion_error)?; let new_provider = PaimonTableProvider::try_new(new_table)?; - let new_source = provider_as_source(Arc::new(new_provider)); - - // Destructure to get alias. - let TableFactor::Table { alias, .. } = relation else { - unreachable!() - }; - - let plan = LogicalPlanBuilder::scan(table_ref, new_source, None)?.build()?; - Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new( - plan, alias, - )))) + plan_scan(relation, table_ref, Arc::new(new_provider)) } } +fn plan_scan( + relation: TableFactor, + table_ref: TableReference, + provider: Arc, +) -> DFResult { + let TableFactor::Table { alias, .. } = relation else { + unreachable!() + }; + let plan = LogicalPlanBuilder::scan(table_ref, provider_as_source(provider), None)?.build()?; + Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new( + plan, alias, + )))) +} + /// Convert a sqlparser `ObjectName` to a DataFusion `TableReference`. fn object_name_to_table_reference( name: &ast::ObjectName, diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index f5f27276b..a81a295d6 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -55,7 +55,7 @@ use datafusion::arrow::compute::cast; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion::common::TableReference; +use datafusion::common::{plan_datafusion_err, TableReference}; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::context::SQLOptions; @@ -86,7 +86,7 @@ use paimon::spec::{ }; use crate::error::to_datafusion_error; -use crate::table_loader::load_table_for_read; +use crate::table::PaimonTableProvider; use crate::{BlobReaderRegistry, DynamicOptions}; /// A SQL context that supports registering multiple Paimon catalogs and executing SQL. @@ -766,20 +766,13 @@ impl SQLContext { } } - /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). - /// - /// DataFusion's default SQL parser does not support these clauses, so we: - /// 1. Extract all table name + version/timestamp pairs (skipping string literals and comments) - /// 2. Strip the time-travel clauses from the SQL - /// 3. For each table, create a `PaimonTableProvider` with the appropriate scan options - /// (merged with session-scoped dynamic options) - /// 4. Register them as UUID-named temp tables, execute the rewritten SQL, then deregister + /// Rewrite time-travel references to temporary snapshot providers. + /// The tracker removes registrations after planning. async fn handle_time_travel_query( &self, sql: &str, sql_options: SQLOptions, ) -> DFResult { - use crate::table::PaimonTableProvider; use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION}; let mut tracker = crate::merge_into::TempTableTracker::new(self); @@ -793,42 +786,21 @@ impl SQLContext { )); } - // Collect all replacements: (clause_range, uuid_name) let mut replacements: Vec<((usize, usize), String)> = Vec::new(); - // Process all VERSION AS OF clauses for info in &version_clauses { let table_ref: datafusion::common::TableReference = info.table_name.as_str().into(); - let (catalog, _catalog_name, identifier) = + let (catalog, catalog_name, identifier) = self.resolve_table_name_from_ref(&table_ref)?; - - let (paimon_table, base_identifier, system_name) = - load_table_for_read(&catalog, &identifier).await?; - - // Merge dynamic options with time-travel options - let mut options = self.dynamic_options.read().unwrap().clone(); - options.insert(SCAN_VERSION_OPTION.to_string(), info.version.clone()); - - let table_with_options = paimon_table - .copy_with_time_travel(options) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; - let provider: Arc = if let Some(system_name) = system_name { - crate::system_tables::provider_for_table( - Arc::clone(&catalog), - base_identifier, - table_with_options, - &system_name, - )? - .ok_or_else(|| { - DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}")) - })? - } else { - Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( - table_with_options, - self.blob_reader_registry.clone(), - )?) - }; + let provider = self + .time_travel_provider( + &catalog, + &catalog_name, + &identifier, + SCAN_VERSION_OPTION, + info.version.clone(), + ) + .await?; let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); self.register_temp_table(uuid_name.as_str(), provider)?; @@ -836,41 +808,20 @@ impl SQLContext { replacements.push((info.clause_range, uuid_name)); } - // Process all TIMESTAMP AS OF clauses for info in ×tamp_clauses { let table_ref: datafusion::common::TableReference = info.table_name.as_str().into(); - let (catalog, _catalog_name, identifier) = + let (catalog, catalog_name, identifier) = self.resolve_table_name_from_ref(&table_ref)?; - - let (paimon_table, base_identifier, system_name) = - load_table_for_read(&catalog, &identifier).await?; - let millis = Self::parse_timestamp_to_millis(&info.timestamp)?; - - // Merge dynamic options with time-travel options - let mut options = self.dynamic_options.read().unwrap().clone(); - options.insert(SCAN_TIMESTAMP_MILLIS_OPTION.to_string(), millis.to_string()); - - let table_with_options = paimon_table - .copy_with_time_travel(options) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; - let provider: Arc = if let Some(system_name) = system_name { - crate::system_tables::provider_for_table( - Arc::clone(&catalog), - base_identifier, - table_with_options, - &system_name, - )? - .ok_or_else(|| { - DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}")) - })? - } else { - Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry( - table_with_options, - self.blob_reader_registry.clone(), - )?) - }; + let provider = self + .time_travel_provider( + &catalog, + &catalog_name, + &identifier, + SCAN_TIMESTAMP_MILLIS_OPTION, + millis.to_string(), + ) + .await?; let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple()); self.register_temp_table(uuid_name.as_str(), provider)?; @@ -878,11 +829,9 @@ impl SQLContext { replacements.push((info.clause_range, uuid_name)); } - // Sort replacements by position (descending) so that replacements - // from right to left don't shift indices of earlier ones + // Replace from right to left to preserve byte offsets. replacements.sort_by_key(|r| std::cmp::Reverse(r.0 .0)); - // Build the rewritten SQL by replacing each clause from right to left let mut rewritten_sql = sql.to_string(); for ((start, end), uuid_name) in &replacements { rewritten_sql = format!( @@ -893,7 +842,6 @@ impl SQLContext { ); } - // Execute the rewritten SQL; tracker auto-deregisters on drop. let current_catalog = self.current_catalog_name(); let current_database = self .ctx @@ -912,6 +860,89 @@ impl SQLContext { self.ctx.sql_with_options(&expanded, sql_options).await } + async fn time_travel_provider( + &self, + catalog: &Arc, + catalog_name: &str, + identifier: &Identifier, + option: &str, + value: String, + ) -> DFResult> { + let parsed = identifier + .parsed_object_name() + .map_err(to_datafusion_error)?; + let base = Identifier::new( + identifier.database().to_string(), + parsed.table().to_string(), + ); + let loaded = catalog + .load_table(&base) + .await + .map_err(to_datafusion_error)?; + if let paimon::catalog::LoadedTable::External(external) = &loaded { + if parsed.branch().is_some() { + return Err(plan_datafusion_err!( + "branches are not supported for '{}' tables ('{}')", + external.declared(), + identifier.full_name() + )); + } + if parsed.system_table().is_some() || option != paimon::spec::SCAN_VERSION_OPTION { + return Err(plan_datafusion_err!( + "time travel is not supported for routed '{}' tables ('{}'); \ + only VERSION AS OF with an integer snapshot id is supported on data tables", + external.declared(), + identifier.full_name() + )); + } + let snapshot_id = crate::catalog::parse_engine_snapshot_id(&value)?; + let provider = self + .ctx + .catalog(catalog_name) + .ok_or_else(|| plan_datafusion_err!("Unknown catalog '{catalog_name}'"))?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + plan_datafusion_err!("Catalog '{catalog_name}' is not a Paimon catalog") + })?; + return provider + .table_at_snapshot( + crate::EngineTableRequest::new( + base.database().to_string(), + base.object().to_string(), + external.declared(), + ), + snapshot_id, + ) + .await; + } + + let (table, base, system_name) = + crate::table_loader::loaded_table_for_read(loaded, identifier).await?; + let mut options = self.dynamic_options.read().unwrap().clone(); + options.insert(option.to_string(), value); + let table = table + .copy_with_time_travel(options) + .await + .map_err(to_datafusion_error)?; + if let Some(system_name) = system_name { + crate::system_tables::provider_for_table( + Arc::clone(catalog), + base, + table, + &system_name, + )? + .ok_or_else(|| plan_datafusion_err!("Unknown Paimon system table: {system_name}")) + } else { + Ok(Arc::new( + PaimonTableProvider::try_new_with_blob_reader_registry( + table, + self.blob_reader_registry.clone(), + )?, + )) + } + } + /// Parse a timestamp string to milliseconds since epoch (using local timezone). fn parse_timestamp_to_millis(ts: &str) -> DFResult { use chrono::{Local, NaiveDateTime, TimeZone}; diff --git a/crates/integrations/datafusion/src/table_loader.rs b/crates/integrations/datafusion/src/table_loader.rs index 664a58c02..748853a8f 100644 --- a/crates/integrations/datafusion/src/table_loader.rs +++ b/crates/integrations/datafusion/src/table_loader.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use datafusion::error::{DataFusionError, Result as DFResult}; -use paimon::catalog::{Catalog, Identifier}; +use paimon::catalog::{Catalog, Identifier, LoadedTable}; use paimon::table::Table; use crate::error::to_datafusion_error; @@ -58,6 +58,24 @@ pub(crate) fn ensure_paimon_served(table: &Table, identifier: &Identifier) -> DF pub(crate) async fn load_table_for_read( catalog: &Arc, identifier: &Identifier, +) -> DFResult<(Table, Identifier, Option)> { + let parsed = identifier + .parsed_object_name() + .map_err(to_datafusion_error)?; + let base = Identifier::new( + identifier.database().to_string(), + parsed.table().to_string(), + ); + let loaded = catalog + .load_table(&base) + .await + .map_err(to_datafusion_error)?; + loaded_table_for_read(loaded, identifier).await +} + +pub(crate) async fn loaded_table_for_read( + loaded: LoadedTable, + identifier: &Identifier, ) -> DFResult<(Table, Identifier, Option)> { let parsed = identifier .parsed_object_name() @@ -66,32 +84,29 @@ pub(crate) async fn load_table_for_read( identifier.database().to_string(), parsed.table().to_string(), ); - let mut table = match catalog.load_table(&base_identifier).await { - Ok(paimon::catalog::LoadedTable::Paimon(table)) => *table, - Ok(paimon::catalog::LoadedTable::Object(_)) => { - // The search UDTFs land here too, so the message names the table - // rather than what the caller asked of it. + let mut table = match loaded { + LoadedTable::Paimon(table) => *table, + LoadedTable::Object(_) => { + // Shared with search UDTFs, so errors name the base table. return Err(DataFusionError::Plan(format!( "table '{}' is declared 'object-table' and cannot be read as a Paimon table", base_identifier.full_name() ))); } - Ok(paimon::catalog::LoadedTable::External(external)) => { + LoadedTable::External(external) => { return Err(DataFusionError::Plan(format!( "table '{}' is declared '{}' and cannot be read as a Paimon table", base_identifier.full_name(), external.declared() ))); } - // `LoadedTable` is non_exhaustive: a variant added upstream is not a - // Paimon table until this path says how to read one. - Ok(_) => { + // Reject unknown variants of non-exhaustive LoadedTable. + _ => { return Err(DataFusionError::Plan(format!( "table '{}' cannot be read as a Paimon table", base_identifier.full_name() ))); } - Err(err) => return Err(to_datafusion_error(err)), }; let system_table = parsed.system_table().map(str::to_string); if let Some(branch) = parsed.branch() { diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index d29a3938a..243140aab 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -306,6 +306,81 @@ struct MetadataEngineResolver { calls: AtomicUsize, } +#[derive(Debug, Default)] +struct SnapshotEngineResolver { + calls: AtomicUsize, + current_calls: AtomicUsize, + current_unavailable: bool, +} + +const LARGE_SNAPSHOT_ID: i64 = 6_251_097_304_546_047_273; + +#[async_trait] +impl TableEngineResolver for SnapshotEngineResolver { + async fn resolve_table( + &self, + request: &EngineTableRequest, + ) -> DFResult>> { + self.current_calls.fetch_add(1, Ordering::Relaxed); + if self.current_unavailable { + return Err(DataFusionError::Execution( + "current table is unavailable".into(), + )); + } + FakeEngineResolver.resolve_table(request).await + } + + async fn resolve_metadata_table( + &self, + request: &EngineTableRequest, + metadata_table: &str, + ) -> DFResult>> { + MetadataEngineResolver::default() + .resolve_metadata_table(request, metadata_table) + .await + } + + async fn resolve_table_snapshot( + &self, + request: &EngineTableRequest, + snapshot_id: i64, + ) -> DFResult>> { + self.calls.fetch_add(1, Ordering::Relaxed); + assert_eq!(request.database, DB); + assert_eq!(request.declared, TableType::IcebergTable); + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + if request.table != "it" { + return Ok(None); + } + let (ids, values) = match snapshot_id { + 10 => (vec![1, 2], vec!["old-a", "old-b"]), + 20 => (vec![1, 4], vec!["new-a", "new-d"]), + LARGE_SNAPSHOT_ID => (vec![9], vec!["large-id"]), + 500 => return Err(DataFusionError::Execution("snapshot backend failed".into())), + _ => { + return Err(DataFusionError::Plan(format!( + "snapshot {snapshot_id} does not exist or has expired" + ))); + } + }; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("historical", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(values)), + ], + )?; + Ok(Some(Arc::new(MemTable::try_new( + schema, + vec![vec![batch]], + )?))) + } +} + #[async_trait] impl TableEngineResolver for MetadataEngineResolver { async fn resolve_table( @@ -418,6 +493,488 @@ fn column_i32(batches: &[RecordBatch]) -> Vec { .collect() } +async fn enable_version_syntax(env: &TestEnv) { + env.ctx + .ctx() + .sql("SET datafusion.sql_parser.dialect = 'databricks'") + .await + .unwrap() + .collect() + .await + .unwrap(); +} + +async fn snapshot_env() -> (TestEnv, Arc) { + let resolver = Arc::new(SnapshotEngineResolver::default()); + let env = setup_with_engine(Some(resolver.clone())).await; + enable_version_syntax(&env).await; + (env, resolver) +} + +async fn routed_sql( + env: &TestEnv, + native: bool, + sql: &str, +) -> DFResult { + if native { + env.ctx.ctx().sql(sql).await + } else { + env.ctx.sql(sql).await + } +} + +#[tokio::test] +async fn engine_snapshot_uses_historical_data_and_schema() { + let (env, _) = snapshot_env().await; + for native in [false, true] { + let batches = routed_sql( + &env, + native, + &format!( + "SELECT h.historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 AS h WHERE h.id = 2" + ), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(string_column_values(&batches, "historical"), ["old-b"]); + + let err = routed_sql( + &env, + native, + &format!("SELECT payload FROM {CATALOG}.{DB}.it VERSION AS OF 10"), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("payload"), "{err}"); + + let batches = routed_sql( + &env, + native, + &format!("SELECT payload FROM {CATALOG}.{DB}.it ORDER BY id"), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(string_column_values(&batches, "payload"), ["x", "y"]); + + let mut versions = vec![LARGE_SNAPSHOT_ID.to_string()]; + if !native { + versions.push(format!("'{LARGE_SNAPSHOT_ID}'")); + } + for version in versions { + let batches = routed_sql( + &env, + native, + &format!("SELECT historical FROM {CATALOG}.{DB}.it VERSION AS OF {version}"), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(string_column_values(&batches, "historical"), ["large-id"]); + } + } +} + +#[tokio::test] +async fn sql_context_snapshot_remains_readable_when_current_resolution_fails() { + let (env, resolver) = current_unavailable_env().await; + let batches = env + .ctx + .sql(&format!( + "SELECT h.historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 AS h WHERE h.id = 2" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(string_column_values(&batches, "historical"), ["old-b"]); + + let err = env + .ctx + .sql(&format!( + "SELECT * FROM {CATALOG}.{DB}.it VERSION AS OF 404 WHERE FALSE" + )) + .await + .unwrap_err(); + assert!(err.to_string().contains("snapshot 404"), "{err}"); + assert_eq!(resolver.current_calls.load(Ordering::Relaxed), 0); + + for native in [false, true] { + for query in [ + format!("SELECT * FROM {CATALOG}.{DB}.it"), + format!("SELECT * FROM {CATALOG}.{DB}.it WHERE FALSE"), + format!("SELECT COUNT(*) FROM {CATALOG}.{DB}.it LIMIT 0"), + format!( + "SELECT h.historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 AS h \ + JOIN {CATALOG}.{DB}.it AS c ON h.id = c.id" + ), + ] { + let err = routed_sql(&env, native, &query).await.unwrap_err(); + assert!( + err.to_string().contains("current table is unavailable"), + "{err}" + ); + } + } +} + +#[tokio::test] +async fn raw_session_snapshot_propagates_current_resolution_errors() { + let (env, resolver) = current_unavailable_env().await; + for suffix in ["", " WHERE FALSE", " LIMIT 0"] { + let err = env + .ctx + .ctx() + .sql(&format!( + "SELECT * FROM {CATALOG}.{DB}.it VERSION AS OF 10{suffix}" + )) + .await + .unwrap_err(); + assert!( + err.to_string().contains("current table is unavailable"), + "{err}" + ); + } + assert_eq!(resolver.calls.load(Ordering::Relaxed), 0); +} + +async fn current_unavailable_env() -> (TestEnv, Arc) { + let resolver = Arc::new(SnapshotEngineResolver { + current_unavailable: true, + ..Default::default() + }); + let env = setup_with_engine(Some(resolver.clone())).await; + enable_version_syntax(&env).await; + (env, resolver) +} + +#[tokio::test] +async fn engine_errors_propagate_to_schema_table() { + let (env, _) = current_unavailable_env().await; + let schema = env.ctx.ctx().catalog(CATALOG).unwrap().schema(DB).unwrap(); + let err = schema.table("it").await.map(|_| ()).unwrap_err(); + assert!( + err.to_string().contains("current table is unavailable"), + "{err}" + ); +} + +#[tokio::test] +async fn engine_errors_propagate_to_context_table_provider() { + let (env, _) = current_unavailable_env().await; + let err = env + .ctx + .ctx() + .table_provider(format!("{CATALOG}.{DB}.it")) + .await + .map(|_| ()) + .unwrap_err(); + assert!( + err.to_string().contains("current table is unavailable"), + "{err}" + ); +} + +#[tokio::test] +async fn engine_errors_propagate_to_context_table() { + let (env, _) = current_unavailable_env().await; + let err = env + .ctx + .ctx() + .table(format!("{CATALOG}.{DB}.it")) + .await + .map(|_| ()) + .unwrap_err(); + assert!( + err.to_string().contains("current table is unavailable"), + "{err}" + ); +} + +#[tokio::test] +async fn engine_snapshots_can_be_joined_with_each_other_and_paimon() { + let (env, _) = snapshot_env().await; + for native in [false, true] { + let batches = routed_sql( + &env, + native, + &format!( + "SELECT a.historical AS old_value, b.historical AS new_value \ + FROM {CATALOG}.{DB}.it VERSION AS OF 10 AS a \ + JOIN {CATALOG}.{DB}.it VERSION AS OF 20 AS b ON a.id = b.id" + ), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+-----------+-----------+", + "| old_value | new_value |", + "+-----------+-----------+", + "| old-a | new-a |", + "+-----------+-----------+", + ], + &batches + ); + + let batches = routed_sql( + &env, + native, + &format!( + "SELECT h.historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 AS h \ + JOIN {CATALOG}.{DB}.pt VERSION AS OF 1 AS p ON h.id = p.id" + ), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(string_column_values(&batches, "historical"), ["old-a"]); + } +} + +#[tokio::test] +async fn engine_snapshot_supports_count_and_empty_results() { + let (env, _) = snapshot_env().await; + for native in [false, true] { + let batches = routed_sql( + &env, + native, + &format!("SELECT COUNT(*) AS n FROM {CATALOG}.{DB}.it VERSION AS OF 10"), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!(["+---+", "| n |", "+---+", "| 2 |", "+---+"], &batches); + + let df = routed_sql( + &env, + native, + &format!("SELECT historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 WHERE FALSE"), + ) + .await + .unwrap(); + assert_eq!(df.schema().field(0).name(), "historical"); + assert_eq!( + df.collect() + .await + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0 + ); + } +} + +#[tokio::test] +async fn engine_snapshot_errors_never_fall_back_to_current_data() { + let (env, _) = snapshot_env().await; + for native in [false, true] { + for (table, version, message) in [ + ("it", 404, "snapshot 404 does not exist or has expired"), + ("it", 500, "snapshot backend failed"), + ("ghost", 10, "not found"), + ] { + let err = routed_sql( + &env, + native, + &format!( + "SELECT * FROM {CATALOG}.{DB}.{table} VERSION AS OF {version} WHERE FALSE" + ), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains(message), "{err}"); + } + } +} + +#[tokio::test] +async fn engine_snapshot_providers_remain_read_only() { + use datafusion::catalog::default_table_source::source_as_provider; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::logical_expr::{dml::InsertOp, LogicalPlan}; + use datafusion::physical_plan::empty::EmptyExec; + + let (env, _) = snapshot_env().await; + for native in [false, true] { + let df = routed_sql( + &env, + native, + &format!("SELECT * FROM {CATALOG}.{DB}.it VERSION AS OF 10"), + ) + .await + .unwrap(); + let mut provider = None; + df.logical_plan() + .apply(|node| { + if let LogicalPlan::TableScan(scan) = node { + provider = Some(source_as_provider(&scan.source)?); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + let provider = provider.unwrap(); + let err = provider + .insert_into( + &env.ctx.ctx().state(), + Arc::new(EmptyExec::new(provider.schema())), + InsertOp::Append, + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("write is not supported"), "{err}"); + } +} + +#[tokio::test] +async fn engine_snapshot_resolution_checks_session_options() { + let (env, resolver) = snapshot_env().await; + for (key, value, message) in [ + ("query-auth.enabled", "true", "query-auth"), + ("scan.version", "20", "time travel"), + ("incremental-between", "1,5", "incremental-between"), + ] { + env.ctx + .sql(&format!("SET 'paimon.{key}' = '{value}'")) + .await + .unwrap(); + for native in [false, true] { + let err = routed_sql( + &env, + native, + &format!("SELECT * FROM {CATALOG}.{DB}.it VERSION AS OF 10"), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains(message), "{err}"); + assert_eq!(resolver.calls.load(Ordering::Relaxed), 0); + } + env.ctx.sql(&format!("RESET 'paimon.{key}'")).await.unwrap(); + } +} + +#[tokio::test] +async fn engine_snapshot_rejects_unsupported_selectors_and_metadata_tables() { + let (env, resolver) = snapshot_env().await; + for native in [false, true] { + let (selector, message) = if native { + ( + "FOR SYSTEM_TIME AS OF '2020-01-01 00:00:00'", + "time travel is not supported", + ) + } else { + ("VERSION AS OF 'tag-name'", "integer snapshot id") + }; + for (table, selector, message) in [ + ("it", selector, message), + ( + "it", + "VERSION AS OF 9223372036854775808", + "integer snapshot id", + ), + ( + "it", + "TIMESTAMP AS OF '2020-01-01 00:00:00'", + "time travel is not supported", + ), + ( + "it$history", + "VERSION AS OF 10", + "time travel is not supported", + ), + ( + "it$branch_b1", + "VERSION AS OF 10", + "branches are not supported", + ), + ] { + let err = routed_sql( + &env, + native, + &format!("SELECT * FROM {CATALOG}.{DB}.\"{table}\" {selector}"), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains(message), "{selector}: {err}"); + assert_eq!(resolver.calls.load(Ordering::Relaxed), 0); + } + } +} + +#[tokio::test] +async fn engine_snapshot_requires_explicit_resolver_support() { + for registered in [false, true] { + let resolver = + registered.then(|| Arc::new(FakeEngineResolver) as Arc); + let env = setup_with_engine(resolver).await; + enable_version_syntax(&env).await; + for native in [false, true] { + let err = routed_sql( + &env, + native, + &format!("SELECT * FROM {CATALOG}.{DB}.it VERSION AS OF 10 WHERE FALSE"), + ) + .await + .unwrap_err(); + let expected = if registered { + "time travel is not supported" + } else { + "no table engine is registered" + }; + assert!(err.to_string().contains(expected), "{err}"); + } + } +} + +#[test] +fn engine_snapshot_resolution_on_runtime_workers() { + let runtime = paimon_datafusion::runtime::runtime(); + runtime.block_on(async { + let (env, _) = snapshot_env().await; + tokio::time::timeout( + std::time::Duration::from_secs(20), + tokio::spawn(async move { + let batches = routed_sql( + &env, + true, + &format!( + "SELECT historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 ORDER BY id" + ), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + string_column_values(&batches, "historical"), + ["old-a", "old-b"] + ); + }), + ) + .await + .expect("snapshot resolution stalled on a runtime worker") + .unwrap(); + }); +} + #[tokio::test] async fn paimon_path_still_serves_paimon_tables() { let env = setup().await; @@ -1385,6 +1942,27 @@ async fn registering_on_a_raw_session_installs_the_planner() { panic!("registration must install the planner, so this cannot read current data"); }; assert!(err.to_string().contains("not supported"), "{err}"); + + register_catalog_table_engine( + &ctx, + CATALOG, + TableType::IcebergTable, + Arc::new(SnapshotEngineResolver::default()), + ) + .unwrap(); + let batches = ctx + .sql(&format!( + "SELECT historical FROM {CATALOG}.{DB}.it VERSION AS OF 10 ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + string_column_values(&batches, "historical"), + ["old-a", "old-b"] + ); } #[tokio::test]