diff --git a/Cargo.lock b/Cargo.lock index 2485a5ceec..e3d0dfb69e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5918,7 +5918,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=7224c2ad6d11db612fbdb621c36135fc37ffce35#7224c2ad6d11db612fbdb621c36135fc37ffce35" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=912f6c21017d5ab6c529568d36e7b8c30a94d84b#912f6c21017d5ab6c529568d36e7b8c30a94d84b" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index ee8d4dcf11..47b04be366 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "7224c2ad6d11db612fbdb621c36135fc37ffce35" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "912f6c21017d5ab6c529568d36e7b8c30a94d84b" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index 6d33d39149..d700d858ac 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -18,9 +18,10 @@ use api::v1::{ArrowIpc, SemanticType}; use bytes::Bytes; use common_grpc::flight::{FlightEncoder, FlightMessage}; use datatypes::arrow::record_batch::RecordBatch; -use snafu::{OptionExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; +use store_api::region_engine::RegionEngine; use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest}; use store_api::storage::RegionId; @@ -71,8 +72,11 @@ impl MetricEngineInner { async fn bulk_insert_physical_region( &self, region_id: RegionId, - request: RegionBulkInsertsRequest, + mut request: RegionBulkInsertsRequest, ) -> Result { + // Simply set the aligned schema to the data region schema version to avoid filling missing columns + // because that schema should be constant and callers have ensured request has the same schema. + request.aligned_schema_version = Some(self.physical_schema_version(region_id).await?); self.data_region .write_data(region_id, RegionRequest::BulkInserts(request)) .await @@ -117,6 +121,8 @@ impl MetricEngineInner { let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?; let partition_expr_version = request.partition_expr_version; + let aligned_schema_version = Some(self.physical_schema_version(data_region_id).await?); + let request = RegionBulkInsertsRequest { region_id: data_region_id, payload: modified_batch, @@ -126,12 +132,22 @@ impl MetricEngineInner { payload, }, partition_expr_version, + aligned_schema_version, }; self.data_region .write_data(data_region_id, RegionRequest::BulkInserts(request)) .await } + async fn physical_schema_version(&self, region_id: RegionId) -> Result { + Ok(self + .mito + .get_metadata(region_id) + .await + .context(error::MitoReadOperationSnafu)? + .schema_version) + } + fn resolve_tag_columns_from_metadata( &self, logical_region_id: RegionId, @@ -282,6 +298,7 @@ mod tests { payload, }, partition_expr_version: None, + aligned_schema_version: None, }) } @@ -319,6 +336,7 @@ mod tests { payload: batch, raw_data: ArrowIpc::default(), partition_expr_version: None, + aligned_schema_version: None, }); let response = env .metric() diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 1d9bf2101b..214060acf8 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -652,6 +652,7 @@ fn build_bulk_insert_request( payload: record_batch.data_body, }, partition_expr_version: None, + aligned_schema_version: None, } } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 8ae05177da..7e071965f2 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -552,7 +552,7 @@ pub(crate) struct SenderBulkRequest { pub(crate) sender: OptionOutputTx, pub(crate) region_id: RegionId, pub(crate) request: BulkPart, - pub(crate) region_metadata: RegionMetadataRef, + pub(crate) region_metadata: Option, pub(crate) partition_expr_version: Option, } diff --git a/src/mito2/src/worker/handle_bulk_insert.rs b/src/mito2/src/worker/handle_bulk_insert.rs index a413271db3..b4a7ae4a18 100644 --- a/src/mito2/src/worker/handle_bulk_insert.rs +++ b/src/mito2/src/worker/handle_bulk_insert.rs @@ -86,11 +86,20 @@ impl RegionWorkerLoop { timestamp_index: ts_index, raw_data: Some(request.raw_data), }; + + let aligned_region_metadata = if request + .aligned_schema_version + .is_some_and(|aligned| aligned == region_metadata.schema_version) + { + Some(region_metadata) + } else { + None + }; pending_bulk_request.push(SenderBulkRequest { sender, request: part, region_id: request.region_id, - region_metadata, + region_metadata: aligned_region_metadata, partition_expr_version: request.partition_expr_version, }); } diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 721415411e..64fdd33edc 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -525,8 +525,10 @@ impl RegionWorkerLoop { } // Double-check the request schema - let need_fill_missing_columns = region_ctx.version().metadata.schema_version - != bulk_req.region_metadata.schema_version; + let need_fill_missing_columns = + !bulk_req.region_metadata.is_some_and(|aligned_schema| { + aligned_schema.schema_version == region_ctx.version().metadata.schema_version + }); // Fill missing columns if needed if need_fill_missing_columns diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 38ced30135..2dbec530eb 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -106,6 +106,7 @@ impl Inserter { region_id: region_id.as_u64(), partition_expr_version: partition_expr_version .map(|value| PartitionExprVersion { value }), + aligned_schema_version: None, body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { schema: schema_bytes.clone(), data_header: raw_flight_data.data_header, @@ -220,6 +221,7 @@ impl Inserter { region_id: region_id.as_u64(), partition_expr_version: partition_expr_version .map(|value| PartitionExprVersion { value }), + aligned_schema_version: None, body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { schema: schema_bytes, data_header: header, diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 307311a4e2..b800ffa05b 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -1380,6 +1380,7 @@ fn transform_logical_batches_to_physical( let modified = { let start = Instant::now(); + // The schema of modified batch is: __primary_key, timestamp, value, other partition columns... let batch = modify_batch_sparse( batch.clone(), table_id, @@ -1535,6 +1536,9 @@ fn encode_region_write_requests( body: Some(region_request::Body::BulkInsert(BulkInsertRequest { region_id: region_id.as_u64(), partition_expr_version: None, + // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any + // column is missing. + aligned_schema_version: None, body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { schema: schema_bytes, data_header, diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index abf96736ec..8ab2e3abd2 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -409,6 +409,7 @@ fn make_region_truncate(truncate: TruncateRequest) -> Result Result> { let region_id = request.region_id.into(); let partition_expr_version = request.partition_expr_version.map(|v| v.value); + let aligned_schema_version = request.aligned_schema_version.map(|v| v.schema_version); let Some(Body::ArrowIpc(request)) = request.body else { return Ok(vec![]); }; @@ -429,6 +430,7 @@ fn make_region_bulk_inserts(request: BulkInsertRequest) -> Result, + pub aligned_schema_version: Option, } impl RegionBulkInsertsRequest { diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index d638926fff..d9bf15a64c 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -77,6 +77,35 @@ mod test { query_and_expect(db.frontend().as_ref(), sql, expected).await; } + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_flight_do_put_missing_nullable_columns() { + common_telemetry::init_default_ut_logging(); + + let (db, server) = setup_grpc_server( + StorageType::File, + "test_standalone_flight_do_put_missing_nullable_columns", + ) + .await; + let addr = server.bind_addr().unwrap().to_string(); + + let client = Client::with_urls(vec![addr]); + let client = Database::new_with_dbname("greptime-public", client); + + create_table(&client).await; + + let record_batches = create_record_batches_without_nullable_column(1); + test_put_record_batches(&client, record_batches).await; + + let sql = "select count(*) from foo where `B` is null"; + let expected = "\ ++----------+ +| count(*) | ++----------+ +| 9 | ++----------+"; + query_and_expect(db.frontend().as_ref(), sql, expected).await; + } + #[tokio::test(flavor = "multi_thread")] async fn test_distributed_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -281,6 +310,41 @@ mod test { assert_eq!(requests_count + 1, responses_count); } + fn create_record_batches_without_nullable_column(start: i64) -> Vec { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false), + ])); + + let mut record_batches = Vec::with_capacity(3); + for chunk in &(start..start + 9).chunks(3) { + let vs = chunk.collect_vec(); + let x1 = vs[0]; + let x2 = vs[1]; + let x3 = vs[2]; + + record_batches.push( + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![x1, x2, x3])) + as VectorRef, + Arc::new(Int32Vector::from_vec(vec![ + -x1 as i32, -x2 as i32, -x3 as i32, + ])), + ], + ) + .unwrap(), + ); + } + record_batches + } + fn create_record_batches(start: i64) -> Vec { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new(