Support CopyTo::partition_by in datafusion proto (apache#9306)

* add support for CopyTo::partition_by in proto Signed-off-by: Hoang Pham <[email protected]> * simplify partition_by logic Signed-off-by: Hoang Pham <[email protected]> --------- Signed-off-by: Hoang Pham <[email protected]>
andygrove · Feb 25, 2024 · 148b4d2 · 148b4d2
1 parent ff36f6d
commit 148b4d2
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 7 deletions.
diff --git a/datafusion/expr/src/logical_plan/dml.rs b/datafusion/expr/src/logical_plan/dml.rs
@@ -36,7 +36,7 @@ pub struct CopyTo {
     pub output_url: String,
     /// The file format to output (explicitly defined or inferred from file extension)
     pub file_format: FileType,
-    /// Detmines which, if any, columns should be used for hive-style partitioned writes
+    /// Determines which, if any, columns should be used for hive-style partitioned writes
     pub partition_by: Vec<String>,
     /// Arbitrary options as tuples
     pub copy_options: CopyOptions,

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -327,6 +327,7 @@ message CopyToNode {
         FileTypeWriterOptions writer_options = 5;
     }
     string file_type = 6;
+    repeated string partition_by = 7;
 }
 
 message SQLOptions {

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
@@ -913,12 +913,13 @@ impl AsLogicalPlan for LogicalPlanNode {
                     }
                     None => return Err(proto_error("CopyTo missing CopyOptions")),
                 };
+
                 Ok(datafusion_expr::LogicalPlan::Copy(
                     datafusion_expr::dml::CopyTo {
                         input: Arc::new(input),
                         output_url: copy.output_url.clone(),
                         file_format: FileType::from_str(&copy.file_type)?,
-                        partition_by: vec![],
+                        partition_by: copy.partition_by.clone(),
                         copy_options,
                     },
                 ))
@@ -1642,7 +1643,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                 output_url,
                 file_format,
                 copy_options,
-                partition_by: _,
+                partition_by,
             }) => {
                 let input = protobuf::LogicalPlanNode::try_from_logical_plan(
                     input,
@@ -1726,6 +1727,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                             output_url: output_url.to_string(),
                             file_type: file_format.to_string(),
                             copy_options: copy_options_proto,
+                            partition_by: partition_by.clone(),
                         },
                     ))),
                 })

diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -324,7 +324,7 @@ async fn roundtrip_logical_plan_copy_to_sql_options() -> Result<()> {
         input: Arc::new(input),
         output_url: "test.csv".to_string(),
         file_format: FileType::CSV,
-        partition_by: vec![],
+        partition_by: vec!["a".to_string(), "b".to_string(), "c".to_string()],
         copy_options: CopyOptions::SQLOptions(StatementOptions::from(&options)),
     });
 
@@ -355,7 +355,7 @@ async fn roundtrip_logical_plan_copy_to_writer_options() -> Result<()> {
         input: Arc::new(input),
         output_url: "test.parquet".to_string(),
         file_format: FileType::PARQUET,
-        partition_by: vec![],
+        partition_by: vec!["a".to_string(), "b".to_string(), "c".to_string()],
         copy_options: CopyOptions::WriterOptions(Box::new(
             FileTypeWriterOptions::Parquet(ParquetWriterOptions::new(writer_properties)),
         )),
@@ -369,6 +369,7 @@ async fn roundtrip_logical_plan_copy_to_writer_options() -> Result<()> {
         LogicalPlan::Copy(copy_to) => {
             assert_eq!("test.parquet", copy_to.output_url);
             assert_eq!(FileType::PARQUET, copy_to.file_format);
+            assert_eq!(vec!["a", "b", "c"], copy_to.partition_by);
             match &copy_to.copy_options {
                 CopyOptions::WriterOptions(y) => match y.as_ref() {
                     FileTypeWriterOptions::Parquet(p) => {
@@ -404,7 +405,7 @@ async fn roundtrip_logical_plan_copy_to_arrow() -> Result<()> {
         input: Arc::new(input),
         output_url: "test.arrow".to_string(),
         file_format: FileType::ARROW,
-        partition_by: vec![],
+        partition_by: vec!["a".to_string(), "b".to_string(), "c".to_string()],
         copy_options: CopyOptions::WriterOptions(Box::new(FileTypeWriterOptions::Arrow(
             ArrowWriterOptions::new(),
         ))),
@@ -418,6 +419,7 @@ async fn roundtrip_logical_plan_copy_to_arrow() -> Result<()> {
         LogicalPlan::Copy(copy_to) => {
             assert_eq!("test.arrow", copy_to.output_url);
             assert_eq!(FileType::ARROW, copy_to.file_format);
+            assert_eq!(vec!["a", "b", "c"], copy_to.partition_by);
             match &copy_to.copy_options {
                 CopyOptions::WriterOptions(y) => match y.as_ref() {
                     FileTypeWriterOptions::Arrow(_) => {}
@@ -450,7 +452,7 @@ async fn roundtrip_logical_plan_copy_to_csv() -> Result<()> {
         input: Arc::new(input),
         output_url: "test.csv".to_string(),
         file_format: FileType::CSV,
-        partition_by: vec![],
+        partition_by: vec!["a".to_string(), "b".to_string(), "c".to_string()],
         copy_options: CopyOptions::WriterOptions(Box::new(FileTypeWriterOptions::CSV(
             CsvWriterOptions::new(
                 writer_properties,
@@ -467,6 +469,7 @@ async fn roundtrip_logical_plan_copy_to_csv() -> Result<()> {
         LogicalPlan::Copy(copy_to) => {
             assert_eq!("test.csv", copy_to.output_url);
             assert_eq!(FileType::CSV, copy_to.file_format);
+            assert_eq!(vec!["a", "b", "c"], copy_to.partition_by);
             match &copy_to.copy_options {
                 CopyOptions::WriterOptions(y) => match y.as_ref() {
                     FileTypeWriterOptions::CSV(p) => {