Support Reading and Writing Extension FileTypes #8667

devinjdangelo · 2023-12-28T15:30:42Z

Which issue does this PR close?

Closes #8657
Closes #8637
Closes #8345

Rationale for this change

DataFusion is built to be as extensible as possible. We currently have a FileType enum which prevents external crates adding support for reading/writing new FileFormats in a ListingTable.

What changes are included in this PR?

Creates a new FileType::Exension(Box<dyn ExtensionFileType>) variant.

Are these changes tested?

Partially by existing tests, additional likely needed

Are there any user-facing changes?

Ability to extend DataFusion to support reading/writing additional file formats more easily

devinjdangelo

I have left built in FileTypes as concrete enum variants, while supporting externally defined FileTypes via a boxed trait object contained in a FileType variant. I think this is the path of least resistance towards supporting externally defined FileFormats along the read and write path in a ListingTable.

We could potentially also refactor the built in types (parquet, csv, ect) to be externally defined structs which implement ExtensionFileType and remove the FileType enum entirely. This would entail a more significant refactor of FileTypeWriterOptions as well.

devinjdangelo · 2023-12-28T16:01:06Z

datafusion/common/src/file_options/file_type.rs

@@ -41,7 +43,8 @@ pub trait GetExt {
 }

 /// Readable file type
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[allow(clippy::derived_hash_with_manual_eq)]


I think the manual and auto traits should be compatible, so silencing this is OK.

You could probably also manually implement Hash pretty easily as well, but I agree it isn't a bad thing

devinjdangelo · 2023-12-28T16:01:53Z

datafusion/common/src/file_options/mod.rs

+    /// 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>),


FileTypeWriterOptions does not handle any parsing or validation for externally defined FileTypes. It simply acts as a container for the options passed in.

devinjdangelo · 2023-12-28T16:03:35Z

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(_) => {


We should be able to support serializing this since it is ultimately just a Vec<(String, String)>

We could also punt it into the FileType::Extension (e.g. add a serialize(&self) -> Result<Vec<u8>> type method)

devinjdangelo · 2023-12-28T16:04:46Z

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.")


The built in Copy plans can't support external types (since we can't initialize them in DataFusion), but it should be not too difficult to build a custom COPY PhysicalPlan that uses a custom FileFormat, but otherwise is the same and relies on the build in Copy logical plan.

Maybe we could have the FileTypeExtension have a method that returned Arc<dyn FileFormat> ?

I tried this (as well as combining ExtensionFileType trait into FileFormat), but that would require adding lots of dependencies between the DataFusion crates. FileType is in common and used in most of DataFusion while FileFormat is only in core.

This could work but I think the consequence would be DataFusion becomes one big crate. I think the only way to keep the independent crates is to have at least two File related abstractions. One simple one for logical planning in Common and one heavy duty one for Execution plans in core.

Even linking them is a challenge because then common depends on core 🤔

alamb · 2023-12-29T12:35:32Z

We could potentially also refactor the built in types (parquet, csv, ect) to be externally defined structs which implement ExtensionFileType and remove the FileType enum entirely. This would entail a more significant refactor of FileTypeWriterOptions as well.

I think this would make an excellent follow on PR.

alamb

Thank you @devinjdangelo -- this is a very neat PR.

As you say, I think making FileType itself a trait would be the ideal solution, but doing so would require refactoring a bunch more refactoring (likely FileTypeWriterOptions a trait for example)

I wonder if you have time to try that alternate approach and see how invasive it turns out to be 🤔

alamb · 2023-12-29T12:35:54Z

.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


Why is the new stack size needed?

I'm not sure. Locally I did not need this. To get CI tests to pass I needed to boost stack size.

alamb · 2023-12-29T12:37:20Z

datafusion/common/src/file_options/file_type.rs

@@ -41,7 +43,8 @@ pub trait GetExt {
 }

 /// Readable file type
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[allow(clippy::derived_hash_with_manual_eq)]


You could probably also manually implement Hash pretty easily as well, but I agree it isn't a bad thing

alamb · 2023-12-29T12:38:40Z

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

@@ -1716,6 +1719,9 @@ mod tests {
                    )
                    .await?;
            }
+            FileType::Extension(_) => {
+                panic!("Extension file type not implemented in write path.")


Can we instead return an NotImplement exception?

alamb · 2023-12-29T12:41:23Z

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(_) => {


We could also punt it into the FileType::Extension (e.g. add a serialize(&self) -> Result<Vec<u8>> type method)

alamb · 2023-12-29T12:43:47Z

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.")


Maybe we could have the FileTypeExtension have a method that returned Arc<dyn FileFormat> ?

devinjdangelo · 2023-12-29T15:42:39Z

Thank you for the review @alamb. It should be achievable to move the built in FileTypes to use the same trait interface and retire the concrete enum types. This would serve as a built in example of how to implement the trait for extension types as well and would even allow pulling out the "built in" files supported to their own crates.

I'll try to find some time within the next week to try this.

alamb · 2023-12-31T13:19:55Z

well and would even allow pulling out the "built in" files supported to their own crates.

I think this would be an ideal long term outcome (e.g. have a datafusion-parquet and datafusion-json, datafusion-arrow crates that came with the specific format support

devinjdangelo · 2024-01-04T03:03:11Z

I did not get back to this when Id hoped. It is still on my list, but I probably will not get to it until this weekend.

devinjdangelo · 2024-01-06T21:42:27Z

I've been thinking about this more, and I think we need a similar interface to how we register a custom TableProvider at the SessionContext level. The idea being we'd have a SessionContext::register_file_type which accepts an Arc<dyn FileType> and stores it in a DashMap<String, Arc<dyn FileType>>. At planning and execution time, all the places we are doing a match FileType currently, would be replaced with a lookup on the map of registered FileTypes + a method call on the trait object.

The above is based on my understanding having studied a bit how we handle external TableProviders. @alamb, does the above sound reasonable to you?

My main uncertainty to the above is how we would handle serialization. Perhaps the FileType trait could have a method which returns a Vec<u8> based serialization of itself. I'm not sure how that would work on the receiving end exactly...

alamb · 2024-01-08T19:11:12Z

The above is based on my understanding having studied a bit how we handle external TableProviders. @alamb, does the above sound reasonable to you?

Yes.

My main uncertainty to the above is how we would handle serialization. Perhaps the FileType trait could have a method which returns a Vec<u8> based serialization of itself. I'm not sure how that would work on the receiving end exactly...

I think the way serialization is handled today is via "ExtensionCodec"

Which do exactly as you suggest and serialize table providers to/from bytes

maybe @thinkharderdev has some more thoughts to share

thinkharderdev · 2024-01-08T22:24:26Z

https://docs.rs/datafusion-proto/latest/datafusion_proto/physical_plan/trait.PhysicalExtensionCodec.html

Yeah, the extension codecs seem like a logical place to handle this.

github-actions · 2024-04-14T01:59:52Z

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

squash

9dce12f

github-actions bot added the core Core DataFusion crate label Dec 28, 2023

devinjdangelo added 2 commits December 28, 2023 10:53

pass through StatementOptions

b93f6f8

lint and doc update

4783a1a

devinjdangelo commented Dec 28, 2023

View reviewed changes

devinjdangelo marked this pull request as draft December 28, 2023 16:10

devinjdangelo changed the title ~~WIP: Support Reading and Writing Extension FileTypes~~ Support Reading and Writing Extension FileTypes Dec 28, 2023

devinjdangelo mentioned this pull request Dec 28, 2023

Make/Remove FileType enum and replace with a trait #8657

Closed

try bump test stack size

0559aa9

alamb mentioned this pull request Dec 28, 2023

DataFusion weekly project plan (Andrew Lamb) - Dec 25, 2023 #8655

Closed

7 tasks

alamb reviewed Dec 29, 2023

View reviewed changes

devinjdangelo mentioned this pull request Dec 29, 2023

Implement serde for CSV and Parquet FileSinkExec #8646

Merged

starting to rip out enums...

4ad7f6e

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions labels Jan 6, 2024

github-actions bot added the Stale PR has not had any activity for some time label Apr 14, 2024

github-actions bot closed this Apr 22, 2024

devinjdangelo mentioned this pull request Jun 22, 2024

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

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support Reading and Writing Extension FileTypes #8667

Support Reading and Writing Extension FileTypes #8667

devinjdangelo commented Dec 28, 2023 •

edited

Loading

devinjdangelo left a comment

devinjdangelo Dec 28, 2023

alamb Dec 29, 2023

devinjdangelo Dec 28, 2023

devinjdangelo Dec 28, 2023

alamb Dec 29, 2023

devinjdangelo Dec 28, 2023

alamb Dec 29, 2023

devinjdangelo Dec 29, 2023

alamb commented Dec 29, 2023

alamb left a comment

alamb Dec 29, 2023

devinjdangelo Dec 29, 2023

alamb Dec 29, 2023

alamb Dec 29, 2023

alamb Dec 29, 2023

alamb Dec 29, 2023

devinjdangelo commented Dec 29, 2023

alamb commented Dec 31, 2023

devinjdangelo commented Jan 4, 2024

devinjdangelo commented Jan 6, 2024

alamb commented Jan 8, 2024

thinkharderdev commented Jan 8, 2024

github-actions bot commented Apr 14, 2024

Support Reading and Writing Extension FileTypes #8667

Support Reading and Writing Extension FileTypes #8667

Conversation

devinjdangelo commented Dec 28, 2023 • 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?

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

alamb commented Dec 29, 2023

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

devinjdangelo commented Dec 29, 2023

alamb commented Dec 31, 2023

devinjdangelo commented Jan 4, 2024

devinjdangelo commented Jan 6, 2024

alamb commented Jan 8, 2024

thinkharderdev commented Jan 8, 2024

github-actions bot commented Apr 14, 2024

devinjdangelo commented Dec 28, 2023 •

edited

Loading