feat: add json data type

Cargo.toml

@@ -124,6 +124,7 @@ greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", r
 humantime = "2.1"
 humantime-serde = "1.1"
 itertools = "0.10"
+jsonb = { git = "https://github.com/CookiePieWw/jsonb.git", branch = "trigger-preserve-order-feature" }
 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

@@ -236,7 +236,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,
@@ -443,7 +443,7 @@ pub fn push_vals(column: &mut Column, origin_count: usize, vector: VectorRef) {
                 .push(convert_i128_to_interval(val.to_i128())),
         },
         Value::Decimal128(val) => values.decimal128_values.push(convert_to_pb_decimal128(val)),
-        Value::List(_) | Value::Duration(_) => unreachable!(),
+        Value::List(_) | Value::Duration(_) | Value::Json(_) => unreachable!(),
     });
     column.null_mask = null_mask.into_vec();
 }
@@ -649,7 +649,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 +814,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 +833,12 @@ 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)
+                || (datatype == ConcreteDataType::binary_datatype()
+                    && *expect_type == ConcreteDataType::json_datatype())
+        })
         .unwrap_or(false)
 }
 
@@ -928,6 +935,9 @@ pub fn to_proto_value(value: Value) -> Option<v1::Value> {
         Value::Decimal128(v) => v1::Value {
             value_data: Some(ValueData::Decimal128Value(convert_to_pb_decimal128(v))),
         },
+        Value::Json(v) => v1::Value {
+            value_data: Some(ValueData::BinaryValue(v.to_vec())),
+        },
         Value::List(_) | Value::Duration(_) => return None,
     };
 
@@ -1023,6 +1033,7 @@ pub fn value_to_grpc_value(value: Value) -> GrpcValue {
                 }
             }),
             Value::Decimal128(v) => Some(ValueData::Decimal128Value(convert_to_pb_decimal128(v))),
+            Value::Json(v) => Some(ValueData::BinaryValue(v.to_vec())),
             Value::List(_) | Value::Duration(_) => unreachable!(),
         },
     }

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/function_registry.rs

@@ -22,6 +22,7 @@ use crate::function::{AsyncFunctionRef, FunctionRef};
 use crate::scalars::aggregate::{AggregateFunctionMetaRef, AggregateFunctions};
 use crate::scalars::date::DateFunction;
 use crate::scalars::expression::ExpressionFunction;
+use crate::scalars::json::JsonFunction;
 use crate::scalars::matches::MatchesFunction;
 use crate::scalars::math::MathFunction;
 use crate::scalars::numpy::NumpyFunction;
@@ -116,6 +117,9 @@ pub static FUNCTION_REGISTRY: Lazy<Arc<FunctionRegistry>> = Lazy::new(|| {
     SystemFunction::register(&function_registry);
     TableFunction::register(&function_registry);
 
+    // JSON functions
+    JsonFunction::register(&function_registry);
+
     Arc::new(function_registry)
 });
 

src/common/function/src/scalars.rs

@@ -15,9 +15,11 @@
 pub mod aggregate;
 pub(crate) mod date;
 pub mod expression;
+pub mod json;
 pub mod matches;
 pub mod math;
 pub mod numpy;
+
 #[cfg(test)]
 pub(crate) mod test;
 pub(crate) mod timestamp;

src/common/function/src/scalars/json.rs

@@ -0,0 +1,28 @@
+// Copyright 2023 Greptime Team
+//
+// 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::sync::Arc;
+mod jq;
+
+use jq::JqFunction;
+
+use crate::function_registry::FunctionRegistry;
+
+pub(crate) struct JsonFunction;
+
+impl JsonFunction {
+    pub fn register(registry: &FunctionRegistry) {
+        registry.register(Arc::new(JqFunction));
+    }
+}

src/common/function/src/scalars/json/jq.rs

@@ -0,0 +1,174 @@
+// Copyright 2023 Greptime Team
+//
+// 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::fmt::{self, Display};
+
+use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
+use common_query::prelude::Signature;
+use datafusion::logical_expr::Volatility;
+use datatypes::data_type::ConcreteDataType;
+use datatypes::prelude::VectorRef;
+use datatypes::scalars::ScalarVectorBuilder;
+use datatypes::vectors::{MutableVector, StringVectorBuilder};
+
+use crate::function::{Function, FunctionContext};
+
+#[derive(Clone, Debug, Default)]
+pub struct JqFunction;
+
+const NAME: &str = "jq";
+
+impl Function for JqFunction {
+    fn name(&self) -> &str {
+        "jq"
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::exact(
+            vec![
+                ConcreteDataType::string_datatype(),
+                ConcreteDataType::json_datatype(),
+            ],
+            Volatility::Immutable,
+        )
+    }
+
+    fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        let jsons = &columns[1];
+        let paths = &columns[0];
+
+        let size = jsons.len();
+        let datatype = jsons.data_type();
+        let mut results = StringVectorBuilder::with_capacity(size);
+
+        match datatype {
+            ConcreteDataType::Binary(_) | ConcreteDataType::Json(_) => {
+                for i in 0..size {
+                    let json = jsons.get_ref(i);
+                    let json = json.as_binary().unwrap();
+                    let path = paths.get_ref(i);
+                    let path = path.as_string().unwrap();
+                    let result = match (json, path) {
+                        (Some(json), Some(path)) => {
+                            let json_path = match jsonb::jsonpath::parse_json_path(path.as_bytes())
+                            {
+                                Ok(json_path) => json_path,
+                                Err(_) => {
+                                    return InvalidFuncArgsSnafu {
+                                        err_msg: format!("Invalid JSON path: {}", path),
+                                    }
+                                    .fail();
+                                }
+                            };
+                            let mut sub_jsonb = Vec::new();
+                            let mut sub_offsets = Vec::new();
+                            match jsonb::get_by_path(
+                                json,
+                                json_path,
+                                &mut sub_jsonb,
+                                &mut sub_offsets,
+                            ) {
+                                Ok(_) => Some(jsonb::to_string(&sub_jsonb)),
+                                Err(_) => None,
+                            }
+                        }
+                        _ => None,
+                    };
+
+                    results.push(result.as_deref());
+                }
+            }
+            _ => {
+                return UnsupportedInputDataTypeSnafu {
+                    function: NAME,
+                    datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
+                }
+                .fail();
+            }
+        }
+
+        Ok(results.to_vector())
+    }
+}
+
+impl Display for JqFunction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "JQ")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use common_query::prelude::TypeSignature;
+    use datatypes::scalars::ScalarVector;
+    use datatypes::vectors::{BinaryVector, StringVector};
+
+    use super::*;
+
+    #[test]
+    fn test_jq_function() {
+        let jq = JqFunction;
+
+        assert_eq!("jq", jq.name());
+        assert_eq!(
+            ConcreteDataType::string_datatype(),
+            jq.return_type(&[
+                ConcreteDataType::json_datatype(),
+                ConcreteDataType::string_datatype()
+            ])
+            .unwrap()
+        );
+
+        assert!(matches!(jq.signature(),
+                         Signature {
+                             type_signature: TypeSignature::Exact(valid_types),
+                             volatility: Volatility::Immutable
+                         } if  valid_types == vec![ConcreteDataType::json_datatype(), ConcreteDataType::string_datatype()]
+        ));
+
+        let json_strings = [
+            r#"{"a": {"b": 2}, "b": 2, "c": 3}"#,
+            r#"{"a": 4, "b": {"c": 6}, "c": 6}"#,
+            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
+        ];
+        let paths = vec!["a", "b", "c"];
+        let results = [r#"{"b":2}"#, r#"{"c":6}"#, r#"{"a":7}"#];
+
+        let jsonbs = json_strings
+            .iter()
+            .map(|s| {
+                let value = jsonb::parse_value(s.as_bytes()).unwrap();
+                value.to_vec()
+            })
+            .collect::<Vec<_>>();
+
+        let json_vector = BinaryVector::from_vec(jsonbs);
+        let path_vector = StringVector::from_vec(paths);
+        let args: Vec<VectorRef> = vec![Arc::new(json_vector), Arc::new(path_vector)];
+        let vector = jq.eval(FunctionContext::default(), &args).unwrap();
+
+        assert_eq!(3, vector.len());
+        for (i, gt) in results.iter().enumerate() {
+            let result = vector.get_ref(i);
+            let result = result.as_string().unwrap().unwrap();
+            assert_eq!(*gt, result);
+        }
+    }
+}

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),
         }
     }};
 }