feat(cubestore): Support custom CSV delimiters (#6597)

Author: paveltiunov

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -103,6 +103,10 @@ pub fn sql_tests() -> Vec<(&'static str, TestFn)> {
             "create_table_with_csv_no_header",
             create_table_with_csv_no_header,
         ),
+        t(
+            "create_table_with_csv_no_header_and_delimiter",
+            create_table_with_csv_no_header_and_delimiter,
+        ),
         t("create_table_with_url", create_table_with_url),
         t("create_table_fail_and_retry", create_table_fail_and_retry),
         t("empty_crash", empty_crash),
@@ -2155,6 +2159,34 @@ async fn create_table_with_csv_no_header(service: Box<dyn SqlClient>) {
     );
 }
 
+async fn create_table_with_csv_no_header_and_delimiter(service: Box<dyn SqlClient>) {
+    let file = write_tmp_file(indoc! {"
+        apple\u{0001}2
+        banana\u{0001}3
+    "})
+    .unwrap();
+    let path = file.path().to_string_lossy();
+    let _ = service
+        .exec_query("CREATE SCHEMA IF NOT EXISTS test")
+        .await
+        .unwrap();
+    let _ = service
+        .exec_query(format!("CREATE TABLE test.table (`fruit` text, `number` int) WITH (input_format = 'csv_no_header', delimiter = '^A') LOCATION '{}'", path).as_str())
+        .await
+        .unwrap();
+    let result = service
+        .exec_query("SELECT * FROM test.table")
+        .await
+        .unwrap();
+    assert_eq!(
+        to_rows(&result),
+        vec![
+            vec![TableValue::String("apple".to_string()), TableValue::Int(2)],
+            vec![TableValue::String("banana".to_string()), TableValue::Int(3)]
+        ]
+    );
+}
+
 async fn create_table_with_url(service: Box<dyn SqlClient>) {
     let url = "https://data.wprdc.org/dataset/0b584c84-7e35-4f4d-a5a2-b01697470c0f/resource/e95dd941-8e47-4460-9bd8-1e51c194370b/download/bikepghpublic.csv";
 

rust/cubestore/cubestore/src/import/mod.rs

@@ -67,25 +67,40 @@ impl ImportFormat {
     ) -> Result<Pin<Box<dyn Stream<Item = Result<Option<Row>, CubeError>> + Send + 'a>>, CubeError>
     {
         match self {
-            ImportFormat::CSV | ImportFormat::CSVNoHeader => {
+            ImportFormat::CSV | ImportFormat::CSVNoHeader | ImportFormat::CSVOptions { .. } => {
                 let lines_stream: Pin<Box<dyn Stream<Item = Result<String, CubeError>> + Send>> =
                     Box::pin(CsvLineStream::new(reader));
 
                 let mut header_mapping = match self {
-                    ImportFormat::CSV => None,
-                    ImportFormat::CSVNoHeader => Some(
+                    ImportFormat::CSVNoHeader
+                    | ImportFormat::CSVOptions {
+                        has_header: false, ..
+                    } => Some(
                         columns
                             .iter()
                             .enumerate()
                             .map(|(i, c)| (i, c.clone()))
                             .collect(),
                     ),
+                    _ => None,
                 };
 
+                let delimiter = match self {
+                    ImportFormat::CSV | ImportFormat::CSVNoHeader => ',',
+                    ImportFormat::CSVOptions { delimiter, .. } => delimiter.unwrap_or(','),
+                };
+
+                if delimiter as u16 > 255 {
+                    return Err(CubeError::user(format!(
+                        "Non ASCII delimiters are unsupported: '{}'",
+                        delimiter
+                    )));
+                }
+
                 let rows = lines_stream.map(move |line| -> Result<Option<Row>, CubeError> {
                     let str = line?;
 
-                    let mut parser = CsvLineParser::new(str.as_str());
+                    let mut parser = CsvLineParser::new(delimiter as u8, str.as_str());
 
                     if header_mapping.is_none() {
                         let mut mapping = Vec::new();
@@ -257,13 +272,15 @@ fn parse_binary_data(value: &str) -> Result<Vec<u8>, CubeError> {
 }
 
 struct CsvLineParser<'a> {
+    delimiter: u8,
     line: &'a str,
     remaining: &'a str,
 }
 
 impl<'a> CsvLineParser<'a> {
-    fn new(line: &'a str) -> Self {
+    fn new(delimiter: u8, line: &'a str) -> Self {
         Self {
+            delimiter,
             line,
             remaining: line,
         }
@@ -308,7 +325,7 @@ impl<'a> CsvLineParser<'a> {
                     .remaining
                     .as_bytes()
                     .iter()
-                    .position(|c| *c == b',')
+                    .position(|c| *c == self.delimiter)
                     .unwrap_or(self.remaining.len());
                 let res = &self.remaining[0..next_comma];
                 self.remaining = self.remaining[next_comma..].as_ref();
@@ -318,8 +335,10 @@ impl<'a> CsvLineParser<'a> {
     }
 
     fn advance(&mut self) -> Result<(), CubeError> {
-        if let Some(b',') = self.remaining.as_bytes().iter().nth(0) {
-            self.remaining = self.remaining[1..].as_ref()
+        if let Some(c) = self.remaining.as_bytes().iter().nth(0) {
+            if *c == self.delimiter {
+                self.remaining = self.remaining[1..].as_ref()
+            }
         }
         Ok(())
     }

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -552,6 +552,12 @@ impl fmt::Display for Column {
 pub enum ImportFormat {
     CSV,
     CSVNoHeader,
+    CSVOptions {
+        delimiter: Option<char>,
+        escape: Option<char>,
+        quote: Option<char>,
+        has_header: bool,
+    },
 }
 
 data_frame_from! {

rust/cubestore/cubestore/src/sql/mod.rs

@@ -969,7 +969,7 @@ impl SqlService for SqlServiceImpl {
                 }
                 let schema_name = &nv[0].value;
                 let table_name = &nv[1].value;
-                let import_format = with_options
+                let mut import_format = with_options
                     .iter()
                     .find(|&opt| opt.name.value == "input_format")
                     .map_or(Result::Ok(ImportFormat::CSV), |option| {
@@ -990,6 +990,49 @@ impl SqlService for SqlServiceImpl {
                             ))),
                         }
                     })?;
+
+                let delimiter = with_options
+                    .iter()
+                    .find(|&opt| opt.name.value == "delimiter")
+                    .map_or(Ok(None), |option| match &option.value {
+                        Value::SingleQuotedString(delimiter) => match delimiter.as_str() {
+                            "tab" => Ok(Some('\t')),
+                            "^A" => Ok(Some('\u{0001}')),
+                            s if s.len() != 1 => {
+                                Err(CubeError::user(format!("Bad delimiter {}", option.value)))
+                            }
+                            s => Ok(Some(s.chars().next().unwrap())),
+                        },
+                        _ => Err(CubeError::user(format!("Bad delimiter {}", option.value))),
+                    })?;
+
+                if let Some(delimiter) = delimiter {
+                    import_format = match import_format {
+                        ImportFormat::CSV => ImportFormat::CSVOptions {
+                            delimiter: Some(delimiter),
+                            has_header: true,
+                            escape: None,
+                            quote: None,
+                        },
+                        ImportFormat::CSVNoHeader => ImportFormat::CSVOptions {
+                            delimiter: Some(delimiter),
+                            has_header: false,
+                            escape: None,
+                            quote: None,
+                        },
+                        ImportFormat::CSVOptions {
+                            has_header,
+                            escape,
+                            quote,
+                            ..
+                        } => ImportFormat::CSVOptions {
+                            delimiter: Some(delimiter),
+                            has_header,
+                            escape,
+                            quote,
+                        },
+                    }
+                }
                 let build_range_end = with_options
                     .iter()
                     .find(|&opt| opt.name.value == "build_range_end")