refactor(query): refactor json functions (#16840)

Author: b41sh

Cargo.lock

@@ -633,6 +633,7 @@ backtrace = { git = "https://github.com/rust-lang/backtrace-rs.git", rev = "7226
 color-eyre = { git = "https://github.com/eyre-rs/eyre.git", rev = "e5d92c3" }
 deltalake = { git = "https://github.com/delta-io/delta-rs", rev = "c149502" }
 ethnum = { git = "https://github.com/datafuse-extras/ethnum-rs", rev = "4cb05f1" }
+jsonb = { git = "https://github.com/databendlabs/jsonb", rev = "37d07f0" }
 map-api = { git = "https://github.com/databendlabs/map-api", tag = "v0.2.3" }
 openai_api_rust = { git = "https://github.com/datafuse-extras/openai-api", rev = "819a0ed" }
 openraft = { git = "https://github.com/databendlabs/openraft", tag = "v0.10.0-alpha.9" }

Cargo.toml

@@ -53,8 +53,31 @@ pub fn new_test_column() -> Vec<Column> {
         Column::Binary(BinaryColumn::from_iter(
             ["abcdefg", "mn", "11", "", "3456", "xyz"].iter(),
         )),
+        // use binary jsonb format values to test
         Column::Variant(BinaryColumn::from_iter(
-            ["abcdefg", "mn", "11", "", "3456", "xyz"].iter(),
+            [
+                // null
+                vec![0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
+                // "abc"
+                vec![
+                    0x20, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x03, 0x61, 0x62, 0x63,
+                ],
+                // 123
+                vec![0x20, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x02, 0x50, 0x7B],
+                // []
+                vec![0x80, 0x00, 0x00, 0x00],
+                // [1,2]
+                vec![
+                    0x80, 0x00, 0x00, 0x02, 0x20, 0x00, 0x00, 0x02, 0x20, 0x00, 0x00, 0x02, 0x50,
+                    0x01, 0x50, 0x02,
+                ],
+                // {"k":"v"}
+                vec![
+                    0x40, 0x00, 0x00, 0x01, 0x10, 0x00, 0x00, 0x01, 0x10, 0x00, 0x00, 0x01, 0x6B,
+                    0x76,
+                ],
+            ]
+            .iter(),
         )),
     ]
 }

src/common/native/tests/it/native/io.rs

@@ -20,6 +20,8 @@ use databend_common_io::deserialize_bitmap;
 use geozero::wkb::Ewkb;
 use geozero::ToJson;
 use jiff::tz::TimeZone;
+use jsonb::OwnedJsonb;
+use jsonb::RawJsonb;
 use jsonb::Value;
 
 use super::binary::BinaryColumn;
@@ -182,7 +184,9 @@ impl ValueType for VariantType {
 
     #[inline(always)]
     fn compare(lhs: Self::ScalarRef<'_>, rhs: Self::ScalarRef<'_>) -> Ordering {
-        jsonb::compare(lhs, rhs).expect("unable to parse jsonb value")
+        let left_jsonb = RawJsonb::new(lhs);
+        let right_jsonb = RawJsonb::new(rhs);
+        left_jsonb.cmp(&right_jsonb)
     }
 }
 
@@ -235,7 +239,9 @@ pub fn cast_scalar_to_variant(scalar: ScalarRef, tz: &TimeZone, buf: &mut Vec<u8
         ScalarRef::Interval(i) => interval_to_string(&i).to_string().into(),
         ScalarRef::Array(col) => {
             let items = cast_scalars_to_variants(col.iter(), tz);
-            jsonb::build_array(items.iter(), buf).expect("failed to build jsonb array");
+            let owned_jsonb = OwnedJsonb::build_array(items.iter().map(RawJsonb::new))
+                .expect("failed to build jsonb array");
+            buf.extend_from_slice(owned_jsonb.as_ref());
             return;
         }
         ScalarRef::Map(col) => {
@@ -255,8 +261,10 @@ pub fn cast_scalar_to_variant(scalar: ScalarRef, tz: &TimeZone, buf: &mut Vec<u8
                 cast_scalar_to_variant(v, tz, &mut val);
                 kvs.insert(key, val);
             }
-            jsonb::build_object(kvs.iter().map(|(k, v)| (k, &v[..])), buf)
-                .expect("failed to build jsonb object from map");
+            let owned_jsonb =
+                OwnedJsonb::build_object(kvs.iter().map(|(k, v)| (k, RawJsonb::new(&v[..]))))
+                    .expect("failed to build jsonb object from map");
+            buf.extend_from_slice(owned_jsonb.as_ref());
             return;
         }
         ScalarRef::Bitmap(b) => {
@@ -272,14 +280,14 @@ pub fn cast_scalar_to_variant(scalar: ScalarRef, tz: &TimeZone, buf: &mut Vec<u8
         }
         ScalarRef::Tuple(fields) => {
             let values = cast_scalars_to_variants(fields, tz);
-            jsonb::build_object(
+            let owned_jsonb = OwnedJsonb::build_object(
                 values
                     .iter()
                     .enumerate()
-                    .map(|(i, bytes)| (format!("{}", i + 1), bytes)),
-                buf,
+                    .map(|(i, bytes)| (format!("{}", i + 1), RawJsonb::new(bytes))),
             )
             .expect("failed to build jsonb object from tuple");
+            buf.extend_from_slice(owned_jsonb.as_ref());
             return;
         }
         ScalarRef::Variant(bytes) => {

src/query/expression/src/types/variant.rs

@@ -28,6 +28,7 @@ use databend_common_io::geo_to_ewkt;
 use geozero::wkb::Ewkb;
 use itertools::Itertools;
 use jiff::tz::TimeZone;
+use jsonb::RawJsonb;
 use num_traits::FromPrimitive;
 use rust_decimal::Decimal;
 use rust_decimal::RoundingStrategy;
@@ -262,7 +263,8 @@ impl Display for ScalarRef<'_> {
                 write!(f, ")")
             }
             ScalarRef::Variant(s) => {
-                let value = jsonb::to_string(s);
+                let raw_jsonb = RawJsonb::new(s);
+                let value = raw_jsonb.to_string();
                 write!(f, "'{value}'")
             }
             ScalarRef::Geometry(s) => {

src/query/expression/src/utils/display.rs

@@ -15,7 +15,7 @@
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use jsonb::parse_value;
-use jsonb::to_string;
+use jsonb::RawJsonb;
 
 use crate::types::AnyType;
 use crate::types::DataType;
@@ -96,7 +96,8 @@ fn transform_scalar(scalar: ScalarRef<'_>, decode: bool) -> Result<Scalar> {
         }
         ScalarRef::Variant(data) => {
             if decode {
-                Scalar::Variant(to_string(data).into_bytes())
+                let raw_jsonb = RawJsonb::new(data);
+                Scalar::Variant(raw_jsonb.to_string().into_bytes())
             } else {
                 let value = parse_value(data).map_err(|err| {
                     ErrorCode::UDFDataError(format!("parse json value error: {err}"))

src/query/expression/src/utils/variant_transform.rs

@@ -39,6 +39,7 @@ use geo::Point;
 use geozero::CoordDimensions;
 use geozero::ToWkb;
 use itertools::Itertools;
+use jsonb::RawJsonb;
 use roaring::RoaringTreemap;
 use serde::de::Visitor;
 use serde::Deserialize;
@@ -776,7 +777,9 @@ impl PartialOrd for Scalar {
             (Scalar::Bitmap(b1), Scalar::Bitmap(b2)) => b1.partial_cmp(b2),
             (Scalar::Tuple(t1), Scalar::Tuple(t2)) => t1.partial_cmp(t2),
             (Scalar::Variant(v1), Scalar::Variant(v2)) => {
-                jsonb::compare(v1.as_slice(), v2.as_slice()).ok()
+                let left_jsonb = RawJsonb::new(v1);
+                let right_jsonb = RawJsonb::new(v2);
+                left_jsonb.partial_cmp(&right_jsonb)
             }
             (Scalar::Geometry(g1), Scalar::Geometry(g2)) => compare_geometry(g1, g2),
             (Scalar::Geography(g1), Scalar::Geography(g2)) => g1.partial_cmp(g2),
@@ -814,7 +817,11 @@ impl<'b> PartialOrd<ScalarRef<'b>> for ScalarRef<'_> {
             (ScalarRef::Map(m1), ScalarRef::Map(m2)) => m1.partial_cmp(m2),
             (ScalarRef::Bitmap(b1), ScalarRef::Bitmap(b2)) => b1.partial_cmp(b2),
             (ScalarRef::Tuple(t1), ScalarRef::Tuple(t2)) => t1.partial_cmp(t2),
-            (ScalarRef::Variant(v1), ScalarRef::Variant(v2)) => jsonb::compare(v1, v2).ok(),
+            (ScalarRef::Variant(v1), ScalarRef::Variant(v2)) => {
+                let left_jsonb = RawJsonb::new(v1);
+                let right_jsonb = RawJsonb::new(v2);
+                left_jsonb.partial_cmp(&right_jsonb)
+            }
             (ScalarRef::Geometry(g1), ScalarRef::Geometry(g2)) => compare_geometry(g1, g2),
             (ScalarRef::Geography(g1), ScalarRef::Geography(g2)) => g1.partial_cmp(g2),
             (ScalarRef::Interval(i1), ScalarRef::Interval(i2)) => i1.partial_cmp(i2),
@@ -913,9 +920,13 @@ impl PartialOrd for Column {
                 col1.iter().partial_cmp(col2.iter())
             }
             (Column::Tuple(fields1), Column::Tuple(fields2)) => fields1.partial_cmp(fields2),
-            (Column::Variant(col1), Column::Variant(col2)) => col1
-                .iter()
-                .partial_cmp_by(col2.iter(), |v1, v2| jsonb::compare(v1, v2).ok()),
+            (Column::Variant(col1), Column::Variant(col2)) => {
+                col1.iter().partial_cmp_by(col2.iter(), |v1, v2| {
+                    let left_jsonb = RawJsonb::new(v1);
+                    let right_jsonb = RawJsonb::new(v2);
+                    left_jsonb.partial_cmp(&right_jsonb)
+                })
+            }
             (Column::Geometry(col1), Column::Geometry(col2)) => {
                 col1.iter().partial_cmp_by(col2.iter(), compare_geometry)
             }
@@ -1423,7 +1434,7 @@ impl Column {
             DataType::Variant => {
                 let mut data = Vec::with_capacity(len);
                 for _ in 0..len {
-                    let val = jsonb::rand_value();
+                    let val = jsonb::Value::rand_value();
                     data.push(val.to_vec());
                 }
                 VariantType::from_data(data)

src/query/expression/src/values.rs

@@ -28,8 +28,8 @@ use databend_common_expression::RowConverter;
 use databend_common_expression::SortField;
 use ethnum::i256;
 use itertools::Itertools;
-use jsonb::convert_to_comparable;
 use jsonb::parse_value;
+use jsonb::RawJsonb;
 use rand::distributions::Alphanumeric;
 use rand::distributions::Standard;
 use rand::prelude::Distribution;
@@ -375,7 +375,9 @@ fn test_variant() {
             validity.push(true);
             let val = parse_value(value.as_bytes()).unwrap();
             let buf = val.to_vec();
-            convert_to_comparable(&buf, &mut builder.data);
+            let raw_jsonb = RawJsonb::new(&buf);
+            let compare_buf = raw_jsonb.convert_to_comparable();
+            builder.put_slice(&compare_buf);
         } else {
             validity.push(false);
         }

src/query/expression/tests/it/row.rs

@@ -21,6 +21,7 @@ use databend_common_io::constants::NULL_BYTES_LOWER;
 use databend_common_io::constants::TRUE_BYTES_LOWER;
 use geozero::wkb::Ewkb;
 use geozero::ToJson;
+use jsonb::RawJsonb;
 
 use crate::field_encoder::helpers::write_json_string;
 use crate::field_encoder::FieldEncoderValues;
@@ -81,7 +82,7 @@ impl FieldEncoderJSON {
 
             Column::Variant(c) => {
                 let v = unsafe { c.index_unchecked(row_index) };
-                out_buf.extend_from_slice(jsonb::to_string(v).as_bytes());
+                out_buf.extend_from_slice(RawJsonb::new(v).to_string().as_bytes());
             }
             Column::Geometry(c) => {
                 let v = unsafe { c.index_unchecked(row_index) };

src/query/formats/src/field_encoder/json.rs

@@ -45,6 +45,7 @@ use databend_common_io::geo_to_wkt;
 use databend_common_io::GeometryDataType;
 use geozero::wkb::Ewkb;
 use jiff::tz::TimeZone;
+use jsonb::RawJsonb;
 use lexical_core::ToLexical;
 use micromarshal::Marshal;
 use micromarshal::Unmarshal;
@@ -322,7 +323,7 @@ impl FieldEncoderValues {
         in_nested: bool,
     ) {
         let v = unsafe { column.index_unchecked(row_index) };
-        let s = jsonb::to_string(v);
+        let s = RawJsonb::new(v).to_string();
         self.write_string_inner(s.as_bytes(), out_buf, in_nested);
     }