diff --git a/Cargo.lock b/Cargo.lock index 0a2d21cfd23..abce62f2d26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4303,6 +4303,7 @@ dependencies = [ "datafusion-common", "datafusion-expr", "datatypes", + "either", "file-engine", "futures", "futures-util", diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 8747c585776..54dc1721028 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -41,6 +41,7 @@ datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datatypes.workspace = true +either.workspace = true file-engine.workspace = true futures.workspace = true futures-util.workspace = true diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 26ea247e5cf..1a153e5a56a 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -44,6 +44,7 @@ use common_telemetry::{debug, error, info, warn}; use dashmap::DashMap; use datafusion::datasource::TableProvider; use datafusion_common::tree_node::TreeNode; +use either::Either; use futures_util::future::try_join_all; use metric_engine::engine::MetricEngine; use mito2::engine::{MITO_ENGINE_NAME, MitoEngine}; @@ -453,35 +454,46 @@ impl RegionServer { ) -> Result { let requests = RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; - let tracing_context = TracingContext::from_current_span(); - let join_tasks = requests.into_iter().map(|(region_id, req)| { - let self_to_move = self; - let span = tracing_context.attach(info_span!( - "RegionServer::handle_region_request", - region_id = region_id.to_string() - )); - async move { - self_to_move - .handle_request(region_id, req) - .trace(span) - .await + // Try to optimize batch Put requests for metric engine + // Returns either Some(response) or None(requests_back) + match self.try_handle_metric_batch_puts(requests).await? { + Either::Left(response) => Ok(response), + Either::Right(requests) => { + // Fallback: original parallel processing + let tracing_context = TracingContext::from_current_span(); + let join_tasks = + requests + .into_iter() + .map(|(region_id, req): (RegionId, RegionRequest)| { + let self_to_move = self; + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = region_id.to_string() + )); + async move { + self_to_move + .handle_request(region_id, req) + .trace(span) + .await + } + }); + + let results = try_join_all(join_tasks).await?; + let mut affected_rows = 0; + let mut extensions = HashMap::new(); + for result in results { + affected_rows += result.affected_rows; + extensions.extend(result.extensions); + } + + Ok(RegionResponse { + affected_rows, + extensions, + metadata: Vec::new(), + }) } - }); - - let results = try_join_all(join_tasks).await?; - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - for result in results { - affected_rows += result.affected_rows; - extensions.extend(result.extensions); } - - Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - }) } async fn handle_requests_in_serial( @@ -494,8 +506,6 @@ impl RegionServer { let mut affected_rows = 0; let mut extensions = HashMap::new(); - // FIXME(jeremy, ruihang): Once the engine supports merged calls, we should immediately - // modify this part to avoid inefficient serial loop calls. for (region_id, req) in requests { let span = tracing_context.attach(info_span!( "RegionServer::handle_region_request", @@ -514,6 +524,104 @@ impl RegionServer { }) } + /// Attempts to optimize batch Put requests for metric engine. + /// + /// Returns Either::Left(response) if optimization succeeded, + /// or Either::Right(original_requests) to fall back to parallel processing. + /// + /// This avoids cloning requests when optimization cannot be applied. + async fn try_handle_metric_batch_puts( + &self, + requests: Vec<(RegionId, RegionRequest)>, + ) -> Result>> { + if requests.is_empty() { + return Ok(Either::Right(requests)); + } + + // Quick check: verify first request is Put and is metric engine + if !matches!(requests[0].1, RegionRequest::Put(_)) { + return Ok(Either::Right(requests)); + } + let first_region_id = requests[0].0; + let request_type = requests[0].1.request_type(); + + // SAFETY: If the first request belongs to metric engine, then ALL requests + // in this batch are guaranteed to belong to metric engine. This invariant + // is maintained by the request batching logic upstream. + let engine = match self + .inner + .get_engine(first_region_id, &RegionChange::None)? + { + CurrentEngine::Engine(e) => e, + _ => return Ok(Either::Right(requests)), + }; + + if engine.name() != METRIC_ENGINE_NAME { + return Ok(Either::Right(requests)); + } + + // Check if ALL requests are Put (now we know it's worth checking) + let mut all_puts = true; + for (_, req) in &requests { + if !matches!(req, RegionRequest::Put(_)) { + all_puts = false; + break; + } + } + + if !all_puts { + return Ok(Either::Right(requests)); + } + + // Now extract Put requests by consuming ownership (zero clone!) + let put_requests = requests.into_iter().map(|(region_id, req)| { + if let RegionRequest::Put(put) = req { + (region_id, put) + } else { + unreachable!("Already checked all are Put") + } + }); + + // Downcast to MetricEngine and call batch API + let metric_engine = + engine + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + violated: "Failed to downcast to MetricEngine", + })?; + + let tracing_context = TracingContext::from_current_span(); + let batch_size = put_requests.len(); + let span = tracing_context.attach(info_span!( + "RegionServer::handle_metric_batch_puts", + batch_size = batch_size, + )); + let result = metric_engine + .put_regions_batch(put_requests) + .trace(span) + .await + .map_err(BoxedError::new) + .context(HandleRegionRequestSnafu { + region_id: first_region_id, + }); + + match result { + Ok(total_affected) => { + crate::metrics::REGION_CHANGED_ROW_COUNT + .with_label_values(&[request_type]) + .inc_by(total_affected as u64); + Ok(Either::Left(RegionResponse::new(total_affected))) + } + Err(err) => { + crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT + .with_label_values(&[request_type]) + .inc_by(batch_size as u64); + Err(err) + } + } + } + async fn handle_sync_region_request(&self, request: &SyncRequest) -> Result { let region_id = RegionId::from_u64(request.region_id); let manifest_info = request diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index cf50c01830f..f84906960c1 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -49,7 +49,8 @@ use store_api::region_engine::{ SyncRegionFromResponse, }; use store_api::region_request::{ - BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionRequest, + AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, }; use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; @@ -496,6 +497,16 @@ impl MetricEngine { self.inner.mito.clone() } + /// Batch put operation for multiple logical regions. + /// Requests are grouped by physical region; a failure can leave earlier + /// physical-region groups committed. + pub async fn put_regions_batch( + &self, + requests: impl ExactSizeIterator, + ) -> Result { + self.inner.put_regions_batch(requests).await + } + /// Returns all logical regions associated with the physical region. pub async fn logical_regions(&self, physical_region_id: RegionId) -> Result> { self.inner diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index c8c513f48c0..f6b6512d82a 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -12,8 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::{Rows, WriteHint}; +use std::collections::HashMap; + +use api::v1::{ + ColumnSchema, PrimaryKeyEncoding as PrimaryKeyEncodingProto, Row, Rows, Value, WriteHint, +}; use common_telemetry::{error, info}; +use fxhash::FxHashMap; use snafu::{OptionExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::region_request::{ @@ -27,7 +32,7 @@ use crate::error::{ PhysicalRegionNotFoundSnafu, Result, UnsupportedRegionRequestSnafu, }; use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED}; -use crate::row_modifier::RowsIter; +use crate::row_modifier::{RowsIter, TableIdInput}; use crate::utils::to_data_region_id; impl MetricEngineInner { @@ -52,6 +57,291 @@ impl MetricEngineInner { } } + /// Batch write multiple logical regions to the same physical region. + /// + /// Dispatch region put requests in batch. + /// + /// Requests may span multiple physical regions. We group them by physical + /// region and write sequentially. This method fails fast on validation or + /// preparation errors within a group and stops at the first failure. + /// Writes in earlier physical-region groups are not rolled back if a later + /// group fails. + pub async fn put_regions_batch( + &self, + requests: impl ExactSizeIterator, + ) -> Result { + let len = requests.len(); + + if len == 0 { + return Ok(0); + } + + let _timer = MITO_OPERATION_ELAPSED + .with_label_values(&["put_batch"]) + .start_timer(); + + // Fast path: single request, no batching overhead + if len == 1 { + let (logical_id, req) = requests.into_iter().next().unwrap(); + return self.put_logical_region(logical_id, req).await; + } + + let mut requests_per_physical: HashMap> = + HashMap::new(); + for (logical_region_id, request) in requests { + let physical_region_id = self.find_physical_region_id(logical_region_id)?; + requests_per_physical + .entry(physical_region_id) + .or_default() + .push((logical_region_id, request)); + } + + let mut total_affected_rows: AffectedRows = 0; + for (physical_region_id, requests) in requests_per_physical { + let affected_rows = self + .put_regions_batch_single_physical(physical_region_id, requests) + .await?; + total_affected_rows += affected_rows; + } + + Ok(total_affected_rows) + } + + /// Write a batch of requests that all belong to the same physical region. + /// + /// This function orchestrates the batch write process: + /// 1. Validates all requests + /// 2. Merges requests according to the encoding strategy (sparse or dense) + /// 3. Writes the merged batch to the physical region + async fn put_regions_batch_single_physical( + &self, + physical_region_id: RegionId, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Result { + if requests.is_empty() { + return Ok(0); + } + + let data_region_id = to_data_region_id(physical_region_id); + let primary_key_encoding = self.get_primary_key_encoding(data_region_id)?; + + // Validate all requests + self.validate_batch_requests(physical_region_id, &requests) + .await?; + + // Merge requests according to encoding strategy + let (merged_request, total_affected_rows) = match primary_key_encoding { + PrimaryKeyEncoding::Sparse => self.merge_sparse_batch(physical_region_id, requests)?, + PrimaryKeyEncoding::Dense => self.merge_dense_batch(data_region_id, requests)?, + }; + + // Write once to the physical region + self.data_region + .write_data(data_region_id, RegionRequest::Put(merged_request)) + .await?; + + Ok(total_affected_rows) + } + + /// Get primary key encoding for a data region. + fn get_primary_key_encoding(&self, data_region_id: RegionId) -> Result { + let state = self.state.read().unwrap(); + state + .get_primary_key_encoding(data_region_id) + .context(PhysicalRegionNotFoundSnafu { + region_id: data_region_id, + }) + } + + /// Validates all requests in a batch. + async fn validate_batch_requests( + &self, + physical_region_id: RegionId, + requests: &[(RegionId, RegionPutRequest)], + ) -> Result<()> { + for (logical_region_id, request) in requests { + self.verify_rows(*logical_region_id, physical_region_id, &request.rows) + .await?; + } + Ok(()) + } + + /// Merges multiple requests using sparse primary key encoding. + fn merge_sparse_batch( + &self, + physical_region_id: RegionId, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Result<(RegionPutRequest, AffectedRows)> { + let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum(); + let mut merged_rows = Vec::with_capacity(total_rows); + let mut total_affected_rows: AffectedRows = 0; + let mut output_schema: Option> = None; + + // Modify and collect rows from each request + for (logical_region_id, mut request) in requests { + self.modify_rows( + physical_region_id, + logical_region_id.table_id(), + &mut request.rows, + PrimaryKeyEncoding::Sparse, + )?; + + let row_count = request.rows.rows.len(); + total_affected_rows += row_count as AffectedRows; + + // Capture the output schema from the first modified request + if output_schema.is_none() { + output_schema = Some(request.rows.schema.clone()); + } + + merged_rows.extend(request.rows.rows); + } + + // Safe to unwrap: requests is guaranteed non-empty by caller + let schema = output_schema.unwrap(); + + let merged_request = RegionPutRequest { + rows: Rows { + schema, + rows: merged_rows, + }, + hint: Some(WriteHint { + primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(), + }), + }; + + Ok((merged_request, total_affected_rows)) + } + + /// Merges multiple requests using dense primary key encoding. + /// + /// In dense mode, different requests can have different columns. + /// We merge all schemas into a union schema, align each row to this schema, + /// then batch-modify all rows together (adding __table_id and __tsid). + fn merge_dense_batch( + &self, + data_region_id: RegionId, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Result<(RegionPutRequest, AffectedRows)> { + // Build union schema from all requests + let merged_schema = Self::build_union_schema(&requests); + + // Align all rows to the merged schema and collect table_ids + let (merged_rows, table_ids) = Self::align_requests_to_schema(requests, &merged_schema); + + // Batch-modify all rows (add __table_id and __tsid columns) + let final_rows = { + let state = self.state.read().unwrap(); + let name_to_id = state + .physical_region_states() + .get(&data_region_id) + .with_context(|| PhysicalRegionNotFoundSnafu { + region_id: data_region_id, + })? + .physical_columns(); + + let iter = RowsIter::new( + Rows { + schema: merged_schema, + rows: merged_rows, + }, + name_to_id, + ); + + self.row_modifier.modify_rows( + iter, + TableIdInput::Batch(&table_ids), + PrimaryKeyEncoding::Dense, + )? + }; + + let merged_request = RegionPutRequest { + rows: final_rows, + hint: None, + }; + + Ok((merged_request, table_ids.len() as AffectedRows)) + } + + /// Builds a union schema containing all columns from all requests. + fn build_union_schema(requests: &[(RegionId, RegionPutRequest)]) -> Vec { + let mut schema_map: HashMap<&str, ColumnSchema> = HashMap::new(); + for (_, request) in requests { + for col in &request.rows.schema { + schema_map + .entry(col.column_name.as_str()) + .or_insert_with(|| col.clone()); + } + } + schema_map.into_values().collect() + } + + fn align_requests_to_schema( + requests: Vec<(RegionId, RegionPutRequest)>, + merged_schema: &[ColumnSchema], + ) -> (Vec, Vec) { + // Pre-calculate total capacity + let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum(); + let mut merged_rows = Vec::with_capacity(total_rows); + let mut table_ids = Vec::with_capacity(total_rows); + + let null_value = Value { value_data: None }; + + for (logical_region_id, request) in requests { + let table_id = logical_region_id.table_id(); + + // Build column name to index mapping once per request + let col_name_to_idx: FxHashMap<&str, usize> = request + .rows + .schema + .iter() + .enumerate() + .map(|(idx, col)| (col.column_name.as_str(), idx)) + .collect(); + + // Build column mapping array once per request + // col_mapping[i] = Some(idx) means merged_schema[i] is at request.schema[idx] + // col_mapping[i] = None means merged_schema[i] doesn't exist in request.schema + let col_mapping: Vec> = merged_schema + .iter() + .map(|merged_col| { + col_name_to_idx + .get(merged_col.column_name.as_str()) + .copied() + }) + .collect(); + + // Apply the mapping to all rows + for mut row in request.rows.rows { + let mut aligned_values = Vec::with_capacity(merged_schema.len()); + for &opt_idx in &col_mapping { + aligned_values.push(match opt_idx { + Some(idx) => std::mem::take(&mut row.values[idx]), + None => null_value.clone(), + }); + } + merged_rows.push(Row { + values: aligned_values, + }); + table_ids.push(table_id); + } + } + + (merged_rows, table_ids) + } + + /// Find the physical region id for a logical region. + fn find_physical_region_id(&self, logical_region_id: RegionId) -> Result { + let state = self.state.read().unwrap(); + state + .logical_regions() + .get(&logical_region_id) + .copied() + .context(LogicalRegionNotFoundSnafu { + region_id: logical_region_id, + }) + } + /// Dispatch region delete request pub async fn delete_region( &self, @@ -98,7 +388,7 @@ impl MetricEngineInner { )?; if primary_key_encoding == PrimaryKeyEncoding::Sparse { request.hint = Some(WriteHint { - primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(), }); } self.data_region @@ -131,7 +421,7 @@ impl MetricEngineInner { )?; if primary_key_encoding == PrimaryKeyEncoding::Sparse { request.hint = Some(WriteHint { - primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(), }); } self.data_region @@ -224,7 +514,9 @@ impl MetricEngineInner { .physical_columns(); RowsIter::new(input, name_to_id) }; - let output = self.row_modifier.modify_rows(iter, table_id, encoding)?; + let output = + self.row_modifier + .modify_rows(iter, TableIdInput::Single(table_id), encoding)?; *rows = output; Ok(()) } @@ -232,14 +524,207 @@ impl MetricEngineInner { #[cfg(test)] mod tests { + use std::collections::HashSet; + use common_recordbatch::RecordBatches; + use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, + MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, + }; + use store_api::path_utils::table_dir; use store_api::region_engine::RegionEngine; use store_api::region_request::RegionRequest; use store_api::storage::ScanRequest; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use super::*; use crate::test_util::{self, TestEnv}; + fn assert_merged_schema(rows: &Rows, expect_sparse: bool) { + let column_names: HashSet = rows + .schema + .iter() + .map(|col| col.column_name.clone()) + .collect(); + + if expect_sparse { + assert!( + column_names.contains(PRIMARY_KEY_COLUMN_NAME), + "sparse encoding should include primary key column" + ); + assert!( + !column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME), + "sparse encoding should not include table id column" + ); + assert!( + !column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME), + "sparse encoding should not include tsid column" + ); + assert!( + !column_names.contains("job"), + "sparse encoding should not include tag columns" + ); + assert!( + !column_names.contains("instance"), + "sparse encoding should not include tag columns" + ); + } else { + assert!( + !column_names.contains(PRIMARY_KEY_COLUMN_NAME), + "dense encoding should not include primary key column" + ); + assert!( + column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME), + "dense encoding should include table id column" + ); + assert!( + column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME), + "dense encoding should include tsid column" + ); + assert!( + column_names.contains("job"), + "dense encoding should keep tag columns" + ); + assert!( + column_names.contains("instance"), + "dense encoding should keep tag columns" + ); + } + } + + async fn create_logical_region_with_tags( + env: &TestEnv, + physical_region_id: RegionId, + logical_region_id: RegionId, + tags: &[&str], + ) { + let region_create_request = test_util::create_logical_region_request( + tags, + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + ); + env.metric() + .handle_request( + logical_region_id, + RegionRequest::Create(region_create_request), + ) + .await + .unwrap(); + } + + async fn run_batch_write_with_schema_variants( + env: &TestEnv, + physical_region_id: RegionId, + options: Vec<(String, String)>, + expect_sparse: bool, + ) { + env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), options) + .await; + + let logical_region_1 = env.default_logical_region_id(); + let logical_region_2 = RegionId::new(1024, 1); + + create_logical_region_with_tags(env, physical_region_id, logical_region_1, &["job"]).await; + create_logical_region_with_tags( + env, + physical_region_id, + logical_region_2, + &["job", "instance"], + ) + .await; + + let schema_1 = test_util::row_schema_with_tags(&["job"]); + let schema_2 = test_util::row_schema_with_tags(&["job", "instance"]); + + let data_region_id = RegionId::new(physical_region_id.table_id(), 2); + let primary_key_encoding = env + .metric() + .inner + .get_primary_key_encoding(data_region_id) + .unwrap(); + assert_eq!( + primary_key_encoding, + if expect_sparse { + PrimaryKeyEncoding::Sparse + } else { + PrimaryKeyEncoding::Dense + } + ); + + let build_requests = || { + let rows_1 = test_util::build_rows(1, 3); + let rows_2 = test_util::build_rows(2, 2); + + vec![ + ( + logical_region_1, + RegionPutRequest { + rows: Rows { + schema: schema_1.clone(), + rows: rows_1, + }, + hint: None, + }, + ), + ( + logical_region_2, + RegionPutRequest { + rows: Rows { + schema: schema_2.clone(), + rows: rows_2, + }, + hint: None, + }, + ), + ] + }; + + let merged_request = if expect_sparse { + let (merged_request, _) = env + .metric() + .inner + .merge_sparse_batch(physical_region_id, build_requests()) + .unwrap(); + let hint = merged_request + .hint + .as_ref() + .expect("missing sparse write hint"); + assert_eq!( + hint.primary_key_encoding, + PrimaryKeyEncodingProto::Sparse as i32 + ); + merged_request + } else { + let (merged_request, _) = env + .metric() + .inner + .merge_dense_batch(data_region_id, build_requests()) + .unwrap(); + assert!(merged_request.hint.is_none()); + merged_request + }; + + assert_merged_schema(&merged_request.rows, expect_sparse); + + let affected_rows = env + .metric() + .inner + .put_regions_batch(build_requests().into_iter()) + .await + .unwrap(); + assert_eq!(affected_rows, 5); + + let request = ScanRequest::default(); + let stream = env + .mito() + .scan_to_stream(data_region_id, request) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 5); + } + #[tokio::test] async fn test_write_logical_region() { let env = TestEnv::new().await; @@ -374,4 +859,249 @@ mod tests { .await .unwrap_err(); } + + #[tokio::test] + async fn test_batch_write_multiple_logical_regions() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + + // Create two additional logical regions + let physical_region_id = env.default_physical_region_id(); + let logical_region_1 = env.default_logical_region_id(); + let logical_region_2 = RegionId::new(1024, 1); + let logical_region_3 = RegionId::new(1024, 2); + + env.create_logical_region(physical_region_id, logical_region_2) + .await; + env.create_logical_region(physical_region_id, logical_region_3) + .await; + + // Prepare batch requests with non-overlapping timestamps + let schema = test_util::row_schema_with_tags(&["job"]); + + // Use build_rows_with_ts to create non-overlapping timestamps + // logical_region_1: ts 0, 1, 2 + // logical_region_2: ts 10, 11 (offset to avoid overlap) + // logical_region_3: ts 20, 21, 22, 23, 24 (offset to avoid overlap) + let rows1 = test_util::build_rows(1, 3); + let mut rows2 = test_util::build_rows(1, 2); + let mut rows3 = test_util::build_rows(1, 5); + + // Adjust timestamps to avoid conflicts + use api::v1::value::ValueData; + for (i, row) in rows2.iter_mut().enumerate() { + if let Some(ValueData::TimestampMillisecondValue(ts)) = + row.values.get_mut(0).and_then(|v| v.value_data.as_mut()) + { + *ts = (10 + i) as i64; + } + } + for (i, row) in rows3.iter_mut().enumerate() { + if let Some(ValueData::TimestampMillisecondValue(ts)) = + row.values.get_mut(0).and_then(|v| v.value_data.as_mut()) + { + *ts = (20 + i) as i64; + } + } + + let requests = vec![ + ( + logical_region_1, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: rows1, + }, + hint: None, + }, + ), + ( + logical_region_2, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: rows2, + }, + hint: None, + }, + ), + ( + logical_region_3, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: rows3, + }, + hint: None, + }, + ), + ]; + + // Batch write + let affected_rows = engine + .inner + .put_regions_batch(requests.into_iter()) + .await + .unwrap(); + assert_eq!(affected_rows, 10); + + // Verify physical region contains data from all logical regions + let request = ScanRequest::default(); + let stream = env + .metric() + .scan_to_stream(physical_region_id, request) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + + // Should have 3 + 2 + 5 = 10 rows total + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 10); + } + + #[tokio::test] + async fn test_batch_write_with_partial_failure() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + + let physical_region_id = env.default_physical_region_id(); + let logical_region_1 = env.default_logical_region_id(); + let logical_region_2 = RegionId::new(1024, 1); + let nonexistent_region = RegionId::new(9999, 9999); + + env.create_logical_region(physical_region_id, logical_region_2) + .await; + + // Prepare batch with one invalid region + let schema = test_util::row_schema_with_tags(&["job"]); + let requests = vec![ + ( + logical_region_1, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: test_util::build_rows(1, 3), + }, + hint: None, + }, + ), + ( + nonexistent_region, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: test_util::build_rows(1, 2), + }, + hint: None, + }, + ), + ( + logical_region_2, + RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: test_util::build_rows(1, 5), + }, + hint: None, + }, + ), + ]; + + // Batch write + let result = engine.inner.put_regions_batch(requests.into_iter()).await; + assert!(result.is_err()); + + // Invalid region is detected before any write, so the physical region remains empty. + // Fail-fast is per physical-region group; cross-group partial success is possible. + let request = ScanRequest::default(); + let stream = env + .metric() + .scan_to_stream(physical_region_id, request) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + } + + #[tokio::test] + async fn test_batch_write_single_request_fast_path() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + + let logical_region_id = env.default_logical_region_id(); + let schema = test_util::row_schema_with_tags(&["job"]); + + // Single request should use fast path + let requests = vec![( + logical_region_id, + RegionPutRequest { + rows: Rows { + schema, + rows: test_util::build_rows(1, 5), + }, + hint: None, + }, + )]; + + let affected_rows = engine + .inner + .put_regions_batch(requests.into_iter()) + .await + .unwrap(); + assert_eq!(affected_rows, 5); + } + + #[tokio::test] + async fn test_batch_write_empty_requests() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + + // Empty batch should return zero affected rows + let requests = vec![]; + let affected_rows = engine + .inner + .put_regions_batch(requests.into_iter()) + .await + .unwrap(); + + assert_eq!(affected_rows, 0); + } + + #[tokio::test] + async fn test_batch_write_sparse_encoding() { + let env = TestEnv::new().await; + let physical_region_id = env.default_physical_region_id(); + + run_batch_write_with_schema_variants( + &env, + physical_region_id, + vec![( + MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING.to_string(), + "sparse".to_string(), + )], + true, + ) + .await; + } + + #[tokio::test] + async fn test_batch_write_dense_encoding() { + let env = TestEnv::new().await; + let physical_region_id = env.default_physical_region_id(); + + run_batch_write_with_schema_variants( + &env, + physical_region_id, + vec![( + MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING.to_string(), + "dense".to_string(), + )], + false, + ) + .await; + } } diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 73662324589..842cf18ee32 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -240,6 +240,14 @@ pub enum Error { location: Location, }, + #[snafu(display("Table id count mismatch. Expect {}, got {}", expected, actual))] + TableIdCountMismatch { + expected: usize, + actual: usize, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Alter request to physical region is forbidden"))] ForbiddenPhysicalAlter { #[snafu(implicit)] @@ -368,6 +376,7 @@ impl ErrorExt for Error { | MissingRegionOption { .. } | ConflictRegionOption { .. } | ColumnTypeMismatch { .. } + | TableIdCountMismatch { .. } | PhysicalRegionBusy { .. } | MultipleFieldColumn { .. } | NoFieldColumn { .. } diff --git a/src/metric-engine/src/row_modifier.rs b/src/metric-engine/src/row_modifier.rs index 0732359c391..d5ed1cc9b0c 100644 --- a/src/metric-engine/src/row_modifier.rs +++ b/src/metric-engine/src/row_modifier.rs @@ -29,7 +29,7 @@ use store_api::metric_engine_consts::{ use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId}; use store_api::storage::{ColumnId, TableId}; -use crate::error::{EncodePrimaryKeySnafu, Result}; +use crate::error::{EncodePrimaryKeySnafu, Result, TableIdCountMismatchSnafu}; /// A row modifier modifies [`Rows`]. /// @@ -42,6 +42,22 @@ pub struct RowModifier { codec: SparsePrimaryKeyCodec, } +/// Table id input for row modification. +#[derive(Clone, Copy)] +pub(crate) enum TableIdInput<'a> { + Single(TableId), + Batch(&'a [TableId]), +} + +impl<'a> TableIdInput<'a> { + fn table_id_for_row(&self, row_idx: usize) -> TableId { + match self { + TableIdInput::Single(table_id) => *table_id, + TableIdInput::Batch(table_ids) => table_ids[row_idx], + } + } +} + impl Default for RowModifier { fn default() -> Self { Self { @@ -51,22 +67,24 @@ impl Default for RowModifier { } impl RowModifier { - /// Modify rows with the given primary key encoding. + /// Modify rows with the given primary key encoding and table ids. pub(crate) fn modify_rows( &self, iter: RowsIter, - table_id: TableId, + table_ids: TableIdInput<'_>, encoding: PrimaryKeyEncoding, ) -> Result { + let row_count = iter.rows.rows.len(); + Self::validate_table_id_count(table_ids, row_count)?; match encoding { - PrimaryKeyEncoding::Sparse => self.modify_rows_sparse(iter, table_id), - PrimaryKeyEncoding::Dense => self.modify_rows_dense(iter, table_id), + PrimaryKeyEncoding::Sparse => self.modify_rows_sparse(iter, table_ids), + PrimaryKeyEncoding::Dense => self.modify_rows_dense(iter, table_ids), } } /// Modifies rows with sparse primary key encoding. /// It replaces the primary key columns with the encoded primary key column(`__primary_key`). - fn modify_rows_sparse(&self, mut iter: RowsIter, table_id: TableId) -> Result { + fn modify_rows_sparse(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result { let num_column = iter.rows.schema.len(); let num_primary_key_column = iter.index.num_primary_key_column; // num_output_column = remaining columns(fields columns + timestamp column) + 1 (encoded primary key column) @@ -74,14 +92,15 @@ impl RowModifier { let mut buffer = vec![]; - for mut iter in iter.iter_mut() { - let (table_id, tsid) = Self::fill_internal_columns(table_id, &iter); + for (row_index, mut row_iter) in iter.iter_mut().enumerate() { + let table_id = table_ids.table_id_for_row(row_index); + let (table_id_value, tsid) = Self::fill_internal_columns(table_id, &row_iter); let mut values = Vec::with_capacity(num_output_column); buffer.clear(); let internal_columns = [ ( ReservedColumnId::table_id(), - api::helper::pb_value_to_value_ref(&table_id, None), + api::helper::pb_value_to_value_ref(&table_id_value, None), ), ( ReservedColumnId::tsid(), @@ -92,13 +111,13 @@ impl RowModifier { .encode_to_vec(internal_columns.into_iter(), &mut buffer) .context(EncodePrimaryKeySnafu)?; self.codec - .encode_to_vec(iter.primary_keys(), &mut buffer) + .encode_to_vec(row_iter.primary_keys(), &mut buffer) .context(EncodePrimaryKeySnafu)?; values.push(ValueData::BinaryValue(buffer.clone()).into()); - values.extend(iter.remaining()); + values.extend(row_iter.remaining()); // Replace the row with the encoded row - *iter.row = Row { values }; + *row_iter.row = Row { values }; } // Update the schema @@ -118,7 +137,7 @@ impl RowModifier { /// Modifies rows with dense primary key encoding. /// It adds two columns(`__table_id`, `__tsid`) to the row. - fn modify_rows_dense(&self, mut iter: RowsIter, table_id: TableId) -> Result { + fn modify_rows_dense(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result { // add table_name column iter.rows.schema.push(ColumnSchema { column_name: DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), @@ -135,15 +154,29 @@ impl RowModifier { datatype_extension: None, options: None, }); - for iter in iter.iter_mut() { - let (table_id, tsid) = Self::fill_internal_columns(table_id, &iter); - iter.row.values.push(table_id); - iter.row.values.push(tsid); + for (row_index, row_iter) in iter.iter_mut().enumerate() { + let table_id = table_ids.table_id_for_row(row_index); + let (table_id_value, tsid) = Self::fill_internal_columns(table_id, &row_iter); + row_iter.row.values.push(table_id_value); + row_iter.row.values.push(tsid); } Ok(iter.rows) } + fn validate_table_id_count(table_ids: TableIdInput<'_>, row_count: usize) -> Result<()> { + if let TableIdInput::Batch(table_ids) = table_ids + && table_ids.len() != row_count + { + return TableIdCountMismatchSnafu { + expected: row_count, + actual: table_ids.len(), + } + .fail(); + } + Ok(()) + } + /// Fills internal columns of a row with table name and a hash of tag values. pub fn fill_internal_columns(table_id: TableId, iter: &RowIter<'_>) -> (Value, Value) { let ts_id = if !iter.has_null_labels() { @@ -389,8 +422,10 @@ mod tests { use std::collections::HashMap; use api::v1::{Row, Rows}; + use store_api::codec::PrimaryKeyEncoding; use super::*; + use crate::error::Error; fn test_schema() -> Vec { vec![ @@ -436,7 +471,13 @@ mod tests { rows: vec![row], }; let rows_iter = RowsIter::new(rows, &name_to_column_id); - let result = encoder.modify_rows_sparse(rows_iter, table_id).unwrap(); + let result = encoder + .modify_rows( + rows_iter, + TableIdInput::Single(table_id), + PrimaryKeyEncoding::Sparse, + ) + .unwrap(); assert_eq!(result.rows[0].values.len(), 1); let encoded_primary_key = vec![ 128, 0, 0, 4, 1, 0, 0, 4, 1, 128, 0, 0, 3, 1, 37, 196, 242, 181, 117, 224, 7, 137, 0, @@ -505,7 +546,13 @@ mod tests { rows: vec![row], }; let rows_iter = RowsIter::new(rows, &name_to_column_id); - let result = encoder.modify_rows_dense(rows_iter, table_id).unwrap(); + let result = encoder + .modify_rows( + rows_iter, + TableIdInput::Single(table_id), + PrimaryKeyEncoding::Dense, + ) + .unwrap(); assert_eq!( result.rows[0].values[0], ValueData::StringValue("greptimedb".to_string()).into() @@ -522,6 +569,34 @@ mod tests { assert_eq!(result.schema, expected_dense_schema()); } + #[test] + fn test_table_id_count_mismatch() { + let name_to_column_id = test_name_to_column_id(); + let encoder = RowModifier::default(); + let schema = test_schema(); + let rows = Rows { + schema, + rows: vec![test_row("a", "b"), test_row("c", "d")], + }; + let rows_iter = RowsIter::new(rows, &name_to_column_id); + let table_ids = [1025]; + let err = encoder + .modify_rows( + rows_iter, + TableIdInput::Batch(&table_ids), + PrimaryKeyEncoding::Dense, + ) + .unwrap_err(); + assert!(matches!( + err, + Error::TableIdCountMismatch { + expected: 2, + actual: 1, + .. + } + )); + } + #[test] fn test_fill_internal_columns() { let name_to_column_id = test_name_to_column_id();