From b6326257d1babcc8cf016af59b92ebfcdee862ae Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 21:30:34 +0800 Subject: [PATCH] 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 (cherry picked from commit 1f4a8ad71cdfc135cf571ea7bca2cb87ed90ad72) (cherry picked from commit 7fae38e667a17b3cd3de3f0540cb3a7b6d5a277a) --- src/mito2/src/engine/alter_test.rs | 214 +++++++++++++++++++++++++++- src/mito2/src/memtable/bulk/part.rs | 64 +++++++++ src/mito2/src/region/opener.rs | 13 +- 3 files changed, 284 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index be4daa1d88e..2461f2ca0cd 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/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}" + ); +} diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index fcba4519478..20482794079 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/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()); + } } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 9388f943ceb..5d2d36df107 100644 --- a/src/mito2/src/region/opener.rs +++ b/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(®ion_metadata)?; + } rows_replayed += part.num_rows(); // During replay, we should adopt the sequence from WAL. let bulk_sequence_from_wal = part.sequence;