greptimedb/src/metric-engine/src/engine/put.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;

use api::helper::ColumnDataTypeWrapper;
use api::v1::{
    ColumnSchema, PrimaryKeyEncoding as PrimaryKeyEncodingProto, Row, Rows, SemanticType, Value,
    WriteHint,
};
use common_telemetry::{error, info};
use fxhash::FxHashMap;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::ColumnMetadata;
use store_api::region_request::{
    AffectedRows, RegionDeleteRequest, RegionPutRequest, RegionRequest,
};
use store_api::storage::{RegionId, TableId};

use crate::engine::MetricEngineInner;
use crate::error::{
    ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu,
    LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu,
    UnsupportedRegionRequestSnafu,
};
use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED};
use crate::row_modifier::{RowsIter, TableIdInput};
use crate::utils::to_data_region_id;

impl MetricEngineInner {
    /// Dispatch region put request
    pub async fn put_region(
        &self,
        region_id: RegionId,
        request: RegionPutRequest,
    ) -> Result<AffectedRows> {
        let is_putting_physical_region =
            self.state.read().unwrap().exist_physical_region(region_id);

        if is_putting_physical_region {
            info!(
                "Metric region received put request {request:?} on physical region {region_id:?}"
            );
            FORBIDDEN_OPERATION_COUNT.inc();

            ForbiddenPhysicalWriteSnafu.fail()
        } else {
            self.put_logical_region(region_id, request).await
        }
    }

    /// 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 (region_id, req) = requests.into_iter().next().unwrap();
            let is_putting_physical_region =
                self.state.read().unwrap().exist_physical_region(region_id);
            if is_putting_physical_region {
                FORBIDDEN_OPERATION_COUNT.inc();
                return ForbiddenPhysicalWriteSnafu.fail();
            }

            return self.put_logical_region(region_id, req).await;
        }

        let mut requests_per_physical: HashMap<RegionId, Vec<(RegionId, RegionPutRequest)>> =
            HashMap::new();
        for (region_id, request) in requests {
            let is_putting_physical_region =
                self.state.read().unwrap().exist_physical_region(region_id);
            if is_putting_physical_region {
                FORBIDDEN_OPERATION_COUNT.inc();
                return ForbiddenPhysicalWriteSnafu.fail();
            }
            let physical_region_id = self.find_physical_region_id(region_id)?;
            requests_per_physical
                .entry(physical_region_id)
                .or_default()
                .push((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,
        mut 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
        let partition_expr_version = self
            .validate_batch_requests(physical_region_id, &mut requests)
            .await?;

        // Keep the common uniform-policy path allocation-free beyond the merge.
        if requests
            .iter()
            .all(|(_, request)| request.skip_wal == requests[0].1.skip_wal)
        {
            let (request, affected_rows) = self.merge_batch(
                physical_region_id,
                primary_key_encoding,
                requests,
                partition_expr_version,
            )?;
            self.data_region
                .write_data(data_region_id, RegionRequest::Put(request))
                .await?;
            return Ok(affected_rows);
        }

        // Only merge consecutive requests with the same WAL policy, preserving
        // write order even when the same logical region occurs more than once.
        // Prepare every batch before writing so a merge error cannot partially
        // apply this physical-region group.
        let batches = Self::split_batch_by_wal_policy(requests);
        let mut merged_requests = Vec::with_capacity(batches.len());
        let mut total_affected_rows = 0;
        for requests in batches {
            let (request, affected_rows) = self.merge_batch(
                physical_region_id,
                primary_key_encoding,
                requests,
                partition_expr_version,
            )?;
            merged_requests.push(request);
            total_affected_rows += affected_rows;
        }

        for request in merged_requests {
            self.data_region
                .write_data(data_region_id, RegionRequest::Put(request))
                .await?;
        }

        Ok(total_affected_rows)
    }

    fn split_batch_by_wal_policy(
        requests: Vec<(RegionId, RegionPutRequest)>,
    ) -> Vec<Vec<(RegionId, RegionPutRequest)>> {
        let run_count = requests
            .chunk_by(|a, b| a.1.skip_wal == b.1.skip_wal)
            .count();
        let mut batches = Vec::with_capacity(run_count);
        let mut requests = requests.into_iter();
        while let Some(first) = requests.next() {
            let remaining_run_len = requests
                .as_slice()
                .iter()
                .take_while(|(_, request)| request.skip_wal == first.1.skip_wal)
                .count();
            let mut batch = Vec::with_capacity(remaining_run_len + 1);
            batch.push(first);
            batch.extend(requests.by_ref().take(remaining_run_len));
            batches.push(batch);
        }
        batches
    }

    /// 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 and returns the physical batch's explicit partition version.
    async fn validate_batch_requests(
        &self,
        physical_region_id: RegionId,
        requests: &mut [(RegionId, RegionPutRequest)],
    ) -> Result<Option<u64>> {
        // Preserve the original physical-batch version boundary before splitting
        // by WAL policy. A conflict must fail before any data batch is written.
        let mut merged_version = None;
        for (_, request) in requests.iter() {
            if let Some(version) = request.partition_expr_version {
                ensure!(
                    merged_version.is_none_or(|merged| merged == version),
                    InvalidRequestSnafu {
                        region_id: physical_region_id,
                        reason: "inconsistent partition expr version in batch"
                    }
                );
                merged_version = Some(version);
            }
        }
        for (logical_region_id, request) in requests {
            self.verify_rows(
                *logical_region_id,
                physical_region_id,
                &mut request.rows,
                true,
            )
            .await?;
        }
        Ok(merged_version)
    }

    /// Merges one WAL-policy run and attaches the physical batch's write options.
    fn merge_batch(
        &self,
        physical_region_id: RegionId,
        encoding: PrimaryKeyEncoding,
        requests: Vec<(RegionId, RegionPutRequest)>,
        partition_expr_version: Option<u64>,
    ) -> Result<(RegionPutRequest, AffectedRows)> {
        let skip_wal = requests
            .first()
            .is_some_and(|(_, request)| request.skip_wal);
        ensure!(
            requests
                .iter()
                .all(|(_, request)| request.skip_wal == skip_wal),
            InvalidRequestSnafu {
                region_id: physical_region_id,
                reason: "inconsistent WAL policy in batch"
            }
        );
        let (rows, hint) = match encoding {
            PrimaryKeyEncoding::Sparse => (
                self.merge_sparse_batch(physical_region_id, requests)?,
                Some(WriteHint {
                    primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
                }),
            ),
            PrimaryKeyEncoding::Dense => (
                self.merge_dense_batch(to_data_region_id(physical_region_id), requests)?,
                None,
            ),
        };
        let affected_rows = rows.rows.len() as AffectedRows;
        Ok((
            RegionPutRequest {
                rows,
                hint,
                skip_wal,
                partition_expr_version,
            },
            affected_rows,
        ))
    }

    /// Merges multiple requests using sparse primary key encoding.
    fn merge_sparse_batch(
        &self,
        physical_region_id: RegionId,
        requests: Vec<(RegionId, RegionPutRequest)>,
    ) -> Result<Rows> {
        let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
        let mut modified_requests = Vec::with_capacity(requests.len());

        for (logical_region_id, mut request) in requests {
            self.modify_rows(
                physical_region_id,
                logical_region_id.table_id(),
                &mut request.rows,
                PrimaryKeyEncoding::Sparse,
            )?;

            modified_requests.push(request.rows);
        }

        let schema =
            Self::build_union_schema(modified_requests.iter().map(|rows| rows.schema.as_slice()));
        let mut merged_rows = Vec::with_capacity(total_rows);
        for rows in modified_requests {
            merged_rows.extend(Self::align_rows_to_schema(rows, &schema));
        }

        Ok(Rows {
            schema,
            rows: merged_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<Rows> {
        // Build union schema from all requests
        let merged_schema =
            Self::build_union_schema(requests.iter().map(|(_, req)| req.rows.schema.as_slice()));

        // 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 physical_columns = 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,
                },
                physical_columns,
            );

            self.row_modifier.modify_rows(
                iter,
                TableIdInput::Batch(&table_ids),
                PrimaryKeyEncoding::Dense,
            )?
        };

        Ok(final_rows)
    }

    fn build_union_schema<'a>(
        schemas: impl IntoIterator<Item = &'a [ColumnSchema]>,
    ) -> Vec<ColumnSchema> {
        let mut schema = Vec::new();
        for columns in schemas {
            for col in columns {
                if !schema
                    .iter()
                    .any(|existing: &ColumnSchema| existing.column_name == col.column_name)
                {
                    schema.push(col.clone());
                }
            }
        }
        schema
    }

    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);

        for (logical_region_id, request) in requests {
            let table_id = logical_region_id.table_id();
            let row_count = request.rows.rows.len();
            merged_rows.extend(Self::align_rows_to_schema(request.rows, merged_schema));
            table_ids.extend(std::iter::repeat_n(table_id, row_count));
        }

        (merged_rows, table_ids)
    }

    fn align_rows_to_schema(rows: Rows, merged_schema: &[ColumnSchema]) -> Vec<Row> {
        let Rows { schema, rows } = rows;
        if schema.len() == merged_schema.len()
            && schema
                .iter()
                .zip(merged_schema)
                .all(|(left, right)| left.column_name == right.column_name)
        {
            return rows;
        }

        let col_name_to_idx: FxHashMap<&str, usize> = schema
            .iter()
            .enumerate()
            .map(|(idx, col)| (col.column_name.as_str(), idx))
            .collect();
        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();
        let null_value = Value { value_data: None };

        rows.into_iter()
            .map(|mut row| {
                let values = col_mapping
                    .iter()
                    .map(|opt_idx| match opt_idx {
                        Some(idx) => std::mem::take(&mut row.values[*idx]),
                        None => null_value.clone(),
                    })
                    .collect();
                Row { values }
            })
            .collect()
    }

    /// 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,
        region_id: RegionId,
        request: RegionDeleteRequest,
    ) -> Result<AffectedRows> {
        if self.is_physical_region(region_id) {
            info!(
                "Metric region received delete request {request:?} on physical region {region_id:?}"
            );
            FORBIDDEN_OPERATION_COUNT.inc();

            UnsupportedRegionRequestSnafu {
                request: RegionRequest::Delete(request),
            }
            .fail()
        } else {
            self.delete_logical_region(region_id, request).await
        }
    }

    async fn put_logical_region(
        &self,
        logical_region_id: RegionId,
        mut request: RegionPutRequest,
    ) -> Result<AffectedRows> {
        let _timer = MITO_OPERATION_ELAPSED
            .with_label_values(&["put"])
            .start_timer();

        let (physical_region_id, data_region_id, primary_key_encoding) =
            self.find_data_region_meta(logical_region_id)?;

        self.verify_rows(
            logical_region_id,
            physical_region_id,
            &mut request.rows,
            true,
        )
        .await?;

        // write to data region
        // TODO: retrieve table name
        self.modify_rows(
            physical_region_id,
            logical_region_id.table_id(),
            &mut request.rows,
            primary_key_encoding,
        )?;
        if primary_key_encoding == PrimaryKeyEncoding::Sparse {
            request.hint = Some(WriteHint {
                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
            });
        }
        self.data_region
            .write_data(data_region_id, RegionRequest::Put(request))
            .await
    }

    async fn delete_logical_region(
        &self,
        logical_region_id: RegionId,
        mut request: RegionDeleteRequest,
    ) -> Result<AffectedRows> {
        let _timer = MITO_OPERATION_ELAPSED
            .with_label_values(&["delete"])
            .start_timer();

        let (physical_region_id, data_region_id, primary_key_encoding) =
            self.find_data_region_meta(logical_region_id)?;

        self.verify_rows(
            logical_region_id,
            physical_region_id,
            &mut request.rows,
            false,
        )
        .await?;

        // write to data region
        // TODO: retrieve table name
        self.modify_rows(
            physical_region_id,
            logical_region_id.table_id(),
            &mut request.rows,
            primary_key_encoding,
        )?;
        if primary_key_encoding == PrimaryKeyEncoding::Sparse {
            request.hint = Some(WriteHint {
                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
            });
        }
        self.data_region
            .write_data(data_region_id, RegionRequest::Delete(request))
            .await
    }

    pub(crate) fn find_data_region_meta(
        &self,
        logical_region_id: RegionId,
    ) -> Result<(RegionId, RegionId, PrimaryKeyEncoding)> {
        let state = self.state.read().unwrap();
        let physical_region_id = *state
            .logical_regions()
            .get(&logical_region_id)
            .with_context(|| LogicalRegionNotFoundSnafu {
                region_id: logical_region_id,
            })?;
        let data_region_id = to_data_region_id(physical_region_id);
        let primary_key_encoding = state.get_primary_key_encoding(data_region_id).context(
            PhysicalRegionNotFoundSnafu {
                region_id: data_region_id,
            },
        )?;
        Ok((physical_region_id, data_region_id, primary_key_encoding))
    }

    /// Verifies a request for a logical region against its corresponding metadata region.
    ///
    /// Includes:
    /// - Check if the logical region exists
    /// - Check if every column in the request exists in the physical region
    /// - Check each column's datatype and semantic type match the physical region's schema
    /// - Check the time index column is present
    /// - When `check_fields` is true, check every logical field column is present.
    ///   Set this to `false` for delete requests, which legitimately carry only
    ///   the primary key + timestamp.
    async fn verify_rows(
        &self,
        logical_region_id: RegionId,
        physical_region_id: RegionId,
        rows: &mut Rows,
        check_fields: bool,
    ) -> Result<()> {
        // Check if the region exists
        let data_region_id = to_data_region_id(physical_region_id);
        let (physical_columns, ts_name) = {
            let state = self.state.read().unwrap();
            if !state.is_logical_region_exist(logical_region_id) {
                error!("Trying to write to an nonexistent region {logical_region_id}");
                return LogicalRegionNotFoundSnafu {
                    region_id: logical_region_id,
                }
                .fail();
            }

            let physical_state = state
                .physical_region_states()
                .get(&data_region_id)
                .context(PhysicalRegionNotFoundSnafu {
                    region_id: data_region_id,
                })?;
            (
                physical_state.physical_columns().clone(),
                physical_state.time_index_column_name().to_string(),
            )
        };

        // Type + semantic check on every column in the request schema.
        for col in &rows.schema {
            let info = physical_columns
                .get(&col.column_name)
                .context(ColumnNotFoundSnafu {
                    name: &col.column_name,
                    region_id: logical_region_id,
                })?;

            ensure!(
                api::helper::is_column_type_value_eq(
                    col.datatype,
                    col.datatype_extension.clone(),
                    &info.column_schema.data_type
                ),
                InvalidRequestSnafu {
                    region_id: logical_region_id,
                    reason: format!(
                        "column {} expect type {:?}, given: {}({})",
                        col.column_name,
                        info.column_schema.data_type,
                        api::v1::ColumnDataType::try_from(col.datatype)
                            .map(|v| v.as_str_name())
                            .unwrap_or("Unknown"),
                        col.datatype,
                    ),
                }
            );

            ensure!(
                api::helper::is_semantic_type_eq(col.semantic_type, info.semantic_type),
                InvalidRequestSnafu {
                    region_id: logical_region_id,
                    reason: format!(
                        "column {} expect semantic type {:?}, given: {}({})",
                        col.column_name,
                        info.semantic_type,
                        api::v1::SemanticType::try_from(col.semantic_type)
                            .map(|v| v.as_str_name())
                            .unwrap_or("Unknown"),
                        col.semantic_type,
                    ),
                }
            );
        }

        ensure!(
            rows.schema.iter().any(|col| col.column_name == ts_name),
            InvalidRequestSnafu {
                region_id: logical_region_id,
                reason: format!("missing required time index column {ts_name}"),
            }
        );

        let logical_columns = self
            .load_logical_columns(physical_region_id, logical_region_id)
            .await?;
        let logical_fields = logical_columns
            .iter()
            .filter(|col| col.semantic_type == SemanticType::Field)
            .map(|col| (col.column_schema.name.as_str(), col))
            .collect::<HashMap<_, _>>();

        for col in &rows.schema {
            if api::helper::is_semantic_type_eq(col.semantic_type, SemanticType::Field) {
                ensure!(
                    logical_fields.contains_key(col.column_name.as_str()),
                    InvalidRequestSnafu {
                        region_id: logical_region_id,
                        reason: format!(
                            "field column {} does not belong to logical region {logical_region_id}",
                            col.column_name,
                        ),
                    }
                );
            }
        }

        if check_fields {
            // Sparse logical writes may omit nullable field columns. Fill them
            // before the rows are rewritten for the shared physical table.
            for (field_name, field_meta) in logical_fields {
                if !rows.schema.iter().any(|col| col.column_name == field_name) {
                    Self::fill_missing_field_column(
                        logical_region_id,
                        field_name,
                        field_meta,
                        rows,
                    )?;
                }
            }

            for (field_name, field_meta) in physical_columns
                .iter()
                .filter(|(_, col)| col.semantic_type == SemanticType::Field)
            {
                if !rows.schema.iter().any(|col| col.column_name == *field_name) {
                    Self::fill_missing_field_column(
                        logical_region_id,
                        field_name,
                        field_meta,
                        rows,
                    )?;
                }
            }
        }

        Ok(())
    }

    fn fill_missing_field_column(
        logical_region_id: RegionId,
        field_name: &str,
        field_meta: &ColumnMetadata,
        rows: &mut Rows,
    ) -> Result<()> {
        // This is only for schema columns with a concrete default, usually NULL
        // for field columns from other logical tables sharing this physical table.
        ensure!(
            !field_meta.column_schema.is_default_impure(),
            UnexpectedRequestSnafu {
                reason: format!(
                    "unexpected impure default value with region_id: {logical_region_id}, column: {field_name}, default_value: {:?}",
                    field_meta.column_schema.default_constraint(),
                ),
            }
        );

        let default_value = field_meta
            .column_schema
            .create_default()
            .context(CreateDefaultSnafu {
                region_id: logical_region_id,
                column: field_name,
            })?
            .with_context(|| InvalidRequestSnafu {
                region_id: logical_region_id,
                reason: format!("missing required field column {field_name}"),
            })?;
        let default_value = api::helper::to_grpc_value(default_value);
        let (datatype, datatype_extension) =
            ColumnDataTypeWrapper::try_from(field_meta.column_schema.data_type.clone())
                .map_err(|e| {
                    InvalidRequestSnafu {
                        region_id: logical_region_id,
                        reason: format!(
                            "no protobuf type for field column {field_name} ({:?}): {e}",
                            field_meta.column_schema.data_type
                        ),
                    }
                    .build()
                })?
                .to_parts();

        rows.schema.push(ColumnSchema {
            column_name: field_name.to_string(),
            datatype: datatype as i32,
            semantic_type: SemanticType::Field as i32,
            datatype_extension,
            options: None,
        });

        for row in &mut rows.rows {
            row.values.push(default_value.clone());
        }

        Ok(())
    }

    /// Perform metric engine specific logic to incoming rows.
    /// - Add table_id column
    /// - Generate tsid
    fn modify_rows(
        &self,
        physical_region_id: RegionId,
        table_id: TableId,
        rows: &mut Rows,
        encoding: PrimaryKeyEncoding,
    ) -> Result<()> {
        let input = std::mem::take(rows);
        let iter = {
            let state = self.state.read().unwrap();
            let physical_columns = state
                .physical_region_states()
                .get(&physical_region_id)
                .with_context(|| PhysicalRegionNotFoundSnafu {
                    region_id: physical_region_id,
                })?
                .physical_columns();
            RowsIter::new(input, physical_columns)
        };
        let output =
            self.row_modifier
                .modify_rows(iter, TableIdInput::Single(table_id), encoding)?;
        *rows = output;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use api::v1::value::ValueData;
    use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema};
    use common_error::ext::ErrorExt;
    use common_error::status_code::StatusCode;
    use common_function::utils::partition_expr_version;
    use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
    use common_recordbatch::RecordBatches;
    use datatypes::arrow::array::{Float64Array, TimestampMillisecondArray};
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
    use datatypes::value::Value as PartitionValue;
    use partition::expr::col;
    use store_api::metadata::ColumnMetadata;
    use store_api::metric_engine_consts::{
        DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, METRIC_ENGINE_NAME,
        PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING,
    };
    use store_api::path_utils::table_dir;
    use store_api::region_engine::RegionEngine;
    use store_api::region_request::{
        EnterStagingRequest, PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest,
        StagingPartitionDirective,
    };
    use store_api::storage::ScanRequest;
    use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;

    use super::*;
    use crate::engine::MetricEngine;
    use crate::test_util::{self, TestEnv};

    async fn scan_timestamp_values(engine: &MetricEngine, region_id: RegionId) -> Vec<(i64, f64)> {
        let stream = engine
            .scan_to_stream(region_id, ScanRequest::default())
            .await
            .unwrap();
        let batches = RecordBatches::try_collect(stream).await.unwrap();
        let mut rows = Vec::new();
        for batch in batches.iter() {
            let batch = batch.df_record_batch();
            let timestamp_index = batch.schema().index_of(greptime_timestamp()).unwrap();
            let value_index = batch.schema().index_of(greptime_value()).unwrap();
            let timestamps = batch
                .column(timestamp_index)
                .as_any()
                .downcast_ref::<TimestampMillisecondArray>()
                .unwrap();
            let values = batch
                .column(value_index)
                .as_any()
                .downcast_ref::<Float64Array>()
                .unwrap();
            rows.extend(
                timestamps
                    .values()
                    .iter()
                    .copied()
                    .zip(values.values().iter().copied()),
            );
        }
        rows.sort_unstable_by_key(|(timestamp, _)| *timestamp);
        rows
    }

    #[tokio::test]
    async fn test_mixed_wal_batch_partition_versions() {
        for encoding in ["sparse", "dense"] {
            let env = TestEnv::new().await;
            let physical_region_id = env.default_physical_region_id();
            let logical_region_id = env.default_logical_region_id();
            env.create_physical_region(
                physical_region_id,
                &TestEnv::default_table_dir(),
                vec![(PRIMARY_KEY_ENCODING.to_string(), encoding.to_string())],
            )
            .await;
            create_logical_region_with_tags(&env, physical_region_id, logical_region_id, &["job"])
                .await;
            let build_requests = |versions: [Option<u64>; 3]| {
                versions
                    .into_iter()
                    .zip([false, true, false])
                    .map(|(partition_expr_version, skip_wal)| {
                        (
                            logical_region_id,
                            RegionPutRequest {
                                skip_wal,
                                rows: Rows {
                                    schema: test_util::row_schema_with_tags(&["job"]),
                                    rows: test_util::build_rows(1, 1),
                                },
                                hint: None,
                                partition_expr_version,
                            },
                        )
                    })
                    .collect::<Vec<_>>()
            };
            // The first run has no version and could otherwise be written
            // before a later run reveals the conflicting explicit versions.
            let err = env
                .metric()
                .inner
                .put_regions_batch_single_physical(
                    physical_region_id,
                    build_requests([None, Some(10), Some(11)]),
                )
                .await
                .unwrap_err();
            assert!(
                err.to_string()
                    .contains("inconsistent partition expr version")
            );
            assert!(
                scan_timestamp_values(&env.metric(), logical_region_id)
                    .await
                    .is_empty()
            );

            for (versions, expected) in [
                ([None, None, None], None),
                ([None, Some(7), None], Some(7)),
                ([Some(7), None, Some(7)], Some(7)),
            ] {
                let mut requests = build_requests(versions);
                let actual = env
                    .metric()
                    .inner
                    .validate_batch_requests(physical_region_id, &mut requests)
                    .await
                    .unwrap();
                assert_eq!(actual, expected);
                for ((_, request), original_version) in requests.iter().zip(versions) {
                    assert_eq!(request.partition_expr_version, original_version);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_put_skip_wal_mixed_batch_recovery() {
        for encoding in ["sparse", "dense"] {
            // Paired runs differ only in the middle request's WAL policy.
            for middle_skip_wal in [false, true] {
                let env = TestEnv::new().await;
                let engine = env.metric();
                engine.inner.flush_task.stop().await.unwrap();
                let physical_region_id = env.default_physical_region_id();
                let logical_region_id = env.default_logical_region_id();
                env.create_physical_region(
                    physical_region_id,
                    &TestEnv::default_table_dir(),
                    vec![(PRIMARY_KEY_ENCODING.to_string(), encoding.to_string())],
                )
                .await;
                create_logical_region_with_tags(
                    &env,
                    physical_region_id,
                    logical_region_id,
                    &["job"],
                )
                .await;
                let metadata_before = engine.get_metadata(logical_region_id).await.unwrap();

                let requests = [false, middle_skip_wal, false].into_iter().enumerate().map(
                    |(index, skip_wal)| {
                        let timestamp = index as i64 + 1;
                        let value = timestamp as f64 * 10.0;
                        // Every request updates the same key at timestamp zero and
                        // also inserts a distinct key, exposing both reordering and
                        // accidental WAL-policy propagation to neighboring requests.
                        let rows = [0, timestamp]
                            .into_iter()
                            .map(|timestamp| Row {
                                values: vec![
                                    Value {
                                        value_data: Some(ValueData::TimestampMillisecondValue(
                                            timestamp,
                                        )),
                                    },
                                    Value {
                                        value_data: Some(ValueData::F64Value(value)),
                                    },
                                    Value {
                                        value_data: Some(ValueData::StringValue(
                                            "tag_0".to_string(),
                                        )),
                                    },
                                ],
                            })
                            .collect();
                        (
                            logical_region_id,
                            RegionPutRequest {
                                rows: Rows {
                                    schema: test_util::row_schema_with_tags(&["job"]),
                                    rows,
                                },
                                hint: None,
                                partition_expr_version: None,
                                skip_wal,
                            },
                        )
                    },
                );
                let affected_rows = engine.inner.put_regions_batch(requests).await.unwrap();
                assert_eq!(affected_rows, 6);
                assert_eq!(
                    scan_timestamp_values(&engine, logical_region_id).await,
                    vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)]
                );

                // Neither data nor metadata has an SST to hide missing WAL.
                for region_id in [
                    to_data_region_id(physical_region_id),
                    crate::utils::to_metadata_region_id(physical_region_id),
                ] {
                    let stat = env.mito().region_statistic(region_id).unwrap();
                    assert!(stat.memtable_size > 0);
                    assert_eq!(stat.sst_num, 0);
                }
                engine
                    .handle_request(
                        physical_region_id,
                        RegionRequest::Close(RegionCloseRequest {
                            flush_on_close: false,
                        }),
                    )
                    .await
                    .unwrap();

                // Recreate the wrapper as well, discarding its metadata cache.
                let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap();
                reopened.inner.flush_task.stop().await.unwrap();
                reopened
                    .handle_request(
                        physical_region_id,
                        RegionRequest::Open(RegionOpenRequest {
                            engine: METRIC_ENGINE_NAME.to_string(),
                            table_dir: TestEnv::default_table_dir(),
                            path_type: PathType::Bare,
                            options: [
                                (PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new()),
                                (PRIMARY_KEY_ENCODING.to_string(), encoding.to_string()),
                            ]
                            .into_iter()
                            .collect(),
                            skip_wal_replay: false,
                            checkpoint: None,
                            requirements: Default::default(),
                        }),
                    )
                    .await
                    .unwrap();
                let recovered_metadata = reopened.get_metadata(logical_region_id).await.unwrap();
                assert_eq!(
                    metadata_before.column_metadatas,
                    recovered_metadata.column_metadatas
                );
                let expected = if middle_skip_wal {
                    vec![(0, 30.0), (1, 10.0), (3, 30.0)]
                } else {
                    vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)]
                };
                assert_eq!(
                    scan_timestamp_values(&reopened, logical_region_id).await,
                    expected,
                    "encoding={encoding}, middle_skip_wal={middle_skip_wal}"
                );
            }
        }
    }

    #[test]
    fn test_split_batch_by_wal_policy_preserves_order() {
        let region_id = RegionId::new(1024, 1);
        let requests = [false, false, true, false, true, true]
            .into_iter()
            .enumerate()
            .map(|(index, skip_wal)| {
                (
                    region_id,
                    RegionPutRequest {
                        rows: Rows::default(),
                        hint: None,
                        partition_expr_version: Some(index as u64),
                        skip_wal,
                    },
                )
            })
            .collect();
        let batches = MetricEngineInner::split_batch_by_wal_policy(requests);
        let actual: Vec<_> = batches
            .iter()
            .map(|batch| {
                (
                    batch[0].1.skip_wal,
                    batch
                        .iter()
                        .map(|(_, request)| request.partition_expr_version.unwrap())
                        .collect::<Vec<_>>(),
                )
            })
            .collect();
        assert_eq!(
            actual,
            vec![
                (false, vec![0, 1]),
                (true, vec![2]),
                (false, vec![3]),
                (true, vec![4, 5])
            ]
        );
        assert!(MetricEngineInner::split_batch_by_wal_policy(vec![]).is_empty());
    }

    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"
            );
        }
    }

    fn job_partition_expr_json() -> String {
        let expr = col("job")
            .gt_eq(PartitionValue::String("job-0".into()))
            .and(col("job").lt(PartitionValue::String("job-9".into())));
        expr.as_json_str().unwrap()
    }

    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();
    }

    fn column_index(rows: &Rows, name: &str) -> usize {
        rows.schema
            .iter()
            .position(|col| col.column_name == name)
            .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 {
                        skip_wal: false,
                        rows: Rows {
                            schema: schema_1.clone(),
                            rows: rows_1,
                        },
                        hint: None,
                        partition_expr_version: None,
                    },
                ),
                (
                    logical_region_2,
                    RegionPutRequest {
                        skip_wal: false,
                        rows: Rows {
                            schema: schema_2.clone(),
                            rows: rows_2,
                        },
                        hint: None,
                        partition_expr_version: None,
                    },
                ),
            ]
        };

        let encoding = if expect_sparse {
            PrimaryKeyEncoding::Sparse
        } else {
            PrimaryKeyEncoding::Dense
        };
        let (merged_request, _) = env
            .metric()
            .inner
            .merge_batch(physical_region_id, encoding, build_requests(), None)
            .unwrap();
        if expect_sparse {
            assert_eq!(
                merged_request.hint.as_ref().unwrap().primary_key_encoding,
                PrimaryKeyEncodingProto::Sparse as i32
            );
        } else {
            assert!(merged_request.hint.is_none());
        }
        assert_merged_schema(&merged_request.rows, expect_sparse);
        assert!(!merged_request.skip_wal);

        for skip_wal in [false, true] {
            let mut requests = build_requests();
            for (_, request) in &mut requests {
                request.skip_wal = skip_wal;
            }
            let (merged, affected_rows) = env
                .metric()
                .inner
                .merge_batch(physical_region_id, encoding, requests, Some(7))
                .unwrap();
            assert_eq!(merged.skip_wal, skip_wal);
            assert_eq!(merged.partition_expr_version, Some(7));
            assert_eq!(affected_rows, 5);
        }

        let mut mixed_requests = build_requests();
        mixed_requests[1].1.skip_wal = true;
        assert!(
            env.metric()
                .inner
                .merge_batch(physical_region_id, encoding, mixed_requests, None)
                .is_err()
        );

        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);
    }

    #[test]
    fn test_sparse_batch_aligns_mixed_field_order() {
        let primary_key = PbColumnSchema {
            column_name: PRIMARY_KEY_COLUMN_NAME.to_string(),
            datatype: ColumnDataType::Binary as i32,
            semantic_type: SemanticType::Tag as _,
            datatype_extension: None,
            options: None,
        };
        let timestamp = PbColumnSchema {
            column_name: greptime_timestamp().to_string(),
            datatype: ColumnDataType::TimestampMillisecond as i32,
            semantic_type: SemanticType::Timestamp as _,
            datatype_extension: None,
            options: None,
        };
        let value = PbColumnSchema {
            column_name: greptime_value().to_string(),
            datatype: ColumnDataType::Float64 as i32,
            semantic_type: SemanticType::Field as _,
            datatype_extension: None,
            options: None,
        };
        let histogram = PbColumnSchema {
            column_name: greptime_native_histogram().to_string(),
            datatype: ColumnDataType::Struct as i32,
            semantic_type: SemanticType::Field as _,
            datatype_extension: None,
            options: None,
        };

        let sample_rows = Rows {
            schema: vec![
                primary_key.clone(),
                timestamp.clone(),
                value.clone(),
                histogram.clone(),
            ],
            rows: vec![Row {
                values: vec![
                    ValueData::BinaryValue(vec![1]).into(),
                    ValueData::TimestampMillisecondValue(0).into(),
                    ValueData::F64Value(1.0).into(),
                    Value { value_data: None },
                ],
            }],
        };
        let histogram_rows = Rows {
            schema: vec![primary_key, timestamp, histogram, value],
            rows: vec![Row {
                values: vec![
                    ValueData::BinaryValue(vec![2]).into(),
                    ValueData::TimestampMillisecondValue(0).into(),
                    ValueData::StructValue(api::v1::StructValue { items: vec![] }).into(),
                    Value { value_data: None },
                ],
            }],
        };

        let schema = MetricEngineInner::build_union_schema([
            sample_rows.schema.as_slice(),
            histogram_rows.schema.as_slice(),
        ]);
        let merged_rows = MetricEngineInner::align_rows_to_schema(sample_rows, &schema)
            .into_iter()
            .chain(MetricEngineInner::align_rows_to_schema(
                histogram_rows,
                &schema,
            ))
            .collect();
        let merged_request = Rows {
            schema,
            rows: merged_rows,
        };

        let value_idx = column_index(&merged_request, greptime_value());
        let histogram_idx = column_index(&merged_request, greptime_native_histogram());
        assert!(matches!(
            merged_request.rows[0].values[value_idx].value_data,
            Some(ValueData::F64Value(_))
        ));
        assert!(
            merged_request.rows[0].values[histogram_idx]
                .value_data
                .is_none()
        );
        assert!(
            merged_request.rows[1].values[value_idx]
                .value_data
                .is_none()
        );
        assert!(matches!(
            merged_request.rows[1].values[histogram_idx].value_data,
            Some(ValueData::StructValue(_))
        ));
    }

    #[tokio::test]
    async fn test_write_logical_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        // prepare data
        let schema = test_util::row_schema_with_tags(&["job"]);
        let rows = test_util::build_rows(1, 5);
        let request = RegionRequest::Put(RegionPutRequest {
            skip_wal: false,
            rows: Rows { schema, rows },
            hint: None,
            partition_expr_version: None,
        });

        // write data
        let logical_region_id = env.default_logical_region_id();
        let result = env
            .metric()
            .handle_request(logical_region_id, request)
            .await
            .unwrap();
        assert_eq!(result.affected_rows, 5);

        // read data from physical region
        let physical_region_id = env.default_physical_region_id();
        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();
        let expected = "\
+-------------------------+----------------+------------+---------------------+-------+
| greptime_timestamp      | greptime_value | __table_id | __tsid              | job   |
+-------------------------+----------------+------------+---------------------+-------+
| 1970-01-01T00:00:00     | 0.0            | 3          | 2955007454552897459 | tag_0 |
| 1970-01-01T00:00:00.001 | 1.0            | 3          | 2955007454552897459 | tag_0 |
| 1970-01-01T00:00:00.002 | 2.0            | 3          | 2955007454552897459 | tag_0 |
| 1970-01-01T00:00:00.003 | 3.0            | 3          | 2955007454552897459 | tag_0 |
| 1970-01-01T00:00:00.004 | 4.0            | 3          | 2955007454552897459 | tag_0 |
+-------------------------+----------------+------------+---------------------+-------+";
        assert_eq!(expected, batches.pretty_print().unwrap(), "physical region");

        // read data from logical region
        let request = ScanRequest::default();
        let stream = env
            .metric()
            .scan_to_stream(logical_region_id, request)
            .await
            .unwrap();
        let batches = RecordBatches::try_collect(stream).await.unwrap();
        let expected = "\
+-------------------------+----------------+-------+
| greptime_timestamp      | greptime_value | job   |
+-------------------------+----------------+-------+
| 1970-01-01T00:00:00     | 0.0            | tag_0 |
| 1970-01-01T00:00:00.001 | 1.0            | tag_0 |
| 1970-01-01T00:00:00.002 | 2.0            | tag_0 |
| 1970-01-01T00:00:00.003 | 3.0            | tag_0 |
| 1970-01-01T00:00:00.004 | 4.0            | tag_0 |
+-------------------------+----------------+-------+";
        assert_eq!(expected, batches.pretty_print().unwrap(), "logical region");
    }

    #[tokio::test]
    async fn test_write_logical_region_row_count() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let engine = env.metric();

        // add columns
        let logical_region_id = env.default_logical_region_id();
        let columns = &["odd", "even", "Ev_En"];
        let alter_request = test_util::alter_logical_region_add_tag_columns(123456, columns);
        engine
            .handle_request(logical_region_id, RegionRequest::Alter(alter_request))
            .await
            .unwrap();

        // prepare data
        let schema = test_util::row_schema_with_tags(columns);
        let rows = test_util::build_rows(3, 100);
        let request = RegionRequest::Put(RegionPutRequest {
            skip_wal: false,
            rows: Rows { schema, rows },
            hint: None,
            partition_expr_version: None,
        });

        // write data
        let result = engine
            .handle_request(logical_region_id, request)
            .await
            .unwrap();
        assert_eq!(100, result.affected_rows);
    }

    #[tokio::test]
    async fn test_write_physical_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let engine = env.metric();

        let physical_region_id = env.default_physical_region_id();
        let schema = test_util::row_schema_with_tags(&["abc"]);
        let rows = test_util::build_rows(1, 100);
        let request = RegionRequest::Put(RegionPutRequest {
            skip_wal: false,
            rows: Rows { schema, rows },
            hint: None,
            partition_expr_version: None,
        });

        engine
            .handle_request(physical_region_id, request)
            .await
            .unwrap_err();
    }

    #[tokio::test]
    async fn test_write_nonexist_logical_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let engine = env.metric();

        let logical_region_id = RegionId::new(175, 8345);
        let schema = test_util::row_schema_with_tags(&["def"]);
        let rows = test_util::build_rows(1, 100);
        let request = RegionRequest::Put(RegionPutRequest {
            skip_wal: false,
            rows: Rows { schema, rows },
            hint: None,
            partition_expr_version: None,
        });

        engine
            .handle_request(logical_region_id, request)
            .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 {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: rows1,
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
            (
                logical_region_2,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: rows2,
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
            (
                logical_region_3,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: rows3,
                    },
                    hint: None,
                    partition_expr_version: 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 {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: test_util::build_rows(1, 3),
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
            (
                nonexistent_region,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: test_util::build_rows(1, 2),
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
            (
                logical_region_2,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: test_util::build_rows(1, 5),
                    },
                    hint: None,
                    partition_expr_version: 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_physical_region_forbidden() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let engine = env.metric();

        let physical_region_id = env.default_physical_region_id();
        let schema = test_util::row_schema_with_tags(&["job"]);
        let requests = vec![(
            physical_region_id,
            RegionPutRequest {
                skip_wal: false,
                rows: Rows {
                    schema,
                    rows: test_util::build_rows(1, 1),
                },
                hint: None,
                partition_expr_version: None,
            },
        )];

        let err = engine
            .inner
            .put_regions_batch(requests.into_iter())
            .await
            .unwrap_err();

        assert!(matches!(
            err,
            crate::error::Error::ForbiddenPhysicalWrite { .. }
        ));
    }

    #[tokio::test]
    async fn test_batch_write_physical_region_forbidden() {
        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_id = env.default_logical_region_id();
        let schema = test_util::row_schema_with_tags(&["job"]);
        let requests = vec![
            (
                logical_region_id,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema: schema.clone(),
                        rows: test_util::build_rows(1, 1),
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
            (
                physical_region_id,
                RegionPutRequest {
                    skip_wal: false,
                    rows: Rows {
                        schema,
                        rows: test_util::build_rows(1, 1),
                    },
                    hint: None,
                    partition_expr_version: None,
                },
            ),
        ];

        let err = engine
            .inner
            .put_regions_batch(requests.into_iter())
            .await
            .unwrap_err();

        assert!(matches!(
            err,
            crate::error::Error::ForbiddenPhysicalWrite { .. }
        ));
    }

    #[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 {
                skip_wal: false,
                rows: Rows {
                    schema,
                    rows: test_util::build_rows(1, 5),
                },
                hint: None,
                partition_expr_version: 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![(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![(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string())],
            false,
        )
        .await;
    }

    #[tokio::test]
    async fn test_metric_put_rejects_bad_partition_expr_version() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();
        let rows = Rows {
            schema: test_util::row_schema_with_tags(&["job"]),
            rows: test_util::build_rows(1, 3),
        };

        let err = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows,
                    hint: None,
                    partition_expr_version: Some(1),
                }),
            )
            .await
            .unwrap_err();

        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
    }

    #[tokio::test]
    async fn test_metric_put_respects_staging_partition_expr_version() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();
        let physical_region_id = env.default_physical_region_id();
        let partition_expr = job_partition_expr_json();
        env.metric()
            .handle_request(
                physical_region_id,
                RegionRequest::EnterStaging(EnterStagingRequest {
                    partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
                        partition_expr.clone(),
                    ),
                }),
            )
            .await
            .unwrap();

        let expected_version = partition_expr_version(Some(&partition_expr));
        let rows = Rows {
            schema: test_util::row_schema_with_tags(&["job"]),
            rows: test_util::build_rows(1, 3),
        };

        let err = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows: rows.clone(),
                    hint: None,
                    partition_expr_version: Some(expected_version.wrapping_add(1)),
                }),
            )
            .await
            .unwrap_err();
        assert_eq!(err.status_code(), StatusCode::InvalidArguments);

        let response = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows: rows.clone(),
                    hint: None,
                    partition_expr_version: None,
                }),
            )
            .await
            .unwrap();
        assert_eq!(response.affected_rows, 3);

        let response = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows,
                    hint: None,
                    partition_expr_version: Some(expected_version),
                }),
            )
            .await
            .unwrap();
        assert_eq!(response.affected_rows, 3);
    }

    /// Regression test for issue #7990: the metric engine must reject a row
    /// whose timestamp column carries a non-timestamp datatype, rather than
    /// letting it panic inside mito's `ValueBuilder::push`.
    #[tokio::test]
    async fn test_verify_rows_rejects_wrong_type() {
        use api::v1::value::ValueData;
        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
        use common_query::prelude::{greptime_timestamp, greptime_value};

        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();

        // Timestamp column is declared as String — the very payload that
        // caused #7990. It should surface a typed error rather than panic.
        let schema = vec![
            PbColumnSchema {
                column_name: greptime_timestamp().to_string(),
                datatype: ColumnDataType::String as i32,
                semantic_type: SemanticType::Timestamp as _,
                datatype_extension: None,
                options: None,
            },
            PbColumnSchema {
                column_name: greptime_value().to_string(),
                datatype: ColumnDataType::Float64 as i32,
                semantic_type: SemanticType::Field as _,
                datatype_extension: None,
                options: None,
            },
            PbColumnSchema {
                column_name: "job".to_string(),
                datatype: ColumnDataType::String as i32,
                semantic_type: SemanticType::Tag as _,
                datatype_extension: None,
                options: None,
            },
        ];
        let rows = vec![Row {
            values: vec![
                Value {
                    value_data: Some(ValueData::StringValue("not-a-timestamp".to_string())),
                },
                Value {
                    value_data: Some(ValueData::F64Value(1.0)),
                },
                Value {
                    value_data: Some(ValueData::StringValue("tag_0".to_string())),
                },
            ],
        }];

        let err = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows: Rows { schema, rows },
                    hint: None,
                    partition_expr_version: None,
                }),
            )
            .await
            .unwrap_err();
        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
    }

    /// The completeness check must reject requests that omit the time index
    /// column, since mito cannot default-fill a `TimeIndex` column and would
    /// previously panic on the empty builder.
    #[tokio::test]
    async fn test_verify_rows_rejects_missing_time_index() {
        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
        use common_query::prelude::greptime_value;

        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();

        // Payload only carries the field and a tag — no timestamp column.
        let schema = vec![
            PbColumnSchema {
                column_name: greptime_value().to_string(),
                datatype: ColumnDataType::Float64 as i32,
                semantic_type: SemanticType::Field as _,
                datatype_extension: None,
                options: None,
            },
            PbColumnSchema {
                column_name: "job".to_string(),
                datatype: ColumnDataType::String as i32,
                semantic_type: SemanticType::Tag as _,
                datatype_extension: None,
                options: None,
            },
        ];
        let rows = vec![Row {
            values: vec![
                Value {
                    value_data: Some(api::v1::value::ValueData::F64Value(1.0)),
                },
                Value {
                    value_data: Some(api::v1::value::ValueData::StringValue("tag_0".to_string())),
                },
            ],
        }];

        let err = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows: Rows { schema, rows },
                    hint: None,
                    partition_expr_version: None,
                }),
            )
            .await
            .unwrap_err();
        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
    }

    #[tokio::test]
    async fn test_verify_rows_rejects_missing_field() {
        use api::v1::value::ValueData;
        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
        use common_query::prelude::greptime_timestamp;

        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();

        // Schema has timestamp + tag but no field column.
        let schema = vec![
            PbColumnSchema {
                column_name: greptime_timestamp().to_string(),
                datatype: ColumnDataType::TimestampMillisecond as i32,
                semantic_type: SemanticType::Timestamp as _,
                datatype_extension: None,
                options: None,
            },
            PbColumnSchema {
                column_name: "job".to_string(),
                datatype: ColumnDataType::String as i32,
                semantic_type: SemanticType::Tag as _,
                datatype_extension: None,
                options: None,
            },
        ];
        let rows = vec![Row {
            values: vec![
                Value {
                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
                },
                Value {
                    value_data: Some(ValueData::StringValue("tag_0".to_string())),
                },
            ],
        }];

        let err = env
            .metric()
            .handle_request(
                logical_region_id,
                RegionRequest::Put(RegionPutRequest {
                    skip_wal: false,
                    rows: Rows { schema, rows },
                    hint: None,
                    partition_expr_version: None,
                }),
            )
            .await
            .unwrap_err();
        let message = err.to_string();
        assert!(
            message.contains("missing required field column"),
            "expected field-completeness rejection, got: {message}"
        );
        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
    }

    #[test]
    fn test_fill_missing_field_column_nullable_no_default() {
        let field_meta = ColumnMetadata {
            column_id: 1,
            semantic_type: SemanticType::Field,
            column_schema: ColumnSchema::new(
                "greptime_value".to_string(),
                ConcreteDataType::float64_datatype(),
                true, // nullable, no default
            ),
        };
        let mut rows = Rows {
            schema: vec![PbColumnSchema {
                column_name: "ts".to_string(),
                datatype: ColumnDataType::TimestampMillisecond as i32,
                semantic_type: SemanticType::Timestamp as _,
                datatype_extension: None,
                options: None,
            }],
            rows: vec![Row {
                values: vec![Value {
                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
                }],
            }],
        };

        MetricEngineInner::fill_missing_field_column(
            RegionId::new(1, 1),
            "greptime_value",
            &field_meta,
            &mut rows,
        )
        .unwrap();

        assert_eq!(rows.schema.len(), 2);
        assert_eq!(rows.schema[1].column_name, "greptime_value");
        assert_eq!(rows.rows[0].values.len(), 2);
        assert!(
            rows.rows[0].values[1].value_data.is_none(),
            "missing nullable field should be filled with null"
        );
    }

    #[test]
    fn test_fill_missing_field_column_rejects_impure_default() {
        let field_meta = ColumnMetadata {
            column_id: 1,
            semantic_type: SemanticType::Field,
            column_schema: ColumnSchema::new(
                "greptime_value".to_string(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_default_constraint(Some(ColumnDefaultConstraint::Function("now()".to_string())))
            .unwrap(),
        };
        let mut rows = Rows {
            schema: vec![PbColumnSchema {
                column_name: "ts".to_string(),
                datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
                semantic_type: SemanticType::Timestamp as _,
                datatype_extension: None,
                options: None,
            }],
            rows: vec![Row {
                values: vec![Value {
                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
                }],
            }],
        };

        let err = MetricEngineInner::fill_missing_field_column(
            RegionId::new(1, 1),
            "greptime_value",
            &field_meta,
            &mut rows,
        )
        .unwrap_err();
        assert!(
            err.to_string().contains("impure default value"),
            "expected impure-default rejection, got: {err}"
        );
    }
}