diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index b7ad7646e3c..a9080745e2e 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/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, extra: Option| { + 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( diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 47752adbe75..3036c61e7ff 100644 --- a/src/mito2/src/engine.rs +++ b/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; diff --git a/src/mito2/src/engine/bulk_insert_validation_test.rs b/src/mito2/src/engine/bulk_insert_validation_test.rs new file mode 100644 index 00000000000..9652152031f --- /dev/null +++ b/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::(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}"); + } +} diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index fad746927d4..1fa6cb24b15 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/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::(); 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(®ion_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(); diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index 8c306effd1d..b0b781eb43c 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/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. diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index c7ec359d6b0..45e9041006d 100644 --- a/src/mito2/src/request.rs +++ b/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; + + /// 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( + region_id: RegionId, + metadata: &RegionMetadata, + mut columns: HashMap<&str, C>, + mut check_missing: impl FnMut(&ColumnMetadata) -> Result<()>, +) -> Result { + 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 { + 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 { + 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 { + 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) = ®ion_metadata + && region_bulk_inserts_request.payload.num_rows() > 0 + && check_bulk_insert_batch( + region_metadata, + ®ion_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(), diff --git a/src/mito2/src/worker/handle_bulk_insert.rs b/src/mito2/src/worker/handle_bulk_insert.rs index 9ecbf537661..d48f06178c6 100644 --- a/src/mito2/src/worker/handle_bulk_insert.rs +++ b/src/mito2/src/worker/handle_bulk_insert.rs @@ -71,10 +71,12 @@ impl RegionWorkerLoop { 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(); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 2b9551308f0..5de0439e26e 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/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 RegionWorkerLoop { 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(®ion_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 = ®ion_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.