Support Reading and Writing Extension FileTypes

.github/workflows/rust.yml

@@ -98,7 +98,7 @@ jobs:
         with:
           rust-version: stable
       - name: Run tests (excluding doctests)
-        run: cargo test --lib --tests --bins --features avro,json,backtrace
+        run: RUST_MIN_STACK=504857600 cargo test --lib --tests --bins --features avro,json,backtrace
       - name: Verify Working Directory Clean
         run: git diff --exit-code
 

datafusion/common/src/file_options/file_type.rs

@@ -18,9 +18,11 @@
 //! File type abstraction
 
 use crate::error::{DataFusionError, Result};
+use crate::parsers::CompressionTypeVariant;
 
 use core::fmt;
 use std::fmt::Display;
+use std::hash::Hasher;
 use std::str::FromStr;
 
 /// The default file extension of arrow files
@@ -41,7 +43,8 @@ pub trait GetExt {
 }
 
 /// Readable file type
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[allow(clippy::derived_hash_with_manual_eq)]
+#[derive(Debug, Clone, Hash)]
 pub enum FileType {
     /// Apache Arrow file
     ARROW,
@@ -54,8 +57,65 @@ pub enum FileType {
     CSV,
     /// JSON file
     JSON,
+    /// FileType Implemented Outside of DataFusion
+    Extension(Box<dyn ExtensionFileType>),
 }
 
+/// A trait to enable externally implementing the functionality of a [FileType].
+pub trait ExtensionFileType:
+    std::fmt::Debug + ExtensionFileTypeClone + Send + Sync
+{
+    /// Returns the default file extension for this type, e.g. CSV would return ".csv".to_owned()
+    /// The default_extension is also used to uniquely identify a specific FileType::Extension variant,
+    /// so ensure this String is unique from any built in FileType and any other ExtensionFileTypes
+    /// defined.
+    fn default_extension(&self) -> String;
+
+    /// Returns the file extension when it is compressed with a given [CompressionTypeVariant]
+    fn extension_with_compression(
+        &self,
+        compression: CompressionTypeVariant,
+    ) -> Result<String>;
+}
+
+pub trait ExtensionFileTypeClone {
+    fn clone_box(&self) -> Box<dyn ExtensionFileType>;
+}
+
+impl Clone for Box<dyn ExtensionFileType> {
+    fn clone(&self) -> Box<dyn ExtensionFileType> {
+        self.clone_box()
+    }
+}
+
+impl std::hash::Hash for Box<dyn ExtensionFileType> {
+    fn hash<H>(&self, state: &mut H)
+    where
+        H: Hasher,
+    {
+        self.default_extension().hash(state)
+    }
+}
+
+impl PartialEq for FileType {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (FileType::ARROW, FileType::ARROW) => true,
+            (FileType::AVRO, FileType::AVRO) => true,
+            #[cfg(feature = "parquet")]
+            (FileType::PARQUET, FileType::PARQUET) => true,
+            (FileType::CSV, FileType::CSV) => true,
+            (FileType::JSON, FileType::JSON) => true,
+            (FileType::Extension(ext_self), FileType::Extension(ext_other)) => {
+                ext_self.default_extension() == ext_other.default_extension()
+            }
+            _ => false,
+        }
+    }
+}
+
+impl Eq for FileType {}
+
 impl GetExt for FileType {
     fn get_ext(&self) -> String {
         match self {
@@ -65,6 +125,7 @@ impl GetExt for FileType {
             FileType::PARQUET => DEFAULT_PARQUET_EXTENSION.to_owned(),
             FileType::CSV => DEFAULT_CSV_EXTENSION.to_owned(),
             FileType::JSON => DEFAULT_JSON_EXTENSION.to_owned(),
+            FileType::Extension(ext) => ext.default_extension(),
         }
     }
 }
@@ -78,6 +139,7 @@ impl Display for FileType {
             FileType::PARQUET => "parquet",
             FileType::AVRO => "avro",
             FileType::ARROW => "arrow",
+            FileType::Extension(ext) => return ext.fmt(f),
         };
         write!(f, "{}", out)
     }

datafusion/common/src/file_options/mod.rs

@@ -154,6 +154,10 @@ pub enum FileTypeWriterOptions {
     JSON(JsonWriterOptions),
     Avro(AvroWriterOptions),
     Arrow(ArrowWriterOptions),
+    /// For extension [FileType]s, FileTypeWriterOptions simply stores
+    /// any passed StatementOptions to be handled later by any custom
+    /// physical plans (Such as a FileFormat::create_writer_physical_plan)
+    Extension(Option<StatementOptions>),
 }
 
 impl FileTypeWriterOptions {
@@ -184,14 +188,17 @@ impl FileTypeWriterOptions {
             FileType::ARROW => {
                 FileTypeWriterOptions::Arrow(ArrowWriterOptions::try_from(options)?)
             }
+            FileType::Extension(_) => {
+                FileTypeWriterOptions::Extension(Some(statement_options.clone()))
+            }
         };
 
         Ok(file_type_write_options)
     }
 
     /// Constructs a FileTypeWriterOptions from session defaults only.
     pub fn build_default(
-        file_type: &FileType,
+        file_type: &mut FileType,
         config_defaults: &ConfigOptions,
     ) -> Result<Self> {
         let empty_statement = StatementOptions::new(vec![]);
@@ -214,6 +221,7 @@ impl FileTypeWriterOptions {
             FileType::ARROW => {
                 FileTypeWriterOptions::Arrow(ArrowWriterOptions::try_from(options)?)
             }
+            FileType::Extension(_) => FileTypeWriterOptions::Extension(None),
         };
 
         Ok(file_type_write_options)
@@ -290,6 +298,7 @@ impl Display for FileTypeWriterOptions {
             FileTypeWriterOptions::JSON(_) => "JsonWriterOptions",
             #[cfg(feature = "parquet")]
             FileTypeWriterOptions::Parquet(_) => "ParquetWriterOptions",
+            FileTypeWriterOptions::Extension(_) => "ExensionWriterOptions",
         };
         write!(f, "{}", name)
     }

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

@@ -250,6 +250,7 @@ impl FileTypeExt for FileType {
                     "FileCompressionType can be specified for CSV/JSON FileType.".into(),
                 )),
             },
+            FileType::Extension(ext) => ext.extension_with_compression(c.variant),
         }
     }
 }

datafusion/core/src/datasource/listing/table.rs

@@ -150,6 +150,9 @@ impl ListingTableConfig {
             ),
             #[cfg(feature = "parquet")]
             FileType::PARQUET => Arc::new(ParquetFormat::default()),
+            FileType::Extension(_) => {
+                unreachable!("FileType::from_str cannot return Extension variant!")
+            }
         };
 
         Ok((file_format, ext))
@@ -768,7 +771,7 @@ impl TableProvider for ListingTable {
         let file_type_writer_options = match &self.options().file_type_write_options {
             Some(opt) => opt.clone(),
             None => FileTypeWriterOptions::build_default(
-                &file_format.file_type(),
+                &mut file_format.file_type(),
                 state.config_options(),
             )?,
         };
@@ -1716,6 +1719,9 @@ mod tests {
                     )
                     .await?;
             }
+            FileType::Extension(_) => {
+                panic!("Extension file type not implemented in write path.")
+            }
         }
 
         // Create and register the source table with the provided schema and inserted data

datafusion/core/src/datasource/listing_table_factory.rs

@@ -88,6 +88,9 @@ impl TableProviderFactory for ListingTableFactory {
                 JsonFormat::default().with_file_compression_type(file_compression_type),
             ),
             FileType::ARROW => Arc::new(ArrowFormat),
+            FileType::Extension(_) => {
+                unreachable!("FileType::from_str cannot return Extension variant!")
+            }
         };
 
         let (provided_schema, table_partition_cols) = if cmd.schema.fields().is_empty() {
@@ -181,6 +184,9 @@ impl TableProviderFactory for ListingTableFactory {
             FileType::PARQUET => file_type_writer_options,
             FileType::ARROW => file_type_writer_options,
             FileType::AVRO => file_type_writer_options,
+            FileType::Extension(_) => {
+                unreachable!("FileType::from_str cannot return Extension variant!")
+            }
         };
 
         let table_path = ListingTableUrl::parse(&cmd.location)?;

datafusion/core/src/datasource/physical_plan/mod.rs

@@ -95,7 +95,9 @@ pub struct FileSinkConfig {
     pub single_file_output: bool,
     /// Controls whether existing data should be overwritten by this sink
     pub overwrite: bool,
-    /// Contains settings specific to writing a given FileType, e.g. parquet max_row_group_size
+    /// Contains settings specific to writing a given FileType, e.g. parquet max_row_group_size.
+    /// Note that for externally defined FileTypes, FileTypeWriterOptions contains arbitrary
+    /// config tuples that must be handled within the physical plan.
     pub file_type_writer_options: FileTypeWriterOptions,
 }
 

datafusion/core/src/physical_planner.rs

@@ -605,6 +605,7 @@ impl DefaultPhysicalPlanner {
                         FileType::JSON => Arc::new(JsonFormat::default()),
                         FileType::AVRO => Arc::new(AvroFormat {} ),
                         FileType::ARROW => Arc::new(ArrowFormat {}),
+                        FileType::Extension(_ext) => return not_impl_err!("Extension FileTypes not supported in Copy To statements.")
                     };
 
                     sink_format.create_writer_physical_plan(input_exec, session_state, config, None).await

datafusion/proto/src/physical_plan/to_proto.rs

@@ -891,6 +891,9 @@ impl TryFrom<&FileTypeWriterOptions> for protobuf::FileTypeWriterOptions {
             FileTypeWriterOptions::Arrow(ArrowWriterOptions {}) => {
                 return not_impl_err!("Arrow file sink protobuf serialization")
             }
+            FileTypeWriterOptions::Extension(_) => {
+                return not_impl_err!("Extension file sink protobuf serialization")
+            }
         };
         Ok(Self {
             file_type: Some(file_type),