refactor how we create listing tables (#4227)

* refactor how we create listing tables

* run linting

* attempt to register default table factories

* pass csv as catalog type

* enure that listingschemaprovider can keep working for csv with header

* no need to registery default table factories

* remove unused imports

* use builder for listingoptions

* remove comment

* fix roundtrip test

* Update datafusion/core/src/datasource/listing_table_factory.rs

Co-authored-by: Andy Grove <[email protected]>

* register a listingtablefactory for NDJSON

* always verify that a factory is available

* allow NDJSON as an alias for JSON

Co-authored-by: Andy Grove <[email protected]>

Author: timvw

datafusion-cli/src/main.rs

@@ -16,9 +16,6 @@
 // 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;
@@ -29,7 +26,6 @@ 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;
@@ -147,31 +143,11 @@ 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))
-        .with_table_factories(table_factories);
+        .with_object_store_registry(Arc::new(object_store_registry));
     RuntimeEnv::new(rn_config)
 }
 

datafusion/core/src/catalog/listing_schema.rs

@@ -49,6 +49,8 @@ pub struct ListingSchemaProvider {
     factory: Arc<dyn TableProviderFactory>,
     store: Arc<dyn ObjectStore>,
     tables: Arc<Mutex<HashMap<String, Arc<dyn TableProvider>>>>,
+    format: String,
+    has_header: bool,
 }
 
 impl ListingSchemaProvider {
@@ -59,18 +61,24 @@ impl ListingSchemaProvider {
     /// `path`: The root path that contains subfolders which represent tables
     /// `factory`: The `TableProviderFactory` to use to instantiate tables for each subfolder
     /// `store`: The `ObjectStore` containing the table data
+    /// `format`: The `FileFormat` of the tables
+    /// `has_header`: Indicates whether the created external table has the has_header flag enabled
     pub fn new(
         authority: String,
         path: object_store::path::Path,
         factory: Arc<dyn TableProviderFactory>,
         store: Arc<dyn ObjectStore>,
+        format: String,
+        has_header: bool,
     ) -> Self {
         Self {
             authority,
             path,
             factory,
             store,
             tables: Arc::new(Mutex::new(HashMap::new())),
+            format,
+            has_header,
         }
     }
 
@@ -118,8 +126,8 @@ impl ListingSchemaProvider {
                             schema: Arc::new(DFSchema::empty()),
                             name: table_name.to_string(),
                             location: table_url,
-                            file_type: "".to_string(),
-                            has_header: false,
+                            file_type: self.format.clone(),
+                            has_header: self.has_header,
                             delimiter: ',',
                             table_partition_cols: vec![],
                             if_not_exists: false,

datafusion/core/src/datasource/file_format/file_type.rs

@@ -158,7 +158,7 @@ impl FromStr for FileType {
             "AVRO" => Ok(FileType::AVRO),
             "PARQUET" => Ok(FileType::PARQUET),
             "CSV" => Ok(FileType::CSV),
-            "JSON" => Ok(FileType::JSON),
+            "JSON" | "NDJSON" => Ok(FileType::JSON),
             _ => Err(DataFusionError::NotImplemented(format!(
                 "Unknown FileType: {}",
                 s

datafusion/core/src/datasource/listing_table_factory.rs

@@ -20,7 +20,7 @@
 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::file_type::{FileCompressionType, FileType};
 use crate::datasource::file_format::json::JsonFormat;
 use crate::datasource::file_format::parquet::ParquetFormat;
 use crate::datasource::file_format::FileFormat;
@@ -30,18 +30,24 @@ use crate::datasource::listing::{
 use crate::datasource::TableProvider;
 use crate::execution::context::SessionState;
 use async_trait::async_trait;
+use datafusion_common::DataFusionError;
 use datafusion_expr::CreateExternalTable;
+use std::str::FromStr;
 use std::sync::Arc;
 
 /// A `TableProviderFactory` capable of creating new `ListingTable`s
-pub struct ListingTableFactory {
-    file_type: FileType,
-}
+pub struct ListingTableFactory {}
 
 impl ListingTableFactory {
     /// Creates a new `ListingTableFactory`
-    pub fn new(file_type: FileType) -> Self {
-        Self { file_type }
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl Default for ListingTableFactory {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
@@ -52,24 +58,59 @@ impl TableProviderFactory for ListingTableFactory {
         state: &SessionState,
         cmd: &CreateExternalTable,
     ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
-        let file_extension = self.file_type.get_ext();
+        let file_compression_type = FileCompressionType::from_str(
+            cmd.file_compression_type.as_str(),
+        )
+        .map_err(|_| {
+            DataFusionError::Execution(format!(
+                "Unknown FileCompressionType {}",
+                cmd.file_compression_type.as_str()
+            ))
+        })?;
+        let file_type = FileType::from_str(cmd.file_type.as_str()).map_err(|_| {
+            DataFusionError::Execution(format!("Unknown FileType {}", cmd.file_type))
+        })?;
 
-        let file_format: Arc<dyn FileFormat> = match self.file_type {
-            FileType::CSV => Arc::new(CsvFormat::default()),
+        let file_extension =
+            file_type.get_ext_with_compression(file_compression_type.to_owned())?;
+
+        let file_format: Arc<dyn FileFormat> = match file_type {
+            FileType::CSV => Arc::new(
+                CsvFormat::default()
+                    .with_has_header(cmd.has_header)
+                    .with_delimiter(cmd.delimiter as u8)
+                    .with_file_compression_type(file_compression_type),
+            ),
             FileType::PARQUET => Arc::new(ParquetFormat::default()),
             FileType::AVRO => Arc::new(AvroFormat::default()),
-            FileType::JSON => Arc::new(JsonFormat::default()),
+            FileType::JSON => Arc::new(
+                JsonFormat::default().with_file_compression_type(file_compression_type),
+            ),
+        };
+
+        let provided_schema = if cmd.schema.fields().is_empty() {
+            None
+        } else {
+            Some(Arc::new(cmd.schema.as_ref().to_owned().into()))
         };
 
-        let options =
-            ListingOptions::new(file_format).with_file_extension(file_extension);
+        let options = ListingOptions::new(file_format)
+            .with_collect_stat(state.config.collect_statistics)
+            .with_file_extension(file_extension)
+            .with_target_partitions(state.config.target_partitions)
+            .with_table_partition_cols(cmd.table_partition_cols.clone())
+            .with_file_sort_order(None);
 
         let table_path = ListingTableUrl::parse(&cmd.location)?;
-        let resolved_schema = options.infer_schema(state, &table_path).await?;
+        let resolved_schema = match provided_schema {
+            None => options.infer_schema(state, &table_path).await?,
+            Some(s) => s,
+        };
         let config = ListingTableConfig::new(table_path)
             .with_listing_options(options)
             .with_schema(resolved_schema);
-        let table = ListingTable::try_new(config)?;
+        let table =
+            ListingTable::try_new(config)?.with_definition(cmd.definition.clone());
         Ok(Arc::new(table))
     }
 }