From cd76d6842a8ec3828cf651bec87f8b75e3abaccd Mon Sep 17 00:00:00 2001 From: WenyXu Date: Wed, 16 Sep 2026 09:07:53 +0000 Subject: [PATCH] fix: validate prepared batch values and clarify write contracts Signed-off-by: WenyXu --- src/api/src/helper.rs | 25 +++ src/mito2/src/request.rs | 10 +- src/operator/src/bulk_insert.rs | 3 + src/operator/src/error.rs | 10 -- .../src/req_convert/insert/row_to_batch.rs | 155 +++++++++++++++++- 5 files changed, 182 insertions(+), 21 deletions(-) diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index d8e3004aa2b..40f9bdfb6d6 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/mito2/src/request.rs b/src/mito2/src/request.rs index 16b7eb4a6be..10002570008 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/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 9355c17686d..51a08d2f324 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -43,6 +43,9 @@ 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, diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index c6b84543c47..8caf829e1e7 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::any::Any; -use std::sync::Arc; use common_datasource::file_format::Format; use common_error::define_into_tonic_status; @@ -102,13 +101,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to flush pending batch: {source}"))] - BatchFlush { - source: Arc, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("external error"))] External { #[snafu(implicit)] @@ -1059,7 +1051,6 @@ impl ErrorExt for Error { | Error::DescribeStatement { source, .. } => source.status_code(), Error::AlterExprToRequest { source, .. } => source.status_code(), Error::External { source, .. } => source.status_code(), - Error::BatchFlush { source, .. } => source.status_code(), Error::FindTablePartitionRule { source, .. } | Error::SplitInsert { source, .. } | Error::SplitDelete { source, .. } @@ -1165,7 +1156,6 @@ impl ErrorExt for Error { Error::Catalog { source, .. } => source.retry_hint(), Error::SubstraitCodec { source, .. } => source.retry_hint(), Error::External { source, .. } => source.retry_hint(), - Error::BatchFlush { source, .. } => source.retry_hint(), Error::BuildRecordBatch { source, .. } => source.retry_hint(), Error::DecodeFlightData { source, .. } => source.retry_hint(), Error::SqlCommon { source, .. } => source.retry_hint(), diff --git a/src/operator/src/req_convert/insert/row_to_batch.rs b/src/operator/src/req_convert/insert/row_to_batch.rs index df5456c4cae..91f2954a922 100644 --- a/src/operator/src/req_convert/insert/row_to_batch.rs +++ b/src/operator/src/req_convert/insert/row_to_batch.rs @@ -14,8 +14,12 @@ use std::collections::HashMap; -use api::helper::{ColumnDataTypeWrapper, pb_value_to_value_ref}; -use api::v1::{Rows, SemanticType}; +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; @@ -29,7 +33,12 @@ use crate::error::{self, Result}; /// 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() { @@ -106,6 +115,16 @@ pub fn rows_to_record_batch(rows: &Rows, table_info: &TableInfo) -> Result Result, +) -> 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; @@ -376,4 +438,93 @@ mod tests { ); } } + + #[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)); + } }