From b93f6f8111171a88161fc62a069c5708cce6ebe6 Mon Sep 17 00:00:00 2001
From: Devin D'Angelo <devinjdangelo@gmail.com>
Date: Thu, 28 Dec 2023 10:53:28 -0500
Subject: [PATCH] pass through StatementOptions

---
 datafusion/common/src/file_options/mod.rs         | 15 +++++++++------
 .../core/src/datasource/physical_plan/mod.rs      |  4 +++-
 datafusion/proto/src/physical_plan/to_proto.rs    |  2 +-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/datafusion/common/src/file_options/mod.rs b/datafusion/common/src/file_options/mod.rs
index c1066b0b1e9e..4c045551f6b6 100644
--- 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
index 5583991355c6..312183588c45 100644
--- 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
index 731898499239..203f8380776a 100644
--- 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")
             }
         };