From 809283e6b4ed780df6c2fa79be0cd97937cf5ea5 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 16 Sep 2026 14:34:22 +0000 Subject: [PATCH] feat: add prepared batch write primitives (#9186) * refactor: expose shared operator test fixtures Signed-off-by: WenyXu * feat: convert prepared table rows to Arrow batches Signed-off-by: WenyXu * feat: write prepared table batches through bulk insert Signed-off-by: WenyXu * feat: define prepared batch submission interface Signed-off-by: WenyXu * fix: validate prepared batch values and clarify write contracts Signed-off-by: WenyXu * fix: reuse JSON2 array schema alignment for prepared writes Signed-off-by: WenyXu * ci: refresh compatibility window for v1.2.1 Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/api/src/helper.rs | 25 + src/datatypes/src/extension/json.rs | 24 +- src/mito2/src/memtable/bulk/part.rs | 22 +- src/mito2/src/request.rs | 10 +- src/operator/src/batcher.rs | 41 ++ src/operator/src/bulk_insert.rs | 95 ++- src/operator/src/error.rs | 2 +- src/operator/src/insert.rs | 2 +- src/operator/src/lib.rs | 5 +- .../src/req_convert/delete/table_to_region.rs | 2 +- src/operator/src/req_convert/insert.rs | 2 + .../src/req_convert/insert/row_to_batch.rs | 591 ++++++++++++++++++ .../src/req_convert/insert/row_to_region.rs | 2 +- .../src/req_convert/insert/table_to_region.rs | 2 +- src/operator/src/{tests.rs => test_util.rs} | 4 +- .../src/{tests => test_util}/kv_backend.rs | 0 .../{tests => test_util}/partition_manager.rs | 10 +- 17 files changed, 789 insertions(+), 50 deletions(-) create mode 100644 src/operator/src/batcher.rs create mode 100644 src/operator/src/req_convert/insert/row_to_batch.rs rename src/operator/src/{tests.rs => test_util.rs} (79%) rename src/operator/src/{tests => test_util}/kv_backend.rs (100%) rename src/operator/src/{tests => test_util}/partition_manager.rs (98%) diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index d8e3004aa2..40f9bdfb6d 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -1033,6 +1033,17 @@ pub fn proto_value_type(value: &v1::Value) -> Option { Some(value_type) } +/// Checks protobuf value types using the write-path compatibility rules. +/// Binary values also represent JSON and vector columns. +pub fn proto_value_type_match(column_type: ColumnDataType, value_type: ColumnDataType) -> bool { + match (column_type, value_type) { + (ct, vt) if ct == vt => true, + (ColumnDataType::Vector, ColumnDataType::Binary) => true, + (ColumnDataType::Json, ColumnDataType::Binary) => true, + _ => false, + } +} + pub fn vectors_to_rows<'a>( columns: impl Iterator, row_count: usize, @@ -2085,4 +2096,18 @@ mod tests { let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); } + + #[test] + fn test_proto_value_type_match() { + for (column, value, expected) in [ + (ColumnDataType::Int32, ColumnDataType::Int32, true), + (ColumnDataType::Json, ColumnDataType::Binary, true), + (ColumnDataType::Vector, ColumnDataType::Binary, true), + (ColumnDataType::Float64, ColumnDataType::List, false), + (ColumnDataType::Float64, ColumnDataType::Struct, false), + (ColumnDataType::Binary, ColumnDataType::Json, false), + ] { + assert_eq!(expected, proto_value_type_match(column, value)); + } + } } diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index b912a0a3f4..3aa0a4c979 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -15,10 +15,11 @@ use std::collections::HashMap; use std::sync::Arc; +use arrow_array::ArrayRef; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; -use arrow_schema::{ArrowError, DataType, Field, FieldRef}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaRef}; use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; @@ -27,6 +28,27 @@ use crate::error::InvalidJson2LayoutSnafu; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; +/// Aligns JSON2 field types with their built arrays while preserving field and schema metadata. +pub fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef { + if schema.fields().iter().all(|f| !is_json2_extension_type(f)) { + return schema; + } + + let mut fields = Vec::with_capacity(schema.fields().len()); + for (field, array) in schema.fields().iter().zip(columns) { + if !is_json2_extension_type(field) { + fields.push(field.clone()); + continue; + } + + let mut field = field.as_ref().clone(); + field.set_data_type(array.data_type().clone()); + fields.push(Arc::new(field)); + } + + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; const JSON2_LAYOUT_V1: u8 = 1; const JSON2_LAYOUT_V2: u8 = 2; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index c3591cd4c0..738bae45f4 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -39,7 +39,7 @@ use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type, }; use datatypes::data_type::DataType; -use datatypes::extension::json::is_json2_extension_type; +use datatypes::extension::json::align_schema_with_json_array; use datatypes::prelude::{MutableVector, Vector}; use datatypes::value::ValueRef; use datatypes::vectors::Helper; @@ -741,26 +741,6 @@ impl BulkPartConverter { } } -fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef { - if schema.fields().iter().all(|f| !is_json2_extension_type(f)) { - return schema; - } - - let mut fields = Vec::with_capacity(schema.fields().len()); - for (field, array) in schema.fields().iter().zip(columns) { - if !is_json2_extension_type(field) { - fields.push(field.clone()); - continue; - } - - let mut field = field.as_ref().clone(); - field.set_data_type(array.data_type().clone()); - fields.push(Arc::new(field)); - } - - Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) -} - fn new_primary_key_column_builders( metadata: &RegionMetadata, capacity: usize, diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 16b7eb4a6b..1000257000 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -20,6 +20,7 @@ use std::time::Instant; use api::helper::{ ColumnDataTypeWrapper, is_column_type_value_eq, is_semantic_type_eq, proto_value_type, + proto_value_type_match, }; use api::v1::column_def::options_from_column_schema; use api::v1::{ColumnDataType, ColumnSchema, OpType, Rows, SemanticType, Value, WriteHint}; @@ -472,15 +473,6 @@ pub(crate) fn validate_proto_value( Ok(()) } -fn proto_value_type_match(column_type: ColumnDataType, value_type: ColumnDataType) -> bool { - match (column_type, value_type) { - (ct, vt) if ct == vt => true, - (ColumnDataType::Vector, ColumnDataType::Binary) => true, - (ColumnDataType::Json, ColumnDataType::Binary) => true, - _ => false, - } -} - /// Oneshot output result sender. #[derive(Debug)] pub struct OutputTx(Sender>); diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs new file mode 100644 index 0000000000..4c5e9892c0 --- /dev/null +++ b/src/operator/src/batcher.rs @@ -0,0 +1,41 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use session::context::QueryContextRef; +use table::metadata::TableInfoRef; +use tokio::sync::OwnedSemaphorePermit; + +use crate::error::Result; + +/// Accepts prepared table writes without coupling the inserter to a batcher implementation. +/// Schema creation, alteration and default evaluation remain the caller's responsibility. +#[async_trait] +pub trait PendingRowsBatcher: Send + Sync { + /// Acquires one slot per original request, shared by all of its table submissions. + async fn acquire(&self) -> Result>; + + /// Waits for the submitted rows to be written, retaining the slot through completion. + /// Cancelling the response wait does not retract an already enqueued submission. + async fn submit( + &self, + table_info: TableInfoRef, + batch: RecordBatch, + ctx: QueryContextRef, + permit: Arc, + ) -> Result; +} diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index c47236d7db..51a08d2f32 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -20,23 +20,106 @@ use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; +use arrow::compute::filter_record_batch; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; +use common_error::ext::BoxedError; use common_grpc::FlightData; -use common_grpc::flight::{FlightEncoder, FlightMessage}; +use common_grpc::flight::{FlightEncoder, FlightMessage, record_batch_to_ipc}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; +use futures::future::{join_all, try_join_all}; +use session::context::QueryContextRef; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; +use crate::error::Result; use crate::insert::Inserter; use crate::req_convert::insert::extract_timestamps; use crate::{error, metrics}; impl Inserter { + /// Routes and writes a prepared table batch, returning the affected row count. + /// + /// Callers must exclude instant-TTL tables and handle Flow notifications after + /// successful writes. This execution helper does not perform either step. + pub async fn flush_bulk_batch( + &self, + table_info: TableInfoRef, + batch: RecordBatch, + ctx: QueryContextRef, + ) -> Result { + let (rule, versions) = self + .partition_manager + .find_table_partition_rule(&table_info) + .await + .context(error::InvalidPartitionSnafu)?; + let masks = rule + .split_record_batch(&batch) + .context(error::SplitInsertSnafu)?; + let mut writes = Vec::with_capacity(masks.len()); + for (region_number, mask) in masks { + if mask.select_none() { + continue; + } + let region_id = RegionId::new(table_info.table_id(), region_number); + let selected = if mask.select_all() { + batch.clone() + } else { + filter_record_batch(&batch, mask.array()).context(error::ComputeArrowSnafu)? + }; + let (schema, data_header, payload) = record_batch_to_ipc(selected) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + let peer = self + .partition_manager + .find_region_leader(region_id) + .await + .context(error::FindRegionLeaderSnafu)?; + let request = RegionRequest { + header: Some(RegionRequestHeader { + dbname: ctx.get_db_string(), + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal: ctx.skip_wal(), + region_id: region_id.as_u64(), + partition_expr_version: versions + .get(®ion_number) + .copied() + .flatten() + .map(|value| PartitionExprVersion { value }), + // Let the datanode revalidate against its current schema. + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema, + data_header, + payload, + })), + })), + }; + writes.push((peer, request)); + } + let results = join_all(writes.into_iter().map(|(peer, request)| async move { + self.node_manager + .datanode(&peer) + .await + .handle(request) + .await + .context(error::RequestInsertsSnafu) + })) + .await; + let affected_rows = results + .into_iter() + .map(|result| result.map(|response| response.affected_rows)) + .sum::>()?; + Ok(affected_rows) + } + /// Handle bulk insert request. pub async fn handle_bulk_insert( &self, @@ -45,7 +128,7 @@ impl Inserter { record_batch: RecordBatch, schema_bytes: Bytes, skip_wal: bool, - ) -> error::Result { + ) -> Result { let table_info = table.table_info(); let table_id = table_info.table_id(); let db_name = table_info.get_db_string(); @@ -175,7 +258,7 @@ impl Inserter { } else { None }; - let handle: common_runtime::JoinHandle> = + let handle: common_runtime::JoinHandle> = common_runtime::spawn_global(async move { let (header, payload) = if mask.select_all() { // SAFETY: raw data must be present, we can avoid re-encoding. @@ -184,7 +267,7 @@ impl Inserter { let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED .with_label_values(&["filter"]) .start_timer(); - let batch = arrow::compute::filter_record_batch(&rb, mask.array()) + let batch = filter_record_batch(&rb, mask.array()) .context(error::ComputeArrowSnafu)?; filter_timer.observe_duration(); metrics::BULK_REQUEST_ROWS @@ -243,9 +326,7 @@ impl Inserter { } } - let region_responses = futures::future::try_join_all(handles) - .await - .context(error::JoinTaskSnafu)?; + let region_responses = try_join_all(handles).await.context(error::JoinTaskSnafu)?; wait_all_datanode_timer.observe_duration(); let mut rows_inserted: usize = 0; for res in region_responses { diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 3c1a6ae360..8caf829e1e 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -1183,7 +1183,7 @@ define_into_tonic_status!(Error); #[cfg(test)] mod tests { - use super::*; + use crate::error::*; #[test] fn admin_function_preserves_external_error_metadata() { diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a4a69ef8ba..c84dda1999 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -1559,7 +1559,7 @@ mod tests { use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType}; use super::*; - use crate::tests::{create_partition_rule_manager, prepare_mocked_backend}; + use crate::test_util::{create_partition_rule_manager, prepare_mocked_backend}; fn make_table_ref_with_schema( ts_name: &str, diff --git a/src/operator/src/lib.rs b/src/operator/src/lib.rs index 5e723faeb5..a971ca41f7 100644 --- a/src/operator/src/lib.rs +++ b/src/operator/src/lib.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod batcher; mod bulk_insert; pub mod delete; pub mod error; @@ -25,6 +26,6 @@ pub mod req_convert; pub mod request; pub mod statement; pub mod table; -#[cfg(test)] -pub(crate) mod tests; +#[cfg(any(test, feature = "testing"))] +pub mod test_util; pub mod utils; diff --git a/src/operator/src/req_convert/delete/table_to_region.rs b/src/operator/src/req_convert/delete/table_to_region.rs index b6fde9b105..0d40e30350 100644 --- a/src/operator/src/req_convert/delete/table_to_region.rs +++ b/src/operator/src/req_convert/delete/table_to_region.rs @@ -62,7 +62,7 @@ mod tests { use store_api::storage::RegionId; use super::*; - use crate::tests::{ + use crate::test_util::{ create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, }; diff --git a/src/operator/src/req_convert/insert.rs b/src/operator/src/req_convert/insert.rs index 72d6693c6b..e36caf2f8c 100644 --- a/src/operator/src/req_convert/insert.rs +++ b/src/operator/src/req_convert/insert.rs @@ -14,6 +14,7 @@ mod column_to_row; mod fill_impure_default; +mod row_to_batch; mod row_to_region; mod stmt_to_region; mod table_to_region; @@ -22,6 +23,7 @@ mod timestamps; use api::v1::SemanticType; pub use column_to_row::ColumnToRow; pub use fill_impure_default::fill_reqs_with_impure_default; +pub use row_to_batch::rows_to_record_batch; pub use row_to_region::RowToRegion; use snafu::{OptionExt, ResultExt}; pub use stmt_to_region::StatementToRegion; diff --git a/src/operator/src/req_convert/insert/row_to_batch.rs b/src/operator/src/req_convert/insert/row_to_batch.rs new file mode 100644 index 0000000000..ccf313737e --- /dev/null +++ b/src/operator/src/req_convert/insert/row_to_batch.rs @@ -0,0 +1,591 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use api::helper::{ + ColumnDataTypeWrapper, pb_value_to_value_ref, proto_value_type, proto_value_type_match, +}; +use api::v1::column_data_type_extension::TypeExt; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, ColumnDataTypeExtension, Rows, SemanticType, Value}; +use arrow::record_batch::RecordBatch; +use common_error::ext::BoxedError; +use datatypes::data_type::ConcreteDataType; +use datatypes::extension::json::align_schema_with_json_array; +use snafu::{OptionExt, ResultExt, ensure}; +use table::metadata::TableInfo; + +use crate::error::{self, Result}; + +/// Converts prepared rows in target-schema order without Prom-specific renaming. +/// +/// Missing columns evaluate their defaults once per conversion. Explicit nulls +/// remain nulls. Semantic types are checked against the table time index and +/// primary-key indices. +/// +/// # Panics +/// +/// Panics if `rows.rows` is empty. Callers must skip empty writes before conversion. +pub fn rows_to_record_batch(rows: &Rows, table_info: &TableInfo) -> Result { + assert!(!rows.rows.is_empty(), "prepared rows must not be empty"); + let schema = &table_info.meta.schema; + let mut source_columns = HashMap::with_capacity(rows.schema.len()); + for (index, source) in rows.schema.iter().enumerate() { + ensure!( + source_columns + .insert(source.column_name.as_str(), index) + .is_none(), + error::InvalidInsertRequestSnafu { + reason: format!("Duplicate input column {}", source.column_name), + } + ); + let target = schema + .column_schema_by_name(&source.column_name) + .with_context(|| error::InvalidInsertRequestSnafu { + reason: format!("Unknown input column {}", source.column_name), + })?; + let data_type = + ColumnDataTypeWrapper::try_new(source.datatype, source.datatype_extension.clone()) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + ensure!( + ConcreteDataType::from(data_type) == target.data_type, + error::InvalidInsertRequestSnafu { + reason: format!("Input datatype differs for column {}", source.column_name), + } + ); + let semantic = SemanticType::try_from(source.semantic_type) + .ok() + .with_context(|| error::InvalidInsertRequestSnafu { + reason: format!("Invalid semantic type for column {}", source.column_name), + })?; + let is_tag = table_info.meta.primary_key_indices.iter().any(|&index| { + schema + .column_schemas() + .get(index) + .is_some_and(|column| column.name == source.column_name) + }); + let expected_semantic = if target.is_time_index() { + SemanticType::Timestamp + } else if is_tag { + SemanticType::Tag + } else { + SemanticType::Field + }; + ensure!( + semantic == expected_semantic, + error::InvalidInsertRequestSnafu { + reason: format!("Input semantics differ for column {}", source.column_name), + } + ); + } + for row in &rows.rows { + ensure!( + row.values.len() == rows.schema.len(), + error::InvalidInsertRequestSnafu { + reason: format!( + "Expected {} values, got {}", + rows.schema.len(), + row.values.len() + ), + } + ); + } + + let mut arrays = Vec::with_capacity(schema.num_columns()); + for column in schema.column_schemas() { + let vector = if let Some(&index) = source_columns.get(column.name.as_str()) { + let mut builder = column.create_mutable_vector(rows.rows.len()); + for row in &rows.rows { + let value = &row.values[index]; + ensure!( + value.value_data.is_some() || column.is_nullable(), + error::InvalidInsertRequestSnafu { + reason: format!("Null supplied for non-nullable column {}", column.name), + } + ); + ensure!( + value_matches_type( + value, + rows.schema[index].datatype, + rows.schema[index].datatype_extension.as_ref(), + ), + error::InvalidInsertRequestSnafu { + reason: format!("Value datatype differs for column {}", column.name), + } + ); + builder + .try_push_value_ref(&pb_value_to_value_ref( + value, + rows.schema[index].datatype_extension.as_ref(), + )) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + } + builder.to_vector() + } else { + column + .create_default_vector(rows.rows.len()) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)? + .with_context(|| error::InvalidInsertRequestSnafu { + reason: format!("Missing required column {}", column.name), + })? + }; + arrays.push(vector.to_arrow_array()); + } + let arrow_schema = align_schema_with_json_array(schema.arrow_schema().clone(), &arrays); + RecordBatch::try_new(arrow_schema, arrays).context(error::ComputeArrowSnafu) +} + +// Validate nested values before the infallible protobuf conversion reads type extensions. +fn value_matches_type( + value: &Value, + datatype: i32, + extension: Option<&ColumnDataTypeExtension>, +) -> bool { + let Some(value_type) = proto_value_type(value) else { + return true; + }; + let Ok(column_type) = ColumnDataType::try_from(datatype) else { + return false; + }; + if !proto_value_type_match(column_type, value_type) { + return false; + } + match value.value_data.as_ref() { + Some(ValueData::ListValue(list)) => { + let Some(TypeExt::ListType(item)) = extension.and_then(|ext| ext.type_ext.as_ref()) + else { + return false; + }; + list.items.iter().all(|value| { + value_matches_type(value, item.datatype, item.datatype_extension.as_deref()) + }) + } + Some(ValueData::StructValue(value)) => { + let Some(TypeExt::StructType(schema)) = extension.and_then(|ext| ext.type_ext.as_ref()) + else { + return false; + }; + value.items.len() == schema.fields.len() + && value + .items + .iter() + .zip(&schema.fields) + .all(|(value, field)| { + value_matches_type(value, field.datatype, field.datatype_extension.as_ref()) + }) + } + _ => true, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::thread::yield_now; + + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, ColumnSchema as ProtoColumnSchema, Row, Value}; + use arrow::array::{Array, Int32Array, StringArray}; + use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema}; + use datatypes::value::Value as DtValue; + use table::metadata::{TableInfoBuilder, TableMetaBuilder}; + + use crate::req_convert::insert::row_to_batch::*; + + fn source(name: &str, datatype: ColumnDataType) -> ProtoColumnSchema { + ProtoColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + } + } + + fn table_info(schema: Schema) -> TableInfo { + let next_column_id = schema.num_columns() as u32; + TableInfoBuilder::default() + .table_id(1) + .table_version(0) + .name("test") + .meta( + TableMetaBuilder::empty() + .schema(Arc::new(schema)) + .primary_key_indices(vec![]) + .next_column_id(next_column_id) + .engine("mito") + .build() + .unwrap(), + ) + .build() + .unwrap() + } + + fn fixture() -> (Rows, Schema) { + let schema = Schema::new(vec![ + ColumnSchema::new("count", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new("label", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("fallback", ConcreteDataType::int32_datatype(), true) + .with_default_constraint(Some(ColumnDefaultConstraint::Value(DtValue::Int32(7)))) + .unwrap(), + ]); + let rows = Rows { + schema: vec![ + source("label", ColumnDataType::String), + source("count", ColumnDataType::Int32), + ], + rows: vec![Row { + values: vec![ + Value::default(), + Value { + value_data: Some(ValueData::I32Value(3)), + }, + ], + }], + }; + (rows, schema) + } + + #[test] + fn test_reorder_multifield_and_defaults() { + let (rows, schema) = fixture(); + let batch = rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap(); + assert_eq!(batch.schema(), schema.arrow_schema().clone()); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3 + ); + assert!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .is_null(0) + ); + assert_eq!( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7 + ); + } + + #[test] + fn test_explicit_null_does_not_use_default() { + let (mut rows, schema) = fixture(); + rows.schema.push(source("fallback", ColumnDataType::Int32)); + rows.rows[0].values.push(Value::default()); + assert!( + rows_to_record_batch(&rows, &table_info(schema.clone())) + .unwrap() + .column(2) + .is_null(0) + ); + } + + #[test] + fn test_timestamp_precisions() { + for (datatype, value) in [ + ( + ColumnDataType::TimestampSecond, + ValueData::TimestampSecondValue(42), + ), + ( + ColumnDataType::TimestampMillisecond, + ValueData::TimestampMillisecondValue(42), + ), + ( + ColumnDataType::TimestampMicrosecond, + ValueData::TimestampMicrosecondValue(42), + ), + ( + ColumnDataType::TimestampNanosecond, + ValueData::TimestampNanosecondValue(42), + ), + ] { + let target_type = ConcreteDataType::from(ColumnDataTypeWrapper::new(datatype, None)); + let schema = Schema::new(vec![ + ColumnSchema::new("ts", target_type, false).with_time_index(true), + ]); + let mut column = source("ts", datatype); + column.semantic_type = SemanticType::Timestamp as i32; + let rows = Rows { + schema: vec![column], + rows: vec![Row { + values: vec![Value { + value_data: Some(value), + }], + }], + }; + let batch = rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap(); + assert_eq!(batch.schema(), schema.arrow_schema().clone()); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.column(0).null_count(), 0); + } + } + + #[test] + fn test_json2_expanded_schema() { + use api::helper::to_grpc_value; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; + use datatypes::schema::SchemaBuilder; + use datatypes::types::json_type::JsonNativeType; + + let settings = JsonSettings::default(); + let mut column_schema = ColumnSchema::new( + "data", + ConcreteDataType::json2(JsonNativeType::object()), + true, + ); + column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings.clone(), + )))); + let schema = SchemaBuilder::try_from(vec![column_schema]) + .unwrap() + .add_metadata("test", "metadata") + .build() + .unwrap(); + let arrow_schema = schema.arrow_schema().clone(); + let field = arrow_schema.field(0); + let datatype = + ColumnDataTypeWrapper::try_from(schema.column_schemas()[0].data_type.clone()).unwrap(); + let (kind, extension) = datatype.to_parts(); + let mut column = source("data", kind); + column.datatype_extension = extension; + let rows = Rows { + schema: vec![column], + rows: vec![Row { + values: vec![to_grpc_value( + settings.encode(serde_json::json!({"id": 3})).unwrap(), + )], + }], + }; + let batch = rows_to_record_batch(&rows, &table_info(schema)).unwrap(); + // The same arrays fail with the static table schema used before alignment. + assert!(RecordBatch::try_new(arrow_schema.clone(), batch.columns().to_vec()).is_err()); + let actual_schema = batch.schema(); + assert_eq!(actual_schema.metadata(), arrow_schema.metadata()); + assert_eq!(actual_schema.field(0).metadata(), field.metadata()); + assert_eq!(actual_schema.field(0).name(), field.name()); + assert_eq!(actual_schema.field(0).is_nullable(), field.is_nullable()); + assert_eq!( + actual_schema.field(0).data_type(), + batch.column(0).data_type() + ); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(array.column_by_name("id").is_some()); + assert!(array.column_by_name("!__remainder__!").is_some()); + assert_eq!(batch.num_rows(), 1); + } + + #[test] + fn test_dynamic_default_is_evaluated_per_conversion() { + use std::time::{Duration, Instant}; + + use arrow::array::TimestampMillisecondArray; + + let (mut rows, base_schema) = fixture(); + let mut columns = base_schema.column_schemas().to_vec(); + columns.push( + ColumnSchema::new( + "created", + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ) + .with_default_constraint(Some(ColumnDefaultConstraint::Function( + "current_timestamp()".to_string(), + ))) + .unwrap(), + ); + let schema = Schema::new(columns); + let timestamp = |batch: &RecordBatch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + let first = timestamp(&rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap()); + // The default uses wall-clock milliseconds, not Tokio's controllable + // clock. Poll its observable result with a bound instead of sleeping. + let deadline = Instant::now() + Duration::from_secs(1); + loop { + let next = + timestamp(&rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap()); + if next != first { + break; + } + assert!( + Instant::now() < deadline, + "dynamic default was not refreshed" + ); + yield_now(); + } + rows.schema + .push(source("created", ColumnDataType::TimestampMillisecond)); + rows.rows[0].values.push(Value::default()); + assert!( + rows_to_record_batch(&rows, &table_info(schema.clone())) + .unwrap() + .column(3) + .is_null(0) + ); + } + + #[test] + fn test_tag_and_field_semantics() { + let (mut rows, schema) = fixture(); + let mut table = table_info(schema); + table.meta.primary_key_indices = vec![1]; + // The "label" column is a tag in table metadata, not a field. + assert!(rows_to_record_batch(&rows, &table).is_err()); + rows.schema[0].semantic_type = SemanticType::Tag as i32; + assert!(rows_to_record_batch(&rows, &table).is_ok()); + // Conversely the count field cannot be submitted as a tag. + rows.schema[1].semantic_type = SemanticType::Tag as i32; + assert!(rows_to_record_batch(&rows, &table).is_err()); + } + + #[test] + fn test_invalid_input() { + for case in 0..7 { + let (mut rows, schema) = fixture(); + match case { + 0 => rows.schema[0].column_name = "unknown".to_string(), + 1 => rows.schema[0].column_name = "count".to_string(), + 2 => { + rows.rows[0].values.pop(); + } + 3 => rows.schema[1].datatype = ColumnDataType::Float64 as i32, + 4 => rows.rows[0].values[1] = Value::default(), + 5 => rows.schema[1].semantic_type = SemanticType::Timestamp as i32, + _ => { + rows.rows[0].values[1].value_data = + Some(ValueData::StringValue("wrong".to_string())) + } + } + assert!( + rows_to_record_batch(&rows, &table_info(schema.clone())).is_err(), + "case {case}" + ); + } + } + + #[test] + #[should_panic(expected = "prepared rows must not be empty")] + fn test_empty_rows_contract() { + let (mut rows, schema) = fixture(); + rows.rows.clear(); + let _ = rows_to_record_batch(&rows, &table_info(schema)); + } + + #[test] + fn test_reject_nested_values_for_scalar_column() { + for value in [ + ValueData::ListValue(api::v1::ListValue { items: vec![] }), + ValueData::StructValue(api::v1::StructValue { items: vec![] }), + ] { + let (mut rows, schema) = fixture(); + rows.rows[0].values[1].value_data = Some(value); + assert!(matches!( + rows_to_record_batch(&rows, &table_info(schema)), + Err(error::Error::InvalidInsertRequest { .. }) + )); + } + } + + #[test] + fn test_nested_value_validation() { + let int = Value { + value_data: Some(ValueData::I32Value(1)), + }; + let list = Value { + value_data: Some(ValueData::ListValue(api::v1::ListValue { + items: vec![int.clone()], + })), + }; + let structure = Value { + value_data: Some(ValueData::StructValue(api::v1::StructValue { + items: vec![int.clone()], + })), + }; + let types = [ + ( + ColumnDataTypeWrapper::list_datatype(ColumnDataTypeWrapper::int32_datatype()), + list.clone(), + ), + ( + ColumnDataTypeWrapper::struct_datatype(vec![( + "count".to_string(), + ColumnDataTypeWrapper::int32_datatype(), + )]), + structure.clone(), + ), + ]; + for (datatype, valid) in types { + let (kind, extension) = datatype.to_parts(); + let schema = Schema::new(vec![ColumnSchema::new("nested", datatype.into(), true)]); + let mut column = source("nested", kind); + column.datatype_extension = extension; + let mut rows = Rows { + schema: vec![column], + rows: vec![Row { + values: vec![valid], + }], + }; + let table = table_info(schema); + assert_eq!(rows_to_record_batch(&rows, &table).unwrap().num_rows(), 1); + for invalid in [list.clone(), structure.clone()] { + match rows.rows[0].values[0].value_data.as_mut().unwrap() { + ValueData::ListValue(v) => v.items = vec![invalid], + ValueData::StructValue(v) => v.items = vec![invalid], + _ => unreachable!(), + } + assert!(matches!( + rows_to_record_batch(&rows, &table), + Err(error::Error::InvalidInsertRequest { .. }) + )); + } + } + let datatype = ColumnDataTypeWrapper::list_datatype(ColumnDataTypeWrapper::list_datatype( + ColumnDataTypeWrapper::int32_datatype(), + )); + let (kind, extension) = datatype.to_parts(); + let value = Value { + value_data: Some(ValueData::ListValue(api::v1::ListValue { + items: vec![list], + })), + }; + assert!(value_matches_type(&value, kind as i32, extension.as_ref())); + assert!(!value_matches_type(&value, kind as i32, None)); + } +} diff --git a/src/operator/src/req_convert/insert/row_to_region.rs b/src/operator/src/req_convert/insert/row_to_region.rs index f512a6c9ac..08b960094b 100644 --- a/src/operator/src/req_convert/insert/row_to_region.rs +++ b/src/operator/src/req_convert/insert/row_to_region.rs @@ -92,7 +92,7 @@ mod tests { use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value}; use super::*; - use crate::tests::{ + use crate::test_util::{ create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, }; diff --git a/src/operator/src/req_convert/insert/table_to_region.rs b/src/operator/src/req_convert/insert/table_to_region.rs index 4fbbbaa02c..adc40118ec 100644 --- a/src/operator/src/req_convert/insert/table_to_region.rs +++ b/src/operator/src/req_convert/insert/table_to_region.rs @@ -78,7 +78,7 @@ mod tests { use store_api::storage::RegionId; use super::*; - use crate::tests::{ + use crate::test_util::{ create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, }; diff --git a/src/operator/src/tests.rs b/src/operator/src/test_util.rs similarity index 79% rename from src/operator/src/tests.rs rename to src/operator/src/test_util.rs index 9263fa4cca..aa4a2f30ab 100644 --- a/src/operator/src/tests.rs +++ b/src/operator/src/test_util.rs @@ -15,5 +15,5 @@ mod kv_backend; mod partition_manager; -pub(crate) use kv_backend::prepare_mocked_backend; -pub(crate) use partition_manager::{create_partition_rule_manager, new_test_table_info}; +pub use crate::test_util::kv_backend::prepare_mocked_backend; +pub use crate::test_util::partition_manager::{create_partition_rule_manager, new_test_table_info}; diff --git a/src/operator/src/tests/kv_backend.rs b/src/operator/src/test_util/kv_backend.rs similarity index 100% rename from src/operator/src/tests/kv_backend.rs rename to src/operator/src/test_util/kv_backend.rs diff --git a/src/operator/src/tests/partition_manager.rs b/src/operator/src/test_util/partition_manager.rs similarity index 98% rename from src/operator/src/tests/partition_manager.rs rename to src/operator/src/test_util/partition_manager.rs index ecf87aedac..e51ac49051 100644 --- a/src/operator/src/tests/partition_manager.rs +++ b/src/operator/src/test_util/partition_manager.rs @@ -15,9 +15,13 @@ use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +#[cfg(test)] use api::v1::helper::tag_column_schema; +#[cfg(test)] use api::v1::value::ValueData; +#[cfg(test)] use api::v1::{ColumnDataType, Row, Rows}; +#[cfg(test)] use common_base::hash::partition_expr_version; use common_meta::cache::{TableRouteCacheRef, new_table_route_cache}; use common_meta::key::TableMetadataManager; @@ -93,10 +97,12 @@ fn new_test_table_info_with_columns( .unwrap() } +#[cfg(test)] fn new_physical_test_table_info(table_id: u32, table_name: &str) -> TableInfo { new_test_table_info_with_columns(table_id, table_name, test_column_schemas(true), vec![0, 2]) } +#[cfg(test)] fn new_logical_test_table_info(table_id: u32, table_name: &str) -> TableInfo { new_test_table_info_with_columns(table_id, table_name, test_column_schemas(false), vec![0]) } @@ -141,9 +147,7 @@ fn test_new_partition_info_cache(table_route_cache: TableRouteCacheRef) -> Parti /// PARTITION r2 VALUES LESS THAN (50, 'sh'), /// PARTITION r3 VALUES LESS THAN (MAXVALUE, MAXVALUE), /// ) -pub(crate) async fn create_partition_rule_manager( - kv_backend: KvBackendRef, -) -> PartitionRuleManagerRef { +pub async fn create_partition_rule_manager(kv_backend: KvBackendRef) -> PartitionRuleManagerRef { let table_metadata_manager = TableMetadataManager::new(kv_backend.clone()); let table_route_cache = test_new_table_route_cache(kv_backend.clone()); let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone());