feat(datafusion): support snapshot time travel for external engines

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 <snapshot_id>` 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`.

crates/integrations/datafusion/src/catalog.rs

@@ -49,9 +49,7 @@ pub(crate) type SessionStateProvider = Arc<dyn Fn() -> Option<SessionState> + Se
 /// providers, so registrations stay visible to schemas obtained earlier.
 type TableEngines = Arc<RwLock<HashMap<PaimonTableType, Arc<dyn TableEngineResolver>>>>;
 
-/// 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<Option<Arc<dyn TableProvider>>>;
 
+    /// 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<Option<Arc<dyn TableProvider>>> {
+        Err(plan_datafusion_err!(
+            "time travel is not supported for routed '{}' tables ('{}.{}')",
+            request.declared,
+            request.database,
+            request.table
+        ))
+    }
+
     /// Resolves `<table>$<metadata_table>`; 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<dyn TableProvider>,
     pub(crate) declared: PaimonTableType,
     pub(crate) table_name: String,
     metadata_table: bool,
+    request: EngineTableRequest,
+    resolver: Arc<dyn TableEngineResolver>,
+    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<Arc<dyn TableProvider>> {
+        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<i64> {
+    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<dyn TableEngineResolver>,
+    dynamic_options: DynamicOptions,
+    snapshot_id: i64,
+) -> DFResult<Arc<dyn TableProvider>> {
+    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<dyn TableProvider>)
+    })
+    .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<Arc<dyn TableProvider>> {
+        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<Arc<dyn SchemaProvider>> {
         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<dyn TableProvider>
                     }))
                 }

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 <integer or string>` → sets `scan.version` option on the table.
-///   At scan time, the version is resolved: tag name (if exists) → snapshot id → error.
-/// - `TIMESTAMP AS OF <timestamp string>` → 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<RelationPlanning> {
-        // 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::<UnavailableEngineTableProvider>() {
+            return Err(unavailable.unavailable_error());
+        }
+
         if let Some(routed) = provider.downcast_ref::<ReadOnlyTableProvider>() {
-            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::<PaimonTableProvider>() 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::<ObjectTableProvider>()
                 || 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,30 +124,32 @@ 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 },
             "paimon time travel resolution thread panicked",
         )
         .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<dyn TableProvider>,
+) -> DFResult<RelationPlanning> {
+    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,