fix(mito2): validate bulk insert batches against the region schema (#9698)

* fix(mito2): validate bulk insert batches against the region schema

Bulk inserts were only checked for the time index column, unlike row
writes, which WriteRequest::check_schema checks against the region
schema. A batch with columns of another type, nulls in the time index or
in non-nullable columns, duplicate or unknown columns reached the region
worker: a time index that is not a timestamp or holds only nulls made
the worker fail, and other mismatches failed after the WAL append or
were stored as written, depending on the SST format.

Move the column checks of check_schema into a core shared by row writes
and bulk inserts, and check bulk insert batches at the same points as
rows: when the engine builds the worker request, and again in the worker
if the region schema changed since. check_schema keeps its behavior.
Bulk insert columns must have the column's Arrow type, except that tags
may use any string encoding, as the memtables already convert them.
Sparse batches must carry one encoded primary key column of binary
without nulls. A batch that misses columns gets them filled by the
worker even if it claims to match the current schema.

The metric engine checks logical bulk inserts against the logical region
before rewriting them into the physical region, as the rewrite copies
columns without checking them and would encode a duplicated tag into the
primary key twice.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): report an error when a write is interrupted by a panic

WriteNotify answers its waiter when it is dropped and reported success
unless an error had been set, also when it was dropped while unwinding
from a panic in the middle of the write. Report an error in that case.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): tolerate a poisoned series lock when dropping a series map

SeriesMap::drop reads every series to update the memtable metrics and
unwrapped the lock, so dropping the map while unwinding from a panic
that poisoned a series lock panicked again. Read through the poisoned
lock instead.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

src/metric-engine/src/engine/bulk_insert.rs

@@ -16,9 +16,11 @@ use std::collections::HashSet;
 
 use api::v1::{ArrowIpc, SemanticType};
 use bytes::Bytes;
+use common_error::ext::BoxedError;
 use common_grpc::flight::{FlightEncoder, FlightMessage};
 use datatypes::arrow::record_batch::RecordBatch;
-use snafu::{OptionExt, ensure};
+use mito2::request::check_bulk_insert_batch;
+use snafu::{OptionExt, ResultExt, ensure};
 use store_api::codec::PrimaryKeyEncoding;
 use store_api::metadata::RegionMetadataRef;
 use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest};
@@ -111,6 +113,11 @@ impl MetricEngineInner {
             &batch,
             &logical_metadata,
         )?;
+        // The rewrite below copies columns without checking them, and would encode a
+        // duplicated tag into the primary key twice.
+        check_bulk_insert_batch(&logical_metadata, &batch)
+            .map_err(BoxedError::new)
+            .context(error::MitoWriteOperationSnafu)?;
         let modified_batch = modify_batch_sparse(
             batch.clone(),
             region_id.table_id(),
@@ -307,6 +314,92 @@ mod tests {
         })
     }
 
+    #[tokio::test]
+    async fn test_bulk_insert_rejects_duplicate_tag() {
+        use common_error::status_code::StatusCode;
+
+        let env = TestEnv::new().await;
+        env.init_metric_region().await;
+        let logical_region_id = env.default_logical_region_id();
+
+        let batch = build_logical_batch(0, 2);
+        let mut fields = batch.schema().fields().to_vec();
+        let mut columns = batch.columns().to_vec();
+        fields.push(fields[2].clone());
+        columns.push(columns[2].clone());
+        let batch = RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap();
+        let err = env
+            .metric()
+            .handle_request(
+                logical_region_id,
+                build_bulk_request(logical_region_id, batch, false),
+            )
+            .await
+            .unwrap_err();
+        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err:?}");
+    }
+
+    #[tokio::test]
+    async fn test_bulk_insert_physical_region_rejects_malformed_primary_key() {
+        use common_error::status_code::StatusCode;
+        use datatypes::arrow::array::{ArrayRef, BinaryArray, Int64Array};
+        use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
+
+        let env = TestEnv::new().await;
+        env.init_metric_region().await;
+        let physical_region_id = env.default_physical_region_id();
+        let logical_region_id = env.default_logical_region_id();
+        let tag_columns = vec![TagColumnInfo {
+            name: "job".to_string(),
+            index: 2,
+            column_id: 2,
+        }];
+        let batch = modify_batch_sparse(
+            build_logical_batch(0, 2),
+            logical_region_id.table_id(),
+            &tag_columns,
+            &[0, 1],
+        )
+        .unwrap();
+        let primary_key_index = batch.schema().index_of(PRIMARY_KEY_COLUMN_NAME).unwrap();
+        let with_primary_key = |primary_key: Option<ArrayRef>, extra: Option<ArrayRef>| {
+            let mut fields = batch.schema().fields().to_vec();
+            let mut columns = batch.columns().to_vec();
+            fields.remove(primary_key_index);
+            columns.remove(primary_key_index);
+            for array in [primary_key, extra].into_iter().flatten() {
+                fields.push(Arc::new(Field::new(
+                    PRIMARY_KEY_COLUMN_NAME,
+                    array.data_type().clone(),
+                    true,
+                )));
+                columns.push(array);
+            }
+            RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap()
+        };
+        let encoded = batch.column(primary_key_index).clone();
+        let rejected = [
+            with_primary_key(None, None),
+            with_primary_key(Some(Arc::new(Int64Array::from(vec![1, 2]))), None),
+            with_primary_key(
+                Some(Arc::new(BinaryArray::from_opt_vec(vec![Some(b"a"), None]))),
+                None,
+            ),
+            with_primary_key(Some(encoded.clone()), Some(encoded)),
+        ];
+        for batch in rejected {
+            let err = env
+                .metric()
+                .handle_request(
+                    physical_region_id,
+                    build_bulk_request(physical_region_id, batch, false),
+                )
+                .await
+                .unwrap_err();
+            assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err:?}");
+        }
+    }
+
     async fn init_dense_metric_region(env: &TestEnv) -> RegionId {
         let physical_region_id = env.default_physical_region_id();
         env.create_physical_region(

src/mito2/src/engine.rs

@@ -25,6 +25,8 @@ mod batch_catchup_test;
 #[cfg(test)]
 mod batch_open_test;
 #[cfg(test)]
+mod bulk_insert_validation_test;
+#[cfg(test)]
 mod bump_committed_sequence_test;
 #[cfg(test)]
 mod catchup_test;

src/mito2/src/engine/bulk_insert_validation_test.rs

@@ -0,0 +1,329 @@
+// 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.
+
+//! Bulk inserts checked against the region schema.
+
+use std::sync::Arc;
+
+use api::v1::{ArrowIpc, Rows};
+use common_error::ext::ErrorExt;
+use common_error::status_code::StatusCode;
+use common_recordbatch::{DfRecordBatch, RecordBatches};
+use common_test_util::flight::encode_to_flight_data;
+use datatypes::arrow::array::{
+    ArrayRef, Float64Array, Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray,
+    TimestampMillisecondArray, TimestampSecondArray,
+};
+use datatypes::arrow::datatypes::{Field, Int64Type, Schema};
+use datatypes::prelude::ConcreteDataType;
+use store_api::region_engine::RegionEngine;
+use store_api::region_request::{
+    AlterKind, ModifyColumnType, RegionAlterRequest, RegionBulkInsertsRequest, RegionPutRequest,
+    RegionRequest,
+};
+use store_api::storage::{RegionId, ScanRequest};
+
+use crate::config::MitoConfig;
+use crate::engine::MitoEngine;
+use crate::request::WorkerRequest;
+use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, reopen_region, rows_schema};
+
+const SST_FORMATS: [&str; 2] = ["flat", "primary_key"];
+
+fn column(name: &str, array: ArrayRef) -> (Field, ArrayRef) {
+    (Field::new(name, array.data_type().clone(), true), array)
+}
+
+fn tag() -> (Field, ArrayRef) {
+    column("tag_0", Arc::new(StringArray::from(vec!["a", "b"])))
+}
+
+fn field() -> (Field, ArrayRef) {
+    column("field_0", Arc::new(Float64Array::from(vec![1.0, 2.0])))
+}
+
+fn ts() -> (Field, ArrayRef) {
+    column(
+        "ts",
+        Arc::new(TimestampMillisecondArray::from(vec![0, 1000])),
+    )
+}
+
+fn bulk_request(region_id: RegionId, columns: Vec<(Field, ArrayRef)>) -> RegionRequest {
+    RegionRequest::BulkInserts(bulk_insert(region_id, columns))
+}
+
+fn bulk_insert(region_id: RegionId, columns: Vec<(Field, ArrayRef)>) -> RegionBulkInsertsRequest {
+    let (fields, arrays): (Vec<_>, Vec<_>) = columns.into_iter().unzip();
+    let payload = DfRecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap();
+    let (schema, record_batch) = encode_to_flight_data(payload.clone());
+    RegionBulkInsertsRequest {
+        skip_wal: false,
+        region_id,
+        payload,
+        raw_data: ArrowIpc {
+            schema: schema.data_header,
+            data_header: record_batch.data_header,
+            payload: record_batch.data_body,
+        },
+        partition_expr_version: None,
+        aligned_schema_version: None,
+    }
+}
+
+async fn scan_num_rows(engine: &MitoEngine, region_id: RegionId) -> usize {
+    let stream = engine
+        .scan_to_stream(region_id, ScanRequest::default())
+        .await
+        .unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    batches.iter().map(|batch| batch.num_rows()).sum()
+}
+
+/// Batches that do not match the region schema are rejected before they are
+/// appended to the WAL, and the region keeps accepting writes.
+#[tokio::test]
+async fn test_bulk_insert_rejects_batches_not_matching_schema() {
+    common_telemetry::init_default_ut_logging();
+    let list = ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
+        Some(vec![Some(1)]),
+        Some(vec![Some(2)]),
+    ]);
+    let rejected = vec![
+        vec![
+            tag(),
+            field(),
+            column("ts", Arc::new(Int64Array::from(vec![0, 1000]))),
+        ],
+        vec![
+            tag(),
+            field(),
+            column(
+                "ts",
+                Arc::new(TimestampMillisecondArray::from(vec![Some(0), None])),
+            ),
+        ],
+        vec![
+            tag(),
+            field(),
+            column(
+                "ts",
+                Arc::new(TimestampMillisecondArray::from(vec![None, None])),
+            ),
+        ],
+        vec![
+            tag(),
+            field(),
+            column("ts", Arc::new(TimestampSecondArray::from(vec![0, 1]))),
+        ],
+        vec![tag(), field()],
+        vec![
+            tag(),
+            column("field_0", Arc::new(StringArray::from(vec!["1", "2"]))),
+            ts(),
+        ],
+        vec![tag(), column("field_0", Arc::new(list)), ts()],
+        vec![
+            tag(),
+            column(
+                "field_0",
+                Arc::new(Float64Array::from(vec![Some(1.0), None])),
+            ),
+            ts(),
+        ],
+        vec![tag(), field(), field(), ts()],
+        vec![
+            tag(),
+            field(),
+            ts(),
+            column("unknown", Arc::new(Int64Array::from(vec![1, 2]))),
+        ],
+    ];
+    for (i, sst_format) in SST_FORMATS.into_iter().enumerate() {
+        let mut env = TestEnv::new().await;
+        let engine = env.create_engine(MitoConfig::default()).await;
+        let region_id = RegionId::new(1, i as u32 + 1);
+        let request = CreateRequestBuilder::new()
+            .all_not_null(true)
+            .insert_option("sst_format", sst_format)
+            .build();
+        let column_schemas = rows_schema(&request);
+        let table_dir = request.table_dir.clone();
+        let options = request.options.clone();
+        engine
+            .handle_request(region_id, RegionRequest::Create(request))
+            .await
+            .unwrap();
+
+        for columns in &rejected {
+            let err = engine
+                .handle_request(region_id, bulk_request(region_id, columns.clone()))
+                .await
+                .unwrap_err();
+            assert_eq!(
+                StatusCode::InvalidArguments,
+                err.status_code(),
+                "{sst_format}: {err:?}"
+            );
+        }
+
+        let put = RegionRequest::Put(RegionPutRequest {
+            rows: Rows {
+                schema: column_schemas,
+                rows: build_rows(0, 3),
+            },
+            hint: None,
+            skip_wal: false,
+            partition_expr_version: None,
+        });
+        engine.handle_request(region_id, put).await.unwrap();
+        // Reopening replays the WAL, which holds none of the rejected batches.
+        reopen_region(&engine, region_id, table_dir, true, options).await;
+        assert_eq!(3, scan_num_rows(&engine, region_id).await, "{sst_format}");
+    }
+}
+
+/// Tags may use any string encoding, which the memtables convert.
+#[tokio::test]
+async fn test_bulk_insert_accepts_tag_encodings() {
+    common_telemetry::init_default_ut_logging();
+    let tags: [ArrayRef; 3] = [
+        Arc::new(StringArray::from(vec!["a", "b"])),
+        Arc::new(LargeStringArray::from(vec!["a", "b"])),
+        Arc::new(StringViewArray::from(vec!["a", "b"])),
+    ];
+    for (i, sst_format) in SST_FORMATS.into_iter().enumerate() {
+        let mut env = TestEnv::new().await;
+        let engine = env.create_engine(MitoConfig::default()).await;
+        for (j, tag) in tags.iter().enumerate() {
+            let region_id = RegionId::new(i as u32 + 1, j as u32 + 1);
+            let request = CreateRequestBuilder::new()
+                .insert_option("sst_format", sst_format)
+                .build();
+            let table_dir = request.table_dir.clone();
+            let options = request.options.clone();
+            engine
+                .handle_request(region_id, RegionRequest::Create(request))
+                .await
+                .unwrap();
+
+            let columns = vec![column("tag_0", tag.clone()), field(), ts()];
+            let response = engine
+                .handle_request(region_id, bulk_request(region_id, columns))
+                .await
+                .unwrap();
+            assert_eq!(2, response.affected_rows, "{sst_format}: {tag:?}");
+            reopen_region(&engine, region_id, table_dir, true, options).await;
+            assert_eq!(
+                2,
+                scan_num_rows(&engine, region_id).await,
+                "{sst_format}: {tag:?}"
+            );
+        }
+    }
+}
+
+/// The worker checks a batch again if the region schema changed after the
+/// request was checked.
+#[tokio::test]
+async fn test_bulk_insert_rechecks_batch_after_schema_change() {
+    common_telemetry::init_default_ut_logging();
+    for (i, sst_format) in SST_FORMATS.into_iter().enumerate() {
+        let mut env = TestEnv::new().await;
+        let engine = env.create_engine(MitoConfig::default()).await;
+        let region_id = RegionId::new(1, i as u32 + 1);
+        let request = CreateRequestBuilder::new()
+            .insert_option("sst_format", sst_format)
+            .build();
+        let table_dir = request.table_dir.clone();
+        let options = request.options.clone();
+        engine
+            .handle_request(region_id, RegionRequest::Create(request))
+            .await
+            .unwrap();
+        let metadata = engine.get_region(region_id).unwrap().metadata();
+
+        // The batch matches the schema before `field_0` becomes a string.
+        let alter = RegionAlterRequest {
+            kind: AlterKind::ModifyColumnTypes {
+                columns: vec![ModifyColumnType {
+                    column_name: "field_0".to_string(),
+                    target_type: ConcreteDataType::string_datatype(),
+                }],
+            },
+        };
+        engine
+            .handle_request(region_id, RegionRequest::Alter(alter))
+            .await
+            .unwrap();
+        let (request, receiver) = WorkerRequest::try_from_region_request(
+            region_id,
+            bulk_request(region_id, vec![tag(), field(), ts()]),
+            Some(metadata),
+        )
+        .unwrap();
+        engine
+            .inner
+            .workers
+            .submit_to_worker(region_id, request)
+            .await
+            .unwrap();
+        let err = receiver.await.unwrap().unwrap_err();
+        assert_eq!(
+            StatusCode::InvalidArguments,
+            err.status_code(),
+            "{sst_format}: {err:?}"
+        );
+
+        reopen_region(&engine, region_id, table_dir, true, options).await;
+        assert_eq!(0, scan_num_rows(&engine, region_id).await, "{sst_format}");
+    }
+}
+
+/// A batch that claims to match the current schema but misses a column still
+/// gets the column filled.
+#[tokio::test]
+async fn test_bulk_insert_fills_columns_missing_from_aligned_batch() {
+    common_telemetry::init_default_ut_logging();
+    for (i, sst_format) in SST_FORMATS.into_iter().enumerate() {
+        let mut env = TestEnv::new().await;
+        let engine = env.create_engine(MitoConfig::default()).await;
+        let region_id = RegionId::new(1, i as u32 + 1);
+        let request = CreateRequestBuilder::new()
+            .insert_option("sst_format", sst_format)
+            .build();
+        let table_dir = request.table_dir.clone();
+        let options = request.options.clone();
+        engine
+            .handle_request(region_id, RegionRequest::Create(request))
+            .await
+            .unwrap();
+
+        let mut request = bulk_insert(region_id, vec![tag(), ts()]);
+        request.aligned_schema_version = Some(
+            engine
+                .get_region(region_id)
+                .unwrap()
+                .metadata()
+                .schema_version,
+        );
+        let response = engine
+            .handle_request(region_id, RegionRequest::BulkInserts(request))
+            .await
+            .unwrap();
+        assert_eq!(2, response.affected_rows, "{sst_format}");
+        reopen_region(&engine, region_id, table_dir, true, options).await;
+        assert_eq!(2, scan_num_rows(&engine, region_id).await, "{sst_format}");
+    }
+}

src/mito2/src/memtable/time_series.rs

@@ -17,7 +17,7 @@ use std::collections::{BTreeMap, Bound, HashSet};
 use std::fmt::{Debug, Formatter};
 use std::iter;
 use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering};
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, PoisonError, RwLock};
 use std::time::{Duration, Instant};
 
 use api::v1::OpType;
@@ -474,7 +474,13 @@ impl Drop for SeriesMap {
         let num_field_builders = self
             .0
             .values()
-            .map(|v| v.read().unwrap().active.num_field_builders())
+            // The map may be dropped while unwinding from a panic that poisoned a lock.
+            .map(|v| {
+                v.read()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .active
+                    .num_field_builders()
+            })
             .sum::<usize>();
         MEMTABLE_ACTIVE_SERIES_COUNT.sub(num_series as i64);
         MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.sub(num_field_builders as i64);
@@ -1493,6 +1499,20 @@ mod tests {
         assert_eq!(expect, &read);
     }
 
+    #[test]
+    fn test_series_map_drop_with_poisoned_lock() {
+        let region_metadata = schema_for_test();
+        let series = Arc::new(RwLock::new(Series::new(&region_metadata)));
+        let map = SeriesMap(BTreeMap::from([(vec![1], series.clone())]));
+        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            let _guard = series.write().unwrap();
+            panic!("poison the series lock");
+        }))
+        .unwrap_err();
+        assert!(series.is_poisoned());
+        drop(map);
+    }
+
     #[test]
     fn test_series() {
         let region_metadata = schema_for_test();

src/mito2/src/region_write_ctx.rs

@@ -23,7 +23,7 @@ use store_api::logstore::LogStore;
 use store_api::logstore::provider::Provider;
 use store_api::storage::{RegionId, SequenceNumber};
 
-use crate::error::{Error, Result, WriteGroupSnafu};
+use crate::error::{Error, Result, UnexpectedSnafu, WriteGroupSnafu};
 use crate::memtable::KeyValues;
 use crate::memtable::bulk::part::BulkPart;
 use crate::metrics;
@@ -66,6 +66,15 @@ impl WriteNotify {
 
 impl Drop for WriteNotify {
     fn drop(&mut self) {
+        // A panic may have interrupted the write, so it must not be reported as written.
+        if self.err.is_none() && std::thread::panicking() {
+            self.err = Some(Arc::new(
+                UnexpectedSnafu {
+                    reason: "the write was interrupted by a panic",
+                }
+                .build(),
+            ));
+        }
         self.notify_result();
     }
 }
@@ -583,6 +592,18 @@ mod tests {
     use crate::memtable::bulk::part::BulkPart;
     use crate::test_util::version_util::VersionControlBuilder;
 
+    #[test]
+    fn test_write_notify_dropped_by_panic_reports_error() {
+        let (tx, rx) = oneshot::channel();
+        let notify = WriteNotify::new(OptionOutputTx::from(tx), 3);
+        std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
+            let _notify = notify;
+            panic!("interrupt the write");
+        }))
+        .unwrap_err();
+        assert!(rx.blocking_recv().unwrap().is_err());
+    }
+
     #[test]
     fn test_request_skip_wal_preserves_sequences_and_other_writes() {
         // Ablate only the request flag: the workload and sequence allocation stay identical.

src/mito2/src/request.rs

@@ -25,6 +25,9 @@ use api::helper::{
 use api::v1::column_def::options_from_column_schema;
 use api::v1::{ColumnDataType, ColumnSchema, OpType, Rows, SemanticType, Value, WriteHint};
 use common_telemetry::info;
+use datatypes::arrow::array::{Array, ArrayRef};
+use datatypes::arrow::datatypes::DataType as ArrowDataType;
+use datatypes::arrow::record_batch::RecordBatch;
 use datatypes::prelude::DataType;
 use partition::expr::PartitionExpr;
 use prometheus::HistogramTimer;
@@ -44,6 +47,7 @@ use store_api::region_request::{
     RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest,
     StagingPartitionDirective,
 };
+use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
 use store_api::storage::{FileId, RegionId, SequenceNumber};
 use tokio::sync::oneshot::{self, Receiver, Sender};
 
@@ -194,85 +198,22 @@ impl WriteRequest {
 
         let region_id = self.region_id;
         // Index all columns in rows.
-        let mut rows_columns: HashMap<_, _> = self
+        let rows_columns = self
             .rows
             .schema
             .iter()
-            .map(|column| (&column.column_name, column))
+            .map(|schema| {
+                let column = RowColumn {
+                    request: self,
+                    schema,
+                };
+                (schema.column_name.as_str(), column)
+            })
             .collect();
-
-        let mut need_fill_default = false;
-        // Checks all columns in this region.
-        for column in &metadata.column_metadatas {
-            if let Some(input_col) = rows_columns.remove(&column.column_schema.name) {
-                // Check data type.
-                ensure!(
-                    is_column_type_value_eq(
-                        input_col.datatype,
-                        input_col.datatype_extension.clone(),
-                        &column.column_schema.data_type
-                    ),
-                    InvalidRequestSnafu {
-                        region_id,
-                        reason: format!(
-                            "column {} expect type {:?}, given: {}({})",
-                            column.column_schema.name,
-                            column.column_schema.data_type,
-                            ColumnDataType::try_from(input_col.datatype)
-                                .map(|v| v.as_str_name())
-                                .unwrap_or("Unknown"),
-                            input_col.datatype,
-                        )
-                    }
-                );
-
-                // Check semantic type.
-                ensure!(
-                    is_semantic_type_eq(input_col.semantic_type, column.semantic_type),
-                    InvalidRequestSnafu {
-                        region_id,
-                        reason: format!(
-                            "column {} has semantic type {:?}, given: {}({})",
-                            column.column_schema.name,
-                            column.semantic_type,
-                            api::v1::SemanticType::try_from(input_col.semantic_type)
-                                .map(|v| v.as_str_name())
-                                .unwrap_or("Unknown"),
-                            input_col.semantic_type
-                        ),
-                    }
-                );
-
-                // Check nullable.
-                // Safety: `rows_columns` ensures this column exists.
-                let has_null = self.has_null[self.name_to_index[&column.column_schema.name]];
-                ensure!(
-                    !has_null || column.column_schema.is_nullable(),
-                    InvalidRequestSnafu {
-                        region_id,
-                        reason: format!(
-                            "column {} is not null but input has null",
-                            column.column_schema.name
-                        ),
-                    }
-                );
-            } else {
-                // Rows don't have this column.
-                self.check_missing_column(column)?;
-
-                need_fill_default = true;
-            }
-        }
-
-        // Checks all columns in rows exist in the region.
-        if !rows_columns.is_empty() {
-            let names: Vec<_> = rows_columns.into_keys().collect();
-            return InvalidRequestSnafu {
-                region_id,
-                reason: format!("unknown columns: {:?}", names),
-            }
-            .fail();
-        }
+        let need_fill_default =
+            check_request_columns(region_id, metadata, rows_columns, |column| {
+                self.check_missing_column(column)
+            })?;
 
         // If we need to fill default values, return a special error.
         ensure!(!need_fill_default, FillDefaultSnafu { region_id });
@@ -369,16 +310,7 @@ impl WriteRequest {
         }
 
         // Not a delete request. Checks whether they have default value.
-        ensure!(
-            column.column_schema.is_nullable()
-                || column.column_schema.default_constraint().is_some(),
-            InvalidRequestSnafu {
-                region_id: self.region_id,
-                reason: format!("missing column {}", column.column_schema.name),
-            }
-        );
-
-        Ok(())
+        check_column_can_be_filled(self.region_id, column)
     }
 
     /// Returns the default value for specific column.
@@ -442,6 +374,242 @@ impl WriteRequest {
     }
 }
 
+/// A column of a write request, as [check_request_columns] sees it.
+trait RequestColumn {
+    /// Returns whether the column has the data type of `column`.
+    fn has_type_of(&self, column: &ColumnMetadata) -> bool;
+
+    /// Describes the data type of the column in errors.
+    fn type_name(&self) -> String;
+
+    /// Returns the semantic type of the column if the request carries one.
+    fn semantic_type(&self) -> Option<i32>;
+
+    /// Returns whether the column contains nulls.
+    fn has_null(&self) -> bool;
+}
+
+/// Checks the columns of a write request against the region schema. Each column
+/// must be a column of the region, have its data type and, if the request carries
+/// one, its semantic type, and contain nulls only if the region column is
+/// nullable. `check_missing` checks the region columns that the request doesn't
+/// have. Returns whether there are such columns.
+fn check_request_columns<C: RequestColumn>(
+    region_id: RegionId,
+    metadata: &RegionMetadata,
+    mut columns: HashMap<&str, C>,
+    mut check_missing: impl FnMut(&ColumnMetadata) -> Result<()>,
+) -> Result<bool> {
+    let mut has_missing = false;
+    for column in &metadata.column_metadatas {
+        let Some(input_col) = columns.remove(column.column_schema.name.as_str()) else {
+            check_missing(column)?;
+            has_missing = true;
+            continue;
+        };
+        // Check data type.
+        ensure!(
+            input_col.has_type_of(column),
+            InvalidRequestSnafu {
+                region_id,
+                reason: format!(
+                    "column {} expect type {:?}, given: {}",
+                    column.column_schema.name,
+                    column.column_schema.data_type,
+                    input_col.type_name(),
+                )
+            }
+        );
+
+        // Check semantic type.
+        if let Some(semantic_type) = input_col.semantic_type() {
+            ensure!(
+                is_semantic_type_eq(semantic_type, column.semantic_type),
+                InvalidRequestSnafu {
+                    region_id,
+                    reason: format!(
+                        "column {} has semantic type {:?}, given: {}({})",
+                        column.column_schema.name,
+                        column.semantic_type,
+                        api::v1::SemanticType::try_from(semantic_type)
+                            .map(|v| v.as_str_name())
+                            .unwrap_or("Unknown"),
+                        semantic_type
+                    ),
+                }
+            );
+        }
+
+        // Check nullable.
+        ensure!(
+            !input_col.has_null() || column.column_schema.is_nullable(),
+            InvalidRequestSnafu {
+                region_id,
+                reason: format!(
+                    "column {} is not null but input has null",
+                    column.column_schema.name
+                ),
+            }
+        );
+    }
+
+    // Checks all columns in the request exist in the region.
+    if !columns.is_empty() {
+        let names: Vec<_> = columns.into_keys().collect();
+        return InvalidRequestSnafu {
+            region_id,
+            reason: format!("unknown columns: {:?}", names),
+        }
+        .fail();
+    }
+
+    Ok(has_missing)
+}
+
+/// Checks that a write may leave out `column`, which is then filled with its
+/// default value or null.
+fn check_column_can_be_filled(region_id: RegionId, column: &ColumnMetadata) -> Result<()> {
+    ensure!(
+        column.column_schema.is_nullable() || column.column_schema.default_constraint().is_some(),
+        InvalidRequestSnafu {
+            region_id,
+            reason: format!("missing column {}", column.column_schema.name),
+        }
+    );
+    Ok(())
+}
+
+/// A column of the rows of a [WriteRequest].
+struct RowColumn<'a> {
+    request: &'a WriteRequest,
+    schema: &'a ColumnSchema,
+}
+
+impl RequestColumn for RowColumn<'_> {
+    fn has_type_of(&self, column: &ColumnMetadata) -> bool {
+        is_column_type_value_eq(
+            self.schema.datatype,
+            self.schema.datatype_extension.clone(),
+            &column.column_schema.data_type,
+        )
+    }
+
+    fn type_name(&self) -> String {
+        format!(
+            "{}({})",
+            ColumnDataType::try_from(self.schema.datatype)
+                .map(|v| v.as_str_name())
+                .unwrap_or("Unknown"),
+            self.schema.datatype,
+        )
+    }
+
+    fn semantic_type(&self) -> Option<i32> {
+        Some(self.schema.semantic_type)
+    }
+
+    fn has_null(&self) -> bool {
+        // Safety: `name_to_index` has every column of the request.
+        self.request.has_null[self.request.name_to_index[&self.schema.column_name]]
+    }
+}
+
+/// A column of the batch of a bulk insert.
+struct BulkColumn<'a> {
+    array: &'a ArrayRef,
+}
+
+impl RequestColumn for BulkColumn<'_> {
+    fn has_type_of(&self, column: &ColumnMetadata) -> bool {
+        let actual = self.array.data_type();
+        let data_type = &column.column_schema.data_type;
+        // The memtables convert tags of any string encoding.
+        if column.semantic_type == SemanticType::Tag && data_type.is_string() {
+            return match actual {
+                ArrowDataType::Dictionary(key, value) => key.is_integer() && value.is_string(),
+                actual => actual.is_string(),
+            };
+        }
+        *actual == data_type.as_arrow_type()
+    }
+
+    fn type_name(&self) -> String {
+        self.array.data_type().to_string()
+    }
+
+    fn semantic_type(&self) -> Option<i32> {
+        None
+    }
+
+    fn has_null(&self) -> bool {
+        self.array.logical_null_count() > 0
+    }
+}
+
+/// Checks the batch of a bulk insert against the region schema with the rules
+/// [WriteRequest::check_schema] applies to rows: column names are unique, every
+/// column is a column of the region with its Arrow type and contains nulls only
+/// if the region column is nullable, and a missing column is nullable or has a
+/// default value. Tags may use any string encoding. Sparse batches carry the
+/// tags in one encoded primary key column instead. Returns whether the batch
+/// misses columns that the writer fills.
+pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) -> Result<bool> {
+    let region_id = metadata.region_id;
+    let sparse = metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse;
+    let mut primary_key = None;
+    let mut columns = HashMap::with_capacity(batch.num_columns());
+    for (field, array) in batch.schema_ref().fields().iter().zip(batch.columns()) {
+        let duplicate = if sparse && field.name() == PRIMARY_KEY_COLUMN_NAME {
+            primary_key.replace(array).is_some()
+        } else {
+            columns
+                .insert(field.name().as_str(), BulkColumn { array })
+                .is_some()
+        };
+        ensure!(
+            !duplicate,
+            InvalidRequestSnafu {
+                region_id,
+                reason: format!("duplicate column {}", field.name()),
+            }
+        );
+    }
+    if sparse {
+        let primary_key = primary_key.with_context(|| InvalidRequestSnafu {
+            region_id,
+            reason: format!("missing column {PRIMARY_KEY_COLUMN_NAME}"),
+        })?;
+        // The writer casts the encoded primary key to a dictionary of binary.
+        let type_matches = match primary_key.data_type() {
+            ArrowDataType::Binary => true,
+            ArrowDataType::Dictionary(key, value) => {
+                key.is_integer() && **value == ArrowDataType::Binary
+            }
+            _ => false,
+        };
+        ensure!(
+            type_matches && primary_key.logical_null_count() == 0,
+            InvalidRequestSnafu {
+                region_id,
+                reason: format!(
+                    "column {PRIMARY_KEY_COLUMN_NAME} must be binary without nulls, given: {}",
+                    primary_key.data_type()
+                ),
+            }
+        );
+    }
+    let mut needs_fill = false;
+    check_request_columns(region_id, metadata, columns, |column| {
+        if sparse && column.semantic_type == SemanticType::Tag {
+            Ok(())
+        } else {
+            needs_fill = true;
+            check_column_can_be_filled(region_id, column)
+        }
+    })?;
+    Ok(needs_fill)
+}
+
 /// Validate proto value schema.
 pub(crate) fn validate_proto_value(
     region_id: RegionId,
@@ -762,7 +930,18 @@ impl WorkerRequest {
                 sender: sender.into(),
                 request: DdlRequest::EnterStaging(v),
             }),
-            RegionRequest::BulkInserts(region_bulk_inserts_request) => {
+            RegionRequest::BulkInserts(mut region_bulk_inserts_request) => {
+                // The worker answers empty batches without writing them.
+                if let Some(region_metadata) = &region_metadata
+                    && region_bulk_inserts_request.payload.num_rows() > 0
+                    && check_bulk_insert_batch(
+                        region_metadata,
+                        &region_bulk_inserts_request.payload,
+                    )?
+                {
+                    // Lets the worker fill the missing columns.
+                    region_bulk_inserts_request.aligned_schema_version = None;
+                }
                 WorkerRequest::BulkInserts(BulkInsertRequest {
                     metadata: region_metadata,
                     sender: sender.into(),

src/mito2/src/worker/handle_bulk_insert.rs

@@ -71,10 +71,12 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             return;
         }
 
-        // safety: ts data type must be a timestamp type.
+        // safety: `WorkerRequest::try_from_region_request` checked the batch against
+        // `region_metadata`, so the time index is a timestamp array without nulls.
         let (ts_primitive, _) = datatypes::timestamp::timestamp_array_to_primitive(ts).unwrap();
 
-        // safety: we've checked ts.len() == batch.num_rows() and batch is not empty
+        // safety: we've checked ts.len() == batch.num_rows(), batch is not empty and
+        // the time index has no nulls.
         let min_ts = arrow::compute::min(&ts_primitive).unwrap();
         let max_ts = arrow::compute::max(&ts_primitive).unwrap();
 

src/mito2/src/worker/handle_write.rs

@@ -34,7 +34,9 @@ use crate::metrics::{
 };
 use crate::region::{RegionLeaderState, RegionRoleState};
 use crate::region_write_ctx::{RegionWriteCtx, WriteSource};
-use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest};
+use crate::request::{
+    SenderBulkRequest, SenderWriteRequest, WriteRequest, check_bulk_insert_batch,
+};
 use crate::wal::Wal;
 use crate::worker::RegionWorkerLoop;
 
@@ -552,14 +554,16 @@ impl<S> RegionWorkerLoop<S> {
                     aligned_schema.schema_version == region_ctx.version().metadata.schema_version
                 });
 
-            // Fill missing columns if needed
-            if need_fill_missing_columns
-                && let Err(e) = bulk_req
-                    .request
-                    .fill_missing_columns(&region_ctx.version().metadata)
-            {
-                bulk_req.sender.send(Err(e));
-                continue;
+            // Fill missing columns if needed. The batch is checked again as the region
+            // schema may have changed since the request was checked.
+            if need_fill_missing_columns {
+                let metadata = &region_ctx.version().metadata;
+                if let Err(e) = check_bulk_insert_batch(metadata, &bulk_req.request.batch)
+                    .and_then(|_| bulk_req.request.fill_missing_columns(metadata))
+                {
+                    bulk_req.sender.send(Err(e));
+                    continue;
+                }
             }
 
             // Collect requests by region.