chore: batch requests in metric engine (#7535)

* chore: batch requests in metric engine Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: fix clippy Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: merge multiple schemas Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: add tracing span Signed-off-by: jeremyhi <fengjiachun@gmail.com> * feat: add metrics for batch put Signed-off-by: jeremyhi <fengjiachun@gmail.com> * feat: add sparse and dense encoding test case Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: avoid allocation of vec Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: table_id_for_row Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: by comment Signed-off-by: jeremyhi <fengjiachun@gmail.com> * feat: algorithm to reduce hash lookups by using array indexing Signed-off-by: jeremyhi <fengjiachun@gmail.com> --------- Signed-off-by: jeremyhi <fengjiachun@gmail.com>
2026-10-27 22:26:15 +00:00 · 2026-01-16 02:57:57 +00:00
parent e0285209cb
commit 007e6cc860
7 changed files with 988 additions and 53 deletions
@@ -4303,6 +4303,7 @@ dependencies = [
 "datafusion-common",
 "datafusion-expr",
 "datatypes",
+ "either",
 "file-engine",
 "futures",
 "futures-util",
@@ -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
@@ -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<RegionResponse> {
        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<Either<RegionResponse, Vec<(RegionId, RegionRequest)>>> {
+        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::<MetricEngine>()
+                .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<RegionResponse> {
        let region_id = RegionId::from_u64(request.region_id);
        let manifest_info = request
@@ -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<Item = (RegionId, RegionPutRequest)>,
+    ) -> Result<AffectedRows> {
+        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<Vec<RegionId>> {
        self.inner
@@ -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<Item = (RegionId, RegionPutRequest)>,
+    ) -> Result<AffectedRows> {
+        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<RegionId, Vec<(RegionId, RegionPutRequest)>> =
+            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<AffectedRows> {
+        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<PrimaryKeyEncoding> {
+        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<Vec<ColumnSchema>> = 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<ColumnSchema> {
+        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<Row>, Vec<TableId>) {
+        // 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<Option<usize>> = 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<RegionId> {
+        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<String> = 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::<usize>(), 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::<usize>(), 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::<usize>(), 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;
+    }
 }
@@ -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 { .. }
@@ -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<Rows> {
+        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<Rows> {
+    fn modify_rows_sparse(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result<Rows> {
        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<Rows> {
+    fn modify_rows_dense(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result<Rows> {
        // 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<ColumnSchema> {
        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();