mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-13 22:36:57 +00:00
fix: support native JSON2 row inserts over gRPC
Signed-off-by: luofucong <luofc@foxmail.com>
This commit is contained in:
luofucong
+34
-1
@@ -449,7 +449,14 @@ impl TryFrom<ConcreteDataType> 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!(
|
||||
|
||||
@@ -79,12 +79,28 @@ fn prepare_rows(rows: &mut Option<Rows>) -> 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::<Vec<_>>();
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
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![
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user