feat: support to insert json data via grpc protocol (#4908)

* feat: support to insert json data via grpc protocol

* chore: handle error

* feat: introduce `prepare_rows`

* chore: fmt toml

* test: add row deletion test

* test: fix unit test

* chore: remove log

* chore: apply suggestions from CR
This commit is contained in:
Weny Xu
2024-11-04 16:55:47 +08:00
committed by GitHub
parent 191755fc42
commit 4ab6dc2825
9 changed files with 248 additions and 75 deletions

View File

@@ -36,6 +36,7 @@ datatypes.workspace = true
file-engine.workspace = true
futures = "0.3"
futures-util.workspace = true
jsonb.workspace = true
lazy_static.workspace = true
meta-client.workspace = true
meter-core.workspace = true

View File

@@ -35,6 +35,7 @@ use crate::error::{
MissingTimeIndexColumnSnafu, RequestDeletesSnafu, Result, TableNotFoundSnafu,
};
use crate::region_req_factory::RegionRequestFactory;
use crate::req_convert::common::preprocess_row_delete_requests;
use crate::req_convert::delete::{ColumnToRow, RowToRegion, TableToRegion};
pub struct Deleter {
@@ -72,6 +73,7 @@ impl Deleter {
mut requests: RowDeleteRequests,
ctx: QueryContextRef,
) -> Result<Output> {
preprocess_row_delete_requests(&mut requests.deletes)?;
// remove empty requests
requests.deletes.retain(|req| {
req.rows

View File

@@ -770,6 +770,13 @@ pub enum Error {
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Invalid json text: {}", json))]
InvalidJsonFormat {
#[snafu(implicit)]
location: Location,
json: String,
},
}
pub type Result<T> = std::result::Result<T, Error>;
@@ -808,7 +815,8 @@ impl ErrorExt for Error {
| Error::BuildAdminFunctionArgs { .. }
| Error::FunctionArityMismatch { .. }
| Error::InvalidPartition { .. }
| Error::PhysicalExpr { .. } => StatusCode::InvalidArguments,
| Error::PhysicalExpr { .. }
| Error::InvalidJsonFormat { .. } => StatusCode::InvalidArguments,
Error::TableAlreadyExists { .. } | Error::ViewAlreadyExists { .. } => {
StatusCode::TableAlreadyExists

View File

@@ -54,6 +54,7 @@ use crate::error::{
};
use crate::expr_factory::CreateExprFactory;
use crate::region_req_factory::RegionRequestFactory;
use crate::req_convert::common::preprocess_row_insert_requests;
use crate::req_convert::insert::{ColumnToRow, RowToRegion, StatementToRegion, TableToRegion};
use crate::statement::StatementExecutor;
@@ -119,10 +120,11 @@ impl Inserter {
/// Handles row inserts request and creates a physical table on demand.
pub async fn handle_row_inserts(
&self,
requests: RowInsertRequests,
mut requests: RowInsertRequests,
ctx: QueryContextRef,
statement_executor: &StatementExecutor,
) -> Result<Output> {
preprocess_row_insert_requests(&mut requests.inserts)?;
self.handle_row_inserts_with_create_type(
requests,
ctx,

View File

@@ -17,9 +17,13 @@ pub(crate) mod partitioner;
use std::collections::HashMap;
use api::helper::ColumnDataTypeWrapper;
use api::v1::column_data_type_extension::TypeExt;
use api::v1::column_def::options_from_column_schema;
use api::v1::value::ValueData;
use api::v1::{Column, ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value};
use api::v1::{
Column, ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row,
RowDeleteRequest, RowInsertRequest, Rows, SemanticType, Value,
};
use common_base::BitVec;
use datatypes::vectors::VectorRef;
use snafu::prelude::*;
@@ -27,10 +31,77 @@ use snafu::ResultExt;
use table::metadata::TableInfo;
use crate::error::{
ColumnDataTypeSnafu, ColumnNotFoundSnafu, InvalidInsertRequestSnafu,
MissingTimeIndexColumnSnafu, Result,
ColumnDataTypeSnafu, ColumnNotFoundSnafu, InvalidInsertRequestSnafu, InvalidJsonFormatSnafu,
MissingTimeIndexColumnSnafu, Result, UnexpectedSnafu,
};
/// Encodes a string value as JSONB binary data if the value is of `StringValue` type.
fn encode_string_to_jsonb_binary(value_data: ValueData) -> Result<ValueData> {
if let ValueData::StringValue(json) = &value_data {
let binary = jsonb::parse_value(json.as_bytes())
.map_err(|_| InvalidJsonFormatSnafu { json }.build())
.map(|jsonb| jsonb.to_vec())?;
Ok(ValueData::BinaryValue(binary))
} else {
UnexpectedSnafu {
violated: "Expected to value data to be a string.",
}
.fail()
}
}
/// Prepares row insertion requests by converting any JSON values to binary JSONB format.
pub fn preprocess_row_insert_requests(requests: &mut Vec<RowInsertRequest>) -> Result<()> {
for request in requests {
prepare_rows(&mut request.rows)?;
}
Ok(())
}
/// Prepares row deletion requests by converting any JSON values to binary JSONB format.
pub fn preprocess_row_delete_requests(requests: &mut Vec<RowDeleteRequest>) -> Result<()> {
for request in requests {
prepare_rows(&mut request.rows)?;
}
Ok(())
}
fn prepare_rows(rows: &mut Option<Rows>) -> Result<()> {
if let Some(rows) = rows {
let indexes = rows
.schema
.iter()
.enumerate()
.filter_map(|(idx, schema)| {
if schema.datatype() == ColumnDataType::Json {
Some(idx)
} else {
None
}
})
.collect::<Vec<_>>();
for idx in &indexes {
let column = &mut rows.schema[*idx];
column.datatype_extension = Some(ColumnDataTypeExtension {
type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())),
});
column.datatype = ColumnDataType::Binary.into();
}
for idx in &indexes {
for row in &mut rows.rows {
if let Some(value_data) = row.values[*idx].value_data.take() {
row.values[*idx].value_data = Some(encode_string_to_jsonb_binary(value_data)?);
}
}
}
}
Ok(())
}
pub fn columns_to_rows(columns: Vec<Column>, row_count: u32) -> Result<Rows> {
let row_count = row_count as usize;
let column_count = columns.len();