fix: validate prepared batch values and clarify write contracts

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/api/src/helper.rs

@@ -1033,6 +1033,17 @@ pub fn proto_value_type(value: &v1::Value) -> Option<ColumnDataType> {
     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<Item = &'a VectorRef>,
     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));
+        }
+    }
 }

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<Result<AffectedRows>>);

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,

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<Error>,
-        #[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(),

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<RecordBatch> {
+    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<Recor
                         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,
@@ -129,6 +148,49 @@ pub fn rows_to_record_batch(rows: &Rows, table_info: &TableInfo) -> Result<Recor
     RecordBatch::try_new(schema.arrow_schema().clone(), 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;
@@ -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));
+    }
 }