add compression type to csv serde

andygrove · Dec 29, 2023 · d02043c · d02043c
1 parent 859e531
commit d02043c
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 22 deletions.
diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -1222,20 +1222,22 @@ message ParquetWriterOptions {
 }
 
 message CsvWriterOptions {
+  // Compression type
+  CompressionTypeVariant compression = 1;
   // Optional column delimiter. Defaults to `b','`
-  string delimiter = 1;
+  string delimiter = 2;
   // Whether to write column names as file headers. Defaults to `true`
-  bool has_header = 2;
+  bool has_header = 3;
   // Optional date format for date arrays
-  string date_format = 3;
+  string date_format = 4;
   // Optional datetime format for datetime arrays
-  string datetime_format = 4;
+  string datetime_format = 5;
   // Optional timestamp format for timestamp arrays
-  string timestamp_format = 5;
+  string timestamp_format = 6;
   // Optional time format for time arrays
-  string time_format = 6;
+  string time_format = 7;
   // Optional value to represent null
-  string null_value = 7;
+  string null_value = 8;
 }
 
 message WriterProperties {

diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs
@@ -1648,8 +1648,10 @@ impl AsLogicalPlan for LogicalPlanNode {
                             match opt.as_ref() {
                                 FileTypeWriterOptions::CSV(csv_opts) => {
                                     let csv_options = &csv_opts.writer_options;
-                                    let csv_writer_options =
-                                        csv_writer_options_to_proto(csv_options);
+                                    let csv_writer_options = csv_writer_options_to_proto(
+                                        csv_options,
+                                        &CompressionTypeVariant::UNCOMPRESSED,
+                                    );
                                     let csv_options =
                                         file_type_writer_options::FileType::CsvOptions(
                                             csv_writer_options,
@@ -1706,8 +1708,11 @@ impl AsLogicalPlan for LogicalPlanNode {
 
 pub(crate) fn csv_writer_options_to_proto(
     csv_options: &WriterBuilder,
+    compression: &CompressionTypeVariant,
 ) -> protobuf::CsvWriterOptions {
+    let compression: protobuf::CompressionTypeVariant = compression.into();
     protobuf::CsvWriterOptions {
+        compression: compression.into(),
         delimiter: (csv_options.delimiter() as char).to_string(),
         has_header: csv_options.header(),
         date_format: csv_options.date_format().unwrap_or("").to_owned(),

diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -908,12 +908,9 @@ impl TryFrom<&FileTypeWriterOptions> for protobuf::FileTypeWriterOptions {
             FileTypeWriterOptions::CSV(CsvWriterOptions {
                 writer_options,
                 compression,
-            }) => {
-                let _compression: protobuf::CompressionTypeVariant = compression.into();
-                protobuf::file_type_writer_options::FileType::CsvOptions(
-                    csv_writer_options_to_proto(writer_options),
-                )
-            }
+            }) => protobuf::file_type_writer_options::FileType::CsvOptions(
+                csv_writer_options_to_proto(writer_options, compression),
+            ),
             FileTypeWriterOptions::JSON(JsonWriterOptions { compression }) => {
                 let compression: protobuf::CompressionTypeVariant = compression.into();
                 protobuf::file_type_writer_options::FileType::JsonOptions(