Merge pull request #4571 from b41sh/variant-array-object

feat(data type): Support Semi-structured array, object data type

common/datavalues/src/columns/column.rs

@@ -177,6 +177,8 @@ where A: AsRef<dyn Array>
             Struct => Arc::new(StructColumn::from_arrow_array(self.as_ref())),
             String => Arc::new(StringColumn::from_arrow_array(self.as_ref())),
             Variant => Arc::new(JsonColumn::from_arrow_array(self.as_ref())),
+            VariantArray => Arc::new(JsonColumn::from_arrow_array(self.as_ref())),
+            VariantObject => Arc::new(JsonColumn::from_arrow_array(self.as_ref())),
         }
     }
 

common/datavalues/src/types/data_type.rs

@@ -138,6 +138,8 @@ pub fn from_arrow_type(dt: &ArrowType) -> DataTypePtr {
         }
         ArrowType::Extension(custom_name, _, _) => match custom_name.as_str() {
             "Variant" => Arc::new(VariantType::default()),
+            "VariantArray" => Arc::new(VariantArrayType::default()),
+            "VariantObject" => Arc::new(VariantObjectType::default()),
             _ => unimplemented!("data_type: {:?}", dt),
         },
 
@@ -166,6 +168,8 @@ pub fn from_arrow_field(f: &ArrowField) -> DataTypePtr {
             },
             "Interval" => return IntervalType::arc(metadata.unwrap().into()),
             "Variant" => return VariantType::arc(),
+            "VariantArray" => return VariantArrayType::arc(),
+            "VariantObject" => return VariantObjectType::arc(),
             _ => {}
         }
     }

common/datavalues/src/types/eq.rs

@@ -48,7 +48,8 @@ pub fn equal(lhs: &dyn DataType, rhs: &dyn DataType) -> bool {
     use crate::prelude::TypeID::*;
     match lhs.data_type_id() {
         Boolean | UInt8 | UInt16 | UInt32 | UInt64 | Int8 | Int16 | Int32 | Int64 | Float32
-        | Float64 | String | Date16 | Date32 | Interval | DateTime32 | Null | Variant => true,
+        | Float64 | String | Date16 | Date32 | Interval | DateTime32 | Null | Variant
+        | VariantArray | VariantObject => true,
 
         DateTime64 => {
             let lhs: &DateTime64Type = lhs.as_any().downcast_ref().unwrap();

common/datavalues/src/types/mod.rs

@@ -29,6 +29,8 @@ pub mod type_string;
 pub mod type_struct;
 pub mod type_traits;
 pub mod type_variant;
+pub mod type_variant_array;
+pub mod type_variant_object;
 
 pub mod eq;
 pub mod type_id;
@@ -61,3 +63,5 @@ pub use type_string::*;
 pub use type_struct::*;
 pub use type_traits::*;
 pub use type_variant::*;
+pub use type_variant_array::*;
+pub use type_variant_object::*;

common/datavalues/src/types/type_factory.rs

@@ -50,6 +50,8 @@ static TYPE_FACTORY: Lazy<Arc<TypeFactory>> = Lazy::new(|| {
     type_factory.register(DateTime32Type::arc(None));
     type_factory.register(DateTime64Type::arc(3, None));
     type_factory.register(VariantType::arc());
+    type_factory.register(VariantArrayType::arc());
+    type_factory.register(VariantObjectType::arc());
 
     type_factory.add_array_wrapper();
     type_factory.add_nullable_wrapper();

common/datavalues/src/types/type_id.rs

@@ -71,6 +71,8 @@ pub enum TypeID {
     /// Variant is a tagged universal type, which can store values of any other type,
     /// including Object and Array, up to a maximum size of 16 MB.
     Variant,
+    VariantArray,
+    VariantObject,
 }
 
 impl TypeID {
@@ -216,7 +218,7 @@ impl TypeID {
             String => PhysicalTypeID::String,
             Array => PhysicalTypeID::Array,
             Struct => PhysicalTypeID::Struct,
-            Variant => PhysicalTypeID::Variant,
+            Variant | VariantArray | VariantObject => PhysicalTypeID::Variant,
         }
     }
 }

common/datavalues/src/types/type_variant.rs

@@ -49,7 +49,7 @@ impl DataType for VariantType {
     }
 
     fn default_value(&self) -> DataValue {
-        DataValue::String("Null".as_bytes().to_vec())
+        DataValue::Json(JsonValue::Null)
     }
 
     fn create_constant_column(&self, data: &DataValue, size: usize) -> Result<ColumnRef> {

common/datavalues/src/types/type_variant_array.rs

@@ -0,0 +1,101 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use common_arrow::arrow::datatypes::DataType as ArrowType;
+use common_exception::Result;
+use serde_json::Value as JsonValue;
+
+use super::data_type::DataType;
+use super::data_type::ARROW_EXTENSION_NAME;
+use super::type_id::TypeID;
+use crate::prelude::*;
+
+#[derive(Default, Clone, serde::Deserialize, serde::Serialize)]
+pub struct VariantArrayType {}
+
+impl VariantArrayType {
+    pub fn arc() -> DataTypePtr {
+        Arc::new(Self {})
+    }
+}
+
+#[typetag::serde]
+impl DataType for VariantArrayType {
+    fn data_type_id(&self) -> TypeID {
+        TypeID::VariantArray
+    }
+
+    #[inline]
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "Array"
+    }
+
+    fn default_value(&self) -> DataValue {
+        DataValue::Json(JsonValue::Array(vec![]))
+    }
+
+    fn create_constant_column(&self, data: &DataValue, size: usize) -> Result<ColumnRef> {
+        let value: JsonValue = DFTryFrom::try_from(data)?;
+        let column = Series::from_data(vec![value]);
+        Ok(Arc::new(ConstColumn::new(column, size)))
+    }
+
+    fn create_column(&self, data: &[DataValue]) -> Result<ColumnRef> {
+        let values: Vec<JsonValue> = data
+            .iter()
+            .map(DFTryFrom::try_from)
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(Series::from_data(values))
+    }
+
+    fn arrow_type(&self) -> ArrowType {
+        ArrowType::Extension(
+            "VariantArray".to_owned(),
+            Box::new(ArrowType::LargeBinary),
+            None,
+        )
+    }
+
+    fn custom_arrow_meta(&self) -> Option<BTreeMap<String, String>> {
+        let mut mp = BTreeMap::new();
+        mp.insert(ARROW_EXTENSION_NAME.to_string(), "VariantArray".to_string());
+        Some(mp)
+    }
+
+    fn create_serializer(&self) -> Box<dyn TypeSerializer> {
+        Box::new(VariantSerializer {})
+    }
+
+    fn create_deserializer(&self, capacity: usize) -> Box<dyn TypeDeserializer> {
+        Box::new(VariantDeserializer::with_capacity(capacity))
+    }
+
+    fn create_mutable(&self, capacity: usize) -> Box<dyn MutableColumn> {
+        Box::new(MutableObjectColumn::<JsonValue>::with_capacity(capacity))
+    }
+}
+
+impl std::fmt::Debug for VariantArrayType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}

common/datavalues/src/types/type_variant_object.rs

@@ -0,0 +1,105 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use common_arrow::arrow::datatypes::DataType as ArrowType;
+use common_exception::Result;
+use serde_json::Map;
+use serde_json::Value as JsonValue;
+
+use super::data_type::DataType;
+use super::data_type::ARROW_EXTENSION_NAME;
+use super::type_id::TypeID;
+use crate::prelude::*;
+
+#[derive(Default, Clone, serde::Deserialize, serde::Serialize)]
+pub struct VariantObjectType {}
+
+impl VariantObjectType {
+    pub fn arc() -> DataTypePtr {
+        Arc::new(Self {})
+    }
+}
+
+#[typetag::serde]
+impl DataType for VariantObjectType {
+    fn data_type_id(&self) -> TypeID {
+        TypeID::VariantObject
+    }
+
+    #[inline]
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "Object"
+    }
+
+    fn default_value(&self) -> DataValue {
+        DataValue::Json(JsonValue::Object(Map::new()))
+    }
+
+    fn create_constant_column(&self, data: &DataValue, size: usize) -> Result<ColumnRef> {
+        let value: JsonValue = DFTryFrom::try_from(data)?;
+        let column = Series::from_data(vec![value]);
+        Ok(Arc::new(ConstColumn::new(column, size)))
+    }
+
+    fn create_column(&self, data: &[DataValue]) -> Result<ColumnRef> {
+        let values: Vec<JsonValue> = data
+            .iter()
+            .map(DFTryFrom::try_from)
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(Series::from_data(values))
+    }
+
+    fn arrow_type(&self) -> ArrowType {
+        ArrowType::Extension(
+            "VariantObject".to_owned(),
+            Box::new(ArrowType::LargeBinary),
+            None,
+        )
+    }
+
+    fn custom_arrow_meta(&self) -> Option<BTreeMap<String, String>> {
+        let mut mp = BTreeMap::new();
+        mp.insert(
+            ARROW_EXTENSION_NAME.to_string(),
+            "VariantObject".to_string(),
+        );
+        Some(mp)
+    }
+
+    fn create_serializer(&self) -> Box<dyn TypeSerializer> {
+        Box::new(VariantSerializer {})
+    }
+
+    fn create_deserializer(&self, capacity: usize) -> Box<dyn TypeDeserializer> {
+        Box::new(VariantDeserializer::with_capacity(capacity))
+    }
+
+    fn create_mutable(&self, capacity: usize) -> Box<dyn MutableColumn> {
+        Box::new(MutableObjectColumn::<JsonValue>::with_capacity(capacity))
+    }
+}
+
+impl std::fmt::Debug for VariantObjectType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}

common/datavalues/tests/it/types/create_column.rs

@@ -129,6 +129,20 @@ fn test_create_constant() -> Result<()> {
             size: 2,
             column_expected: Series::from_data(vec![json!("hello"), json!("hello")]),
         },
+        Test {
+            name: "variant_array",
+            data_type: VariantArrayType::arc(),
+            value: DataValue::Json(json!([1, 2, 3])),
+            size: 2,
+            column_expected: Series::from_data(vec![json!([1, 2, 3]), json!([1, 2, 3])]),
+        },
+        Test {
+            name: "variant_object",
+            data_type: VariantObjectType::arc(),
+            value: DataValue::Json(json!({"a":1,"b":2})),
+            size: 2,
+            column_expected: Series::from_data(vec![json!({"a":1,"b":2}), json!({"a":1,"b":2})]),
+        },
     ];
 
     for test in tests {

common/functions/src/scalars/expressions/cast_with_type.rs

@@ -178,15 +178,51 @@ pub fn cast_with_type(
 pub fn cast_to_variant(
     column: &ColumnRef,
     from_type: &DataTypePtr,
+    data_type: &DataTypePtr,
 ) -> Result<(ColumnRef, Option<Bitmap>)> {
+    let column = Series::remove_nullable(column);
+    let size = column.len();
+
+    if data_type.data_type_id() == TypeID::VariantArray {
+        if from_type.data_type_id() == TypeID::Variant {
+            let col: &JsonColumn = Series::check_get(&column)?;
+            let mut builder = ColumnBuilder::<JsonValue>::with_capacity(size);
+            for v in col.iter() {
+                if v.is_array() {
+                    builder.append(v);
+                } else {
+                    let arr = JsonValue::Array(vec![v.clone()]);
+                    builder.append(&arr);
+                }
+            }
+            return Ok((builder.build(size), None));
+        }
+        return Err(ErrorCode::BadDataValueType(format!(
+            "Expression type does not match column data type, expecting ARRAY but got {}",
+            from_type.data_type_id()
+        )));
+    } else if data_type.data_type_id() == TypeID::VariantObject {
+        if from_type.data_type_id() == TypeID::Variant {
+            let col: &JsonColumn = Series::check_get(&column)?;
+            for v in col.iter() {
+                if !v.is_object() {
+                    return Err(ErrorCode::BadDataValueType(format!(
+                        "Failed to cast variant value {} to OBJECT",
+                        v
+                    )));
+                }
+            }
+            return Ok((column.clone(), None));
+        }
+        return Err(ErrorCode::BadDataValueType(format!(
+            "Expression type does not match column data type, expecting OBJECT but got {}",
+            from_type.data_type_id()
+        )));
+    }
     if from_type.data_type_id() == TypeID::Variant {
         return Ok((column.clone(), None));
     }
-
-    let column = Series::remove_nullable(column);
-    let size = column.len();
     let mut builder = ColumnBuilder::<JsonValue>::with_capacity(size);
-
     if from_type.data_type_id().is_numeric() || from_type.data_type_id() == TypeID::Boolean {
         let serializer = from_type.create_serializer();
         match serializer.serialize_json_object(&column, None) {
@@ -201,7 +237,7 @@ pub fn cast_to_variant(
     }
     // other data types can't automatically casted to variant
     return Err(ErrorCode::BadDataValueType(format!(
-        "Expression type does not match column data type, expecting VARIANT but got {:?}",
+        "Expression type does not match column data type, expecting VARIANT but got {}",
         from_type.data_type_id()
     )));
 }
@@ -213,8 +249,11 @@ pub fn arrow_cast_compute(
     data_type: &DataTypePtr,
     cast_options: &CastOptions,
 ) -> Result<(ColumnRef, Option<Bitmap>)> {
-    if data_type.data_type_id() == TypeID::Variant {
-        return cast_to_variant(column, from_type);
+    if data_type.data_type_id() == TypeID::Variant
+        || data_type.data_type_id() == TypeID::VariantArray
+        || data_type.data_type_id() == TypeID::VariantObject
+    {
+        return cast_to_variant(column, from_type, data_type);
     }
 
     let arrow_array = column.as_arrow_array();