Introduce the DynamicFileCatalog in datafusion-catalog

datafusion-examples/examples/dynamic_csv_sql.rs

@@ -0,0 +1,47 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+
+/// This example demonstrates executing a simple query against an Arrow data source (CSV) and
+/// fetching results
+#[tokio::main]
+async fn main() -> Result<()> {
+    // create local execution context
+    let ctx = SessionContext::new();
+
+    let testdata = datafusion::test_util::arrow_test_data();
+    let path = &format!("file:///{testdata}/csv/aggregate_test_100.csv");
+    // execute the query
+    let df = ctx
+        .sql(
+            format!(
+                r#"SELECT column_1, MIN(column_12), MAX(column_12)
+        FROM '{}'
+        WHERE column_11 > 0.1 AND column_11 < 0.9
+        GROUP BY column_1"#,
+                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"

datafusion/core/src/catalog/dynamic_file_schema.rs

@@ -0,0 +1,118 @@
+use std::any::Any;
+use std::sync::{Arc, Weak};
+
+use async_trait::async_trait;
+use dirs::home_dir;
+use parking_lot::{Mutex, RwLock};
+
+use datafusion_common::plan_datafusion_err;
+
+use crate::catalog::schema::SchemaProvider;
+use crate::datasource::listing::{ListingTable, ListingTableConfig, ListingTableUrl};
+use crate::datasource::TableProvider;
+use crate::error::Result;
+use crate::execution::context::SessionState;
+
+/// Wraps another schema provider
+pub struct DynamicFileSchemaProvider {
+    inner: Arc<dyn SchemaProvider>,
+    state_store: StateStore,
+}
+
+impl DynamicFileSchemaProvider {
+    pub fn new(inner: Arc<dyn SchemaProvider>) -> Self {
+        Self {
+            inner,
+            state_store: StateStore::new(),
+        }
+    }
+
+    pub fn with_state(&self, state: Weak<RwLock<SessionState>>) {
+        self.state_store.with_state(state);
+    }
+}
+
+#[async_trait]
+impl SchemaProvider for DynamicFileSchemaProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn table_names(&self) -> Vec<String> {
+        self.inner.table_names()
+    }
+
+    fn register_table(
+        &self,
+        name: String,
+        table: Arc<dyn TableProvider>,
+    ) -> Result<Option<Arc<dyn TableProvider>>> {
+        self.inner.register_table(name, table)
+    }
+
+    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 optimized_url = substitute_tilde(name.to_owned());
+        let table_url = ListingTableUrl::parse(optimized_url.as_str())?;
+        let state = &self
+            .state_store
+            .get_state()
+            .upgrade()
+            .ok_or_else(|| plan_datafusion_err!("locking error"))?
+            .read()
+            .clone();
+        let cfg = ListingTableConfig::new(table_url.clone())
+            .infer(state)
+            .await?;
+
+        Ok(Some(Arc::new(ListingTable::try_new(cfg)?)))
+    }
+
+    fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
+        self.inner.deregister_table(name)
+    }
+
+    fn table_exist(&self, name: &str) -> bool {
+        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
+}
+
+pub struct StateStore {
+    state: Arc<Mutex<Option<Weak<RwLock<SessionState>>>>>,
+}
+
+impl StateStore {
+    pub fn new() -> Self {
+        Self {
+            state: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    pub fn with_state(&self, state: Weak<RwLock<SessionState>>) {
+        let mut lock = self.state.lock();
+        *lock = Some(state);
+    }
+
+    pub fn get_state(&self) -> Weak<RwLock<SessionState>> {
+        self.state.lock().clone().unwrap()
+    }
+}
+
+impl Default for StateStore {
+    fn default() -> Self {
+        Self::new()
+    }
+}

datafusion/core/src/catalog/mod.rs

@@ -17,6 +17,7 @@
 
 //! Interfaces and default implementations of catalogs and schemas.
 
+pub mod dynamic_file_schema;
 pub mod information_schema;
 pub mod listing_schema;
 pub mod schema;

datafusion/core/src/execution/context/mod.rs

@@ -72,6 +72,7 @@ use object_store::ObjectStore;
 use parking_lot::RwLock;
 use url::Url;
 
+use crate::catalog::dynamic_file_schema::DynamicFileSchemaProvider;
 pub use datafusion_execution::config::SessionConfig;
 pub use datafusion_execution::TaskContext;
 pub use datafusion_expr::execution_props::ExecutionProps;
@@ -305,10 +306,18 @@ impl SessionContext {
 
     /// Creates a new `SessionContext` using the provided [`SessionState`]
     pub fn new_with_state(state: SessionState) -> Self {
+        let state_ref = Arc::new(RwLock::new(state.clone()));
+        state
+            .schema_for_ref("datafusion.public.xx")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<DynamicFileSchemaProvider>()
+            .unwrap()
+            .with_state(Arc::downgrade(&state_ref));
         Self {
-            session_id: state.session_id().to_string(),
+            session_id: state_ref.clone().read().session_id().to_string(),
             session_start_time: Utc::now(),
-            state: Arc::new(RwLock::new(state)),
+            state: state_ref,
         }
     }
 

datafusion/core/src/execution/session_state.rs

@@ -17,6 +17,7 @@
 
 //! [`SessionState`]: information required to run queries in a session
 
+use crate::catalog::dynamic_file_schema::DynamicFileSchemaProvider;
 use crate::catalog::information_schema::{InformationSchemaProvider, INFORMATION_SCHEMA};
 use crate::catalog::listing_schema::ListingSchemaProvider;
 use crate::catalog::schema::{MemorySchemaProvider, SchemaProvider};
@@ -196,11 +197,13 @@ impl SessionState {
 
         if config.create_default_catalog_and_schema() {
             let default_catalog = MemoryCatalogProvider::new();
+            let schema =
+                DynamicFileSchemaProvider::new(Arc::new(MemorySchemaProvider::new()));
 
             default_catalog
                 .register_schema(
                     &config.options().catalog.default_schema,
-                    Arc::new(MemorySchemaProvider::new()),
+                    Arc::new(schema),
                 )
                 .expect("memory catalog provider can register schema");