feat(json2): encode json2 variant payloads as jsonb (#8435)

* feat: encode json2 variant payloads as jsonb

* feat: encode json2 variant payloads as jsonb

* minor refactor

* fix: cargo check

* fix: unit test

* fix: cr

src/datatypes/src/error.rs

@@ -234,18 +234,21 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
     #[snafu(display("Invalid fulltext option: {}", msg))]
     InvalidFulltextOption {
         msg: String,
         #[snafu(implicit)]
         location: Location,
     },
+
     #[snafu(display("Invalid skipping index option: {}", msg))]
     InvalidSkippingIndexOption {
         msg: String,
         #[snafu(implicit)]
         location: Location,
     },
+
     #[snafu(display("Inconsistent struct field count {field_len} and item count {item_len}"))]
     InconsistentStructFieldsAndItems {
         field_len: usize,
@@ -253,6 +256,7 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
     #[snafu(display("Failed to process JSONB value"))]
     InvalidJsonb {
         error: jsonb::Error,

src/datatypes/src/json/value.rs

@@ -21,11 +21,11 @@ use num_traits::ToPrimitive;
 use ordered_float::OrderedFloat;
 use serde::{Deserialize, Serialize};
 use serde_json::Number;
-use snafu::{OptionExt, ResultExt, ensure};
+use snafu::{OptionExt, ensure};
 
 use crate::Result;
 use crate::data_type::ConcreteDataType;
-use crate::error::{AlignJsonValueSnafu, SerializeSnafu};
+use crate::error::{AlignJsonValueSnafu, InvalidJsonSnafu, InvalidJsonbSnafu};
 use crate::types::json_type::{JsonNativeType, JsonNumberType};
 use crate::types::{JsonType, StructField, StructType};
 use crate::value::{ListValue, ListValueRef, StructValue, StructValueRef, Value, ValueRef};
@@ -283,13 +283,10 @@ impl Display for JsonVariant {
                         .join(", ")
                 )
             }
-            Self::Variant(x) => {
-                let result: serde_json::Result<serde_json::Value> = serde_json::from_slice(x);
-                match result {
-                    Ok(v) => write!(f, "{v}"),
-                    Err(_) => write!(f, "{x:?}"),
-                }
-            }
+            Self::Variant(x) => match decode_json_variant(x) {
+                Ok(v) => write!(f, "{v}"),
+                Err(_) => write!(f, "{x:?}"),
+            },
         }
     }
 }
@@ -483,13 +480,7 @@ impl JsonValue {
                     JsonVariant::Object(object)
                 }
 
-                (v, JsonNativeType::Variant) => {
-                    let json: serde_json::Value =
-                        JsonValue::new(v).try_into().context(SerializeSnafu)?;
-                    serde_json::to_vec(&json)
-                        .map(JsonVariant::Variant)
-                        .context(SerializeSnafu)?
-                }
+                (v, JsonNativeType::Variant) => JsonVariant::Variant(encode_json_variant(v)?),
 
                 (value, expected) => {
                     return AlignJsonValueSnafu {
@@ -555,13 +546,85 @@ impl TryFrom<JsonValue> for serde_json::Value {
                     }
                     serde_json::Value::Object(map)
                 }
-                JsonVariant::Variant(x) => serde_json::from_slice(&x)?,
+                JsonVariant::Variant(x) => decode_json_variant(&x).map_err(|err| {
+                    serde_json::Error::io(std::io::Error::new(
+                        std::io::ErrorKind::InvalidData,
+                        err.to_string(),
+                    ))
+                })?,
             })
         }
         helper(v.json_variant)
     }
 }
 
+pub(crate) fn encode_json_variant(value: JsonVariant) -> Result<Vec<u8>> {
+    match value {
+        JsonVariant::Variant(bytes) => Ok(bytes),
+        value => jsonb::Value::try_from(value).map(|value| value.to_vec()),
+    }
+}
+
+pub(crate) fn decode_json_variant(
+    bytes: &[u8],
+) -> std::result::Result<serde_json::Value, jsonb::Error> {
+    jsonb::from_slice(bytes).map(Into::into)
+}
+
+pub(crate) fn encode_serde_json_as_jsonb(value: serde_json::Value) -> Vec<u8> {
+    jsonb::Value::from(value).to_vec()
+}
+
+impl TryFrom<JsonVariant> for jsonb::Value<'static> {
+    type Error = crate::Error;
+
+    fn try_from(value: JsonVariant) -> Result<Self> {
+        Ok(match value {
+            JsonVariant::Null => jsonb::Value::Null,
+            JsonVariant::Bool(value) => jsonb::Value::Bool(value),
+            JsonVariant::Number(value) => jsonb::Value::Number(value.try_into()?),
+            JsonVariant::String(value) => jsonb::Value::String(value.into()),
+            JsonVariant::Array(values) => jsonb::Value::Array(
+                values
+                    .into_iter()
+                    .map(jsonb::Value::try_from)
+                    .collect::<Result<_>>()?,
+            ),
+            JsonVariant::Object(values) => jsonb::Value::Object(
+                values
+                    .into_iter()
+                    .map(|(key, value)| jsonb::Value::try_from(value).map(|value| (key, value)))
+                    .collect::<Result<_>>()?,
+            ),
+            JsonVariant::Variant(value) => {
+                let value = decode_json_variant(&value)
+                    .map_err(|error| InvalidJsonbSnafu { error }.build())?;
+                jsonb::Value::from(value)
+            }
+        })
+    }
+}
+
+impl TryFrom<JsonNumber> for jsonb::Number {
+    type Error = crate::Error;
+
+    fn try_from(value: JsonNumber) -> Result<Self> {
+        Ok(match value {
+            JsonNumber::PosInt(value) => jsonb::Number::UInt64(value),
+            JsonNumber::NegInt(value) => jsonb::Number::Int64(value),
+            JsonNumber::Float(value) => {
+                ensure!(
+                    !value.0.is_nan(),
+                    InvalidJsonSnafu {
+                        value: "NaN is not a valid JSON number"
+                    }
+                );
+                jsonb::Number::Float64(value.0)
+            }
+        })
+    }
+}
+
 impl Clone for JsonValue {
     fn clone(&self) -> Self {
         let Self {
@@ -890,6 +953,11 @@ mod tests {
     use super::*;
     use crate::types::json_type::JsonObjectType;
 
+    fn jsonb_bytes(json: &str) -> Vec<u8> {
+        let value: serde_json::Value = serde_json::from_str(json).unwrap();
+        encode_json_variant(value.into()).unwrap()
+    }
+
     #[test]
     fn test_align_json_value() -> Result<()> {
         fn parse_json_value(json: &str) -> JsonValue {
@@ -935,7 +1003,7 @@ mod tests {
                         ("note".to_string(), JsonVariant::Null),
                         (
                             "payload".to_string(),
-                            JsonVariant::Variant(br#"{"k":"v"}"#.to_vec()),
+                            JsonVariant::Variant(jsonb_bytes(r#"{"k":"v"}"#)),
                         ),
                     ]))]),
                 ),
@@ -964,7 +1032,9 @@ mod tests {
         value.try_align(&JsonType::new_json2(JsonNativeType::Variant))?;
         assert_eq!(
             value,
-            JsonValue::from(JsonVariant::Variant(br#"{"foo":[1,true,null]}"#.to_vec()))
+            JsonValue::from(JsonVariant::Variant(jsonb_bytes(
+                r#"{"foo":[1,true,null]}"#
+            )))
         );
 
         // Incompatible scalar alignment should fail instead of coercing the value.
@@ -977,6 +1047,15 @@ mod tests {
             r#"Failed to align JSON value, reason: unable to align 'hello' of type "<String>" to type "<Bool>""#
         );
 
+        let mut value = JsonValue::from(f64::NAN);
+        let err = value
+            .try_align(&JsonType::new_json2(JsonNativeType::Variant))
+            .unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Invalid JSON: NaN is not a valid JSON number"
+        );
+
         Ok(())
     }
 }

src/datatypes/src/vectors/json/array.rs

@@ -32,9 +32,10 @@ use crate::arrow_array::{
     MutableBinaryArray, StringViewArray, binary_array_value, string_array_value,
 };
 use crate::error::{
-    AlignJsonArraySnafu, ArrowComputeSnafu, CastTypeSnafu, DeserializeSnafu, InvalidJsonSnafu,
-    Result, SerializeSnafu,
+    AlignJsonArraySnafu, ArrowComputeSnafu, CastTypeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu,
+    Result,
 };
+use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb};
 
 pub struct JsonArray<'a> {
     inner: &'a ArrayRef,
@@ -59,9 +60,7 @@ impl JsonArray<'_> {
             }
             DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
                 let bytes = binary_array_value(array, i);
-                serde_json::from_slice(bytes).with_context(|_| DeserializeSnafu {
-                    json: String::from_utf8_lossy(bytes),
-                })?
+                decode_json_variant(bytes).map_err(|error| InvalidJsonbSnafu { error }.build())?
             }
             DataType::Struct(_) => {
                 let structs = array.as_struct();
@@ -230,7 +229,7 @@ impl JsonArray<'_> {
             if value.is_null() {
                 encoded.push(None);
             } else {
-                let bytes = serde_json::to_vec(&value).context(SerializeSnafu)?;
+                let bytes = encode_serde_json_as_jsonb(value);
                 total_bytes += bytes.len();
                 encoded.push(Some(bytes));
             }
@@ -368,10 +367,12 @@ mod test {
         assert_eq!(JsonArray::from(&strings).try_get_value(0)?, json!("hello"));
         assert_eq!(JsonArray::from(&strings).try_get_value(1)?, Value::Null);
 
-        let binaries: ArrayRef = Arc::new(BinaryArray::from(vec![
-            br#"{"nested":[1,null,"x"]}"#.as_slice(),
-            b"null".as_slice(),
-        ]));
+        let nested = jsonb::parse_value(br#"{"nested":[1,null,"x"]}"#)
+            .unwrap()
+            .to_vec();
+        let null = jsonb::parse_value(b"null").unwrap().to_vec();
+        let binaries: ArrayRef =
+            Arc::new(BinaryArray::from(vec![nested.as_slice(), null.as_slice()]));
         assert_eq!(
             JsonArray::from(&binaries).try_get_value(0)?,
             json!({"nested": [1, null, "x"]})

src/datatypes/src/vectors/json/builder.rs

@@ -138,6 +138,10 @@ mod tests {
             Value::Json(Box::new(value.into()))
         }
 
+        fn jsonb_bytes(json: &str) -> Bytes {
+            Bytes::from(jsonb::parse_value(json.as_bytes()).unwrap().to_vec())
+        }
+
         // Object inputs should merge into a superset schema, preserve null rows,
         // and align conflicting nested values into Variant payloads.
         let mut builder = JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 3);
@@ -166,7 +170,7 @@ mod tests {
                 vec![
                     Value::Null,
                     Value::Int64(1),
-                    Value::Binary(Bytes::from(br#"{"name":"foo"}"#.to_vec())),
+                    Value::Binary(jsonb_bytes(r#"{"name":"foo"}"#)),
                 ],
                 merged_struct_type.clone(),
             ))
@@ -178,7 +182,7 @@ mod tests {
                 vec![
                     Value::Boolean(true),
                     Value::Int64(2),
-                    Value::Binary(Bytes::from(br#""raw""#.to_vec())),
+                    Value::Binary(jsonb_bytes(r#""raw""#)),
                 ],
                 merged_struct_type,
             ))