diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index 0a8c999461a..d8e3004aa2b 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -449,7 +449,14 @@ impl TryFrom for ColumnDataTypeWrapper { }), JsonFormat::Json2(native_type) => { if native_type.is_null() { - None + Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonNativeType(Box::new( + JsonNativeTypeExtension { + datatype: ColumnDataType::Json as i32, + datatype_extension: None, + }, + ))), + }) } else { let concrete_type = ConcreteDataType::from_arrow_type(&native_type.as_arrow_type()); @@ -1291,6 +1298,32 @@ mod tests { assert!(values.bool_values.is_empty()); } + #[test] + fn test_json2_unknown_type_encoding() { + let datatype = ConcreteDataType::json2(JsonNativeType::Null); + let wrapper = ColumnDataTypeWrapper::try_from(datatype.clone()).unwrap(); + assert_eq!( + wrapper.to_parts(), + ( + ColumnDataType::Json, + Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonNativeType(Box::new(JsonNativeTypeExtension { + datatype: ColumnDataType::Json as i32, + datatype_extension: None, + }))), + }), + ) + ); + assert_eq!(ConcreteDataType::from(wrapper), datatype); + + for extension in [None, Some(ColumnDataTypeExtension::default())] { + assert_eq!( + ConcreteDataType::from(ColumnDataTypeWrapper::new(ColumnDataType::Json, extension)), + datatype + ); + } + } + #[test] fn test_concrete_datatype_from_column_datatype() { assert_eq!( diff --git a/src/operator/src/req_convert/common.rs b/src/operator/src/req_convert/common.rs index 63226ef8e48..7b4d7de2bb1 100644 --- a/src/operator/src/req_convert/common.rs +++ b/src/operator/src/req_convert/common.rs @@ -79,12 +79,28 @@ fn prepare_rows(rows: &mut Option) -> Result<()> { .enumerate() .filter_map(|(idx, schema)| { if schema.datatype() == ColumnDataType::Json { - Some(idx) + match schema + .datatype_extension + .as_ref() + .and_then(|x| x.type_ext.as_ref()) + { + None | Some(TypeExt::JsonType(_)) => Some(Ok(idx)), + Some(TypeExt::JsonNativeType(_)) => None, + Some(_) => Some( + InvalidInsertRequestSnafu { + reason: format!( + "Invalid type extension for JSON column '{}'", + schema.column_name + ), + } + .fail(), + ), + } } else { None } }) - .collect::>(); + .collect::>>()?; for idx in &indexes { let column = &mut rows.schema[*idx]; column.datatype_extension = Some(ColumnDataTypeExtension { @@ -390,9 +406,65 @@ mod tests { use api::v1::column::Values; use api::v1::{SemanticType, VectorTypeExtension}; use common_base::bit_vec::prelude::*; + use common_error::ext::WhateverResult; use super::*; + #[test] + fn test_prepare_jsonb_rows() -> WhateverResult<()> { + let binary = Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }); + for extension in [ + None, + Some(ColumnDataTypeExtension::default()), + binary.clone(), + ] { + for value in [Some(ValueData::StringValue(r#"{"key":1}"#.into())), None] { + let expected = value.as_ref().map(|_| { + ValueData::BinaryValue(jsonb::parse_value(br#"{"key":1}"#).unwrap().to_vec()) + }); + let mut rows = Some(Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::Json as i32, + datatype_extension: extension.clone(), + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { value_data: value }], + }], + }); + prepare_rows(&mut rows)?; + let rows = rows.unwrap(); + assert_eq!(rows.schema[0].datatype_extension, binary); + assert_eq!(rows.rows[0].values[0].value_data, expected); + } + } + Ok(()) + } + + #[test] + fn test_preprocess_json_invalid_extension() { + let mut rows = Some(Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::Json as i32, + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: 3 })), + }), + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { value_data: None }], + }], + }); + assert!( + prepare_rows(&mut rows) + .unwrap_err() + .to_string() + .contains("Invalid type extension") + ); + } + #[test] fn test_request_column_to_row() { let columns = vec![ diff --git a/src/pipeline/tests/json_parse.rs b/src/pipeline/tests/json_parse.rs index c374c846cde..b0c32950518 100644 --- a/src/pipeline/tests/json_parse.rs +++ b/src/pipeline/tests/json_parse.rs @@ -23,6 +23,8 @@ use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; +use greptime_proto::v1::column_data_type_extension::TypeExt; +use greptime_proto::v1::{ColumnDataTypeExtension, JsonNativeTypeExtension}; const INPUT_VALUE_OBJ: &str = r#" [ @@ -138,7 +140,15 @@ transform: 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()); + + let expected = Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonNativeType(Box::new(JsonNativeTypeExtension { + datatype: ColumnDataType::Json as _, + datatype_extension: None, + }))), + }); + assert_eq!(output.schema[0].datatype_extension, expected); + assert_eq!( output.schema[0] .options diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 9607d611627..025d56b959a 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -1331,6 +1331,125 @@ async fn test_skip_wal_grpc_unary_stream_and_flight_sql(distributed: bool) { env.shutdown().await; } +#[apply(both_deployment_cases)] +async fn test_grpc_json2_row_inserts(distributed: bool) { + let mut env = MockInstanceImpl::new("grpc_json2_row_inserts", distributed).await; + let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await; + let addr = server.bind_addr().unwrap().to_string(); + let database = Database::new_with_dbname("greptime-public", Client::with_urls(vec![addr])); + database + .sql( + "CREATE TABLE grpc_json2 (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, j JSON2) \ + WITH (append_mode='true', 'memtable.type'='bulk')", + ) + .await + .unwrap(); + + let (datatype, extension) = + api::helper::ColumnDataTypeWrapper::try_from(datatypes::prelude::ConcreteDataType::json2( + datatypes::types::json_type::JsonNativeType::Null, + )) + .unwrap() + .into_parts(); + + // Send the all-NULL JSON2 batch separately to check its type marker. + for (index, (host, payload)) in [ + ( + "host-a", + Some(serde_json::json!({ + "active": true, + "nested": {"items": [1, "two", null, {"ok": false}], "ratio": 1.5}, + "tags": ["api", "prod"] + })), + ), + ( + "host-b", + Some(serde_json::json!({ + "active": false, + "nested": {"items": [-2, "three", null, {"ok": true}], "ratio": -0.25}, + "tags": ["worker"] + })), + ), + ("host-a", None), + ] + .into_iter() + .enumerate() + { + let output = database + .row_inserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "grpc_json2".into(), + rows: Some(Rows { + schema: vec![ + ColumnSchema { + column_name: "ts".into(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "host".into(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "j".into(), + datatype: datatype as i32, + datatype_extension: extension.clone(), + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue( + index as i64, + )), + }, + Value { + value_data: Some(ValueData::StringValue(host.into())), + }, + Value { + value_data: payload.map(|x| { + ValueData::JsonValue(api::helper::encode_json_value( + x.into(), + )) + }), + }, + ], + }], + }), + }], + }) + .await + .unwrap(); + assert_eq!(output, 1); + } + + let output = database + .sql("SELECT host, j FROM grpc_json2 ORDER BY ts") + .await + .unwrap(); + let batches = match output.data { + OutputData::RecordBatches(batches) => batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + let pretty = batches.pretty_print().unwrap(); + let expected = r#"+--------+---------------------------------------------------------------------------------------------------+ +| host | j | ++--------+---------------------------------------------------------------------------------------------------+ +| host-a | {"active":true,"nested":{"items":[1,"two",null,{"ok":false}],"ratio":1.5},"tags":["api","prod"]} | +| host-b | {"active":false,"nested":{"items":[-2,"three",null,{"ok":true}],"ratio":-0.25},"tags":["worker"]} | +| host-a | | ++--------+---------------------------------------------------------------------------------------------------+"#; + assert_eq!(pretty, expected); + server.shutdown().await.unwrap(); + env.shutdown().await; +} + pub async fn test_insert_and_select(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_insert_and_select").await;