From a7d564124da647790854003e3305b4bc7df53d90 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Wed, 8 Jul 2026 19:27:42 +0800 Subject: [PATCH] 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 | 4 + src/datatypes/src/json/value.rs | 117 ++++++++++++++++++---- src/datatypes/src/vectors/json/array.rs | 21 ++-- src/datatypes/src/vectors/json/builder.rs | 8 +- 4 files changed, 119 insertions(+), 31 deletions(-) diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index e897e984df..9136270d1c 100644 --- a/src/datatypes/src/error.rs +++ b/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, diff --git a/src/datatypes/src/json/value.rs b/src/datatypes/src/json/value.rs index b660894c2b..9e41dcd8e9 100644 --- a/src/datatypes/src/json/value.rs +++ b/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::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 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> { + 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 { + jsonb::from_slice(bytes).map(Into::into) +} + +pub(crate) fn encode_serde_json_as_jsonb(value: serde_json::Value) -> Vec { + jsonb::Value::from(value).to_vec() +} + +impl TryFrom for jsonb::Value<'static> { + type Error = crate::Error; + + fn try_from(value: JsonVariant) -> Result { + 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::>()?, + ), + JsonVariant::Object(values) => jsonb::Value::Object( + values + .into_iter() + .map(|(key, value)| jsonb::Value::try_from(value).map(|value| (key, value))) + .collect::>()?, + ), + JsonVariant::Variant(value) => { + let value = decode_json_variant(&value) + .map_err(|error| InvalidJsonbSnafu { error }.build())?; + jsonb::Value::from(value) + } + }) + } +} + +impl TryFrom for jsonb::Number { + type Error = crate::Error; + + fn try_from(value: JsonNumber) -> Result { + 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 { + 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 "" to type """# ); + 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(()) } } diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 37174adbda..3e58af7019 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/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"]}) diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index c3a7474a80..e0226af409 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/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, ))