From da9d314fa4329854bb4dd17e09c98e91906c2603 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 4 Jun 2026 10:55:09 +0800 Subject: [PATCH] fix: align bulk insert schema handling (#8222) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: align bulk insert schema handling Carry `aligned_schema_version` through bulk insert requests so datanode can recheck stale schemas only when needed. Use the physical data region schema version for metric logical bulk inserts, while external bulk insert callers leave alignment unknown. Related files: - `Cargo.toml` - `Cargo.lock` - `src/store-api/src/region_request.rs` - `src/metric-engine/src/engine/bulk_insert.rs` - `src/mito2/src/request.rs` - `src/mito2/src/worker/handle_bulk_insert.rs` - `src/mito2/src/worker/handle_write.rs` - `src/operator/src/bulk_insert.rs` - `src/servers/src/pending_rows_batcher.rs` - `src/mito2/src/engine/edit_region_test.rs` Signed-off-by: Lei, HUANG * test: cover bulk insert missing columns Add Flight bulk insert coverage for batches that omit nullable columns and verify the database fills the missing values with nulls. Related files: - `tests-integration/src/grpc/flight.rs` Signed-off-by: Lei, HUANG * chore: update proto to main branch commit Signed-off-by: Lei, HUANG * fix: align physical metric bulk insert schema Pass the physical schema version through metric bulk insert requests so physical and logical region writes share the same schema alignment behavior. Files: - \`src/metric-engine/src/engine/bulk_insert.rs\` — updates \`bulk_insert_physical_region\` and reuses \`physical_schema_version\` for bulk insert schema alignment. Signed-off-by: Lei, HUANG * chore: add some doc Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/metric-engine/src/engine/bulk_insert.rs | 22 ++++++- src/mito2/src/engine/edit_region_test.rs | 1 + src/mito2/src/request.rs | 2 +- src/mito2/src/worker/handle_bulk_insert.rs | 11 +++- src/mito2/src/worker/handle_write.rs | 6 +- src/operator/src/bulk_insert.rs | 2 + src/servers/src/pending_rows_batcher.rs | 4 ++ src/store-api/src/region_request.rs | 3 + tests-integration/src/grpc/flight.rs | 64 +++++++++++++++++++++ 11 files changed, 111 insertions(+), 8 deletions(-) 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(