fix(mito): re-encode bulk WAL entry after filling missing columns (#8808)

When a bulk insert request carries a stale schema, the worker fills the
missing columns via BulkPart::fill_missing_columns before writing. The
method replaced the batch but kept raw_data (the original Arrow Flight
IPC bytes), while BulkWalEntry::try_from(&BulkPart) prefers raw_data, so
the memtable received the filled batch but the WAL recorded the pre-fill
bytes.

Replaying such an entry restores a batch that misses the filled columns:

- Bulk memtable (flat format): convert_bulk_part fails with
  ColumnNotFound; the error is swallowed by the no-op write notifier and
  the rows are silently lost after restart.
- Time series memtable: BulkPart::to_mutation builds rows shorter than
  the declared schema and the region worker panics with index out of
  bounds during replay, hanging the region open.

Fixes:

- fill_missing_columns clears raw_data so the WAL entry is re-encoded
  from the filled batch.
- replay_memtable fills missing columns for replayed bulk parts of dense
  regions, so entries already written by affected versions replay
  correctly.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
(cherry picked from commit 1f4a8ad71cdfc135cf571ea7bca2cb87ed90ad72)
(cherry picked from commit 7fae38e667a17b3cd3de3f0540cb3a7b6d5a277a)

src/mito2/src/engine/alter_test.rs

@@ -17,21 +17,26 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
+use api::v1::bulk_wal_entry::Body;
 use api::v1::helper::{row, tag_column_schema};
 use api::v1::value::ValueData;
-use api::v1::{ColumnDataType, Row, Rows, SemanticType, Value};
+use api::v1::{ArrowIpc, BulkWalEntry, ColumnDataType, Row, Rows, SemanticType, Value, WalEntry};
 use common_error::ext::ErrorExt;
 use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions};
-use common_recordbatch::RecordBatches;
+use common_recordbatch::{DfRecordBatch, RecordBatches};
+use common_test_util::flight::encode_to_flight_data;
 use datafusion_expr::col;
+use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray};
+use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
 use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions};
+use store_api::logstore::provider::Provider;
 use store_api::metadata::ColumnMetadata;
 use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY;
 use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole};
 use store_api::region_request::{
-    AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, RegionOpenRequest,
-    RegionRequest, SetIndexOption, SetRegionOption,
+    AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest,
+    RegionBulkInsertsRequest, RegionOpenRequest, RegionRequest, SetIndexOption, SetRegionOption,
 };
 use store_api::storage::{ColumnId, RegionId, ScanRequest};
 
@@ -42,9 +47,10 @@ use crate::error;
 use crate::sst::FormatType;
 use crate::test_util::batch_util::sort_batches_and_print;
 use crate::test_util::{
-    CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key,
-    column_metadata_to_column_schema, flush_region, put_rows, rows_schema,
+    CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, build_rows_for_key,
+    column_metadata_to_column_schema, flush_region, put_rows, reopen_region, rows_schema,
 };
+use crate::wal::Wal;
 
 async fn scan_check_after_alter(engine: &MitoEngine, region_id: RegionId, expected: &str) {
     let request = ScanRequest::default();
@@ -1978,3 +1984,199 @@ async fn test_alter_region_append_mode_invalid() {
     // append_mode should still be true
     check_append_mode(&engine, true);
 }
+
+/// Builds a batch against the schema before [add_nullable_field1]
+/// (schema version 0): `(tag_0, field_0, ts)`.
+fn build_schema_v0_batch() -> DfRecordBatch {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("tag_0", DataType::Utf8, true),
+        Field::new("field_0", DataType::Float64, true),
+        Field::new(
+            "ts",
+            DataType::Timestamp(TimeUnit::Millisecond, None),
+            false,
+        ),
+    ]));
+    let tag = Arc::new(StringArray::from_iter_values(
+        (0..3).map(|v| format!("a{v}")),
+    )) as ArrayRef;
+    let field = Arc::new(Float64Array::from_iter_values((0..3).map(|v| v as f64))) as ArrayRef;
+    let ts = Arc::new(TimestampMillisecondArray::from_iter_values(
+        (0..3).map(|v| v as i64 * 1000),
+    )) as ArrayRef;
+    DfRecordBatch::try_new(schema, vec![tag, field, ts]).unwrap()
+}
+
+fn encode_arrow_ipc(batch: &DfRecordBatch) -> ArrowIpc {
+    let (schema, record_batch) = encode_to_flight_data(batch.clone());
+    ArrowIpc {
+        schema: schema.data_header,
+        data_header: record_batch.data_header,
+        payload: record_batch.data_body,
+    }
+}
+
+/// Builds a bulk insert request whose batch was built against the schema
+/// before [add_nullable_field1] (schema version 0): `(tag_0, field_0, ts)`.
+fn build_schema_v0_bulk_request(region_id: RegionId) -> RegionBulkInsertsRequest {
+    let payload = build_schema_v0_batch();
+    let raw_data = encode_arrow_ipc(&payload);
+
+    RegionBulkInsertsRequest {
+        region_id,
+        payload,
+        raw_data,
+        partition_expr_version: None,
+        // The writer only knows schema version 0 while the region is already
+        // at version 1, so the worker has to fill the missing columns.
+        aligned_schema_version: Some(0),
+    }
+}
+
+async fn scan_all_sorted(engine: &MitoEngine, region_id: RegionId) -> String {
+    let stream = engine
+        .scan_to_stream(region_id, ScanRequest::default())
+        .await
+        .unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    sort_batches_and_print(&batches, &["tag_0", "ts"])
+}
+
+const SCHEMA_V0_BULK_ROWS: &str = "\
++-------+---------+---------------------+---------+
+| tag_0 | field_0 | ts                  | field_1 |
++-------+---------+---------------------+---------+
+| a0    | 0.0     | 1970-01-01T00:00:00 |         |
+| a1    | 1.0     | 1970-01-01T00:00:01 |         |
+| a2    | 2.0     | 1970-01-01T00:00:02 |         |
++-------+---------+---------------------+---------+";
+
+#[tokio::test]
+async fn test_bulk_insert_stale_schema_wal_replay() {
+    test_bulk_insert_stale_schema_wal_replay_with_format(false).await;
+    test_bulk_insert_stale_schema_wal_replay_with_format(true).await;
+}
+
+async fn test_bulk_insert_stale_schema_wal_replay_with_format(flat_format: bool) {
+    common_telemetry::init_default_ut_logging();
+
+    let mut env = TestEnv::new().await;
+    let engine = env
+        .create_engine(MitoConfig {
+            default_flat_format: flat_format,
+            ..Default::default()
+        })
+        .await;
+
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+    let table_dir = request.table_dir.clone();
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    // Bumps the region schema version to 1.
+    engine
+        .handle_request(region_id, RegionRequest::Alter(add_nullable_field1()))
+        .await
+        .unwrap();
+
+    // Bulk insert whose batch misses `field_1`. The worker fills the missing
+    // column before writing to the memtable and the WAL.
+    let response = engine
+        .handle_request(
+            region_id,
+            RegionRequest::BulkInserts(build_schema_v0_bulk_request(region_id)),
+        )
+        .await
+        .unwrap();
+    assert_eq!(3, response.affected_rows);
+
+    assert_eq!(
+        SCHEMA_V0_BULK_ROWS,
+        scan_all_sorted(&engine, region_id).await,
+        "flat_format: {flat_format}"
+    );
+
+    // Reopens the region to replay the WAL. The rows must survive the replay.
+    reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await;
+
+    assert_eq!(
+        SCHEMA_V0_BULK_ROWS,
+        scan_all_sorted(&engine, region_id).await,
+        "flat_format: {flat_format}"
+    );
+}
+
+/// Replays a bulk WAL entry that misses columns added by an alter. This
+/// simulates entries written by old versions that keep the stale raw data
+/// when filling missing columns.
+#[tokio::test]
+async fn test_bulk_insert_stale_wal_entry_replay() {
+    test_bulk_insert_stale_wal_entry_replay_with_format(false).await;
+    test_bulk_insert_stale_wal_entry_replay_with_format(true).await;
+}
+
+async fn test_bulk_insert_stale_wal_entry_replay_with_format(flat_format: bool) {
+    common_telemetry::init_default_ut_logging();
+
+    let mut env = TestEnv::new().await;
+    let engine = env
+        .create_engine(MitoConfig {
+            default_flat_format: flat_format,
+            ..Default::default()
+        })
+        .await;
+
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+    let table_dir = request.table_dir.clone();
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    // Bumps the region schema version to 1.
+    engine
+        .handle_request(region_id, RegionRequest::Alter(add_nullable_field1()))
+        .await
+        .unwrap();
+
+    // Writes a bulk entry whose batch misses `field_1` to the WAL directly.
+    let batch = build_schema_v0_batch();
+    let entry = WalEntry {
+        mutations: vec![],
+        bulk_entries: vec![BulkWalEntry {
+            sequence: 1,
+            max_ts: 2000,
+            min_ts: 0,
+            timestamp_index: 2,
+            body: Some(Body::ArrowIpc(encode_arrow_ipc(&batch))),
+        }],
+    };
+    let LogStoreImpl::RaftEngine(log_store) = env.get_log_store().unwrap() else {
+        unreachable!()
+    };
+    let wal = Wal::new(log_store);
+    let mut writer = wal.writer();
+    writer
+        .add_entry(
+            region_id,
+            1,
+            &entry,
+            &Provider::raft_engine_provider(region_id.as_u64()),
+        )
+        .unwrap();
+    writer.write_to_wal().await.unwrap();
+
+    // Reopens the region to replay the WAL. The replay must fill the missing
+    // column instead of losing the rows.
+    reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await;
+
+    assert_eq!(
+        SCHEMA_V0_BULK_ROWS,
+        scan_all_sorted(&engine, region_id).await,
+        "flat_format: {flat_format}"
+    );
+}

src/mito2/src/memtable/bulk/part.rs

@@ -262,6 +262,12 @@ impl BulkPart {
         // Update the batch
         self.batch = new_batch;
 
+        // The raw data no longer matches the batch after filling missing
+        // columns. Clears it so the WAL entry is re-encoded from the filled
+        // batch, otherwise replaying the entry restores a batch without the
+        // filled columns and the memtable rejects it.
+        self.raw_data = None;
+
         Ok(())
     }
 
@@ -2896,4 +2902,62 @@ mod tests {
         let total_rows: usize = reader.map(|r| r.unwrap().num_rows()).sum();
         assert_eq!(total_rows, 6);
     }
+
+    #[test]
+    fn test_fill_missing_columns_resets_raw_data() {
+        let metadata = metadata_for_test();
+
+        // A batch missing the `v1` field column, e.g. built against an older
+        // schema before `v1` was added.
+        let k0_array = Arc::new(arrow::array::StringArray::from(vec!["key1", "key2"]));
+        let k1_array = Arc::new(arrow::array::UInt32Array::from(vec![1u32, 2]));
+        let v0_array = Arc::new(arrow::array::Int64Array::from(vec![100i64, 200]));
+        let ts_array = Arc::new(TimestampMillisecondArray::from(vec![1000i64, 2000]));
+        let input_schema = Arc::new(Schema::new(vec![
+            Field::new("k0", ArrowDataType::Utf8, false),
+            Field::new("k1", ArrowDataType::UInt32, false),
+            Field::new("v0", ArrowDataType::Int64, true),
+            Field::new(
+                "ts",
+                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
+                false,
+            ),
+        ]));
+        let batch =
+            RecordBatch::try_new(input_schema, vec![k0_array, k1_array, v0_array, ts_array])
+                .unwrap();
+
+        // Encodes the batch as raw data like the bulk insert request does.
+        let mut encoder = FlightEncoder::default();
+        let schema_bytes = encoder.encode_schema(batch.schema().as_ref()).data_header;
+        let [flight_data] = encoder
+            .encode(FlightMessage::RecordBatch(batch.clone()))
+            .try_into()
+            .unwrap();
+        let mut part = BulkPart {
+            batch,
+            max_timestamp: 2000,
+            min_timestamp: 1000,
+            sequence: 5,
+            timestamp_index: 3,
+            raw_data: Some(ArrowIpc {
+                schema: schema_bytes,
+                data_header: flight_data.data_header,
+                payload: flight_data.data_body,
+            }),
+        };
+
+        part.fill_missing_columns(&metadata).unwrap();
+        assert!(part.batch.column_by_name("v1").is_some());
+        // The raw data no longer matches the batch after filling missing
+        // columns and must be cleared, otherwise the WAL entry built from
+        // this part loses the filled columns.
+        assert!(part.raw_data.is_none());
+
+        // The WAL entry round trip keeps the filled column.
+        let entry = BulkWalEntry::try_from(&part).unwrap();
+        let replayed = BulkPart::try_from(entry).unwrap();
+        assert_eq!(2, replayed.num_rows());
+        assert!(replayed.batch.column_by_name("v1").is_some());
+    }
 }

src/mito2/src/region/opener.rs

@@ -32,6 +32,7 @@ use object_store::manager::ObjectStoreManagerRef;
 use object_store::util::{is_object_storage, normalize_dir};
 use parquet::file::metadata::PageIndexPolicy;
 use snafu::{OptionExt, ResultExt, ensure};
+use store_api::codec::PrimaryKeyEncoding;
 use store_api::logstore::LogStore;
 use store_api::logstore::provider::Provider;
 use store_api::metadata::{
@@ -832,6 +833,7 @@ where
     // data in the WAL.
     let mut last_entry_id = flushed_entry_id;
     let replay_from_entry_id = flushed_entry_id + 1;
+    let region_metadata = version_control.current().version.metadata.clone();
 
     let mut wal_stream = wal_entry_reader.read(provider, replay_from_entry_id)?;
     while let Some(res) = wal_stream.next().await {
@@ -876,7 +878,16 @@ where
         }
 
         for bulk_entry in entry.bulk_entries {
-            let part = BulkPart::try_from(bulk_entry)?;
+            let mut part = BulkPart::try_from(bulk_entry)?;
+            // The entry may miss columns added by a concurrent alter if it was
+            // written by a writer with a stale schema (older versions kept the
+            // stale raw data when filling missing columns). Fills missing
+            // columns like the write path, otherwise the memtable rejects the
+            // batch. Skips sparse batches as they don't carry tag columns by
+            // design.
+            if region_metadata.primary_key_encoding != PrimaryKeyEncoding::Sparse {
+                part.fill_missing_columns(&region_metadata)?;
+            }
             rows_replayed += part.num_rows();
             // During replay, we should adopt the sequence from WAL.
             let bulk_sequence_from_wal = part.sequence;