extended CSV options in SQL

korowa · Nov 27, 2023 · 2422275 · 2422275
1 parent 234217e
commit 2422275
Show file tree

Hide file tree

Showing 9 changed files with 124 additions and 3 deletions.
diff --git a/datafusion/core/src/catalog/listing_schema.rs b/datafusion/core/src/catalog/listing_schema.rs
@@ -141,6 +141,8 @@ impl ListingSchemaProvider {
                             file_type: self.format.clone(),
                             has_header: self.has_header,
                             delimiter: ',',
+                            quote: '"',
+                            escape: None,
                             table_partition_cols: vec![],
                             if_not_exists: false,
                             definition: None,

diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs
@@ -71,6 +71,8 @@ impl TableProviderFactory for ListingTableFactory {
                 CsvFormat::default()
                     .with_has_header(cmd.has_header)
                     .with_delimiter(cmd.delimiter as u8)
+                    .with_quote(cmd.quote as u8)
+                    .with_escape(cmd.escape.map(|c| c as u8))
                     .with_file_compression_type(file_compression_type),
             ),
             #[cfg(feature = "parquet")]
@@ -262,6 +264,8 @@ mod tests {
             file_type: "csv".to_string(),
             has_header: true,
             delimiter: ',',
+            quote: '"',
+            escape: None,
             schema: Arc::new(DFSchema::empty()),
             table_partition_cols: vec![],
             if_not_exists: false,

diff --git a/datafusion/expr/src/logical_plan/ddl.rs b/datafusion/expr/src/logical_plan/ddl.rs
@@ -178,6 +178,10 @@ pub struct CreateExternalTable {
     pub has_header: bool,
     /// Delimiter for CSV
     pub delimiter: char,
+    /// Quote character for CSV
+    pub quote: char,
+    /// Escape character for CSV
+    pub escape: Option<char>,
     /// Partition Columns
     pub table_partition_cols: Vec<String>,
     /// Option to not error if table already exists

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -216,6 +216,8 @@ message CreateExternalTableNode {
   bool unbounded = 14;
   map<string, string> options = 11;
   Constraints constraints = 15;
+  string quote = 16;
+  string escape = 17;
 }
 
 message PrepareNode {

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
@@ -505,6 +505,13 @@ impl AsLogicalPlan for LogicalPlanNode {
                 } else {
                     None
                 };
+                let escape = if !create_extern_table.escape.is_empty() {
+                    Some(create_extern_table.delimiter.chars().next().ok_or_else(|| {
+                        DataFusionError::Internal(String::from("Protobuf deserialization error, unable to parse CSV escape"))
+                    })?)
+                } else {
+                    None
+                };
 
                 let file_type = create_extern_table.file_type.as_str();
                 if ctx.table_factory(file_type).is_none() {
@@ -530,6 +537,10 @@ impl AsLogicalPlan for LogicalPlanNode {
                     delimiter: create_extern_table.delimiter.chars().next().ok_or_else(|| {
                         DataFusionError::Internal(String::from("Protobuf deserialization error, unable to parse CSV delimiter"))
                     })?,
+                    quote: create_extern_table.quote.chars().next().ok_or_else(|| {
+                        DataFusionError::Internal(String::from("Protobuf deserialization error, unable to parse CSV quote"))
+                    })?,
+                    escape,
                     table_partition_cols: create_extern_table
                         .table_partition_cols
                         .clone(),
@@ -1289,6 +1300,8 @@ impl AsLogicalPlan for LogicalPlanNode {
                     file_type,
                     has_header,
                     delimiter,
+                    quote,
+                    escape,
                     schema: df_schema,
                     table_partition_cols,
                     if_not_exists,
@@ -1323,6 +1336,8 @@ impl AsLogicalPlan for LogicalPlanNode {
                             table_partition_cols: table_partition_cols.clone(),
                             if_not_exists: *if_not_exists,
                             delimiter: String::from(*delimiter),
+                            quote: String::from(*quote),
+                            escape: escape.clone().unwrap_or_default().to_string(),
                             order_exprs: converted_order_exprs,
                             definition: definition.clone().unwrap_or_default(),
                             file_compression_type: file_compression_type.to_string(),