Support COPY TO Externally Defined File Formats, add FileType trait #11060

devinjdangelo · 2024-06-22T02:55:13Z

Which issue does this PR close?

Closes #8345
Revives #8667

Rationale for this change

Previously, COPY TO queries could only be planned/executed for a preset enum of FileTypes. CREATE EXTERNAL TABLE supports arbitrary file formats via the FileFormat and TableProvider traits. This PR removes the FileType enum in favor of a FileType trait and updates COPY TO plans to work with any externally implemented dyn FileFormat.

This PR also enables future work to remove csv, parquet, json, ect... formats into dedicated crates (datafusion-csv, ...).

What changes are included in this PR?

Deletes the FileType enum and replaces with a trait
Creates a wrapper struct for FileFormat that implements FileType to enable conversions between the two
Extends LogicalExtensionCodec to allow serde of dyn FileType

Are these changes tested?

Existing tests verify behavior for currently supported formats remains the same.

Are there any user-facing changes?

Yes, there are API changes relevant for developers using DataFusion.

…pache#10879) * adding new function contains * adding substrait test * adding doc * adding doc * Update docs/source/user-guide/sql/scalar_functions.md Co-authored-by: Alex Huang <[email protected]> * adding entry --------- Co-authored-by: Alex Huang <[email protected]>

alamb · 2024-06-22T11:54:59Z

This looks awesome -- thank you @devinjdangelo -- I plan to review this carefully, but probably not until tomorrow

devinjdangelo

cc @alamb @thinkharderdev I finally had some time to figure out how to make #8667 work.

There are a number of tasks to be done, preferably as follow on PRs given the size/complexity of the changes in this PR already:

end-to-end example of implementing a custom FileFormat and using COPY with it
Finishing LogicalExtensionCodec impls for csv, json, parquet, arrow, and avro FileFormat
Moving FileFormat impls and supporting code out of core and into dedicated crates
(maybe) Moving ListingTable to a dedicated crate

devinjdangelo · 2024-06-22T11:42:50Z

datafusion/common/src/config.rs

+/// These file types have special built in behavior for configuration.
+/// Use TableOptions::Extensions for configuring other file types.
+#[derive(Debug, Clone)]
+pub enum ConfigFileType {


This is the last vestige of the old FileType enum. It can be completely ignored if using a custom format.

🤔 perhaps we should add some Trait to unify the handling of options for built in formats and custom formats 🤔

We could potentially remove the TableOptions code and instead have each FileFormatFactory handle configuration. This is actually mostly the case in this PR already, but TableOptions is a common helper.

devinjdangelo · 2024-06-22T11:44:15Z

datafusion/common/src/file_options/file_type.rs

@@ -40,107 +37,10 @@ pub trait GetExt {
    fn get_ext(&self) -> String;
 }

-/// Readable file type


This is the primary change in the PR. Most other changes are just to recover the original behavior for these file types, but with the new trait.

devinjdangelo · 2024-06-22T11:49:26Z

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

+/// The former trait is a superset of the latter trait, which includes execution time
+/// relevant methods. [FileType] is only used in logical planning and only implements
+/// the subset of methods required during logical planning.
+pub struct DefaultFileFormat {


This design pattern is copied from how the DefaultTableSource allows a TableProvider to be used during logical planning as a TableSource.

In other words, FileType is to FileFormat as TableSource is to TableProvider.

I found the naming confusing -- I would have expected that DefaultFileFormat implemented FileFormatFactory or something

Maybe we could call this FileFormatSource ?

I renamed this to DefaultFileType.

devinjdangelo · 2024-06-22T11:54:36Z

datafusion/proto/src/logical_plan/file_formats.rs

+#[derive(Debug)]
+pub struct CsvLogicalExtensionCodec;
+
+// TODO! This is a placeholder for now and needs to be implemented for real.


Note this comment that I have extended LogicalExtensionCodec and created placeholder impls, but have not actually implemented logic to serde all built in FileFormat yet.

I will file follow on tickets

devinjdangelo · 2024-06-22T11:57:29Z

datafusion/sql/tests/common/mod.rs

 use datafusion_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
 use datafusion_sql::planner::ContextProvider;

+struct MockCsvType {}


This struct could be pulled out of this test as an example of LogicalPlanning a COPY query with a custom FileType

devinjdangelo · 2024-06-22T12:00:41Z

datafusion/expr/src/logical_plan/dml.rs

-    /// File format options.
-    pub format_options: FormatOptions,
+    /// File type trait
+    pub file_type: Arc<dyn FileType>,


If using DataFusion core for execution plans, this FileType can be converted to a FileFormatFactory which will produce a FileFormat at runtime.

What is the reason it can't be FileFormatFactory directly 🤔

This would introduce a dependency between datafusion-expr and datafusion-core. The FileType trait exists to avoid this dependency and so users only using datafusion for logical planning would not have to implement all of the physical execution related FileFormatmethods.

alamb

Thank you @devinjdangelo -- I had a good look at this PR, and I think it looks really nice

I did not quite get through the entire PR yet, but I wanted to provide some early feeback

I think it very much adheres to DataFusion's design goals of permiting extension at all points and makes it easier to support new file formats

One thing I didn't fully grok at first was whyn there was a need for FileFormatFactory -- like why not simply use FileFormat everywhere.

Then I realized that the reason was that the FileFormats have state, specifically default options and that the format itself then carries the specific options

pub struct JsonFormatFactory {
    /// default options
    options: Option<JsonOptions>,  
}

Perhaps we can document this rationale somewhere so others are not confused as well.

Suggestion: Add an example

I think something that would be awesome to add (as a follow on PR perhaps) is an example of creating a custom file format and read/write to it.

#11079

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

alamb · 2024-06-23T14:14:17Z

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

@@ -58,6 +74,15 @@ pub trait FileFormat: Send + Sync + fmt::Debug {
    /// downcast to a specific implementation.
    fn as_any(&self) -> &dyn Any;

+    /// Returns the extension for this FileFormat, e.g. "file.csv" -> csv
+    fn get_ext(&self) -> String;


Minor: Is there any reason this has to return an owned String? Could it return &str ?

I think &str would be too restrictive, but we could avoid requiring a String allocation with a cow instead.

let's keep it String for now then

alamb · 2024-06-23T14:14:32Z

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

+    fn get_ext_with_compression(
+        &self,
+        _file_compression_type: &FileCompressionType,
+    ) -> Result<String>;


Same question about owned String

alamb · 2024-06-23T14:16:31Z

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

+/// The former trait is a superset of the latter trait, which includes execution time
+/// relevant methods. [FileType] is only used in logical planning and only implements
+/// the subset of methods required during logical planning.
+pub struct DefaultFileFormat {


I found the naming confusing -- I would have expected that DefaultFileFormat implemented FileFormatFactory or something

Maybe we could call this FileFormatSource ?

alamb · 2024-06-23T14:18:31Z

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

+    }
+}
+
+impl FileType for DefaultFileFormat {


Rather than this wrapper, I wonder if we could implement FileType for dyn FileFormatFactory

impl FileType for &dyn FileFormatFactory{ ... }

And avoid having to have this struct 🤔

I tried briefly, but ran into various errors. I am not sure if it is possible with the traits defined in different crates to make this work without a wrapper. I think that is why DefaultTableSource exists as well.

alamb · 2024-06-23T14:21:37Z

datafusion/core/src/execution/session_state.rs

@@ -237,6 +248,33 @@ impl SessionState {
            function_factory: None,
        };

+        #[cfg(feature = "parquet")]


datafusion/core/src/execution/session_state.rs

alamb · 2024-06-23T14:24:47Z

datafusion/core/src/physical_planner.rs

@@ -786,32 +779,9 @@ impl DefaultPhysicalPlanner {
                    table_partition_cols,
                    overwrite: false,
                };
-                let mut table_options = session_state.default_table_options();


this is great this code has been removed from the default planner

alamb

Thank you so much @devinjdangelo -- I made it through this PR and I think it looks really nice

I do think it is worth considering how to simplify the traits (possibly unifying the configuration) and there look like there are a few todos, but otherwise this PR looks really nice to me

cc @ozankabak @berkaysynnada @metesynnada

alamb · 2024-06-23T15:31:34Z

datafusion/common/src/config.rs

+/// These file types have special built in behavior for configuration.
+/// Use TableOptions::Extensions for configuring other file types.
+#[derive(Debug, Clone)]
+pub enum ConfigFileType {


🤔 perhaps we should add some Trait to unify the handling of options for built in formats and custom formats 🤔

alamb · 2024-06-23T15:32:20Z

datafusion/common/src/config.rs

-        match value {
-            FileType::ARROW => FormatOptions::ARROW,
-            FileType::AVRO => FormatOptions::AVRO,
-            #[cfg(feature = "parquet")]


it is certainly nice to avoid many of these #[cfg(feature = "parquet")] that were sprinkled throughout the code

alamb · 2024-06-23T15:33:45Z

datafusion/common/src/config.rs

+    /// # Parameters
+    ///
+    /// * `format`: The file format to use (e.g., CSV, Parquet).
+    pub fn set_file_format(&mut self, format: ConfigFileType) {


the parameter is called "format" but the struct is ConfigFileType

Maybe we should call this set_config_type? Or maybe it could just pass in a `Arc~?

I changed the method name for now.

It might make sense for TableOptions to be broken up and its functionality moved into each implementation of FileFormatFactory. I have avoided any reference to FileFormat in datafusion-common, since that would add an indirect dependency on datafusion-core to datafusion-sql and other crates.

alamb · 2024-06-23T15:35:07Z

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

    use bytes::Bytes;
    use futures::StreamExt;

-    #[test]
-    fn get_ext_with_compression() {


Should we port these tests instead of just removing them?

Yes, I moved this functionality to a trait method on each FileFormatFactory. We could add a test for each format which is similar to this test.

alamb · 2024-06-23T15:36:35Z

datafusion/expr/src/logical_plan/dml.rs

-    /// File format options.
-    pub format_options: FormatOptions,
+    /// File type trait
+    pub file_type: Arc<dyn FileType>,


What is the reason it can't be FileFormatFactory directly 🤔

alamb · 2024-06-23T15:36:59Z

datafusion/proto/src/logical_plan/file_formats.rs

+#[derive(Debug)]
+pub struct JsonLogicalExtensionCodec;
+
+// TODO! This is a placeholder for now and needs to be implemented for real.


do we have to complete the TODO?

Users depending on the ability to serialize COPY plans (e.g. ballista) will need this TODO to be completed before upgrading to any version of datafusion including this change.

I think it would be OK to cut an follow up ticket for this.

Co-authored-by: Andrew Lamb <[email protected]>

devinjdangelo

Thank you @alamb for the thorough review as always! I have made a few small changes.

devinjdangelo · 2024-06-24T01:56:11Z

datafusion/common/src/config.rs

+    /// # Parameters
+    ///
+    /// * `format`: The file format to use (e.g., CSV, Parquet).
+    pub fn set_file_format(&mut self, format: ConfigFileType) {


I changed the method name for now.

It might make sense for TableOptions to be broken up and its functionality moved into each implementation of FileFormatFactory. I have avoided any reference to FileFormat in datafusion-common, since that would add an indirect dependency on datafusion-core to datafusion-sql and other crates.

devinjdangelo · 2024-06-24T02:02:18Z

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

    use bytes::Bytes;
    use futures::StreamExt;

-    #[test]
-    fn get_ext_with_compression() {


Yes, I moved this functionality to a trait method on each FileFormatFactory. We could add a test for each format which is similar to this test.

devinjdangelo · 2024-06-24T02:17:29Z

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

@@ -58,6 +74,15 @@ pub trait FileFormat: Send + Sync + fmt::Debug {
    /// downcast to a specific implementation.
    fn as_any(&self) -> &dyn Any;

+    /// Returns the extension for this FileFormat, e.g. "file.csv" -> csv
+    fn get_ext(&self) -> String;


I think &str would be too restrictive, but we could avoid requiring a String allocation with a cow instead.

devinjdangelo · 2024-06-24T02:25:46Z

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

+/// The former trait is a superset of the latter trait, which includes execution time
+/// relevant methods. [FileType] is only used in logical planning and only implements
+/// the subset of methods required during logical planning.
+pub struct DefaultFileFormat {


I renamed this to DefaultFileType.

devinjdangelo · 2024-06-24T02:26:55Z

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

+    }
+}
+
+impl FileType for DefaultFileFormat {


I tried briefly, but ran into various errors. I am not sure if it is possible with the traits defined in different crates to make this work without a wrapper. I think that is why DefaultTableSource exists as well.

devinjdangelo · 2024-06-24T02:31:13Z

datafusion/expr/src/logical_plan/dml.rs

-    /// File format options.
-    pub format_options: FormatOptions,
+    /// File type trait
+    pub file_type: Arc<dyn FileType>,


This would introduce a dependency between datafusion-expr and datafusion-core. The FileType trait exists to avoid this dependency and so users only using datafusion for logical planning would not have to implement all of the physical execution related FileFormatmethods.

devinjdangelo · 2024-06-24T02:33:06Z

datafusion/proto/src/logical_plan/file_formats.rs

+#[derive(Debug)]
+pub struct JsonLogicalExtensionCodec;
+
+// TODO! This is a placeholder for now and needs to be implemented for real.


Users depending on the ability to serialize COPY plans (e.g. ballista) will need this TODO to be completed before upgrading to any version of datafusion including this change.

I think it would be OK to cut an follow up ticket for this.

devinjdangelo · 2024-06-24T02:39:13Z

datafusion/common/src/config.rs

+/// These file types have special built in behavior for configuration.
+/// Use TableOptions::Extensions for configuring other file types.
+#[derive(Debug, Clone)]
+pub enum ConfigFileType {


We could potentially remove the TableOptions code and instead have each FileFormatFactory handle configuration. This is actually mostly the case in this PR already, but TableOptions is a common helper.

berkaysynnada · 2024-06-24T11:41:00Z

Thank you, @devinjdangelo, for proposing this extensibility feature. I would like to clarify a few points to ensure I understand correctly:

We cannot merge FileType and FileFormatFactory due to dependency issues.
We cannot merge FileFormatFactory and FileFormat because of stateful format factories.

If these points are correct, I am considering whether we could move these states (extension options, like CsvOptions) into FileType somehow to remove one of these layers. This is just a thought experiment; I don't want to block the PR, which I believe improves extensibility.

devinjdangelo · 2024-06-24T12:44:17Z

Thank you for taking a look @berkaysynnada

We cannot merge FileType and FileFormatFactory due to dependency issues.

That is correct. Logical planning cannot depend on FileFormatFactory or FileFormat to avoid a dependency on datafusion's execution engine during logical planning.

We cannot merge FileFormatFactory and FileFormat because of stateful format factories.

Yes, SessionState stores a HashMap<String, Arc<dyn FileFormatFactory>>. Each FileFormatFactory could be shared among multiple concurrent queries and must be able to produce many, differently configured, Arc<dyn FileFormat>s. If SessionState stored only a Arc<dyn FileFormat> it would have to have a &mut self method enabled by a Mutex<> and even with that it would possibly lead to unexpected behavior if multiple queries modified the format during execution.

alamb · 2024-06-24T22:02:35Z

Yes, SessionState stores a HashMap<String, Arc>. Each FileFormatFactory could be shared among multiple concurrent queries and must be able to produce many, differently configured,

I wonder if it wuld be possible to store the file format state outside the FileFormat and pass in the format options as part of each query somehow 🤔

alamb

Lets see if we can find some way to avoid the extra level of stateful factories. If not then I think this PR could still be merged

Before we merge I'll file follow on tickets for supporting protobuf encoding

devinjdangelo · 2024-06-25T11:25:27Z

I wonder if it wuld be possible to store the file format state outside the FileFormat and pass in the format options as part of each query somehow 🤔

I think this is possible, but it would also make implementing FileFormat a bit inconvenient. We would have to pass in and parse format_options on every trait method. E.g. something like:

pub trait FileFormat{
        async fn create_physical_plan(
        &self,
        state: &SessionState,
        conf: FileScanConfig,
        filters: Option<&Arc<dyn PhysicalExpr>>,
        format_options: &HashMap<String, String>,
    ) -> Result<Arc<dyn ExecutionPlan>>;
    ...
}

And each method would have to parse the format_options. FileFormatFactory avoids this by allowing FileFormat to be created during physical planning when query specific options are known and can be parsed once/stored in the FileFormat struct.

alamb

I think we should proceed with this PR, even though some of the traits are less than ideal (due to having to handle state)

Also before merging we should file the follow on tickets you suggest in #11060 (review)

Let's leave it open for another day or two to gather additional comments

Thanks again for pushing this part of the code forward

alamb · 2024-06-27T22:21:57Z

I merged up to resolve a conflict

alamb

Filed #11150 to track moving the extension codes

I plan to merge this PR when CI is complete

alamb · 2024-06-28T00:20:52Z

Forward ! Thank you @devinjdangelo and @Weijun-H

…pache#11060) * wip create and register ext file types with session * Add contains function, and support in datafusion substrait consumer (apache#10879) * adding new function contains * adding substrait test * adding doc * adding doc * Update docs/source/user-guide/sql/scalar_functions.md Co-authored-by: Alex Huang <[email protected]> * adding entry --------- Co-authored-by: Alex Huang <[email protected]> * logical planning updated * compiling * removing filetype enum * compiling * working on tests * fix some tests * test fixes * cli fix * cli fmt * Update datafusion/core/src/datasource/file_format/mod.rs Co-authored-by: Andrew Lamb <[email protected]> * Update datafusion/core/src/execution/session_state.rs Co-authored-by: Andrew Lamb <[email protected]> * review comments * review comments * review comments * typo fix * fmt * fix err log style * fmt --------- Co-authored-by: Lordworms <[email protected]> Co-authored-by: Alex Huang <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>

devinjdangelo and others added 12 commits June 15, 2024 13:53

wip create and register ext file types with session

a67d73b

logical planning updated

05d329b

compiling

027131b

removing filetype enum

54f5f8f

compiling

0274eec

working on tests

4319881

fix some tests

487da13

test fixes

9527f82

cli fix

891afc3

cli fmt

128bf2b

Merge branch 'main' into external_file_types

e6d2fa3

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions core Core DataFusion crate labels Jun 22, 2024

alamb mentioned this pull request Jun 22, 2024

DataFusion weekly project plan (Andrew Lamb) - June 17, 2024 #10955

Closed

5 tasks

devinjdangelo commented Jun 22, 2024

View reviewed changes

alamb added the api change Changes the API exposed to users of the crate label Jun 23, 2024

alamb changed the title ~~Support COPY TO Externally Defined File Formats~~ Support COPY TO Externally Defined File Formats, add FileType trait Jun 23, 2024

alamb mentioned this pull request Jun 23, 2024

Add example for writing a FileFormat #11079

Closed

alamb reviewed Jun 23, 2024

View reviewed changes

devinjdangelo and others added 7 commits June 23, 2024 21:49

Update datafusion/core/src/datasource/file_format/mod.rs

b065cd7

Co-authored-by: Andrew Lamb <[email protected]>

Update datafusion/core/src/execution/session_state.rs

485f754

Co-authored-by: Andrew Lamb <[email protected]>

review comments

da8e635

review comments

6dec6ef

review comments

013ef91

typo fix

132af14

fmt

493e82f

devinjdangelo commented Jun 24, 2024

View reviewed changes

fix err log style

bcc3119

alamb reviewed Jun 24, 2024

View reviewed changes

devinjdangelo added 2 commits June 25, 2024 07:17

merge main

77bcc54

fmt

11071ac

alamb approved these changes Jun 26, 2024

View reviewed changes

Merge remote-tracking branch 'apache/main' into external_file_types

826721f

alamb mentioned this pull request Jun 27, 2024

Implement physical plan serialization for COPY plans CsvLogicalExtensionCodec #11150

Closed

alamb approved these changes Jun 27, 2024

View reviewed changes

alamb merged commit 5501e8e into apache:main Jun 28, 2024
23 checks passed

alamb mentioned this pull request Jul 1, 2024

DataFusion weekly project plan (Andrew Lamb) - July 1, 2024 #11190

Closed

10 tasks

Support COPY TO Externally Defined File Formats, add FileType trait #11060

Support COPY TO Externally Defined File Formats, add FileType trait #11060

Conversation

devinjdangelo commented Jun 22, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb commented Jun 22, 2024

devinjdangelo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Suggestion: Add an example

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

devinjdangelo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

berkaysynnada commented Jun 24, 2024

devinjdangelo commented Jun 24, 2024 • edited Loading

alamb commented Jun 24, 2024

alamb left a comment • edited Loading

Choose a reason for hiding this comment

devinjdangelo commented Jun 25, 2024

alamb left a comment

Choose a reason for hiding this comment

alamb commented Jun 27, 2024

alamb left a comment

Choose a reason for hiding this comment

alamb commented Jun 28, 2024

devinjdangelo commented Jun 22, 2024 •

edited

Loading

devinjdangelo commented Jun 24, 2024 •

edited

Loading

alamb left a comment •

edited

Loading