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

crates/integrations/datafusion/README.md

@@ -27,13 +27,16 @@ This crate contains the integration of [Apache DataFusion](https://datafusion.ap
 ## External engine snapshots
 
 Implement `TableEngineResolver::resolve_table_snapshot` to support
-`VERSION AS OF <snapshot_id>` on external data tables (integer IDs only).
+`VERSION AS OF <snapshot_id>` on external data tables (integer IDs only), and
+`resolve_table_at_timestamp` to support `TIMESTAMP AS OF '<yyyy-MM-dd HH:mm:ss>'`.
+The timestamp is read in the local time zone and passed as epoch milliseconds;
+the engine picks the snapshot.
 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
+dialect supporting these clauses; it preloads the current provider and propagates
 loading errors.
 
 ## Reader memory

crates/integrations/datafusion/src/catalog.rs

@@ -93,6 +93,22 @@ pub trait TableEngineResolver: Debug + Send + Sync {
         ))
     }
 
+    /// Resolve the table as of `timestamp_millis` (epoch milliseconds) and its
+    /// schema; the engine picks the snapshot. A time with no snapshot is an error;
+    /// `Ok(None)` means the table is missing.
+    async fn resolve_table_at_timestamp(
+        &self,
+        request: &EngineTableRequest,
+        _timestamp_millis: i64,
+    ) -> DFResult<Option<Arc<dyn TableProvider>>> {
+        Err(plan_datafusion_err!(
+            "timestamp 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.
@@ -117,25 +133,50 @@ pub(crate) struct ReadOnlyTableProvider {
     dynamic_options: DynamicOptions,
 }
 
+/// The historical state requested from an engine.
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum EngineTimeTravel {
+    Snapshot(i64),
+    TimestampMillis(i64),
+}
+
+impl std::fmt::Display for EngineTimeTravel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Snapshot(id) => write!(f, "snapshot {id}"),
+            Self::TimestampMillis(millis) => write!(f, "timestamp-millis {millis}"),
+        }
+    }
+}
+
+pub(crate) fn unsupported_engine_time_travel(
+    declared: PaimonTableType,
+    table_name: &str,
+) -> datafusion::error::DataFusionError {
+    plan_datafusion_err!(
+        "time travel is not supported for routed '{declared}' tables ('{table_name}'); \
+         only VERSION AS OF with an integer snapshot id and TIMESTAMP AS OF \
+         are supported on data tables"
+    )
+}
+
 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
-        )
+        unsupported_engine_time_travel(self.declared, &self.table_name)
     }
 
-    pub(crate) async fn at_snapshot(&self, snapshot_id: i64) -> DFResult<Arc<dyn TableProvider>> {
+    pub(crate) async fn at_time_travel(
+        &self,
+        travel: EngineTimeTravel,
+    ) -> DFResult<Arc<dyn TableProvider>> {
         if self.metadata_table {
             return Err(self.unsupported_time_travel());
         }
-        resolve_engine_snapshot(
+        resolve_engine_time_travel(
             self.request.clone(),
             Arc::clone(&self.resolver),
             Arc::clone(&self.dynamic_options),
-            snapshot_id,
+            travel,
         )
         .await
     }
@@ -149,11 +190,11 @@ pub(crate) fn parse_engine_snapshot_id(version: &str) -> DFResult<i64> {
     })
 }
 
-async fn resolve_engine_snapshot(
+async fn resolve_engine_time_travel(
     request: EngineTableRequest,
     resolver: Arc<dyn TableEngineResolver>,
     dynamic_options: DynamicOptions,
-    snapshot_id: i64,
+    travel: EngineTimeTravel,
 ) -> DFResult<Arc<dyn TableProvider>> {
     await_with_runtime(async move {
         let table_name = format!("{}.{}", request.database, request.table);
@@ -164,12 +205,17 @@ async fn resolve_engine_snapshot(
         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}")
-            })?;
+        let resolved = match travel {
+            EngineTimeTravel::Snapshot(id) => resolver.resolve_table_snapshot(&request, id).await?,
+            EngineTimeTravel::TimestampMillis(millis) => {
+                resolver
+                    .resolve_table_at_timestamp(&request, millis)
+                    .await?
+            }
+        };
+        let inner = resolved.ok_or_else(|| {
+            plan_datafusion_err!("table '{table_name}' was not found at {travel}")
+        })?;
         Ok(Arc::new(ReadOnlyTableProvider {
             inner,
             declared: request.declared,
@@ -423,10 +469,10 @@ impl PaimonCatalogProvider {
         Arc::clone(&self.table_engines)
     }
 
-    pub(crate) async fn table_at_snapshot(
+    pub(crate) async fn table_at_time_travel(
         &self,
         request: EngineTableRequest,
-        snapshot_id: i64,
+        travel: EngineTimeTravel,
     ) -> DFResult<Arc<dyn TableProvider>> {
         let resolver = self
             .table_engines
@@ -442,13 +488,8 @@ impl PaimonCatalogProvider {
                     request.table
                 )
             })?;
-        resolve_engine_snapshot(
-            request,
-            resolver,
-            Arc::clone(&self.dynamic_options),
-            snapshot_id,
-        )
-        .await
+        resolve_engine_time_travel(request, resolver, Arc::clone(&self.dynamic_options), travel)
+            .await
     }
 
     fn paimon_schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {

crates/integrations/datafusion/src/relation_planner.rs

@@ -34,11 +34,12 @@ use datafusion::sql::sqlparser::ast::{self, TableFactor, TableVersion};
 use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION};
 
 use crate::catalog::{
-    parse_engine_snapshot_id, ReadOnlyTableProvider, UnavailableEngineTableProvider,
+    parse_engine_snapshot_id, EngineTimeTravel, ReadOnlyTableProvider,
+    UnavailableEngineTableProvider,
 };
 use crate::table::{ObjectTableProvider, PaimonTableProvider};
 
-/// Resolves Paimon time travel and registered engine snapshot IDs.
+/// Resolves Paimon time travel, and registered engine snapshot IDs and timestamps.
 #[derive(Debug)]
 pub struct PaimonRelationPlanner;
 
@@ -84,15 +85,22 @@ impl RelationPlanner for PaimonRelationPlanner {
         }
 
         if let Some(routed) = provider.downcast_ref::<ReadOnlyTableProvider>() {
-            let TableVersion::VersionAsOf(expr) = version else {
-                return Err(routed.unsupported_time_travel());
+            let travel = match version {
+                TableVersion::VersionAsOf(expr) => {
+                    let options = resolve_version_as_of(expr)?;
+                    EngineTimeTravel::Snapshot(parse_engine_snapshot_id(
+                        &options[SCAN_VERSION_OPTION],
+                    )?)
+                }
+                TableVersion::TimestampAsOf(expr) => {
+                    EngineTimeTravel::TimestampMillis(timestamp_as_of_millis(expr)?)
+                }
+                _ => 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",
+                async move { routed.at_time_travel(travel).await },
+                "engine time travel resolution thread panicked",
             )?;
             return plan_scan(relation, table_ref, new_provider);
         }
@@ -210,14 +218,19 @@ fn resolve_version_as_of(expr: &ast::Expr) -> DFResult<HashMap<String, String>>
 
 /// Resolve `TIMESTAMP AS OF <expr>` into `scan.timestamp-millis` option.
 fn resolve_timestamp_as_of(expr: &ast::Expr) -> DFResult<HashMap<String, String>> {
+    let millis = timestamp_as_of_millis(expr)?;
+    Ok(HashMap::from([(
+        SCAN_TIMESTAMP_MILLIS_OPTION.to_string(),
+        millis.to_string(),
+    )]))
+}
+
+/// Parse the timestamp string of `TIMESTAMP AS OF <expr>` to epoch millis.
+fn timestamp_as_of_millis(expr: &ast::Expr) -> DFResult<i64> {
     match expr {
         ast::Expr::Value(v) => match &v.value {
             ast::Value::SingleQuotedString(s) | ast::Value::DoubleQuotedString(s) => {
-                let millis = parse_timestamp_to_millis(s)?;
-                Ok(HashMap::from([(
-                    SCAN_TIMESTAMP_MILLIS_OPTION.to_string(),
-                    millis.to_string(),
-                )]))
+                parse_timestamp_to_millis(s)
             }
             _ => Err(datafusion::error::DataFusionError::Plan(format!(
                 "Unsupported TIMESTAMP AS OF expression: {expr}. Expected a timestamp string."

crates/integrations/datafusion/src/sql_context.rs

@@ -773,8 +773,6 @@ impl SQLContext {
         sql: &str,
         sql_options: SQLOptions,
     ) -> DFResult<DataFrame> {
-        use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION};
-
         let mut tracker = crate::merge_into::TempTableTracker::new(self);
 
         let version_clauses = extract_all_version_as_of(sql);
@@ -797,8 +795,7 @@ impl SQLContext {
                     &catalog,
                     &catalog_name,
                     &identifier,
-                    SCAN_VERSION_OPTION,
-                    info.version.clone(),
+                    TimeTravelClause::Version(info.version.clone()),
                 )
                 .await?;
 
@@ -818,8 +815,7 @@ impl SQLContext {
                     &catalog,
                     &catalog_name,
                     &identifier,
-                    SCAN_TIMESTAMP_MILLIS_OPTION,
-                    millis.to_string(),
+                    TimeTravelClause::TimestampMillis(millis),
                 )
                 .await?;
 
@@ -865,9 +861,11 @@ impl SQLContext {
         catalog: &Arc<dyn Catalog>,
         catalog_name: &str,
         identifier: &Identifier,
-        option: &str,
-        value: String,
+        clause: TimeTravelClause,
     ) -> DFResult<Arc<dyn TableProvider>> {
+        use crate::catalog::EngineTimeTravel;
+        use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION};
+
         let parsed = identifier
             .parsed_object_name()
             .map_err(to_datafusion_error)?;
@@ -887,15 +885,20 @@ impl SQLContext {
                     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",
+            if parsed.system_table().is_some() {
+                return Err(crate::catalog::unsupported_engine_time_travel(
                     external.declared(),
-                    identifier.full_name()
+                    &identifier.full_name(),
                 ));
             }
-            let snapshot_id = crate::catalog::parse_engine_snapshot_id(&value)?;
+            let travel = match clause {
+                TimeTravelClause::Version(version) => {
+                    EngineTimeTravel::Snapshot(crate::catalog::parse_engine_snapshot_id(&version)?)
+                }
+                TimeTravelClause::TimestampMillis(millis) => {
+                    EngineTimeTravel::TimestampMillis(millis)
+                }
+            };
             let provider = self
                 .ctx
                 .catalog(catalog_name)
@@ -906,19 +909,25 @@ impl SQLContext {
                     plan_datafusion_err!("Catalog '{catalog_name}' is not a Paimon catalog")
                 })?;
             return provider
-                .table_at_snapshot(
+                .table_at_time_travel(
                     crate::EngineTableRequest::new(
                         base.database().to_string(),
                         base.object().to_string(),
                         external.declared(),
                     ),
-                    snapshot_id,
+                    travel,
                 )
                 .await;
         }
 
         let (table, base, system_name) =
             crate::table_loader::loaded_table_for_read(loaded, identifier).await?;
+        let (option, value) = match clause {
+            TimeTravelClause::Version(version) => (SCAN_VERSION_OPTION, version),
+            TimeTravelClause::TimestampMillis(millis) => {
+                (SCAN_TIMESTAMP_MILLIS_OPTION, millis.to_string())
+            }
+        };
         let mut options = self.dynamic_options.read().unwrap().clone();
         options.insert(option.to_string(), value);
         let table = table
@@ -3584,6 +3593,12 @@ struct TimestampAsOfInfo {
     clause_range: (usize, usize),
 }
 
+/// A parsed `VERSION AS OF` or `TIMESTAMP AS OF` selector.
+enum TimeTravelClause {
+    Version(String),
+    TimestampMillis(i64),
+}
+
 /// Check whether a SQL string contains a time-travel keyword (`VERSION AS OF` or
 /// `TIMESTAMP AS OF`) **outside** of single-quoted string literals, `--` line
 /// comments, and `/* */` block comments.