feat: add json data type

Cargo.toml

@@ -120,10 +120,11 @@ etcd-client = { version = "0.13" }
 fst = "0.4.7"
 futures = "0.3"
 futures-util = "0.3"
-greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "c437b55725b7f5224fe9d46db21072b4a682ee4b" }
+greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "c7247fbc2c6058a5cdba47182e02b4943e54e220" }
 humantime = "2.1"
 humantime-serde = "1.1"
 itertools = "0.10"
+jsonb = { git = "https://github.com/CookiePieWw/jsonb.git", rev = "d0166c130fce903bf6c58643417a3173a6172d31", default-features = false }
 lazy_static = "1.4"
 meter-core = { git = "https://github.com/GreptimeTeam/greptime-meter.git", rev = "80eb97c24c88af4dd9a86f8bbaf50e741d4eb8cd" }
 mockall = "0.11.4"

src/api/src/helper.rs

@@ -42,7 +42,8 @@ use greptime_proto::v1::greptime_request::Request;
 use greptime_proto::v1::query_request::Query;
 use greptime_proto::v1::value::ValueData;
 use greptime_proto::v1::{
-    ColumnDataTypeExtension, DdlRequest, DecimalTypeExtension, QueryRequest, Row, SemanticType,
+    ColumnDataTypeExtension, DdlRequest, DecimalTypeExtension, JsonTypeExtension, QueryRequest,
+    Row, SemanticType,
 };
 use paste::paste;
 use snafu::prelude::*;
@@ -103,7 +104,17 @@ impl From<ColumnDataTypeWrapper> for ConcreteDataType {
             ColumnDataType::Uint64 => ConcreteDataType::uint64_datatype(),
             ColumnDataType::Float32 => ConcreteDataType::float32_datatype(),
             ColumnDataType::Float64 => ConcreteDataType::float64_datatype(),
-            ColumnDataType::Binary => ConcreteDataType::binary_datatype(),
+            ColumnDataType::Binary => {
+                if let Some(TypeExt::JsonType(_)) = datatype_wrapper
+                    .datatype_ext
+                    .as_ref()
+                    .and_then(|datatype_ext| datatype_ext.type_ext.as_ref())
+                {
+                    ConcreteDataType::json_datatype()
+                } else {
+                    ConcreteDataType::binary_datatype()
+                }
+            }
             ColumnDataType::String => ConcreteDataType::string_datatype(),
             ColumnDataType::Date => ConcreteDataType::date_datatype(),
             ColumnDataType::Datetime => ConcreteDataType::datetime_datatype(),
@@ -236,7 +247,7 @@ impl TryFrom<ConcreteDataType> for ColumnDataTypeWrapper {
             ConcreteDataType::UInt64(_) => ColumnDataType::Uint64,
             ConcreteDataType::Float32(_) => ColumnDataType::Float32,
             ConcreteDataType::Float64(_) => ColumnDataType::Float64,
-            ConcreteDataType::Binary(_) => ColumnDataType::Binary,
+            ConcreteDataType::Binary(_) | ConcreteDataType::Json(_) => ColumnDataType::Binary,
             ConcreteDataType::String(_) => ColumnDataType::String,
             ConcreteDataType::Date(_) => ColumnDataType::Date,
             ConcreteDataType::DateTime(_) => ColumnDataType::Datetime,
@@ -276,6 +287,16 @@ impl TryFrom<ConcreteDataType> for ColumnDataTypeWrapper {
                         })),
                     })
             }
+            ColumnDataType::Binary => {
+                if datatype == ConcreteDataType::json_datatype() {
+                    // Json is the same as  binary in proto. The extension marks the binary in proto is actually a json.
+                    Some(ColumnDataTypeExtension {
+                        type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())),
+                    })
+                } else {
+                    None
+                }
+            }
             _ => None,
         };
         Ok(Self {
@@ -649,7 +670,8 @@ pub fn pb_values_to_vector_ref(data_type: &ConcreteDataType, values: Values) ->
         ConcreteDataType::Null(_)
         | ConcreteDataType::List(_)
         | ConcreteDataType::Dictionary(_)
-        | ConcreteDataType::Duration(_) => {
+        | ConcreteDataType::Duration(_)
+        | ConcreteDataType::Json(_) => {
             unreachable!()
         }
     }
@@ -813,7 +835,8 @@ pub fn pb_values_to_values(data_type: &ConcreteDataType, values: Values) -> Vec<
         ConcreteDataType::Null(_)
         | ConcreteDataType::List(_)
         | ConcreteDataType::Dictionary(_)
-        | ConcreteDataType::Duration(_) => {
+        | ConcreteDataType::Duration(_)
+        | ConcreteDataType::Json(_) => {
             unreachable!()
         }
     }
@@ -831,7 +854,13 @@ pub fn is_column_type_value_eq(
     expect_type: &ConcreteDataType,
 ) -> bool {
     ColumnDataTypeWrapper::try_new(type_value, type_extension)
-        .map(|wrapper| ConcreteDataType::from(wrapper) == *expect_type)
+        .map(|wrapper| {
+            let datatype = ConcreteDataType::from(wrapper);
+            (datatype == *expect_type)
+            // Json type leverage binary type in pb, so this is valid.
+                || (datatype == ConcreteDataType::binary_datatype()
+                    && *expect_type == ConcreteDataType::json_datatype())
+        })
         .unwrap_or(false)
 }
 

src/common/function/Cargo.toml

@@ -23,6 +23,7 @@ common-time.workspace = true
 common-version.workspace = true
 datafusion.workspace = true
 datatypes.workspace = true
+jsonb.workspace = true
 num = "0.4"
 num-traits = "0.2"
 once_cell.workspace = true

src/common/function/src/scalars.rs

@@ -18,6 +18,7 @@ pub mod expression;
 pub mod matches;
 pub mod math;
 pub mod numpy;
+
 #[cfg(test)]
 pub(crate) mod test;
 pub(crate) mod timestamp;

src/common/grpc/src/select.rs

@@ -70,7 +70,7 @@ macro_rules! convert_arrow_array_to_grpc_vals {
                     return Ok(vals);
                 },
             )+
-            ConcreteDataType::Null(_) | ConcreteDataType::List(_) | ConcreteDataType::Dictionary(_) | ConcreteDataType::Duration(_) => unreachable!("Should not send {:?} in gRPC", $data_type),
+            ConcreteDataType::Null(_) | ConcreteDataType::List(_) | ConcreteDataType::Dictionary(_) | ConcreteDataType::Duration(_) | ConcreteDataType::Json(_) => unreachable!("Should not send {:?} in gRPC", $data_type),
         }
     }};
 }

src/datanode/src/region_server.rs

@@ -405,7 +405,7 @@ impl FlightCraft for RegionServer {
             .trace(tracing_context.attach(info_span!("RegionServer::handle_read")))
             .await?;
 
-        let stream = Box::pin(FlightRecordBatchStream::new(result, tracing_context));
+        let stream = Box::pin(FlightRecordBatchStream::new(result, tracing_context, false));
         Ok(Response::new(stream))
     }
 }

src/datatypes/Cargo.toml

@@ -23,6 +23,7 @@ common-telemetry.workspace = true
 common-time.workspace = true
 datafusion-common.workspace = true
 enum_dispatch = "0.3"
+jsonb.workspace = true
 num = "0.4"
 num-traits = "0.2"
 ordered-float = { version = "3.0", features = ["serde"] }

src/datatypes/src/data_type.rs

@@ -33,8 +33,8 @@ use crate::types::{
     BinaryType, BooleanType, DateTimeType, DateType, Decimal128Type, DictionaryType,
     DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType,
     DurationType, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
-    IntervalDayTimeType, IntervalMonthDayNanoType, IntervalType, IntervalYearMonthType, ListType,
-    NullType, StringType, TimeMillisecondType, TimeType, TimestampMicrosecondType,
+    IntervalDayTimeType, IntervalMonthDayNanoType, IntervalType, IntervalYearMonthType, JsonType,
+    ListType, NullType, StringType, TimeMillisecondType, TimeType, TimestampMicrosecondType,
     TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, TimestampType,
     UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
@@ -81,6 +81,9 @@ pub enum ConcreteDataType {
     // Compound types:
     List(ListType),
     Dictionary(DictionaryType),
+
+    // JSON type:
+    Json(JsonType),
 }
 
 impl fmt::Display for ConcreteDataType {
@@ -128,6 +131,7 @@ impl fmt::Display for ConcreteDataType {
             ConcreteDataType::Decimal128(v) => write!(f, "{}", v.name()),
             ConcreteDataType::List(v) => write!(f, "{}", v.name()),
             ConcreteDataType::Dictionary(v) => write!(f, "{}", v.name()),
+            ConcreteDataType::Json(v) => write!(f, "{}", v.name()),
         }
     }
 }
@@ -162,6 +166,7 @@ impl ConcreteDataType {
                 | ConcreteDataType::Duration(_)
                 | ConcreteDataType::Decimal128(_)
                 | ConcreteDataType::Binary(_)
+                | ConcreteDataType::Json(_)
         )
     }
 
@@ -216,6 +221,10 @@ impl ConcreteDataType {
         matches!(self, ConcreteDataType::Decimal128(_))
     }
 
+    pub fn is_json(&self) -> bool {
+        matches!(self, ConcreteDataType::Json(_))
+    }
+
     pub fn numerics() -> Vec<ConcreteDataType> {
         vec![
             ConcreteDataType::int8_datatype(),
@@ -404,7 +413,7 @@ macro_rules! impl_new_concrete_type_functions {
 
 impl_new_concrete_type_functions!(
     Null, Boolean, UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64, Float32, Float64,
-    Binary, Date, DateTime, String
+    Binary, Date, DateTime, String, Json
 );
 
 impl ConcreteDataType {

src/datatypes/src/schema.rs

@@ -25,6 +25,7 @@ use datafusion_common::DFSchemaRef;
 use snafu::{ensure, ResultExt};
 
 use crate::error::{self, DuplicateColumnSnafu, Error, ProjectArrowSchemaSnafu, Result};
+use crate::prelude::DataType;
 pub use crate::schema::column_schema::{
     ColumnSchema, FulltextAnalyzer, FulltextOptions, Metadata, COMMENT_KEY, FULLTEXT_KEY,
     TIME_INDEX_KEY,
@@ -34,6 +35,8 @@ pub use crate::schema::raw::RawSchema;
 
 /// Key used to store version number of the schema in metadata.
 pub const VERSION_KEY: &str = "greptime:version";
+/// Key used to store actual column type in field metadata.
+pub const TYPE_KEY: &str = "greptime:type";
 
 /// A common schema, should be immutable.
 #[derive(Clone, PartialEq, Eq)]
@@ -256,7 +259,13 @@ fn collect_fields(column_schemas: &[ColumnSchema]) -> Result<FieldsAndIndices> {
         if column_schema.is_time_index() && timestamp_index.is_none() {
             timestamp_index = Some(index);
         }
-        let field = Field::try_from(column_schema)?;
+        let mut field = Field::try_from(column_schema)?;
+
+        // Json column performs the same as binary column in Arrow, so we need to mark it
+        if column_schema.data_type.is_json() {
+            let metadata = HashMap::from([(TYPE_KEY.to_string(), column_schema.data_type.name())]);
+            field = field.with_metadata(metadata);
+        }
         fields.push(field);
         ensure!(
             name_to_index

src/datatypes/src/schema/column_schema.rs

@@ -22,6 +22,7 @@ use snafu::{ensure, ResultExt};
 use crate::data_type::{ConcreteDataType, DataType};
 use crate::error::{self, Error, Result};
 use crate::schema::constraint::ColumnDefaultConstraint;
+use crate::schema::TYPE_KEY;
 use crate::value::Value;
 use crate::vectors::VectorRef;
 
@@ -34,6 +35,8 @@ pub const COMMENT_KEY: &str = "greptime:storage:comment";
 const DEFAULT_CONSTRAINT_KEY: &str = "greptime:default_constraint";
 /// Key used to store fulltext options in arrow field's metadata.
 pub const FULLTEXT_KEY: &str = "greptime:fulltext";
+/// Value used to identify column data type from arrow field's metadata. It should be the same as `ConcreteDataType::json_datatype().name()`.
+pub const JSON_TYPE_VALUE: &str = "Json";
 
 /// Schema of a column, used as an immutable struct.
 #[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -268,7 +271,14 @@ impl TryFrom<&Field> for ColumnSchema {
     type Error = Error;
 
     fn try_from(field: &Field) -> Result<ColumnSchema> {
-        let data_type = ConcreteDataType::try_from(field.data_type())?;
+        let mut data_type = ConcreteDataType::try_from(field.data_type())?;
+        // Override the data type if it is specified in the metadata.
+        if field.metadata().contains_key(TYPE_KEY) {
+            data_type = match field.metadata().get(TYPE_KEY).unwrap().as_str() {
+                JSON_TYPE_VALUE => ConcreteDataType::json_datatype(),
+                _ => data_type,
+            };
+        }
         let mut metadata = field.metadata().clone();
         let default_constraint = match metadata.remove(DEFAULT_CONSTRAINT_KEY) {
             Some(json) => {
@@ -528,4 +538,32 @@ mod tests {
         assert_eq!(formatted_int8, "test_column_1 Int8 null");
         assert_eq!(formatted_int32, "test_column_2 Int32 not null");
     }
+
+    #[test]
+    fn test_from_field_to_column_schema() {
+        let field = Field::new("test", ArrowDataType::Int32, true);
+        let column_schema = ColumnSchema::try_from(&field).unwrap();
+        assert_eq!("test", column_schema.name);
+        assert_eq!(ConcreteDataType::int32_datatype(), column_schema.data_type);
+        assert!(column_schema.is_nullable);
+        assert!(!column_schema.is_time_index);
+        assert!(column_schema.default_constraint.is_none());
+        assert!(column_schema.metadata.is_empty());
+
+        let field = Field::new("test", ArrowDataType::Binary, true);
+        let field = field.with_metadata(Metadata::from([(
+            TYPE_KEY.to_string(),
+            ConcreteDataType::json_datatype().name(),
+        )]));
+        let column_schema = ColumnSchema::try_from(&field).unwrap();
+        assert_eq!("test", column_schema.name);
+        assert_eq!(ConcreteDataType::json_datatype(), column_schema.data_type);
+        assert!(column_schema.is_nullable);
+        assert!(!column_schema.is_time_index);
+        assert!(column_schema.default_constraint.is_none());
+        assert_eq!(
+            column_schema.metadata.get(TYPE_KEY).unwrap(),
+            &ConcreteDataType::json_datatype().name()
+        );
+    }
 }