feat(cubestore): allow to import base64-encoded bytes in CSV (#1891)

cube-js · Jan 26, 2021 · 2f43afa · 2f43afa
1 parent 849f3d8
commit 2f43afa
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 13 deletions.
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
diff --git a/rust/cubestore/Cargo.toml b/rust/cubestore/Cargo.toml
@@ -12,6 +12,7 @@ repository = "https://github.com/cube-js/cube.js"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+base64 = "0.13.0"
 tokio = { version = "1.0", features = ["full", "rt"] }
 warp = "0.2"
 sqlparser = "0.7.0"

diff --git a/rust/cubestore/src/import/mod.rs b/rust/cubestore/src/import/mod.rs
@@ -1,4 +1,5 @@
 use crate::config::ConfigObj;
+use crate::metastore::is_valid_hll;
 use crate::metastore::{Column, ColumnType, ImportFormat, MetaStore};
 use crate::sql::timestamp_from_string;
 use crate::store::{DataFrame, WALDataStore};
@@ -106,8 +107,15 @@ impl ImportFormat {
                                                 .map(|d| TableValue::Decimal(d.to_string()))
                                                 .unwrap_or(TableValue::Null)
                                         }
-                                        ColumnType::Bytes => unimplemented!(),
-                                        ColumnType::HyperLogLog(_) => unimplemented!(),
+                                        ColumnType::Bytes => {
+                                            TableValue::Bytes(base64::decode(value)?)
+                                        }
+                                        ColumnType::HyperLogLog(f) => {
+                                            let data = base64::decode(value)?;
+                                            is_valid_hll(&data, *f)?;
+
+                                            TableValue::Bytes(data)
+                                        }
                                         ColumnType::Timestamp => {
                                             timestamp_from_string(value.as_str())?
                                         }

diff --git a/rust/cubestore/src/lib.rs b/rust/cubestore/src/lib.rs
@@ -335,3 +335,9 @@ impl From<cloud_storage::Error> for CubeError {
         return CubeError::from_error(v);
     }
 }
+
+impl From<base64::DecodeError> for CubeError {
+    fn from(v: base64::DecodeError) -> Self {
+        return CubeError::from_error(v);
+    }
+}
diff --git a/rust/cubestore/src/metastore/mod.rs b/rust/cubestore/src/metastore/mod.rs
@@ -37,6 +37,8 @@ use arrow::datatypes::{DataType, Field};
 use chrono::{DateTime, Utc};
 use chunks::ChunkRocksTable;
 use core::{fmt, mem};
+use cubehll::HllSketch;
+use cubezetasketch::HyperLogLogPlusPlus;
 use futures::future::join_all;
 use futures_timer::Delay;
 use index::{IndexRocksIndex, IndexRocksTable};
@@ -282,6 +284,19 @@ pub enum HllFlavour {
     ZetaSketch, // Compatible with BigQuery.
 }
 
+pub fn is_valid_hll(data: &[u8], f: HllFlavour) -> Result<(), CubeError> {
+    // TODO: do no memory allocations for better performance, this is run on hot path.
+    match f {
+        HllFlavour::Airlift => {
+            HllSketch::read(data)?;
+        }
+        HllFlavour::ZetaSketch => {
+            HyperLogLogPlusPlus::read(data)?;
+        }
+    }
+    return Ok(());
+}
+
 #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
 pub enum ColumnType {
     String,

diff --git a/rust/cubestore/src/sql/mod.rs b/rust/cubestore/src/sql/mod.rs
@@ -7,8 +7,8 @@ use sqlparser::ast::*;
 use sqlparser::dialect::Dialect;
 
 use crate::metastore::{
-    table::Table, HllFlavour, IdRow, ImportFormat, Index, IndexDef, MetaStoreTable, RowKey, Schema,
-    TableId,
+    is_valid_hll, table::Table, HllFlavour, IdRow, ImportFormat, Index, IndexDef, MetaStoreTable,
+    RowKey, Schema, TableId,
 };
 use crate::table::{Row, TableValue, TimestampValue};
 use crate::CubeError;
@@ -547,15 +547,8 @@ fn decode_byte(s: &str) -> Option<u8> {
 
 fn parse_hyper_log_log(v: &Value, f: HllFlavour) -> Result<Vec<u8>, CubeError> {
     let bytes = parse_binary_string(v)?;
-    // TODO: check without memory allocations. this is run on hot path.
-    match f {
-        HllFlavour::Airlift => {
-            cubehll::HllSketch::read(&bytes)?;
-        }
-        HllFlavour::ZetaSketch => {
-            cubezetasketch::HyperLogLogPlusPlus::read(&bytes)?;
-        }
-    }
+    is_valid_hll(&bytes, f)?;
+
     return Ok(bytes);
 }