Allow listing tables to be created via TableFactories

datafusion-cli/src/main.rs

@@ -16,6 +16,9 @@
 // under the License.
 
 use clap::Parser;
+use datafusion::datasource::datasource::TableProviderFactory;
+use datafusion::datasource::file_format::file_type::FileType;
+use datafusion::datasource::listing_table_factory::ListingTableFactory;
 use datafusion::datasource::object_store::ObjectStoreRegistry;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::SessionConfig;
@@ -26,6 +29,7 @@ use datafusion_cli::{
     exec, print_format::PrintFormat, print_options::PrintOptions, DATAFUSION_CLI_VERSION,
 };
 use mimalloc::MiMalloc;
+use std::collections::HashMap;
 use std::env;
 use std::path::Path;
 use std::sync::Arc;
@@ -93,7 +97,7 @@ pub async fn main() -> Result<()> {
 
     if let Some(ref path) = args.data_path {
         let p = Path::new(path);
-        env::set_current_dir(&p).unwrap();
+        env::set_current_dir(p).unwrap();
     };
 
     let mut session_config = SessionConfig::from_env().with_information_schema(true);
@@ -105,6 +109,7 @@ pub async fn main() -> Result<()> {
     let runtime_env = create_runtime_env()?;
     let mut ctx =
         SessionContext::with_config_rt(session_config.clone(), Arc::new(runtime_env));
+    ctx.refresh_catalogs().await?;
 
     let mut print_options = PrintOptions {
         format: args.format,
@@ -142,11 +147,31 @@ pub async fn main() -> Result<()> {
 }
 
 fn create_runtime_env() -> Result<RuntimeEnv> {
+    let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
+        HashMap::new();
+    table_factories.insert(
+        "csv".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::CSV)),
+    );
+    table_factories.insert(
+        "parquet".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::PARQUET)),
+    );
+    table_factories.insert(
+        "avro".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::AVRO)),
+    );
+    table_factories.insert(
+        "json".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::JSON)),
+    );
+
     let object_store_provider = DatafusionCliObjectStoreProvider {};
     let object_store_registry =
         ObjectStoreRegistry::new_with_provider(Some(Arc::new(object_store_provider)));
-    let rn_config =
-        RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry));
+    let rn_config = RuntimeConfig::new()
+        .with_object_store_registry(Arc::new(object_store_registry))
+        .with_table_factories(table_factories);
     RuntimeEnv::new(rn_config)
 }
 

datafusion-cli/src/object_storage.rs

@@ -141,7 +141,13 @@ mod tests {
         assert!(err.to_string().contains("Generic S3 error: Missing region"));
 
         env::set_var("AWS_REGION", "us-east-1");
-        assert!(provider.get_by_url(&Url::from_str(s3).unwrap()).is_ok());
+        let url = Url::from_str(s3).expect("Unable to parse s3 url");
+        let res = provider.get_by_url(&url);
+        let msg = match res {
+            Err(e) => format!("{}", e),
+            Ok(_) => "".to_string()
+        };
+        assert_eq!("".to_string(), msg); // Fail with error message
         env::remove_var("AWS_REGION");
     }
 }

datafusion-cli/src/print_format.rs

@@ -107,9 +107,9 @@ mod tests {
         let batch = RecordBatch::try_new(
             schema,
             vec![
-                Arc::new(Int32Array::from_slice(&[1, 2, 3])),
-                Arc::new(Int32Array::from_slice(&[4, 5, 6])),
-                Arc::new(Int32Array::from_slice(&[7, 8, 9])),
+                Arc::new(Int32Array::from_slice([1, 2, 3])),
+                Arc::new(Int32Array::from_slice([4, 5, 6])),
+                Arc::new(Int32Array::from_slice([7, 8, 9])),
             ],
         )
         .unwrap();
@@ -137,9 +137,9 @@ mod tests {
         let batch = RecordBatch::try_new(
             schema,
             vec![
-                Arc::new(Int32Array::from_slice(&[1, 2, 3])),
-                Arc::new(Int32Array::from_slice(&[4, 5, 6])),
-                Arc::new(Int32Array::from_slice(&[7, 8, 9])),
+                Arc::new(Int32Array::from_slice([1, 2, 3])),
+                Arc::new(Int32Array::from_slice([4, 5, 6])),
+                Arc::new(Int32Array::from_slice([7, 8, 9])),
             ],
         )
         .unwrap();

datafusion/core/src/catalog/listing_schema.rs

@@ -19,8 +19,10 @@
 use crate::catalog::schema::SchemaProvider;
 use crate::datasource::datasource::TableProviderFactory;
 use crate::datasource::TableProvider;
-use datafusion_common::DataFusionError;
+use crate::execution::context::SessionState;
+use datafusion_common::{context, DataFusionError};
 use futures::TryStreamExt;
+use itertools::Itertools;
 use object_store::ObjectStore;
 use std::any::Any;
 use std::collections::{HashMap, HashSet};
@@ -72,7 +74,7 @@ impl ListingSchemaProvider {
     }
 
     /// Reload table information from ObjectStore
-    pub async fn refresh(&self) -> datafusion_common::Result<()> {
+    pub async fn refresh(&self, state: &SessionState) -> datafusion_common::Result<()> {
         let entries: Vec<_> = self
             .store
             .list(Some(&self.path))
@@ -100,13 +102,20 @@ impl ListingSchemaProvider {
                 .ok_or_else(|| {
                     DataFusionError::Internal("Cannot parse file name!".to_string())
                 })?;
-            let table_name = table.to_str().ok_or_else(|| {
+            let table_name = file_name.split('.').collect_vec()[0];
+            let table_path = table.to_str().ok_or_else(|| {
                 DataFusionError::Internal("Cannot parse file name!".to_string())
             })?;
-            if !self.table_exist(file_name) {
-                let table_name = format!("{}/{}", self.authority, table_name);
-                let provider = self.factory.create(table_name.as_str()).await?;
-                let _ = self.register_table(file_name.to_string(), provider.clone())?;
+            if !self.table_exist(table_name) {
+                let table_url = format!("{}/{}", self.authority, table_path);
+                let provider = self
+                    .factory
+                    .create(state, table_url.as_str())
+                    .await
+                    .map_err(|e| {
+                        context!(format!("Could not create table for {}", table_url), e)
+                    })?;
+                let _ = self.register_table(table_name.to_string(), provider.clone())?;
             }
         }
         Ok(())

datafusion/core/src/datasource/datasource.rs

@@ -92,5 +92,9 @@ pub trait TableProvider: Sync + Send {
 #[async_trait]
 pub trait TableProviderFactory: Sync + Send {
     /// Create a TableProvider with the given url
-    async fn create(&self, url: &str) -> Result<Arc<dyn TableProvider>>;
+    async fn create(
+        &self,
+        ctx: &SessionState,
+        url: &str,
+    ) -> Result<Arc<dyn TableProvider>>;
 }

datafusion/core/src/datasource/listing_table_factory.rs

@@ -0,0 +1,79 @@
+// 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.
+
+//! Factory for creating ListingTables with default options
+
+use crate::datasource::datasource::TableProviderFactory;
+use crate::datasource::file_format::avro::AvroFormat;
+use crate::datasource::file_format::csv::CsvFormat;
+use crate::datasource::file_format::file_type::{FileType, GetExt};
+use crate::datasource::file_format::json::JsonFormat;
+use crate::datasource::file_format::parquet::ParquetFormat;
+use crate::datasource::file_format::FileFormat;
+use crate::datasource::listing::{
+    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
+use crate::datasource::TableProvider;
+use crate::execution::context::SessionState;
+use async_trait::async_trait;
+use std::sync::Arc;
+
+/// A `TableProviderFactory` capable of creating new `ListingTable`s
+pub struct ListingTableFactory {
+    file_type: FileType,
+}
+
+impl ListingTableFactory {
+    /// Creates a new `ListingTableFactory`
+    pub fn new(file_type: FileType) -> Self {
+        Self { file_type }
+    }
+}
+
+#[async_trait]
+impl TableProviderFactory for ListingTableFactory {
+    async fn create(
+        &self,
+        state: &SessionState,
+        url: &str,
+    ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
+        let file_extension = self.file_type.get_ext();
+
+        let file_format: Arc<dyn FileFormat> = match self.file_type {
+            FileType::CSV => Arc::new(CsvFormat::default()),
+            FileType::PARQUET => Arc::new(ParquetFormat::default()),
+            FileType::AVRO => Arc::new(AvroFormat::default()),
+            FileType::JSON => Arc::new(JsonFormat::default()),
+        };
+
+        let options = ListingOptions {
+            format: file_format,
+            collect_stat: true,
+            file_extension: file_extension.to_owned(),
+            target_partitions: 1,
+            table_partition_cols: vec![],
+        };
+
+        let table_path = ListingTableUrl::parse(url)?;
+        let resolved_schema = options.infer_schema(state, &table_path).await?;
+        let config = ListingTableConfig::new(table_path)
+            .with_listing_options(options)
+            .with_schema(resolved_schema);
+        let table = ListingTable::try_new(config)?;
+        Ok(Arc::new(table))
+    }
+}

datafusion/core/src/datasource/mod.rs

@@ -23,6 +23,7 @@ pub mod default_table_source;
 pub mod empty;
 pub mod file_format;
 pub mod listing;
+pub mod listing_table_factory;
 pub mod memory;
 pub mod object_store;
 pub mod view;

datafusion/core/src/execution/context.rs

@@ -175,6 +175,26 @@ impl SessionContext {
         Self::with_config(SessionConfig::new())
     }
 
+    /// Finds any ListSchemaProviders and instructs them to reload tables from "disk"
+    pub async fn refresh_catalogs(&self) -> Result<()> {
+        let cat_names = self.catalog_names().clone();
+        for cat_name in cat_names.iter() {
+            let cat = self.catalog(cat_name.as_str()).ok_or_else(|| {
+                DataFusionError::Internal("Catalog not found!".to_string())
+            })?;
+            for schema_name in cat.schema_names() {
+                let schema = cat.schema(schema_name.as_str()).ok_or_else(|| {
+                    DataFusionError::Internal("Schema not found!".to_string())
+                })?;
+                let lister = schema.as_any().downcast_ref::<ListingSchemaProvider>();
+                if let Some(lister) = lister {
+                    lister.refresh(&self.state()).await?;
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Creates a new session context using the provided session configuration.
     pub fn with_config(config: SessionConfig) -> Self {
         let runtime = Arc::new(RuntimeEnv::default());
@@ -486,7 +506,7 @@ impl SessionContext {
                     cmd.file_type
                 ))
             })?;
-        let table = (*factory).create(cmd.location.as_str()).await?;
+        let table = (*factory).create(&state, cmd.location.as_str()).await?;
         self.register_table(cmd.name.as_str(), table)?;
         let plan = LogicalPlanBuilder::empty(false).build()?;
         Ok(Arc::new(DataFrame::new(self.state.clone(), &plan)))
@@ -1764,7 +1784,7 @@ impl ContextProvider for SessionState {
             Ok(schema) => {
                 let provider = schema.table(resolved_ref.table).ok_or_else(|| {
                     DataFusionError::Plan(format!(
-                        "'{}.{}.{}' not found",
+                        "table '{}.{}.{}' not found",
                         resolved_ref.catalog, resolved_ref.schema, resolved_ref.table
                     ))
                 })?;
@@ -2005,11 +2025,12 @@ mod tests {
     use super::*;
     use crate::assert_batches_eq;
     use crate::datasource::datasource::TableProviderFactory;
+    use crate::datasource::listing_table_factory::ListingTableFactory;
     use crate::execution::context::QueryPlanner;
     use crate::execution::runtime_env::RuntimeConfig;
     use crate::physical_plan::expressions::AvgAccumulator;
     use crate::test;
-    use crate::test_util::{parquet_test_data, TestTableFactory};
+    use crate::test_util::parquet_test_data;
     use crate::variable::VarType;
     use arrow::array::ArrayRef;
     use arrow::datatypes::*;
@@ -2267,30 +2288,33 @@ mod tests {
 
         let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
             HashMap::new();
-        table_factories.insert("test".to_string(), Arc::new(TestTableFactory {}));
+        let factory = Arc::new(ListingTableFactory::new(FileType::CSV));
+        table_factories.insert("test".to_string(), factory);
         let rt_cfg = RuntimeConfig::new().with_table_factories(table_factories);
         let runtime = Arc::new(RuntimeEnv::new(rt_cfg).unwrap());
         let cfg = SessionConfig::new()
             .set_str("datafusion.catalog.location", url.as_str())
             .set_str("datafusion.catalog.type", "test");
         let session_state = SessionState::with_config_rt(cfg, runtime);
         let ctx = SessionContext::with_state(session_state);
+        ctx.refresh_catalogs().await?;
 
-        let mut table_count = 0;
-        for cat_name in ctx.catalog_names().iter() {
-            let cat = ctx.catalog(cat_name).unwrap();
-            for s_name in cat.schema_names().iter() {
-                let schema = cat.schema(s_name).unwrap();
-                if let Some(listing) =
-                    schema.as_any().downcast_ref::<ListingSchemaProvider>()
-                {
-                    listing.refresh().await.unwrap();
-                    table_count = schema.table_names().len();
-                }
-            }
-        }
+        let result =
+            plan_and_collect(&ctx, "select c_name from default.customer limit 3;")
+                .await?;
+
+        let actual = arrow::util::pretty::pretty_format_batches(&result)
+            .unwrap()
+            .to_string();
+        let expected = r#"+--------------------+
+| c_name             |
++--------------------+
+| Customer#000000002 |
+| Customer#000000003 |
+| Customer#000000004 |
++--------------------+"#;
+        assert_eq!(actual, expected);
 
-        assert_eq!(table_count, 8);
         Ok(())
     }
 

datafusion/core/src/test_util.rs

@@ -284,6 +284,7 @@ pub struct TestTableFactory {}
 impl TableProviderFactory for TestTableFactory {
     async fn create(
         &self,
+        _state: &SessionState,
         url: &str,
     ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
         Ok(Arc::new(TestTableProvider {

datafusion/core/tests/sql/errors.rs

@@ -80,7 +80,7 @@ async fn query_cte_incorrect() -> Result<()> {
     assert!(plan.is_err());
     assert_eq!(
         format!("{}", plan.unwrap_err()),
-        "Error during planning: 'datafusion.public.t' not found"
+        "Error during planning: table 'datafusion.public.t' not found"
     );
 
     // forward referencing
@@ -89,7 +89,7 @@ async fn query_cte_incorrect() -> Result<()> {
     assert!(plan.is_err());
     assert_eq!(
         format!("{}", plan.unwrap_err()),
-        "Error during planning: 'datafusion.public.u' not found"
+        "Error during planning: table 'datafusion.public.u' not found"
     );
 
     // wrapping should hide u
@@ -98,7 +98,7 @@ async fn query_cte_incorrect() -> Result<()> {
     assert!(plan.is_err());
     assert_eq!(
         format!("{}", plan.unwrap_err()),
-        "Error during planning: 'datafusion.public.u' not found"
+        "Error during planning: table 'datafusion.public.u' not found"
     );
 
     Ok(())