pass through StatementOptions

apache · Dec 28, 2023 · b93f6f8 · b93f6f8
1 parent 9dce12f
commit b93f6f8
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/datafusion/common/src/file_options/mod.rs b/datafusion/common/src/file_options/mod.rs
@@ -154,9 +154,10 @@ pub enum FileTypeWriterOptions {
     JSON(JsonWriterOptions),
     Avro(AvroWriterOptions),
     Arrow(ArrowWriterOptions),
-    /// For extension [FileType]s, FileTypeWriterOptions ignores all
-    /// passed options and returns an empty variant.
-    Extension,
+    /// 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 {
@@ -187,7 +188,9 @@ impl FileTypeWriterOptions {
             FileType::ARROW => {
                 FileTypeWriterOptions::Arrow(ArrowWriterOptions::try_from(options)?)
             }
-            FileType::Extension(_) => FileTypeWriterOptions::Extension,
+            FileType::Extension(_) => {
+                FileTypeWriterOptions::Extension(Some(statement_options.clone()))
+            }
         };
 
         Ok(file_type_write_options)
@@ -218,7 +221,7 @@ impl FileTypeWriterOptions {
             FileType::ARROW => {
                 FileTypeWriterOptions::Arrow(ArrowWriterOptions::try_from(options)?)
             }
-            FileType::Extension(_) => FileTypeWriterOptions::Extension,
+            FileType::Extension(_) => FileTypeWriterOptions::Extension(None),
         };
 
         Ok(file_type_write_options)
@@ -295,7 +298,7 @@ impl Display for FileTypeWriterOptions {
             FileTypeWriterOptions::JSON(_) => "JsonWriterOptions",
             #[cfg(feature = "parquet")]
             FileTypeWriterOptions::Parquet(_) => "ParquetWriterOptions",
-            FileTypeWriterOptions::Extension => "ExensionWriterOptions",
+            FileTypeWriterOptions::Extension(_) => "ExensionWriterOptions",
         };
         write!(f, "{}", name)
     }

diff --git a/datafusion/core/src/datasource/physical_plan/mod.rs b/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,
 }
 

diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -891,7 +891,7 @@ impl TryFrom<&FileTypeWriterOptions> for protobuf::FileTypeWriterOptions {
             FileTypeWriterOptions::Arrow(ArrowWriterOptions {}) => {
                 return not_impl_err!("Arrow file sink protobuf serialization")
             }
-            FileTypeWriterOptions::Extension => {
+            FileTypeWriterOptions::Extension(_) => {
                 return not_impl_err!("Extension file sink protobuf serialization")
             }
         };