From fb86f6573e244ce4cb3cbbce2140f032ef509ff2 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 31 Aug 2026 07:49:05 +0000 Subject: [PATCH] feat(pipeline): support table-aware JSON2 transforms (#8964) * feat(pipeline): support table-aware JSON2 transforms Signed-off-by: shuiyisong * feat(pipeline): support JSON2 type hints in transforms Signed-off-by: shuiyisong * fix(pipeline): default failed JSON2 transforms to null Signed-off-by: shuiyisong * refactor(json2): distinguish invalid settings from layout errors Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/datatypes/src/error.rs | 8 + src/datatypes/src/json.rs | 30 ++- src/pipeline/src/error.rs | 16 ++ src/pipeline/src/etl.rs | 93 +++++-- src/pipeline/src/etl/ctx_req.rs | 15 +- src/pipeline/src/etl/transform.rs | 246 +++++++++++++++++- .../src/etl/transform/transformer/greptime.rs | 217 ++++++++++++--- .../transform/transformer/greptime/coerce.rs | 134 ++++++++-- src/pipeline/src/etl/value.rs | 1 + src/pipeline/src/lib.rs | 3 +- src/pipeline/tests/json_parse.rs | 51 ++++ src/servers/src/pipeline.rs | 45 +++- src/sql/src/lib.rs | 2 +- src/sql/src/parsers/create_parser.rs | 1 + src/sql/src/parsers/create_parser/json.rs | 30 +++ 15 files changed, 776 insertions(+), 116 deletions(-) diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index 1def0d6be8..48ced6a4ac 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -210,6 +210,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid JSON2 settings: {reason}"))] + InvalidJson2Settings { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid Vector: {}", msg))] InvalidVector { msg: String, @@ -334,6 +341,7 @@ impl ErrorExt for Error { | InvalidPrecisionOrScale { .. } | InvalidJson { .. } | InvalidJson2Layout { .. } + | InvalidJson2Settings { .. } | InvalidJsonb { .. } | InvalidVector { .. } | InvalidFulltextOption { .. } diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 79fc10b12e..1ef8c85a22 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -29,7 +29,7 @@ use serde_json::{Map, Value as Json}; use snafu::ResultExt; use crate::data_type::ConcreteDataType; -use crate::error::{self, InvalidJson2LayoutSnafu, Result, UnsupportedJsonTypeSnafu}; +use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; use crate::schema::ColumnDefaultConstraint; use crate::types::json_type::{JsonNativeType, JsonObjectType}; @@ -162,7 +162,7 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { let mut object = JsonObjectType::new(); for hint in type_hints { if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: format!( "JSON2 type hint path cannot exceed {JSON2_MAX_STRUCTURED_DEPTH} segments" ), @@ -174,7 +174,7 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { .first() .is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME) { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: format!( "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" ), @@ -195,12 +195,26 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { | ConcreteDataType::Float64(_) | ConcreteDataType::String(_) => (&hint.data_type).into(), data_type => { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: format!("unsupported JSON2 type hint data type: {data_type}"), } .fail(); } }; + let non_finite_default = match &hint.default_constraint { + Some(ColumnDefaultConstraint::Value(Value::Float32(value))) => !value.0.is_finite(), + Some(ColumnDefaultConstraint::Value(Value::Float64(value))) => !value.0.is_finite(), + _ => false, + }; + if non_finite_default { + return InvalidJson2SettingsSnafu { + reason: format!( + "JSON2 type hint default for '{}' must be finite", + hint.path.join(".") + ), + } + .fail(); + } validate_type_hint(&mut object, &hint.path, data_type)?; } Ok(()) @@ -212,14 +226,14 @@ fn validate_type_hint( data_type: JsonNativeType, ) -> Result<()> { let Some((name, path)) = path.split_first() else { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: "JSON2 type hint path must not be empty".to_string(), } .fail(); }; if path.is_empty() { if object.insert(name.clone(), data_type).is_some() { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: format!("duplicate JSON2 type hint path '{name}'"), } .fail(); @@ -231,7 +245,7 @@ fn validate_type_hint( .entry(name.clone()) .or_insert_with(|| JsonNativeType::Object(JsonObjectType::new())); let JsonNativeType::Object(child) = child else { - return InvalidJson2LayoutSnafu { + return InvalidJson2SettingsSnafu { reason: format!("conflicting JSON2 type hint path at '{name}'"), } .fail(); @@ -712,7 +726,7 @@ mod tests { } #[test] - fn test_json_settings_reject_invalid_type_hint_layout() { + fn test_json_settings_reject_invalid_type_hints() { for type_hints in [ json!([{"path": [], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]), json!([ diff --git a/src/pipeline/src/error.rs b/src/pipeline/src/error.rs index 2da62dda2e..36abcaae9e 100644 --- a/src/pipeline/src/error.rs +++ b/src/pipeline/src/error.rs @@ -389,6 +389,20 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + #[snafu(display("Invalid JSON2 type hint: {reason}"))] + InvalidJson2TypeHint { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid JSON2 type hint path '{path}'"))] + ParseJson2TypeHintPath { + path: String, + #[snafu(source)] + source: sql::error::Error, + #[snafu(implicit)] + location: Location, + }, #[snafu(display("Transform index `type` must be set."))] TransformIndexTypeMustBeSet { #[snafu(implicit)] @@ -952,6 +966,8 @@ impl ErrorExt for Error { | TransformElementMustBeMap { .. } | TransformFieldMustBeSet { .. } | TransformTypeMustBeSet { .. } + | InvalidJson2TypeHint { .. } + | ParseJson2TypeHintPath { .. } | TransformIndexTypeMustBeSet { .. } | TransformIndexUnsupportedField { .. } | TransformIndexOptionMustBeScalar { .. } diff --git a/src/pipeline/src/etl.rs b/src/pipeline/src/etl.rs index 26d3520829..f83b5c2626 100644 --- a/src/pipeline/src/etl.rs +++ b/src/pipeline/src/etl.rs @@ -37,7 +37,9 @@ use crate::error::{ YamlLoadSnafu, YamlParseSnafu, }; use crate::etl::processor::ProcessorKind; -use crate::etl::transform::transformer::greptime::{RowWithTableSuffix, values_to_rows}; +use crate::etl::transform::transformer::greptime::{ + RowWithTableSuffix, values_to_row, values_to_rows, +}; use crate::tablesuffix::TableSuffixTemplate; use crate::{ ContextOpt, GreptimeTransformer, IdentityTimeIndex, PipelineContext, SchemaInfo, @@ -236,6 +238,14 @@ pub enum PipelineExecOutput { Filtered, } +/// The result after processors and dispatcher rules have run. +#[derive(Debug)] +pub enum PipelineProcessOutput { + Processed(VrlValue), + DispatchedTo(DispatchedTo, VrlValue), + Filtered, +} + /// Output from a successful pipeline transformation. /// /// Rows are grouped by their ContextOpt, with each row having its own optional @@ -303,24 +313,42 @@ impl Pipeline { pub fn exec_mut( &self, - mut val: VrlValue, + val: VrlValue, pipeline_ctx: &PipelineContext<'_>, schema_info: &mut SchemaInfo, ) -> Result { - // process + match self.process_mut(val)? { + PipelineProcessOutput::Processed(val) => self + .transform_mut(val, pipeline_ctx, schema_info) + .map(PipelineExecOutput::Transformed), + PipelineProcessOutput::DispatchedTo(dispatched_to, val) => { + Ok(PipelineExecOutput::DispatchedTo(dispatched_to, val)) + } + PipelineProcessOutput::Filtered => Ok(PipelineExecOutput::Filtered), + } + } + + pub fn process_mut(&self, mut val: VrlValue) -> Result { for processor in self.processors.iter() { val = processor.exec_mut(val)?; if val.is_null() { - // line is filtered - return Ok(PipelineExecOutput::Filtered); + return Ok(PipelineProcessOutput::Filtered); } } - // dispatch, fast return if matched if let Some(rule) = self.dispatcher.as_ref().and_then(|d| d.exec(&val)) { - return Ok(PipelineExecOutput::DispatchedTo(rule.into(), val)); + return Ok(PipelineProcessOutput::DispatchedTo(rule.into(), val)); } + Ok(PipelineProcessOutput::Processed(val)) + } + + pub fn transform_mut( + &self, + val: VrlValue, + pipeline_ctx: &PipelineContext<'_>, + schema_info: &mut SchemaInfo, + ) -> Result { let mut val = if val.is_array() { val } else { @@ -356,9 +384,7 @@ impl Pipeline { } }; - Ok(PipelineExecOutput::Transformed(TransformedOutput { - rows_by_context, - })) + Ok(TransformedOutput { rows_by_context }) } pub fn processors(&self) -> &processor::Processors { @@ -369,6 +395,10 @@ impl Pipeline { &self.transformer } + pub fn resolve_table_suffix(&self, value: &VrlValue) -> Option { + ContextOpt::resolve_table_suffix(self.tablesuffix.as_ref(), value) + } + // the method is for test purpose pub fn schemas(&self) -> Option<&Vec> { match &self.transformer { @@ -409,34 +439,43 @@ fn transform_array_elements_by_ctx( ); } - let values = - unwrap_or_continue_if_err!(transformer.transform_mut(element, is_v1), skip_error); + let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, element); + let values = unwrap_or_continue_if_err!( + transformer.transform_mut_with_schema( + element, + is_v1, + schema_info, + table_suffix.as_deref(), + ), + skip_error + ); if is_v1 { // v1 mode: just use transformer output directly - let mut opt = unwrap_or_continue_if_err!( + let opt = unwrap_or_continue_if_err!( ContextOpt::from_pipeline_map_to_opt(element), skip_error ); - let table_suffix = opt.resolve_table_suffix(tablesuffix_template, element); rows_by_context .entry(opt) .or_insert_with(Vec::new) .push((Row { values }, table_suffix)); } else { // v2 mode: combine with auto-transform for remaining fields - let element_rows_map = values_to_rows( - schema_info, - element.clone(), - pipeline_ctx, - Some(values), - false, - tablesuffix_template, - ) - .map_err(Box::new) - .context(TransformArrayElementSnafu { index })?; - for (k, v) in element_rows_map { - rows_by_context.entry(k).or_default().extend(v); - } + let mut value = element.clone(); + let opt = unwrap_or_continue_if_err!( + ContextOpt::from_pipeline_map_to_opt(&mut value), + skip_error + ); + let row = unwrap_or_continue_if_err!( + values_to_row(schema_info, value, pipeline_ctx, Some(values), false,) + .map_err(Box::new) + .context(TransformArrayElementSnafu { index }), + skip_error + ); + rows_by_context + .entry(opt) + .or_default() + .push((row, table_suffix)); } } diff --git a/src/pipeline/src/etl/ctx_req.rs b/src/pipeline/src/etl/ctx_req.rs index 737e6147b2..86d160314f 100644 --- a/src/pipeline/src/etl/ctx_req.rs +++ b/src/pipeline/src/etl/ctx_req.rs @@ -69,10 +69,6 @@ pub struct ContextOpt { // reset the schema in query context schema: Option, - - // pipeline options, not set in query context - // can be removed before the end of the pipeline execution - table_suffix: Option, } impl ContextOpt { @@ -112,9 +108,7 @@ impl ContextOpt { GREPTIME_SKIP_WAL => { opt.skip_wal = Some(v); } - GREPTIME_TABLE_SUFFIX => { - opt.table_suffix = Some(v); - } + GREPTIME_TABLE_SUFFIX => {} _ => {} } } @@ -123,12 +117,13 @@ impl ContextOpt { } pub(crate) fn resolve_table_suffix( - &mut self, table_suffix: Option<&TableSuffixTemplate>, pipeline_map: &VrlValue, ) -> Option { - self.table_suffix - .take() + pipeline_map + .as_object() + .and_then(|map| map.get(GREPTIME_TABLE_SUFFIX)) + .map(|value| value.to_string_lossy().to_string()) .or_else(|| table_suffix.and_then(|s| s.apply(pipeline_map))) } diff --git a/src/pipeline/src/etl/transform.rs b/src/pipeline/src/etl/transform.rs index efcb4765ad..c9a1510735 100644 --- a/src/pipeline/src/etl/transform.rs +++ b/src/pipeline/src/etl/transform.rs @@ -17,17 +17,21 @@ pub mod transformer; use std::collections::HashMap; +use api::helper::ColumnDataTypeWrapper; use api::v1::ColumnDataType; use api::v1::value::ValueData; use chrono::Utc; -use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; +use datatypes::json::{JsonSettings, JsonTypeHint}; +use datatypes::schema::{ColumnDefaultConstraint, FulltextOptions, SkippingIndexOptions}; +use datatypes::value::Value; use snafu::{OptionExt, ResultExt, ensure}; use sql::parsers::utils::{ validate_column_fulltext_create_option, validate_column_skipping_index_create_option, }; use crate::error::{ - Error, FieldMustBeTypeSnafu, KeyMustBeStringSnafu, Result, TransformElementMustBeMapSnafu, + Error, FieldMustBeTypeSnafu, InvalidJson2TypeHintSnafu, KeyMustBeStringSnafu, + ParseJson2TypeHintPathSnafu, Result, TransformElementMustBeMapSnafu, TransformFieldMustBeSetSnafu, TransformIndexOptionMustBeScalarSnafu, TransformIndexOptionSnafu, TransformIndexOptionUnsupportedSnafu, TransformIndexOptionsUnsupportedSnafu, TransformIndexTypeMismatchSnafu, TransformIndexTypeMustBeSetSnafu, @@ -49,6 +53,10 @@ const TRANSFORM_INDEX_OPTIONS_FIELD: &str = "index.options"; const TRANSFORM_TAG: &str = "tag"; const TRANSFORM_DEFAULT: &str = "default"; const TRANSFORM_ON_FAILURE: &str = "on_failure"; +const JSON2_TYPE: &str = "json2"; +const JSON2_TYPE_HINT: &str = "type.json2[]"; +const JSON2_TYPE_HINT_PATH: &str = "path"; +const JSON2_TYPE_HINT_NULLABLE: &str = "nullable"; pub use transformer::greptime::GreptimeTransformer; @@ -141,6 +149,7 @@ impl TryFrom<&Vec> for Transforms { pub struct Transform { pub fields: Fields, pub type_: ColumnDataType, + pub(crate) json_settings: Option, pub default: Option, pub index: Option, pub index_options: Option, @@ -196,7 +205,8 @@ impl TransformIndexOptions { // ColumnDataType::TimestampMicrosecond // ColumnDataType::TimestampMillisecond // ColumnDataType::TimestampSecond -// ColumnDataType::Binary +// ColumnDataType::Binary (JSONB) +// ColumnDataType::Json (JSON2) impl Transform { pub(crate) fn get_default(&self) -> Option<&ValueData> { @@ -260,6 +270,7 @@ fn get_default_for_type(ty: &ColumnDataType) -> Result { ColumnDataType::Float32 => ValueData::F32Value(0.0), ColumnDataType::Float64 => ValueData::F64Value(0.0), ColumnDataType::Binary => ValueData::BinaryValue(jsonb::Value::Null.to_vec()), + ColumnDataType::Json => ValueData::JsonValue(Default::default()), ColumnDataType::String => ValueData::StringValue(String::new()), ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(0), @@ -419,6 +430,162 @@ fn lower_transform_index_options( } } } + +fn parse_transform_type(value: &yaml_rust::Yaml) -> Result<(ColumnDataType, Option)> { + if let Some(type_name) = value.as_str() { + return Ok((parse_str_type(type_name)?, None)); + } + + let config = value.as_hash().context(FieldMustBeTypeSnafu { + field: TRANSFORM_TYPE, + ty: "string or map", + })?; + ensure!( + config.len() == 1, + InvalidJson2TypeHintSnafu { + reason: "transform type map must contain exactly one `json2` field".to_string() + } + ); + let (type_name, hints) = config.iter().next().context(InvalidJson2TypeHintSnafu { + reason: "transform type map must contain a `json2` field".to_string(), + })?; + let type_name = type_name.as_str().with_context(|| KeyMustBeStringSnafu { + k: type_name.clone(), + })?; + ensure!( + type_name.eq_ignore_ascii_case(JSON2_TYPE), + InvalidJson2TypeHintSnafu { + reason: format!("unsupported transform type map `{type_name}`") + } + ); + + let hints = hints.as_vec().context(FieldMustBeTypeSnafu { + field: JSON2_TYPE, + ty: "list", + })?; + let hints = hints + .iter() + .map(parse_json2_type_hint) + .collect::>>()?; + Ok(( + ColumnDataType::Json, + Some(JsonSettings::try_new(hints, None)?), + )) +} + +fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { + let config = value.as_hash().context(FieldMustBeTypeSnafu { + field: JSON2_TYPE_HINT, + ty: "map", + })?; + let mut path = None; + let mut type_name = None; + let mut nullable = true; + let mut default = None; + let mut index = None; + + for (key, value) in config { + let key = key + .as_str() + .with_context(|| KeyMustBeStringSnafu { k: key.clone() })?; + match key { + JSON2_TYPE_HINT_PATH => path = Some(yaml_string(value, JSON2_TYPE_HINT_PATH)?), + TRANSFORM_TYPE => type_name = Some(yaml_string(value, TRANSFORM_TYPE)?), + JSON2_TYPE_HINT_NULLABLE => { + nullable = yaml_bool(value, JSON2_TYPE_HINT_NULLABLE)?; + } + TRANSFORM_DEFAULT => default = Some(value), + TRANSFORM_INDEX => index = Some(value), + _ => { + return InvalidJson2TypeHintSnafu { + reason: format!("unsupported field `{key}`"), + } + .fail(); + } + } + } + + let path = path.context(InvalidJson2TypeHintSnafu { + reason: "`path` must be set".to_string(), + })?; + let path = sql::parse_json2_type_hint_path(&path) + .with_context(|_| ParseJson2TypeHintPathSnafu { path: path.clone() })?; + let type_name = type_name.context(InvalidJson2TypeHintSnafu { + reason: "`type` must be set".to_string(), + })?; + let type_ = parse_str_type(&type_name)?; + ensure!( + matches!( + type_, + ColumnDataType::String + | ColumnDataType::Int64 + | ColumnDataType::Uint64 + | ColumnDataType::Float64 + | ColumnDataType::Boolean + ), + InvalidJson2TypeHintSnafu { + reason: format!("unsupported type `{type_name}`") + } + ); + let data_type = ColumnDataTypeWrapper::new(type_, None).into(); + let default_constraint = default + .map(|value| parse_json2_type_hint_default(value, &type_)) + .transpose()?; + if let Some(default_constraint) = &default_constraint { + default_constraint.validate(&data_type, nullable)?; + } + + let inverted_index = if let Some(value) = index { + let (index, options) = parse_transform_index(value)?; + ensure!( + index == Index::Inverted, + InvalidJson2TypeHintSnafu { + reason: format!("unsupported index `{index}`") + } + ); + lower_transform_index_options(index, &ColumnDataType::Json, options)?; + true + } else { + false + }; + + Ok(JsonTypeHint { + path, + data_type, + nullable, + default_constraint, + inverted_index, + }) +} + +fn parse_json2_type_hint_default( + value: &yaml_rust::Yaml, + type_: &ColumnDataType, +) -> Result { + if value.is_null() { + return Ok(ColumnDefaultConstraint::Value(Value::Null)); + } + + let value = match value { + yaml_rust::Yaml::Real(value) | yaml_rust::Yaml::String(value) => value.clone(), + yaml_rust::Yaml::Integer(value) => value.to_string(), + yaml_rust::Yaml::Boolean(value) => value.to_string(), + _ => { + return FieldMustBeTypeSnafu { + field: TRANSFORM_DEFAULT, + ty: "scalar", + } + .fail(); + } + }; + let value = api::v1::Value { + value_data: Some(parse_str_value(type_, &value)?), + }; + Ok(ColumnDefaultConstraint::Value( + api::helper::pb_value_to_value_ref(&value, None).into(), + )) +} + impl TryFrom<&yaml_rust::yaml::Hash> for Transform { type Error = Error; @@ -430,6 +597,7 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform { let mut on_failure = None; let mut type_ = None; + let mut json_settings = None; for (k, v) in hash { let key = k @@ -445,8 +613,9 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform { } TRANSFORM_TYPE => { - let t = yaml_string(v, TRANSFORM_TYPE)?; - type_ = Some(parse_str_type(&t)?); + let (parsed_type, parsed_json_settings) = parse_transform_type(v)?; + type_ = Some(parsed_type); + json_settings = parsed_json_settings; } TRANSFORM_INDEX => { @@ -507,6 +676,7 @@ impl TryFrom<&yaml_rust::yaml::Hash> for Transform { let builder = Transform { fields, type_, + json_settings, default: final_default, index, index_options, @@ -529,6 +699,72 @@ mod tests { docs[0].as_hash().unwrap().try_into() } + #[test] + fn test_transform_parses_json2_type_hints() { + let transform = parse_transform( + r#" +field: payload +type: + json2: + - path: "user.id" + type: int64 + nullable: false + default: 7 + index: + type: inverted + - path: 'attrs."http.status_code"' + type: string +"#, + ) + .unwrap(); + + assert_eq!(transform.type_, ColumnDataType::Json); + let hints = transform.json_settings.as_ref().unwrap().type_hints(); + assert_eq!(hints.len(), 2); + assert_eq!(hints[0].path, ["user", "id"]); + assert_eq!( + hints[0].data_type, + datatypes::prelude::ConcreteDataType::int64_datatype() + ); + assert!(!hints[0].nullable); + assert_eq!( + hints[0].default_constraint, + Some(ColumnDefaultConstraint::Value(Value::Int64(7))) + ); + assert!(hints[0].inverted_index); + assert_eq!(hints[1].path, ["attrs", "http.status_code"]); + assert!(hints[1].nullable); + } + + #[test] + fn test_transform_rejects_non_finite_json2_default() { + for default in ["NaN", "1e9999"] { + let err = parse_transform(&format!( + r#" +field: payload +type: + json2: + - path: score + type: float64 + default: {default} +"#, + )) + .unwrap_err(); + + assert!( + matches!( + &err, + Error::Datatypes { + source: datatypes::error::Error::InvalidJson2Settings { .. }, + .. + } + ), + "{err:?}" + ); + assert!(err.to_string().contains("must be finite"), "{err}"); + } + } + #[test] fn test_transform_parses_legacy_string_index() { let transform = parse_transform( diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index 3292de749a..ca15808c00 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -150,6 +150,7 @@ impl GreptimeTransformer { let transform = Transform { fields: Fields::one(Field::new(greptime_timestamp().to_string(), None)), type_, + json_settings: None, default, index: Some(Index::Time), index_options: None, @@ -225,6 +226,16 @@ impl GreptimeTransformer { &self, pipeline_map: &mut VrlValue, is_v1: bool, + ) -> Result> { + self.transform_mut_with_schema(pipeline_map, is_v1, &SchemaInfo::default(), None) + } + + pub(crate) fn transform_mut_with_schema( + &self, + pipeline_map: &mut VrlValue, + is_v1: bool, + schema_info: &SchemaInfo, + table_suffix: Option<&str>, ) -> Result> { let mut values = vec![GreptimeValue { value_data: None }; self.schema.len()]; let mut output_index = 0; @@ -236,7 +247,17 @@ impl GreptimeTransformer { // let keep us `get` here to be compatible with v1 match pipeline_map.get(column_name) { Some(v) => { - let value_data = coerce_value(v, transform)?; + let json_settings = if transform.type_ == ColumnDataType::Json + && matches!(v, VrlValue::Array(_) | VrlValue::Object(_)) + { + schema_info.json_settings_for_column( + field.target_or_input_field(), + table_suffix, + )? + } else { + None + }; + let value_data = coerce_value(v, transform, json_settings.as_ref())?; // every transform fields has only one output field values[output_index] = GreptimeValue { value_data }; } @@ -276,6 +297,12 @@ impl GreptimeTransformer { &self.schema } + pub fn has_json_transform(&self) -> bool { + self.transforms + .iter() + .any(|transform| transform.type_ == ColumnDataType::Json) + } + pub fn transforms_mut(&mut self) -> &mut Transforms { &mut self.transforms } @@ -348,8 +375,8 @@ pub struct SchemaInfo { pub schema: Vec, /// index of the column name pub index: HashMap, - /// The pipeline's corresponding table (if already created). Useful to retrieve column schemas. - table: Option>, + /// Tables already looked up, keyed by their resolved suffix. Missing tables are cached as None. + tables: HashMap>>, } impl SchemaInfo { @@ -357,7 +384,7 @@ impl SchemaInfo { Self { schema: Vec::with_capacity(capacity), index: HashMap::with_capacity(capacity), - table: None, + tables: HashMap::new(), } } @@ -369,16 +396,33 @@ impl SchemaInfo { Self { schema: schema_list.into_iter().map(Into::into).collect(), index, - table: None, + tables: HashMap::new(), } } pub fn set_table(&mut self, table: Option>) { - self.table = table; + self.set_table_for_suffix(String::new(), table); + } + + pub fn has_table_for_suffix(&self, table_suffix: &str) -> bool { + self.tables.contains_key(table_suffix) + } + + pub fn set_table_for_suffix(&mut self, table_suffix: String, table: Option>) { + self.tables.insert(table_suffix, table); + } + + fn table_for_suffix(&self, table_suffix: Option<&str>) -> Option<&Arc> { + let table_suffix = table_suffix.unwrap_or_default(); + match self.tables.get(table_suffix) { + Some(table) => table.as_ref(), + None if !table_suffix.is_empty() => self.tables.get("").and_then(Option::as_ref), + None => None, + } } fn find_column_schema_in_table(&self, column_name: &str) -> Option { - if let Some(table) = &self.table + if let Some(table) = self.table_for_suffix(None) && let Some(i) = table.schema_ref().column_index_by_name(column_name) { let column_schema = table.schema_ref().column_schemas()[i].clone(); @@ -400,6 +444,30 @@ impl SchemaInfo { } } + fn json_settings_for_column( + &self, + column_name: &str, + table_suffix: Option<&str>, + ) -> Result> { + let Some(column_schema) = self + .table_for_suffix(table_suffix) + .and_then(|table| table.schema_ref().column_schema_by_name(column_name)) + else { + return Ok(None); + }; + if !column_schema.data_type.is_json2() { + return Ok(None); + } + + if let Some(extension) = column_schema.extension_type::()? { + Ok(Some(extension.metadata().json_settings().clone())) + } else { + Ok(Some( + parse_legacy_json2_settings(column_schema.metadata())?.unwrap_or_default(), + )) + } + } + pub fn column_schemas(&self) -> api::error::Result> { self.schema .iter() @@ -499,12 +567,12 @@ pub(crate) fn values_to_rows( // Single object: extract ContextOpt and table_suffix let mut result = std::collections::HashMap::new(); - let mut opt = match ContextOpt::from_pipeline_map_to_opt(&mut values) { + let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &values); + let opt = match ContextOpt::from_pipeline_map_to_opt(&mut values) { Ok(r) => r, Err(e) => return if skip_error { Ok(result) } else { Err(e) }, }; - let table_suffix = opt.resolve_table_suffix(tablesuffix_template, &values); let row = match values_to_row(schema_info, values, pipeline_ctx, row, need_calc_ts) { Ok(r) => r, Err(e) => return if skip_error { Ok(result) } else { Err(e) }, @@ -528,11 +596,11 @@ pub(crate) fn values_to_rows( } // Extract ContextOpt and table_suffix for this element - let mut opt = unwrap_or_continue_if_err!( + let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &value); + let opt = unwrap_or_continue_if_err!( ContextOpt::from_pipeline_map_to_opt(&mut value), skip_error ); - let table_suffix = opt.resolve_table_suffix(tablesuffix_template, &value); let transformed_row = unwrap_or_continue_if_err!( values_to_row(schema_info, value, pipeline_ctx, row.clone(), need_calc_ts), skip_error @@ -691,35 +759,13 @@ fn resolve_value( } VrlValue::Array(_) | VrlValue::Object(_) => { - let is_json2 = schema_info - .find_column_schema_in_table(&column_name) - // TODO(LFC): Default to JSON2 for auto-created tables. - .is_some_and(|x| { - matches!( - &x.column_schema.data_type, - ConcreteDataType::Json(column_type) if column_type.is_json2() - ) - }); + let json_settings = schema_info.json_settings_for_column(&column_name, None)?; - let value = if is_json2 { + let value = if let Some(json_settings) = json_settings { let value: serde_json::Value = value.try_into().map_err(|e: StdError| { CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build() })?; - let value = - if let Some(column) = schema_info.find_column_schema_in_table(&column_name) { - if let Some(extension) = column - .column_schema - .extension_type::()? - { - extension.metadata().json_settings().encode(value)? - } else { - parse_legacy_json2_settings(column.column_schema.metadata())? - .unwrap_or_default() - .encode(value)? - } - } else { - JsonSettings::default().encode(value)? - }; + let value = json_settings.encode(value)?; resolve_schema( index, @@ -907,7 +953,7 @@ pub fn flatten_object(object: VrlValue, max_nested_levels: usize) -> Result serde_json_crate::Value { +pub(crate) fn vrl_value_to_serde_json(value: &VrlValue) -> serde_json_crate::Value { match value { VrlValue::Null => serde_json_crate::Value::Null, VrlValue::Boolean(b) => serde_json_crate::Value::Bool(*b), @@ -980,10 +1026,107 @@ fn do_flatten_object( #[cfg(test)] mod tests { use api::v1::SemanticType; + use common_recordbatch::RecordBatch; + use datatypes::extension::json::JsonMetadata; + use datatypes::json::JsonTypeHint; + use datatypes::schema::{ColumnSchema as DatatypeColumnSchema, Schema}; + use table::test_util::MemTable; use super::*; use crate::{PipelineDefinition, identity_pipeline}; + #[test] + fn test_transform_json2_uses_destination_table_settings() { + let table = |name: &str, settings: JsonSettings, sample: serde_json::Value| { + let data_type = settings.encode(sample).unwrap().data_type(); + let mut column_schema = DatatypeColumnSchema::new("payload", data_type, true); + column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new( + JsonMetadata::new(settings), + ))); + MemTable::table( + name, + RecordBatch::new_empty(Arc::new(Schema::new(vec![column_schema]))), + ) + }; + let int_settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["age".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: false, + default_constraint: None, + inverted_index: false, + }], + None, + ) + .unwrap(); + let mut schema_info = SchemaInfo::default(); + schema_info.set_table(Some(table( + "events", + int_settings, + serde_json::json!({"age": 42}), + ))); + let string_settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["age".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: false, + default_constraint: None, + inverted_index: false, + }], + None, + ) + .unwrap(); + schema_info.set_table_for_suffix( + "_mobile".to_string(), + Some(table( + "events_mobile", + string_settings, + serde_json::json!({"age": "42"}), + )), + ); + + let pipeline = crate::parse(&crate::Content::Yaml( + r#" +transform: + - field: source, payload + type: json2 +table_suffix: _${device} +"#, + )) + .unwrap(); + let (pipeline, _, pipeline_definition, pipeline_params) = crate::setup_pipeline!(pipeline); + let pipeline_context = + PipelineContext::new(&pipeline_definition, &pipeline_params, Channel::Unknown); + + let error = pipeline + .exec_mut( + serde_json::json!({"source": {"age": "42"}}).into(), + &pipeline_context, + &mut schema_info, + ) + .unwrap_err(); + assert!( + error.to_string().contains("does not match JSON2 type hint"), + "{error:?}" + ); + + let mut rows = pipeline + .exec_mut( + serde_json::json!({"source": {"age": "42"}, "device": "mobile"}).into(), + &pipeline_context, + &mut schema_info, + ) + .unwrap() + .into_transformed() + .unwrap(); + let (row, table_suffix) = rows.swap_remove(0); + assert_eq!(table_suffix.as_deref(), Some("_mobile")); + assert!(matches!( + &row.values[0].value_data, + Some(ValueData::JsonValue(_)) + )); + } + #[test] fn test_identify_pipeline() { let params = GreptimePipelineParams::default(); diff --git a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs index 5eaaa6e691..9f01af708b 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs @@ -12,10 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use api::v1::column_data_type_extension::TypeExt; use api::v1::column_def::{options_from_fulltext, options_from_inverted, options_from_skipping}; use api::v1::{ColumnDataTypeExtension, ColumnOptions, JsonTypeExtension}; +use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, +}; +use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; +use datatypes::json::JsonSettings; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; +use datatypes::value::Value; use greptime_proto::v1::value::ValueData; use greptime_proto::v1::{ColumnDataType, ColumnSchema, SemanticType}; use snafu::{OptionExt, ResultExt, ensure}; @@ -28,7 +36,9 @@ use crate::error::{ UnsupportedTypeInPipelineSnafu, VrlRegexValueSnafu, }; use crate::etl::transform::index::Index; -use crate::etl::transform::transformer::greptime::vrl_value_to_jsonb_value; +use crate::etl::transform::transformer::greptime::{ + vrl_value_to_jsonb_value, vrl_value_to_serde_json, +}; use crate::etl::transform::{OnFailure, Transform, TransformIndexOptions}; pub(crate) fn coerce_columns(transform: &Transform) -> Result> { @@ -128,7 +138,7 @@ fn build_skipping_index_options(transform: &Transform) -> Result Result> { validate_transform_index_state(transform)?; - match transform.index { + let mut options = match transform.index { Some(Index::Fulltext) => { let options = build_fulltext_index_options(transform)?; options_from_fulltext(&options).context(ColumnOptionsSnafu) @@ -139,10 +149,32 @@ fn coerce_options(transform: &Transform) -> Result> { } Some(Index::Inverted) => Ok(Some(options_from_inverted())), _ => Ok(None), + }?; + + if transform.type_ == ColumnDataType::Json { + let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new( + transform.json_settings.clone().unwrap_or_default(), + ))); + let options = options.get_or_insert_default(); + options.options.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + if let Some(metadata) = extension.serialize_metadata() { + options + .options + .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata); + } } + + Ok(options) } -pub(crate) fn coerce_value(val: &VrlValue, transform: &Transform) -> Result> { +pub(crate) fn coerce_value( + val: &VrlValue, + transform: &Transform, + json_settings: Option<&JsonSettings>, +) -> Result> { match val { VrlValue::Null => Ok(None), VrlValue::Integer(n) => coerce_i64_value(*n, transform), @@ -169,7 +201,9 @@ pub(crate) fn coerce_value(val: &VrlValue, transform: &Transform) -> Result coerce_json_value(val, transform), + VrlValue::Array(_) | VrlValue::Object(_) => { + coerce_json_value(val, transform, json_settings) + } VrlValue::Regex(_) => VrlRegexValueSnafu.fail(), } } @@ -205,7 +239,7 @@ fn coerce_bool_value(b: bool, transform: &Transform) -> Result } }, - ColumnDataType::Binary => { + ColumnDataType::Binary | ColumnDataType::Json => { return CoerceJsonTypeToSnafu { ty: transform.type_.as_str_name(), } @@ -294,7 +328,7 @@ fn coerce_i64_value(n: i64, transform: &Transform) -> Result> ColumnDataType::TimestampMillisecond => ValueData::TimestampMillisecondValue(n), ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(n), - ColumnDataType::Binary => { + ColumnDataType::Binary | ColumnDataType::Json => { return CoerceJsonTypeToSnafu { ty: transform.type_.as_str_name(), } @@ -363,7 +397,7 @@ fn coerce_u64_value(n: u64, transform: &Transform) -> Result> Err(_) => return integer_out_of_range(n, transform), }, - ColumnDataType::Binary => { + ColumnDataType::Binary | ColumnDataType::Json => { return CoerceJsonTypeToSnafu { ty: transform.type_.as_str_name(), } @@ -407,7 +441,7 @@ fn coerce_f64_value(n: f64, transform: &Transform) -> Result> } }, - ColumnDataType::Binary => { + ColumnDataType::Binary | ColumnDataType::Json => { return CoerceJsonTypeToSnafu { ty: transform.type_.as_str_name(), } @@ -486,7 +520,7 @@ fn coerce_string_value(s: &str, transform: &Transform) -> Result CoerceUnsupportedEpochTypeSnafu { ty: "String" }.fail(), }, - ColumnDataType::Binary => CoerceStringToTypeSnafu { + ColumnDataType::Binary | ColumnDataType::Json => CoerceStringToTypeSnafu { s, ty: transform.type_.as_str_name(), } @@ -496,23 +530,48 @@ fn coerce_string_value(s: &str, transform: &Transform) -> Result Result> { - match &transform.type_ { - ColumnDataType::Binary => (), +fn coerce_json_value( + v: &VrlValue, + transform: &Transform, + json_settings: Option<&JsonSettings>, +) -> Result> { + let value = match transform.type_ { + ColumnDataType::Binary => { + let data: jsonb::Value = vrl_value_to_jsonb_value(v); + ValueData::BinaryValue(data.to_vec()) + } + ColumnDataType::Json => { + let json = vrl_value_to_serde_json(v); + let encoded = if let Some(settings) = json_settings.or(transform.json_settings.as_ref()) + { + settings.encode(json) + } else { + JsonSettings::default().encode(json) + }; + let value = match encoded { + Ok(value) => value, + Err(error) => return handle_coercion_failure(transform, error.into()), + }; + let Value::Json(value) = value else { + unreachable!() + }; + ValueData::JsonValue(api::helper::encode_json_value(*value)) + } t => { return CoerceTypeToJsonSnafu { ty: t.as_str_name(), } .fail(); } - } - let data: jsonb::Value = vrl_value_to_jsonb_value(v); - Ok(Some(ValueData::BinaryValue(data.to_vec()))) + }; + Ok(Some(value)) } #[cfg(test)] mod tests { + use datatypes::data_type::ConcreteDataType; + use datatypes::json::JsonTypeHint; use datatypes::schema::{FulltextAnalyzer, FulltextBackend, SkippingIndexType}; use vrl::prelude::Bytes; @@ -523,6 +582,7 @@ mod tests { Transform { fields: Fields::default(), type_, + json_settings: None, default: None, index: None, index_options: None, @@ -666,6 +726,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::Int32, + json_settings: None, default: None, index: None, index_options: None, @@ -676,14 +737,14 @@ mod tests { // valid string { let val = VrlValue::Integer(123); - let result = coerce_value(&val, &transform).unwrap(); + let result = coerce_value(&val, &transform, None).unwrap(); assert_eq!(result, Some(ValueData::I32Value(123))); } // invalid string { let val = VrlValue::Bytes(Bytes::from("hello")); - let result = coerce_value(&val, &transform); + let result = coerce_value(&val, &transform, None); assert!(result.is_err()); } } @@ -693,6 +754,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::Int32, + json_settings: None, default: None, index: None, index_options: None, @@ -701,15 +763,43 @@ mod tests { }; let val = VrlValue::Bytes(Bytes::from("hello")); - let result = coerce_value(&val, &transform).unwrap(); + let result = coerce_value(&val, &transform, None).unwrap(); assert_eq!(result, None); } + #[test] + fn test_coerce_json2_with_on_failure() { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["age".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: false, + default_constraint: None, + inverted_index: false, + }], + None, + ) + .unwrap(); + let mut transform = transform(ColumnDataType::Json); + transform.json_settings = Some(settings); + transform.on_failure = Some(OnFailure::Ignore); + let value: VrlValue = serde_json::json!({"age": "42"}).into(); + + assert_eq!(coerce_value(&value, &transform, None).unwrap(), None); + + transform.on_failure = Some(OnFailure::Default); + assert_eq!( + coerce_value(&value, &transform, None).unwrap(), + Some(ValueData::JsonValue(Default::default())) + ); + } + #[test] fn test_coerce_string_with_on_failure_default() { let mut transform = Transform { fields: Fields::default(), type_: ColumnDataType::Int32, + json_settings: None, default: None, index: None, index_options: None, @@ -720,7 +810,7 @@ mod tests { // with no explicit default value { let val = VrlValue::Bytes(Bytes::from("hello")); - let result = coerce_value(&val, &transform).unwrap(); + let result = coerce_value(&val, &transform, None).unwrap(); assert_eq!(result, Some(ValueData::I32Value(0))); } @@ -728,7 +818,7 @@ mod tests { { transform.default = Some(ValueData::I32Value(42)); let val = VrlValue::Bytes(Bytes::from("hello")); - let result = coerce_value(&val, &transform).unwrap(); + let result = coerce_value(&val, &transform, None).unwrap(); assert_eq!(result, Some(ValueData::I32Value(42))); } } @@ -738,6 +828,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::String, + json_settings: None, default: None, index: Some(Index::Fulltext), index_options: Some(TransformIndexOptions::Fulltext( @@ -769,6 +860,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::Int64, + json_settings: None, default: None, index: Some(Index::Skipping), index_options: Some(TransformIndexOptions::Skipping( @@ -792,6 +884,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::String, + json_settings: None, default: None, index: Some(Index::Fulltext), index_options: Some(TransformIndexOptions::Skipping( @@ -809,6 +902,7 @@ mod tests { let transform = Transform { fields: Fields::default(), type_: ColumnDataType::String, + json_settings: None, default: None, index: None, index_options: Some(TransformIndexOptions::Fulltext( diff --git a/src/pipeline/src/etl/value.rs b/src/pipeline/src/etl/value.rs index 2f5560b456..09d63058ff 100644 --- a/src/pipeline/src/etl/value.rs +++ b/src/pipeline/src/etl/value.rs @@ -102,6 +102,7 @@ pub fn parse_str_type(t: &str) -> Result { // We only consider object and array to be json types. and use Map to represent json // TODO(qtang): Needs to be defined with better semantics "json" => Ok(ColumnDataType::Binary), + "json2" => Ok(ColumnDataType::Json), _ => ValueParseTypeSnafu { t }.fail(), } diff --git a/src/pipeline/src/lib.rs b/src/pipeline/src/lib.rs index c657f61342..c3a45d4fe4 100644 --- a/src/pipeline/src/lib.rs +++ b/src/pipeline/src/lib.rs @@ -27,7 +27,8 @@ pub use etl::transform::GreptimeTransformer; pub use etl::transform::transformer::greptime::{GreptimePipelineParams, SchemaInfo}; pub use etl::transform::transformer::identity_pipeline; pub use etl::{ - Content, DispatchedTo, Pipeline, PipelineExecOutput, TransformedOutput, TransformerMode, parse, + Content, DispatchedTo, Pipeline, PipelineExecOutput, PipelineProcessOutput, TransformedOutput, + TransformerMode, parse, }; pub use manager::{ GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME, diff --git a/src/pipeline/tests/json_parse.rs b/src/pipeline/tests/json_parse.rs index 37979745ed..c374c846cd 100644 --- a/src/pipeline/tests/json_parse.rs +++ b/src/pipeline/tests/json_parse.rs @@ -17,7 +17,12 @@ mod common; use std::borrow::Cow; use api::v1::ColumnDataType; +use api::v1::json_value::Value as JsonValue; use api::v1::value::ValueData; +use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, +}; +use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; const INPUT_VALUE_OBJ: &str = r#" [ @@ -113,6 +118,52 @@ transform: assert_eq!(v, jsonb::Value::Object(expected)); } +#[test] +fn test_json2_parse() { + let pipeline_yaml = r#" +--- +processors: + - json_parse: + field: commit + +transform: + - field: commit + type: + json2: + - path: "commitAuthor" + type: string + nullable: false +"#; + + let output = common::parse_and_exec(INPUT_VALUE_OBJ, pipeline_yaml); + + assert_eq!(output.schema[0].datatype, ColumnDataType::Json as i32); + assert!(output.schema[0].datatype_extension.is_none()); + assert_eq!( + output.schema[0] + .options + .as_ref() + .and_then(|options| options.options.get(EXTENSION_TYPE_NAME_KEY)) + .map(String::as_str), + Some(Json2ExtensionType::NAME) + ); + let metadata = output.schema[0] + .options + .as_ref() + .and_then(|options| options.options.get(EXTENSION_TYPE_METADATA_KEY)) + .unwrap(); + let metadata: JsonMetadata = serde_json::from_str(metadata).unwrap(); + assert_eq!( + metadata.json_settings().type_hints()[0].path, + ["commitAuthor"] + ); + + let ValueData::JsonValue(value) = output.rows[0].values[0].value_data.as_ref().unwrap() else { + panic!("expect JSON2 value"); + }; + assert!(matches!(value.value, Some(JsonValue::Object(_)))); +} + #[test] fn test_json_parse_with_simple_extractor() { let pipeline_yaml = r#" diff --git a/src/servers/src/pipeline.rs b/src/servers/src/pipeline.rs index 7f6fa20f1f..a7081cf89d 100644 --- a/src/servers/src/pipeline.rs +++ b/src/servers/src/pipeline.rs @@ -21,7 +21,7 @@ use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value}; use common_time::timestamp::TimeUnit; use pipeline::{ ContextOpt, ContextReq, DispatchedTo, GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, Pipeline, - PipelineContext, PipelineDefinition, PipelineExecOutput, SchemaInfo, TransformedOutput, + PipelineContext, PipelineDefinition, PipelineProcessOutput, SchemaInfo, TransformedOutput, TransformerMode, identity_pipeline, unwrap_or_continue_if_err, }; use session::context::{Channel, QueryContextRef}; @@ -140,10 +140,14 @@ async fn run_custom_pipeline( let table = handler.get_table(&table_name, query_ctx).await?; schema_info.set_table(table); + let needs_json_settings = matches!( + pipeline.transformer(), + TransformerMode::GreptimeTransformer(transformer) if transformer.has_json_transform() + ); for pipeline_map in pipeline_maps { let result = pipeline - .exec_mut(pipeline_map, pipeline_ctx, &mut schema_info) + .process_mut(pipeline_map) .inspect_err(|_| { METRIC_HTTP_LOGS_TRANSFORM_ELAPSED .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE]) @@ -151,9 +155,36 @@ async fn run_custom_pipeline( }) .context(PipelineSnafu); - let r = unwrap_or_continue_if_err!(result, skip_error); - match r { - PipelineExecOutput::Transformed(TransformedOutput { rows_by_context }) => { + match unwrap_or_continue_if_err!(result, skip_error) { + PipelineProcessOutput::Processed(value) => { + if needs_json_settings { + // JSON2 coercion must use settings from the final routed table. + let values = match &value { + VrlValue::Array(values) => values.as_slice(), + value => std::slice::from_ref(value), + }; + for value in values.iter().filter(|value| value.is_object()) { + let table_suffix = pipeline.resolve_table_suffix(value).unwrap_or_default(); + if !schema_info.has_table_for_suffix(&table_suffix) { + let destination = + table_suffix_to_table_name(&table_name, &table_suffix); + let table = handler.get_table(&destination, query_ctx).await?; + schema_info.set_table_for_suffix(table_suffix, table); + } + } + } + + let result = pipeline + .transform_mut(value, pipeline_ctx, &mut schema_info) + .inspect_err(|_| { + METRIC_HTTP_LOGS_TRANSFORM_ELAPSED + .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE]) + .observe(transform_timer.elapsed().as_secs_f64()); + }) + .context(PipelineSnafu); + let TransformedOutput { rows_by_context } = + unwrap_or_continue_if_err!(result, skip_error); + // Process each ContextOpt group separately for (opt, rows_with_suffix) in rows_by_context { let rows_by_suffix = transformed_map.entry(opt).or_default(); @@ -166,10 +197,10 @@ async fn run_custom_pipeline( } } } - PipelineExecOutput::DispatchedTo(dispatched_to, val) => { + PipelineProcessOutput::DispatchedTo(dispatched_to, val) => { push_to_map!(dispatched, dispatched_to, val, arr_len); } - PipelineExecOutput::Filtered => { + PipelineProcessOutput::Filtered => { continue; } } diff --git a/src/sql/src/lib.rs b/src/sql/src/lib.rs index e8c6bdf8ef..b017f798f3 100644 --- a/src/sql/src/lib.rs +++ b/src/sql/src/lib.rs @@ -23,6 +23,6 @@ pub mod partition; pub mod statements; pub mod util; -pub use parsers::create_parser::{ENGINE, MAXVALUE}; +pub use parsers::create_parser::{ENGINE, MAXVALUE, parse_json2_type_hint_path}; pub use parsers::tql_parser::TQL; pub use parsers::with_tql_parser::{CteContent, HybridCteWith}; diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index be7ed0670d..0e0ece1b33 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -24,6 +24,7 @@ use datafusion_common::ScalarValue; use datatypes::arrow::datatypes::{DataType as ArrowDataType, IntervalUnit}; use datatypes::data_type::ConcreteDataType; use itertools::Itertools; +pub use json::parse_json2_type_hint_path; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::{ ColumnOption, ColumnOptionDef, DataType, Expr, KeyOrIndexDisplay, NullsDistinctOption, diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 63399f81a1..2f8a80c79a 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -21,6 +21,7 @@ use sqlparser::parser::Parser; use sqlparser::tokenizer::Token; use crate::ast::Ident; +use crate::dialect::GreptimeDbDialect; use crate::error::{InvalidSqlSnafu, Result, SyntaxSnafu}; use crate::parsers::create_parser::{INVERTED, SKIPPING}; use crate::statements::create::{Json2Options, JsonTypeHint}; @@ -29,6 +30,25 @@ use crate::statements::transform::type_alias::get_type_by_alias; const JSON2_TYPE_NAME: &str = "JSON2"; const MAX_AUTO_EXPANDED_PATHS: &str = "max_auto_expanded_paths"; +/// Parses a JSON2 type hint path with the same grammar used by `CREATE TABLE`. +pub fn parse_json2_type_hint_path(path: &str) -> Result> { + let dialect = GreptimeDbDialect {}; + let mut parser = Parser::new(&dialect) + .try_with_sql(path) + .context(SyntaxSnafu)?; + let path = parse_json2_path(&mut parser)?; + ensure!( + parser.peek_token().token == Token::EOF, + InvalidSqlSnafu { + msg: format!( + "unexpected token '{}' in JSON2 type hint path", + parser.peek_token() + ) + } + ); + Ok(path) +} + pub(super) fn parse_json2_type_and_options( parser: &mut Parser<'_>, ) -> Result)>> { @@ -322,6 +342,7 @@ fn ensure_no_path_conflict(hints: &[JsonTypeHint], path: &[String]) -> Result<() mod tests { use sqlparser::ast::{DataType, ExactNumberInfo}; + use super::parse_json2_type_hint_path; use crate::dialect::GreptimeDbDialect; use crate::parser::{ParseOptions, ParserContext}; use crate::statements::create::Column; @@ -339,6 +360,15 @@ mod tests { create_table.columns.remove(0) } + #[test] + fn test_parse_json2_type_hint_path() { + assert_eq!( + parse_json2_type_hint_path(r#"attrs."http.status_code""#).unwrap(), + vec!["attrs", "http.status_code"] + ); + assert!(parse_json2_type_hint_path("user.id trailing").is_err()); + } + #[test] fn test_parse_json2_type_hints() { let column = parse_json2_column(