refactor: json conversion (#4893)

* refactor: json type update

* test: update test

* fix: convert when needed

* revert: leave sqlness tests unchanged

* fix: fmt

* refactor: just refactor

* Apply suggestions from code review

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* refactor: parse jsonb first

* test: add bad cases

* Update src/datatypes/src/vectors/binary.rs

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* fix: fmt

* fix: fix clippy/check

---------

Co-authored-by: Weny Xu <wenymedia@gmail.com>

src/datatypes/src/error.rs

@@ -189,6 +189,13 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Invalid JSON text: {}", value))]
+    InvalidJson {
+        value: String,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Value exceeds the precision {} bound", precision))]
     ValueExceedsPrecision {
         precision: u8,
@@ -222,7 +229,8 @@ impl ErrorExt for Error {
             | DefaultValueType { .. }
             | DuplicateMeta { .. }
             | InvalidTimestampPrecision { .. }
-            | InvalidPrecisionOrScale { .. } => StatusCode::InvalidArguments,
+            | InvalidPrecisionOrScale { .. }
+            | InvalidJson { .. } => StatusCode::InvalidArguments,
 
             ValueExceedsPrecision { .. }
             | CastType { .. }

src/datatypes/src/lib.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #![feature(let_chains)]
+#![feature(assert_matches)]
 
 pub mod arrow_array;
 pub mod data_type;

src/datatypes/src/vectors/binary.rs

@@ -36,6 +36,36 @@ impl BinaryVector {
     pub(crate) fn as_arrow(&self) -> &dyn Array {
         &self.array
     }
+
+    /// Creates a new binary vector of JSONB from a binary vector.
+    /// The binary vector must contain valid JSON strings.
+    pub fn convert_binary_to_json(&self) -> Result<BinaryVector> {
+        let arrow_array = self.to_arrow_array();
+        let mut vector = vec![];
+        for binary in arrow_array
+            .as_any()
+            .downcast_ref::<BinaryArray>()
+            .unwrap()
+            .iter()
+        {
+            let jsonb = if let Some(binary) = binary {
+                match jsonb::from_slice(binary) {
+                    Ok(jsonb) => Some(jsonb.to_vec()),
+                    Err(_) => {
+                        let s = String::from_utf8_lossy(binary);
+                        return error::InvalidJsonSnafu {
+                            value: s.to_string(),
+                        }
+                        .fail();
+                    }
+                }
+            } else {
+                None
+            };
+            vector.push(jsonb);
+        }
+        Ok(BinaryVector::from(vector))
+    }
 }
 
 impl From<BinaryArray> for BinaryVector {
@@ -233,6 +263,8 @@ vectors::impl_try_from_arrow_array_for_vector!(BinaryArray, BinaryVector);
 
 #[cfg(test)]
 mod tests {
+    use std::assert_matches::assert_matches;
+
     use arrow::datatypes::DataType as ArrowDataType;
     use common_base::bytes::Bytes;
     use serde_json;
@@ -383,4 +415,52 @@ mod tests {
         assert_eq!(b"four", vector.get_data(3).unwrap());
         assert_eq!(builder.len(), 4);
     }
+
+    #[test]
+    fn test_binary_json_conversion() {
+        // json strings
+        let json_strings = vec![
+            b"{\"hello\": \"world\"}".to_vec(),
+            b"{\"foo\": 1}".to_vec(),
+            b"123".to_vec(),
+        ];
+        let json_vector = BinaryVector::from(json_strings.clone())
+            .convert_binary_to_json()
+            .unwrap();
+        let jsonbs = json_strings
+            .iter()
+            .map(|v| jsonb::parse_value(v).unwrap().to_vec())
+            .collect::<Vec<_>>();
+        for i in 0..3 {
+            assert_eq!(
+                json_vector.get_ref(i).as_binary().unwrap().unwrap(),
+                jsonbs.get(i).unwrap().as_slice()
+            );
+        }
+
+        // jsonb
+        let json_vector = BinaryVector::from(jsonbs.clone())
+            .convert_binary_to_json()
+            .unwrap();
+        for i in 0..3 {
+            assert_eq!(
+                json_vector.get_ref(i).as_binary().unwrap().unwrap(),
+                jsonbs.get(i).unwrap().as_slice()
+            );
+        }
+
+        // binary with jsonb header (0x80, 0x40, 0x20)
+        let binary_with_jsonb_header: Vec<u8> = [0x80, 0x23, 0x40, 0x22].to_vec();
+        let error = BinaryVector::from(vec![binary_with_jsonb_header])
+            .convert_binary_to_json()
+            .unwrap_err();
+        assert_matches!(error, error::Error::InvalidJson { .. });
+
+        // invalid json string
+        let json_strings = vec![b"{\"hello\": \"world\"".to_vec()];
+        let error = BinaryVector::from(json_strings)
+            .convert_binary_to_json()
+            .unwrap_err();
+        assert_matches!(error, error::Error::InvalidJson { .. });
+    }
 }

src/datatypes/src/vectors/operations.rs

@@ -18,6 +18,8 @@ mod find_unique;
 mod replicate;
 mod take;
 
+use std::sync::Arc;
+
 use common_base::BitVec;
 
 use crate::error::{self, Result};
@@ -89,6 +91,12 @@ macro_rules! impl_scalar_vector_op {
             }
 
             fn cast(&self, to_type: &ConcreteDataType) -> Result<VectorRef> {
+                if to_type == &ConcreteDataType::json_datatype() {
+                    if let Some(vector) = self.as_any().downcast_ref::<BinaryVector>() {
+                        let json_vector = vector.convert_binary_to_json()?;
+                        return Ok(Arc::new(json_vector) as VectorRef);
+                    }
+                }
                 cast::cast_non_constant!(self, to_type)
             }
 

src/servers/src/postgres/types.rs

@@ -961,7 +961,7 @@ pub(super) fn parameters_to_scalar_values(
                 if let Some(server_type) = &server_type {
                     match server_type {
                         ConcreteDataType::Binary(_) => {
-                            ScalarValue::Binary(data.map(|d| jsonb::Value::from(d).to_vec()))
+                            ScalarValue::Binary(data.map(|d| d.to_string().into_bytes()))
                         }
                         _ => {
                             return Err(invalid_parameter_error(
@@ -971,7 +971,7 @@ pub(super) fn parameters_to_scalar_values(
                         }
                     }
                 } else {
-                    ScalarValue::Binary(data.map(|d| jsonb::Value::from(d).to_vec()))
+                    ScalarValue::Binary(data.map(|d| d.to_string().into_bytes()))
                 }
             }
             _ => Err(invalid_parameter_error(