Introduce the DynamicFileCatalog in datafusion-catalog

datafusion-cli/src/catalog.rs

@@ -18,22 +18,19 @@
 use std::any::Any;
 use std::sync::{Arc, Weak};
 
-use crate::object_storage::{get_object_store, AwsOptions, GcpOptions};
-
+use async_trait::async_trait;
+use datafusion::catalog::dynamic_file_schema::substitute_tilde;
 use datafusion::catalog::schema::SchemaProvider;
 use datafusion::catalog::{CatalogProvider, CatalogProviderList};
 use datafusion::common::plan_datafusion_err;
-use datafusion::datasource::listing::{
-    ListingTable, ListingTableConfig, ListingTableUrl,
-};
+use datafusion::datasource::listing::ListingTableUrl;
 use datafusion::datasource::TableProvider;
 use datafusion::error::Result;
 use datafusion::execution::context::SessionState;
-
-use async_trait::async_trait;
-use dirs::home_dir;
 use parking_lot::RwLock;
 
+use crate::object_storage::{get_object_store, AwsOptions, GcpOptions};
+
 /// Wraps another catalog, automatically creating table providers
 /// for local files if needed
 pub struct DynamicFileCatalog {
@@ -115,7 +112,8 @@ impl CatalogProvider for DynamicFileCatalogProvider {
     }
 }
 
-/// Wraps another schema provider
+/// Wraps another schema provider. [DynamicFileSchemaProvider] is responsible for registering the required
+/// object stores for the file locations.
 struct DynamicFileSchemaProvider {
     inner: Arc<dyn SchemaProvider>,
     state: Weak<RwLock<SessionState>>,
@@ -149,9 +147,11 @@ impl SchemaProvider for DynamicFileSchemaProvider {
     }
 
     async fn table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
-        let inner_table = self.inner.table(name).await?;
-        if inner_table.is_some() {
-            return Ok(inner_table);
+        let inner_table = self.inner.table(name).await;
+        if inner_table.is_ok() {
+            if let Some(inner_table) = inner_table? {
+                return Ok(Some(inner_table));
+            }
         }
 
         // if the inner schema provider didn't have a table by
@@ -195,16 +195,7 @@ impl SchemaProvider for DynamicFileSchemaProvider {
                 state.runtime_env().register_object_store(url, store);
             }
         }
-
-        let config = match ListingTableConfig::new(table_url).infer(&state).await {
-            Ok(cfg) => cfg,
-            Err(_) => {
-                // treat as non-existing
-                return Ok(None);
-            }
-        };
-
-        Ok(Some(Arc::new(ListingTable::try_new(config)?)))
+        self.inner.table(name).await
     }
 
     fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
@@ -215,24 +206,14 @@ impl SchemaProvider for DynamicFileSchemaProvider {
         self.inner.table_exist(name)
     }
 }
-fn substitute_tilde(cur: String) -> String {
-    if let Some(usr_dir_path) = home_dir() {
-        if let Some(usr_dir) = usr_dir_path.to_str() {
-            if cur.starts_with('~') && !usr_dir.is_empty() {
-                return cur.replacen('~', usr_dir, 1);
-            }
-        }
-    }
-    cur
-}
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-
     use datafusion::catalog::schema::SchemaProvider;
     use datafusion::prelude::SessionContext;
 
+    use super::*;
+
     fn setup_context() -> (SessionContext, Arc<dyn SchemaProvider>) {
         let mut ctx = SessionContext::new();
         ctx.register_catalog_list(Arc::new(DynamicFileCatalog::new(
@@ -262,7 +243,7 @@ mod tests {
         let (ctx, schema) = setup_context();
 
         // That's a non registered table so expecting None here
-        let table = schema.table(&location).await.unwrap();
+        let table = schema.table(&location).await?;
         assert!(table.is_none());
 
         // It should still create an object store for the location in the SessionState
@@ -286,7 +267,7 @@ mod tests {
 
         let (ctx, schema) = setup_context();
 
-        let table = schema.table(&location).await.unwrap();
+        let table = schema.table(&location).await?;
         assert!(table.is_none());
 
         let store = ctx
@@ -308,7 +289,7 @@ mod tests {
 
         let (ctx, schema) = setup_context();
 
-        let table = schema.table(&location).await.unwrap();
+        let table = schema.table(&location).await?;
         assert!(table.is_none());
 
         let store = ctx
@@ -330,41 +311,4 @@ mod tests {
 
         assert!(schema.table(location).await.is_err());
     }
-    #[cfg(not(target_os = "windows"))]
-    #[test]
-    fn test_substitute_tilde() {
-        use std::env;
-        use std::path::MAIN_SEPARATOR;
-        let original_home = home_dir();
-        let test_home_path = if cfg!(windows) {
-            "C:\\Users\\user"
-        } else {
-            "/home/user"
-        };
-        env::set_var(
-            if cfg!(windows) { "USERPROFILE" } else { "HOME" },
-            test_home_path,
-        );
-        let input = "~/Code/datafusion/benchmarks/data/tpch_sf1/part/part-0.parquet";
-        let expected = format!(
-            "{}{}Code{}datafusion{}benchmarks{}data{}tpch_sf1{}part{}part-0.parquet",
-            test_home_path,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR,
-            MAIN_SEPARATOR
-        );
-        let actual = substitute_tilde(input.to_string());
-        assert_eq!(actual, expected);
-        match original_home {
-            Some(home_path) => env::set_var(
-                if cfg!(windows) { "USERPROFILE" } else { "HOME" },
-                home_path.to_str().unwrap(),
-            ),
-            None => env::remove_var(if cfg!(windows) { "USERPROFILE" } else { "HOME" }),
-        }
-    }
 }

datafusion-examples/examples/csv_sql.rs

@@ -24,25 +24,39 @@ use datafusion::prelude::*;
 #[tokio::main]
 async fn main() -> Result<()> {
     // create local execution context
-    let ctx = SessionContext::new();
+    let cfg = SessionConfig::new().set_str("datafusion.catalog.has_header", "true");
+    let ctx = SessionContext::new_with_config(cfg);
 
     let testdata = datafusion::test_util::arrow_test_data();
-
+    let path = &format!("{testdata}/csv/aggregate_test_100.csv");
     // register csv file with the execution context
-    ctx.register_csv(
-        "aggregate_test_100",
-        &format!("{testdata}/csv/aggregate_test_100.csv"),
-        CsvReadOptions::new(),
-    )
-    .await?;
+    ctx.register_csv("aggregate_test_100", path, CsvReadOptions::new())
+        .await?;
 
     // execute the query
     let df = ctx
         .sql(
-            "SELECT c1, MIN(c12), MAX(c12) \
-        FROM aggregate_test_100 \
-        WHERE c11 > 0.1 AND c11 < 0.9 \
-        GROUP BY c1",
+            r#"SELECT c1, MIN(c12), MAX(c12)
+        FROM aggregate_test_100
+        WHERE c11 > 0.1 AND c11 < 0.9
+        GROUP BY c1"#,
+        )
+        .await?;
+
+    // print the results
+    df.show().await?;
+
+    // query the file by the path dynamically.
+    let df = ctx
+        .sql(
+            format!(
+                r#"SELECT c1, MIN(c12), MAX(c12)
+        FROM '{}'
+        WHERE c11 > 0.1 AND c11 < 0.9
+        GROUP BY c1"#,
+                path
+            )
+            .as_str(),
         )
         .await?;
 

datafusion-examples/examples/external_dependency/query-aws-s3.rs

@@ -63,5 +63,13 @@ async fn main() -> Result<()> {
     // print the results
     df.show().await?;
 
+    // dynamic query by the file path
+    let df = ctx
+        .sql(format!(r#"SELECT * FROM '{}' LIMIT 10"#, &path).as_str())
+        .await?;
+
+    // print the results
+    df.show().await?;
+
     Ok(())
 }

datafusion/core/Cargo.toml

@@ -108,6 +108,7 @@ datafusion-physical-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
 datafusion-sql = { workspace = true }
+dirs = "4.0.0"
 flate2 = { version = "1.0.24", optional = true }
 futures = { workspace = true }
 glob = "0.3.0"