fix: align bulk insert schema handling (#8222)

* 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 <mrsatangel@gmail.com> * 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 <mrsatangel@gmail.com> * chore: update proto to main branch commit Signed-off-by: Lei, HUANG <mrsatangel@gmail.com> * 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 <mrsatangel@gmail.com> * chore: add some doc Signed-off-by: Lei, HUANG <mrsatangel@gmail.com> --------- Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>
2026-07-06 22:10:39 +00:00 · 2026-06-04 10:55:09 +08:00
parent d304df6e75
commit da9d314fa4
11 changed files with 111 additions and 8 deletions
--- 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",
--- 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"
--- 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<AffectedRows> {
+        // 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<u64> {
+        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()
--- 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,
    }
 }

--- 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<RegionMetadataRef>,
    pub(crate) partition_expr_version: Option<u64>,
 }

--- a/src/mito2/src/worker/handle_bulk_insert.rs
+++ b/src/mito2/src/worker/handle_bulk_insert.rs
@@ -86,11 +86,20 @@ impl<S: LogStore> RegionWorkerLoop<S> {
            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,
        });
    }
--- a/src/mito2/src/worker/handle_write.rs
+++ b/src/mito2/src/worker/handle_write.rs
@@ -525,8 +525,10 @@ impl<S> RegionWorkerLoop<S> {
            }

            // 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
--- 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,
--- 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,
--- 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<Vec<(RegionId, Regi
 fn make_region_bulk_inserts(request: BulkInsertRequest) -> Result<Vec<(RegionId, RegionRequest)>> {
    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<Vec<(RegionId,
            payload,
            raw_data: request,
            partition_expr_version,
+            aligned_schema_version,
        }),
    )])
 }
@@ -1510,6 +1512,7 @@ pub struct RegionBulkInsertsRequest {
    pub payload: DfRecordBatch,
    pub raw_data: ArrowIpc,
    pub partition_expr_version: Option<u64>,
+    pub aligned_schema_version: Option<u64>,
 }

 impl RegionBulkInsertsRequest {
--- 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<RecordBatch> {
+        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<RecordBatch> {
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(