From ed8e43e109afee6ca2a72444f3e377b87b06f197 Mon Sep 17 00:00:00 2001 From: luofucong Date: Thu, 16 Jul 2026 14:55:44 +0800 Subject: [PATCH] x3 Signed-off-by: luofucong --- src/datatypes/src/json.rs | 152 ++++++++++-------- .../src/etl/transform/transformer/greptime.rs | 67 +++++++- 2 files changed, 146 insertions(+), 73 deletions(-) diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 54bd4fe97a2..710ae77babe 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -30,7 +30,7 @@ use snafu::ResultExt; use crate::data_type::ConcreteDataType; use crate::error::{self, Result, UnsupportedJsonTypeSnafu}; -use crate::json::value::{JsonValue, JsonVariant}; +use crate::json::value::{JsonNumber, JsonValue, JsonVariant}; use crate::schema::ColumnDefaultConstraint; use crate::types::json_type::JsonNativeType; use crate::value::{ListValue, StructValue, Value}; @@ -91,11 +91,16 @@ impl JsonSettings { /// Encode a serde_json::Value into a Value::Json using current settings. pub fn encode(&self, json: Json) -> Result { + self.encode_variant(json.into()) + } + + /// Encode a JsonVariant into a Value::Json using current settings. + pub fn encode_variant(&self, json: JsonVariant) -> Result { let mut context = JsonContext { path: Vec::new(), settings: self, }; - encode_json_with_context(json, &mut context).map(|v| Value::Json(Box::new(v))) + encode_json_variant_with_context(json, &mut context).map(|v| Value::Json(Box::new(v))) } } @@ -120,38 +125,40 @@ fn with_key_context( } /// Main encoding function with key path tracking -fn encode_json_with_context(json: Json, context: &mut JsonContext) -> Result { - if context.path.is_empty() && !matches!(json, Json::Object(_)) { +fn encode_json_variant_with_context( + json: JsonVariant, + context: &mut JsonContext, +) -> Result { + if context.path.is_empty() && !matches!(json, JsonVariant::Object(_)) { return UnsupportedJsonTypeSnafu.fail(); } match json { - Json::Object(json_object) => encode_json_object_with_context(json_object, context), - Json::Array(json_array) => encode_json_array_with_context(json_array, context), - _ => encode_json_value_with_context(json, context), + JsonVariant::Object(json_object) => encode_json_object_with_context(json_object, context), + JsonVariant::Array(json_array) => encode_json_array_with_context(json_array, context), + _ => encode_json_variant_value_with_context(json, context), } } fn encode_json_object_with_context<'a>( - json_object: Map, + mut json_object: BTreeMap, context: &mut JsonContext<'a>, ) -> Result { - let mut object = BTreeMap::new(); - for (key, value) in json_object { - let value = with_key_context(context, &key, |context| { + for (key, slot) in &mut json_object { + let value = std::mem::take(slot); + let value = with_key_context(context, key, |context| { if let Some(hint) = context.type_hint() { encode_json_value_with_hint(value, hint, context) } else { - encode_json_value_with_context(value, context) + encode_json_variant_value_with_context(value, context) } })?; - - object.insert(key, value.into_variant()); + *slot = value.into_variant(); } - apply_missing_type_hints(&mut object, context)?; + apply_missing_type_hints(&mut json_object, context)?; - Ok(JsonValue::new(JsonVariant::Object(object))) + Ok(JsonValue::new(JsonVariant::Object(json_object))) } fn apply_missing_type_hints( @@ -217,7 +224,7 @@ fn encode_missing_type_hint_value( if let Some(default_constraint) = &hint.default_constraint { let value = default_constraint.create_default(&hint.data_type, hint.nullable)?; let json = decode_primitive_value(value)?; - return encode_json_value_with_hint(json, hint, context); + return encode_json_value_with_hint(json.into(), hint, context); } if hint.nullable { @@ -234,11 +241,11 @@ fn encode_missing_type_hint_value( } fn encode_json_value_with_hint( - json: Json, + json: JsonVariant, hint: &JsonTypeHint, context: &mut JsonContext, ) -> Result { - if json.is_null() { + if matches!(json, JsonVariant::Null) { return if hint.nullable { Ok(JsonValue::null()) } else { @@ -264,13 +271,13 @@ fn encode_json_value_with_hint( }; match (&hint.data_type, json) { - (ConcreteDataType::String(_), Json::String(v)) => Ok(v.into()), + (ConcreteDataType::String(_), JsonVariant::String(v)) => Ok(v.into()), ( ConcreteDataType::Int8(_) | ConcreteDataType::Int16(_) | ConcreteDataType::Int32(_) | ConcreteDataType::Int64(_), - Json::Number(v), + v @ JsonVariant::Number(_), ) => match v.as_i64() { Some(v) => Ok(v.into()), None => invalid_type(), @@ -280,34 +287,34 @@ fn encode_json_value_with_hint( | ConcreteDataType::UInt16(_) | ConcreteDataType::UInt32(_) | ConcreteDataType::UInt64(_), - Json::Number(v), + v @ JsonVariant::Number(_), ) => match v.as_u64() { Some(v) => Ok(v.into()), None => invalid_type(), }, - (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(v)) => { - match v.as_f64() { - Some(v) => Ok(v.into()), - None => invalid_type(), - } + (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), JsonVariant::Number(v)) => { + let value = match v { + JsonNumber::PosInt(v) => v as f64, + JsonNumber::NegInt(v) => v as f64, + JsonNumber::Float(v) => v.0, + }; + Ok(value.into()) } - (ConcreteDataType::Boolean(_), Json::Bool(v)) => Ok(v.into()), + (ConcreteDataType::Boolean(_), JsonVariant::Bool(v)) => Ok(v.into()), _ => invalid_type(), } } fn encode_json_array_with_context<'a>( - json_array: Vec, + mut json_array: Vec, context: &mut JsonContext<'a>, ) -> Result { - let json_array_len = json_array.len(); - let mut items = Vec::with_capacity(json_array_len); - - for (index, value) in json_array.into_iter().enumerate() { + for (index, slot) in json_array.iter_mut().enumerate() { + let value = std::mem::take(slot); let item_value = with_key_context(context, &index.to_string(), |context| { - encode_json_value_with_context(value, context) + encode_json_variant_value_with_context(value, context) })?; - items.push(item_value); + *slot = item_value.into_variant(); } // In specification, it's valid for a JSON array to have different types of items, for example, @@ -315,54 +322,37 @@ fn encode_json_array_with_context<'a>( // array, which requires all items have exactly the same type. So we merge out the maybe // different item types to a unified type, and align all the item values to it. - let merged_item_type = if let Some((first, rests)) = items.split_first() { - let mut merged = first.json_type().clone(); - for rest in rests.iter().map(|x| x.json_type()) { + let merged_item_type = if let Some((first, rests)) = json_array.split_first() { + let mut merged = first.native_type(); + for rest in rests.iter().map(JsonVariant::native_type) { if matches!(merged, JsonNativeType::Variant) { break; } - merged.merge(rest); + merged.merge(&rest); } Some(merged) } else { None }; if let Some(unified_item_type) = merged_item_type { - for item in &mut items { - item.try_align(&unified_item_type)?; + for item in &mut json_array { + let mut value = JsonValue::new(std::mem::take(item)); + value.try_align(&unified_item_type)?; + *item = value.into_variant(); } } - let items = items - .into_iter() - .map(|x| x.into_variant()) - .collect::>(); - Ok(JsonValue::new(JsonVariant::Array(items))) + Ok(JsonValue::new(JsonVariant::Array(json_array))) } /// Helper function to encode a JSON value to a Value and determine its ConcreteDataType with context -fn encode_json_value_with_context(json: Json, context: &mut JsonContext) -> Result { +fn encode_json_variant_value_with_context( + json: JsonVariant, + context: &mut JsonContext, +) -> Result { match json { - Json::Null => Ok(JsonValue::null()), - Json::Bool(b) => Ok(b.into()), - Json::Number(n) => { - if let Some(i) = n.as_i64() { - Ok(i.into()) - } else if let Some(u) = n.as_u64() { - if u <= i64::MAX as u64 { - Ok((u as i64).into()) - } else { - Ok(u.into()) - } - } else if let Some(f) = n.as_f64() { - Ok(f.into()) - } else { - // Fallback to string representation - Ok(n.to_string().into()) - } - } - Json::String(s) => Ok(s.into()), - Json::Array(arr) => encode_json_array_with_context(arr, context), - Json::Object(obj) => encode_json_object_with_context(obj, context), + JsonVariant::Array(arr) => encode_json_array_with_context(arr, context), + JsonVariant::Object(obj) => encode_json_object_with_context(obj, context), + value => Ok(JsonValue::new(value)), } } @@ -572,6 +562,32 @@ mod tests { } } + #[test] + fn test_encode_json_variant() { + let settings = JsonSettings::new(vec![JsonTypeHint { + path: vec!["count".to_string()], + data_type: ConcreteDataType::uint64_datatype(), + nullable: false, + default_constraint: None, + inverted_index: false, + }]); + let value = JsonVariant::from(json!({ + "count": 42, + "nested": { "enabled": true }, + "values": [1, 2, 3] + })); + + let result = settings + .encode_variant(value) + .unwrap() + .into_json_inner() + .unwrap(); + let Value::Struct(result) = result else { + panic!("Expected Struct value"); + }; + assert_eq!(struct_field_value(&result, "count"), &Value::UInt64(42)); + } + #[test] fn test_encode_json_object() { let json = json!({ diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index 814e3b36c08..36ec7031816 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -29,6 +29,7 @@ use common_query::prelude::{greptime_timestamp, greptime_value}; use common_telemetry::warn; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::JsonExtensionType; +use datatypes::json::value::JsonVariant; use datatypes::value::Value; use greptime_proto::v1::{ColumnSchema, Row, Rows, Value as GreptimeValue}; use itertools::Itertools; @@ -39,7 +40,6 @@ use session::context::Channel; use snafu::OptionExt; use table::Table; use vrl::prelude::{Bytes, VrlValueConvert}; -use vrl::value::value::StdError; use vrl::value::{KeyString, Value as VrlValue}; use crate::error::{ @@ -47,6 +47,7 @@ use crate::error::{ IdentifyPipelineColumnTypeMismatchSnafu, InvalidTimestampSnafu, Result, TimeIndexMustBeNonNullSnafu, TransformColumnNameMustBeUniqueSnafu, TransformMultipleTimestampIndexSnafu, TransformTimestampIndexCountSnafu, ValueMustBeMapSnafu, + VrlRegexValueSnafu, }; use crate::etl::PipelineDocVersion; use crate::etl::ctx_req::ContextOpt; @@ -710,10 +711,7 @@ fn resolve_value( let settings = json_extension_type .and_then(|x| x.metadata().json_settings.clone()) .unwrap_or_default(); - let value: serde_json::Value = value.try_into().map_err(|e: StdError| { - CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build() - })?; - let value = settings.encode(value)?; + let value = settings.encode_variant(vrl_value_into_json_variant(value)?)?; resolve_schema( index, @@ -776,6 +774,33 @@ fn vrl_value_to_jsonb_value<'a>(value: &'a VrlValue) -> jsonb::Value<'a> { } } +fn vrl_value_into_json_variant(value: VrlValue) -> Result { + match value { + VrlValue::Null => Ok(JsonVariant::Null), + VrlValue::Boolean(value) => Ok(value.into()), + VrlValue::Integer(value) => Ok(value.into()), + VrlValue::Float(value) => Ok(value.into_inner().into()), + VrlValue::Bytes(value) => Ok(String::from_utf8_lossy_owned(value.to_vec()).into()), + VrlValue::Array(values) => values + .into_iter() + .map(vrl_value_into_json_variant) + .collect::>() + .map(JsonVariant::Array), + VrlValue::Object(values) => values + .into_iter() + .map(|(key, value)| { + vrl_value_into_json_variant(value).map(|value| (key.to_string(), value)) + }) + .collect::>() + .map(JsonVariant::Object), + VrlValue::Regex(_) => VrlRegexValueSnafu.fail(), + VrlValue::Timestamp(_) => CoerceIncompatibleTypesSnafu { + msg: "timestamp is not supported in JSON2 values", + } + .fail(), + } +} + fn identity_pipeline_inner( pipeline_maps: Vec, pipeline_ctx: &PipelineContext<'_>, @@ -978,6 +1003,38 @@ mod tests { use super::*; use crate::{PipelineDefinition, identity_pipeline}; + #[test] + fn test_vrl_value_into_json_variant() { + let json = serde_json::json!({ + "array": [1, 2.5, null], + "boolean": true, + "integer": 42, + "string": "value" + }); + let value: VrlValue = json.clone().into(); + assert_eq!( + vrl_value_into_json_variant(value).unwrap(), + JsonVariant::from(json) + ); + + let value = VrlValue::Array(vec![VrlValue::Regex( + regex::Regex::new("value").unwrap().into(), + )]); + assert!(matches!( + vrl_value_into_json_variant(value), + Err(crate::error::Error::VrlRegexValue { .. }) + )); + + let value = VrlValue::Object(BTreeMap::from([( + KeyString::from("timestamp"), + VrlValue::Timestamp(chrono::Utc::now()), + )])); + assert!(matches!( + vrl_value_into_json_variant(value), + Err(crate::error::Error::CoerceIncompatibleTypes { .. }) + )); + } + #[test] fn test_identify_pipeline() { let params = GreptimePipelineParams::default();