metric_engine/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)?;

        // TODO(weny): Consolidate validation and merging to avoid redundant request traversals,
        // while ensuring the entire batch is validated before writing.
        // Validate all requests
        self.validate_batch_requests(physical_region_id, &mut 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: &mut [(RegionId, RegionPutRequest)],
    ) -> Result<()> {
        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"
            }
        );
        for (logical_region_id, request) in requests {
            self.verify_rows(
                *logical_region_id,
                physical_region_id,
                &mut request.rows,
                true,
            )
            .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 skip_wal = requests
            .first()
            .is_some_and(|(_, request)| request.skip_wal);
        let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
        let mut modified_requests = Vec::with_capacity(requests.len());
        let mut total_affected_rows: AffectedRows = 0;
        let mut merged_version: Option<u64> = None;

        for (logical_region_id, mut request) in requests {
            if let Some(request_version) = request.partition_expr_version {
                if let Some(merged_version) = merged_version {
                    ensure!(
                        merged_version == request_version,
                        InvalidRequestSnafu {
                            region_id: physical_region_id,
                            reason: "inconsistent partition expr version in batch"
                        }
                    );
                } else {
                    merged_version = Some(request_version);
                }
            }
            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;
            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));
        }

        let merged_request = RegionPutRequest {
            skip_wal,
            rows: Rows {
                schema,
                rows: merged_rows,
            },
            hint: Some(WriteHint {
                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
            }),
            partition_expr_version: merged_version,
        };

        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)> {
        let skip_wal = requests
            .first()
            .is_some_and(|(_, request)| request.skip_wal);
        // 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, merged_version) =
            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,
            )?
        };

        let merged_request = RegionPutRequest {
            skip_wal,
            rows: final_rows,
            hint: None,
            partition_expr_version: merged_version,
        };

        Ok((merged_request, table_ids.len() as AffectedRows))
    }

    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],
    ) -> Result<(Vec<Row>, Vec<TableId>, Option<u64>)> {
        // 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 mut merged_version: Option<u64> = None;

        for (logical_region_id, request) in requests {
            if let Some(request_version) = request.partition_expr_version {
                if let Some(merged_version) = merged_version {
                    ensure!(
                        merged_version == request_version,
                        InvalidRequestSnafu {
                            region_id: logical_region_id,
                            reason: "inconsistent partition expr version in batch"
                        }
                    );
                } else {
                    merged_version = Some(request_version);
                }
            }
            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));
        }

        Ok((merged_rows, table_ids, merged_version))
    }

    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_snapshot(),
                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_batch_partition_versions() {
        check_batch_partition_versions("sparse").await;
        check_batch_partition_versions("dense").await;
    }

    async fn check_batch_partition_versions(encoding: &str) {
        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()
                .map(|partition_expr_version| {
                    (
                        logical_region_id,
                        RegionPutRequest {
                            skip_wal: false,
                            rows: Rows {
                                schema: test_util::row_schema_with_tags(&["job"]),
                                rows: test_util::build_rows(1, 1),
                            },
                            hint: None,
                            partition_expr_version,
                        },
                    )
                })
                .collect::<Vec<_>>()
        };
        // Conflicting explicit versions must fail before any data is written.
        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 engine = env.metric();
            engine
                .inner
                .validate_batch_requests(physical_region_id, &mut requests)
                .await
                .unwrap();
            let (merged, _) = match encoding {
                "sparse" => engine
                    .inner
                    .merge_sparse_batch(physical_region_id, requests),
                "dense" => engine
                    .inner
                    .merge_dense_batch(to_data_region_id(physical_region_id), requests),
                _ => unreachable!(),
            }
            .unwrap();
            assert_eq!(merged.partition_expr_version, expected);
        }
    }

    #[tokio::test]
    async fn test_put_skip_wal_batch_recovery() {
        check_put_skip_wal_batch_recovery("sparse", false).await;
        check_put_skip_wal_batch_recovery("sparse", true).await;
        check_put_skip_wal_batch_recovery("dense", false).await;
        check_put_skip_wal_batch_recovery("dense", true).await;
    }

    async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool) {
        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 = [skip_wal; 3]
            .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 to verify merge order.
                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 skip_wal {
            vec![]
        } 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}, skip_wal={skip_wal}"
        );
    }

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

    fn check_batch_merge_wal_policy(
        env: &TestEnv,
        physical_region_id: RegionId,
        mut requests: Vec<(RegionId, RegionPutRequest)>,
        expect_sparse: bool,
        skip_wal: bool,
    ) {
        for (_, request) in &mut requests {
            request.skip_wal = skip_wal;
        }
        let (merged_request, affected_rows) = if expect_sparse {
            let (merged_request, affected_rows) = env
                .metric()
                .inner
                .merge_sparse_batch(physical_region_id, 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, affected_rows)
        } else {
            let (merged_request, affected_rows) = env
                .metric()
                .inner
                .merge_dense_batch(to_data_region_id(physical_region_id), requests)
                .unwrap();
            assert!(merged_request.hint.is_none());
            (merged_request, affected_rows)
        };
        assert_merged_schema(&merged_request.rows, expect_sparse);
        assert_eq!(merged_request.skip_wal, skip_wal);
        assert_eq!(affected_rows, 5);
    }

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

        check_batch_merge_wal_policy(
            env,
            physical_region_id,
            build_requests(),
            expect_sparse,
            false,
        );
        check_batch_merge_wal_policy(
            env,
            physical_region_id,
            build_requests(),
            expect_sparse,
            true,
        );

        for policies in [[false, true], [true, false]] {
            let mut mixed_requests = build_requests();
            for ((_, request), skip_wal) in mixed_requests.iter_mut().zip(policies) {
                request.skip_wal = skip_wal;
            }
            let err = env
                .metric()
                .inner
                .put_regions_batch(mixed_requests.into_iter())
                .await
                .unwrap_err();
            assert!(err.to_string().contains("inconsistent WAL policy in batch"));
            for logical_region_id in [logical_region_1, logical_region_2] {
                assert!(
                    scan_timestamp_values(&env.metric(), logical_region_id)
                        .await
                        .is_empty()
                );
            }
        }

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