fix(cubestore): Support NULL values in CSV import

cube-js · Jan 26, 2021 · 529e5ac · 529e5ac
1 parent a3e118e
commit 529e5ac
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 26 deletions.
diff --git a/rust/cubestore/src/import/mod.rs b/rust/cubestore/src/import/mod.rs
@@ -61,7 +61,7 @@ impl ImportFormat {
                                         .find(|c| c.get_name() == &next_column)
                                         .cloned()
                                         .ok_or(CubeError::user(format!(
-                                            "Column {} not found during import in {:?}",
+                                            "Column '{}' is not found during import in {:?}",
                                             next_column, columns
                                         )))?,
                                 );
@@ -74,33 +74,37 @@ impl ImportFormat {
                         for column in header_mapping
                             .as_ref()
                             .ok_or(CubeError::user(
-                                "Header required for CSV import".to_string(),
+                                "Header is required for CSV import".to_string(),
                             ))?
                             .iter()
                         {
                             let value = parser.next_value()?;
 
-                            row.push(match column.get_column_type() {
-                                ColumnType::String => TableValue::String(value),
-                                ColumnType::Int => value
-                                    .parse()
-                                    .map(|v| TableValue::Int(v))
-                                    .unwrap_or(TableValue::Null),
-                                ColumnType::Decimal { .. } => {
-                                    BigDecimal::from_str_radix(value.as_str(), 10)
-                                        .map(|d| TableValue::Decimal(d.to_string()))
-                                        .unwrap_or(TableValue::Null)
-                                }
-                                ColumnType::Bytes => unimplemented!(),
-                                ColumnType::HyperLogLog(_) => unimplemented!(),
-                                ColumnType::Timestamp => timestamp_from_string(value.as_str())?,
-                                ColumnType::Float => {
-                                    TableValue::Float(value.parse::<f64>()?.to_string())
-                                }
-                                ColumnType::Boolean => {
-                                    TableValue::Boolean(value.to_lowercase() == "true")
-                                }
-                            });
+                            if &value == "" {
+                                row.push(TableValue::Null);
+                            } else {
+                                row.push(match column.get_column_type() {
+                                    ColumnType::String => TableValue::String(value),
+                                    ColumnType::Int => value
+                                        .parse()
+                                        .map(|v| TableValue::Int(v))
+                                        .unwrap_or(TableValue::Null),
+                                    ColumnType::Decimal { .. } => {
+                                        BigDecimal::from_str_radix(value.as_str(), 10)
+                                            .map(|d| TableValue::Decimal(d.to_string()))
+                                            .unwrap_or(TableValue::Null)
+                                    }
+                                    ColumnType::Bytes => unimplemented!(),
+                                    ColumnType::HyperLogLog(_) => unimplemented!(),
+                                    ColumnType::Timestamp => timestamp_from_string(value.as_str())?,
+                                    ColumnType::Float => {
+                                        TableValue::Float(value.parse::<f64>()?.to_string())
+                                    }
+                                    ColumnType::Boolean => {
+                                        TableValue::Boolean(value.to_lowercase() == "true")
+                                    }
+                                });
+                            }
 
                             parser.advance()?;
                         }

diff --git a/rust/cubestore/src/sql/mod.rs b/rust/cubestore/src/sql/mod.rs
@@ -2078,6 +2078,7 @@ mod tests {
                 file.write_all("id,city,t,arr\n".as_bytes()).unwrap();
                 file.write_all("1,San Francisco,2021-01-24 12:12:23 UTC,\"[\"\"Foo\"\",\"\"Bar\"\",\"\"FooBar\"\"]\"\n".as_bytes()).unwrap();
                 file.write_all("2,New York,2021-01-24 19:12:23 UTC,\"[\"\"\"\"]\"\n".as_bytes()).unwrap();
+                file.write_all("3,New York,2021-01-25 19:12:23 UTC,\n".as_bytes()).unwrap();
 
                 dir
             };
@@ -2089,10 +2090,10 @@ mod tests {
             assert!(error.contains("has data"));
 
             let result = service.exec_query("SELECT count(*) as cnt from Foo.Persons").await.unwrap();
-            assert_eq!(result.get_rows()[0], Row::new(vec![TableValue::Int(2)]));
+            assert_eq!(result.get_rows(), &vec![Row::new(vec![TableValue::Int(3)])]);
 
-            let result = service.exec_query("SELECT count(*) as cnt from Foo.Persons WHERE arr = '[\"Foo\",\"Bar\",\"FooBar\"]' or arr = '[\"\"]'").await.unwrap();
-            assert_eq!(result.get_rows(), &vec![Row::new(vec![TableValue::Int(2)])]);
+            let result = service.exec_query("SELECT count(*) as cnt from Foo.Persons WHERE arr = '[\"Foo\",\"Bar\",\"FooBar\"]' or arr = '[\"\"]' or arr is null").await.unwrap();
+            assert_eq!(result.get_rows(), &vec![Row::new(vec![TableValue::Int(3)])]);
         }).await;
     }