Skip to main content

metric_engine/engine/
put.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16
17use api::helper::ColumnDataTypeWrapper;
18use api::v1::{
19    ColumnSchema, PrimaryKeyEncoding as PrimaryKeyEncodingProto, Row, Rows, SemanticType, Value,
20    WriteHint,
21};
22use common_telemetry::{error, info};
23use fxhash::FxHashMap;
24use snafu::{OptionExt, ResultExt, ensure};
25use store_api::codec::PrimaryKeyEncoding;
26use store_api::metadata::ColumnMetadata;
27use store_api::region_request::{
28    AffectedRows, RegionDeleteRequest, RegionPutRequest, RegionRequest,
29};
30use store_api::storage::{RegionId, TableId};
31
32use crate::engine::MetricEngineInner;
33use crate::error::{
34    ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu,
35    LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu,
36    UnsupportedRegionRequestSnafu,
37};
38use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED};
39use crate::row_modifier::{RowsIter, TableIdInput};
40use crate::utils::to_data_region_id;
41
42impl MetricEngineInner {
43    /// Dispatch region put request
44    pub async fn put_region(
45        &self,
46        region_id: RegionId,
47        request: RegionPutRequest,
48    ) -> Result<AffectedRows> {
49        let is_putting_physical_region =
50            self.state.read().unwrap().exist_physical_region(region_id);
51
52        if is_putting_physical_region {
53            info!(
54                "Metric region received put request {request:?} on physical region {region_id:?}"
55            );
56            FORBIDDEN_OPERATION_COUNT.inc();
57
58            ForbiddenPhysicalWriteSnafu.fail()
59        } else {
60            self.put_logical_region(region_id, request).await
61        }
62    }
63
64    /// Batch write multiple logical regions to the same physical region.
65    ///
66    /// Dispatch region put requests in batch.
67    ///
68    /// Requests may span multiple physical regions. We group them by physical
69    /// region and write sequentially. This method fails fast on validation or
70    /// preparation errors within a group and stops at the first failure.
71    /// Writes in earlier physical-region groups are not rolled back if a later
72    /// group fails.
73    pub async fn put_regions_batch(
74        &self,
75        requests: impl ExactSizeIterator<Item = (RegionId, RegionPutRequest)>,
76    ) -> Result<AffectedRows> {
77        let len = requests.len();
78
79        if len == 0 {
80            return Ok(0);
81        }
82
83        let _timer = MITO_OPERATION_ELAPSED
84            .with_label_values(&["put_batch"])
85            .start_timer();
86
87        // Fast path: single request, no batching overhead
88        if len == 1 {
89            let (region_id, req) = requests.into_iter().next().unwrap();
90            let is_putting_physical_region =
91                self.state.read().unwrap().exist_physical_region(region_id);
92            if is_putting_physical_region {
93                FORBIDDEN_OPERATION_COUNT.inc();
94                return ForbiddenPhysicalWriteSnafu.fail();
95            }
96
97            return self.put_logical_region(region_id, req).await;
98        }
99
100        let mut requests_per_physical: HashMap<RegionId, Vec<(RegionId, RegionPutRequest)>> =
101            HashMap::new();
102        for (region_id, request) in requests {
103            let is_putting_physical_region =
104                self.state.read().unwrap().exist_physical_region(region_id);
105            if is_putting_physical_region {
106                FORBIDDEN_OPERATION_COUNT.inc();
107                return ForbiddenPhysicalWriteSnafu.fail();
108            }
109            let physical_region_id = self.find_physical_region_id(region_id)?;
110            requests_per_physical
111                .entry(physical_region_id)
112                .or_default()
113                .push((region_id, request));
114        }
115
116        let mut total_affected_rows: AffectedRows = 0;
117        for (physical_region_id, requests) in requests_per_physical {
118            let affected_rows = self
119                .put_regions_batch_single_physical(physical_region_id, requests)
120                .await?;
121            total_affected_rows += affected_rows;
122        }
123
124        Ok(total_affected_rows)
125    }
126
127    /// Write a batch of requests that all belong to the same physical region.
128    ///
129    /// This function orchestrates the batch write process:
130    /// 1. Validates all requests
131    /// 2. Merges requests according to the encoding strategy (sparse or dense)
132    /// 3. Writes the merged batch to the physical region
133    async fn put_regions_batch_single_physical(
134        &self,
135        physical_region_id: RegionId,
136        mut requests: Vec<(RegionId, RegionPutRequest)>,
137    ) -> Result<AffectedRows> {
138        if requests.is_empty() {
139            return Ok(0);
140        }
141
142        let data_region_id = to_data_region_id(physical_region_id);
143        let primary_key_encoding = self.get_primary_key_encoding(data_region_id)?;
144
145        // Validate all requests
146        self.validate_batch_requests(physical_region_id, &mut requests)
147            .await?;
148
149        // Merge requests according to encoding strategy
150        let (merged_request, total_affected_rows) = match primary_key_encoding {
151            PrimaryKeyEncoding::Sparse => self.merge_sparse_batch(physical_region_id, requests)?,
152            PrimaryKeyEncoding::Dense => self.merge_dense_batch(data_region_id, requests)?,
153        };
154
155        // Write once to the physical region
156        self.data_region
157            .write_data(data_region_id, RegionRequest::Put(merged_request))
158            .await?;
159
160        Ok(total_affected_rows)
161    }
162
163    /// Get primary key encoding for a data region.
164    fn get_primary_key_encoding(&self, data_region_id: RegionId) -> Result<PrimaryKeyEncoding> {
165        let state = self.state.read().unwrap();
166        state
167            .get_primary_key_encoding(data_region_id)
168            .context(PhysicalRegionNotFoundSnafu {
169                region_id: data_region_id,
170            })
171    }
172
173    /// Validates all requests in a batch.
174    async fn validate_batch_requests(
175        &self,
176        physical_region_id: RegionId,
177        requests: &mut [(RegionId, RegionPutRequest)],
178    ) -> Result<()> {
179        for (logical_region_id, request) in requests {
180            self.verify_rows(
181                *logical_region_id,
182                physical_region_id,
183                &mut request.rows,
184                true,
185            )
186            .await?;
187        }
188        Ok(())
189    }
190
191    /// Merges multiple requests using sparse primary key encoding.
192    fn merge_sparse_batch(
193        &self,
194        physical_region_id: RegionId,
195        requests: Vec<(RegionId, RegionPutRequest)>,
196    ) -> Result<(RegionPutRequest, AffectedRows)> {
197        let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
198        let mut modified_requests = Vec::with_capacity(requests.len());
199        let mut total_affected_rows: AffectedRows = 0;
200        let mut merged_version: Option<u64> = None;
201
202        for (logical_region_id, mut request) in requests {
203            if let Some(request_version) = request.partition_expr_version {
204                if let Some(merged_version) = merged_version {
205                    ensure!(
206                        merged_version == request_version,
207                        InvalidRequestSnafu {
208                            region_id: physical_region_id,
209                            reason: "inconsistent partition expr version in batch"
210                        }
211                    );
212                } else {
213                    merged_version = Some(request_version);
214                }
215            }
216            self.modify_rows(
217                physical_region_id,
218                logical_region_id.table_id(),
219                &mut request.rows,
220                PrimaryKeyEncoding::Sparse,
221            )?;
222
223            let row_count = request.rows.rows.len();
224            total_affected_rows += row_count as AffectedRows;
225            modified_requests.push(request.rows);
226        }
227
228        let schema =
229            Self::build_union_schema(modified_requests.iter().map(|rows| rows.schema.as_slice()));
230        let mut merged_rows = Vec::with_capacity(total_rows);
231        for rows in modified_requests {
232            merged_rows.extend(Self::align_rows_to_schema(rows, &schema));
233        }
234
235        let merged_request = RegionPutRequest {
236            rows: Rows {
237                schema,
238                rows: merged_rows,
239            },
240            hint: Some(WriteHint {
241                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
242            }),
243            partition_expr_version: merged_version,
244        };
245
246        Ok((merged_request, total_affected_rows))
247    }
248
249    /// Merges multiple requests using dense primary key encoding.
250    ///
251    /// In dense mode, different requests can have different columns.
252    /// We merge all schemas into a union schema, align each row to this schema,
253    /// then batch-modify all rows together (adding __table_id and __tsid).
254    fn merge_dense_batch(
255        &self,
256        data_region_id: RegionId,
257        requests: Vec<(RegionId, RegionPutRequest)>,
258    ) -> Result<(RegionPutRequest, AffectedRows)> {
259        // Build union schema from all requests
260        let merged_schema =
261            Self::build_union_schema(requests.iter().map(|(_, req)| req.rows.schema.as_slice()));
262
263        // Align all rows to the merged schema and collect table_ids
264        let (merged_rows, table_ids, merged_version) =
265            Self::align_requests_to_schema(requests, &merged_schema)?;
266
267        // Batch-modify all rows (add __table_id and __tsid columns)
268        let final_rows = {
269            let state = self.state.read().unwrap();
270            let physical_columns = state
271                .physical_region_states()
272                .get(&data_region_id)
273                .with_context(|| PhysicalRegionNotFoundSnafu {
274                    region_id: data_region_id,
275                })?
276                .physical_columns();
277
278            let iter = RowsIter::new(
279                Rows {
280                    schema: merged_schema,
281                    rows: merged_rows,
282                },
283                physical_columns,
284            );
285
286            self.row_modifier.modify_rows(
287                iter,
288                TableIdInput::Batch(&table_ids),
289                PrimaryKeyEncoding::Dense,
290            )?
291        };
292
293        let merged_request = RegionPutRequest {
294            rows: final_rows,
295            hint: None,
296            partition_expr_version: merged_version,
297        };
298
299        Ok((merged_request, table_ids.len() as AffectedRows))
300    }
301
302    fn build_union_schema<'a>(
303        schemas: impl IntoIterator<Item = &'a [ColumnSchema]>,
304    ) -> Vec<ColumnSchema> {
305        let mut schema = Vec::new();
306        for columns in schemas {
307            for col in columns {
308                if !schema
309                    .iter()
310                    .any(|existing: &ColumnSchema| existing.column_name == col.column_name)
311                {
312                    schema.push(col.clone());
313                }
314            }
315        }
316        schema
317    }
318
319    fn align_requests_to_schema(
320        requests: Vec<(RegionId, RegionPutRequest)>,
321        merged_schema: &[ColumnSchema],
322    ) -> Result<(Vec<Row>, Vec<TableId>, Option<u64>)> {
323        // Pre-calculate total capacity
324        let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
325        let mut merged_rows = Vec::with_capacity(total_rows);
326        let mut table_ids = Vec::with_capacity(total_rows);
327        let mut merged_version: Option<u64> = None;
328
329        for (logical_region_id, request) in requests {
330            if let Some(request_version) = request.partition_expr_version {
331                if let Some(merged_version) = merged_version {
332                    ensure!(
333                        merged_version == request_version,
334                        InvalidRequestSnafu {
335                            region_id: logical_region_id,
336                            reason: "inconsistent partition expr version in batch"
337                        }
338                    );
339                } else {
340                    merged_version = Some(request_version);
341                }
342            }
343            let table_id = logical_region_id.table_id();
344            let row_count = request.rows.rows.len();
345            merged_rows.extend(Self::align_rows_to_schema(request.rows, merged_schema));
346            table_ids.extend(std::iter::repeat_n(table_id, row_count));
347        }
348
349        Ok((merged_rows, table_ids, merged_version))
350    }
351
352    fn align_rows_to_schema(rows: Rows, merged_schema: &[ColumnSchema]) -> Vec<Row> {
353        let Rows { schema, rows } = rows;
354        if schema.len() == merged_schema.len()
355            && schema
356                .iter()
357                .zip(merged_schema)
358                .all(|(left, right)| left.column_name == right.column_name)
359        {
360            return rows;
361        }
362
363        let col_name_to_idx: FxHashMap<&str, usize> = schema
364            .iter()
365            .enumerate()
366            .map(|(idx, col)| (col.column_name.as_str(), idx))
367            .collect();
368        let col_mapping: Vec<Option<usize>> = merged_schema
369            .iter()
370            .map(|merged_col| {
371                col_name_to_idx
372                    .get(merged_col.column_name.as_str())
373                    .copied()
374            })
375            .collect();
376        let null_value = Value { value_data: None };
377
378        rows.into_iter()
379            .map(|mut row| {
380                let values = col_mapping
381                    .iter()
382                    .map(|opt_idx| match opt_idx {
383                        Some(idx) => std::mem::take(&mut row.values[*idx]),
384                        None => null_value.clone(),
385                    })
386                    .collect();
387                Row { values }
388            })
389            .collect()
390    }
391
392    /// Find the physical region id for a logical region.
393    fn find_physical_region_id(&self, logical_region_id: RegionId) -> Result<RegionId> {
394        let state = self.state.read().unwrap();
395        state
396            .logical_regions()
397            .get(&logical_region_id)
398            .copied()
399            .context(LogicalRegionNotFoundSnafu {
400                region_id: logical_region_id,
401            })
402    }
403
404    /// Dispatch region delete request
405    pub async fn delete_region(
406        &self,
407        region_id: RegionId,
408        request: RegionDeleteRequest,
409    ) -> Result<AffectedRows> {
410        if self.is_physical_region(region_id) {
411            info!(
412                "Metric region received delete request {request:?} on physical region {region_id:?}"
413            );
414            FORBIDDEN_OPERATION_COUNT.inc();
415
416            UnsupportedRegionRequestSnafu {
417                request: RegionRequest::Delete(request),
418            }
419            .fail()
420        } else {
421            self.delete_logical_region(region_id, request).await
422        }
423    }
424
425    async fn put_logical_region(
426        &self,
427        logical_region_id: RegionId,
428        mut request: RegionPutRequest,
429    ) -> Result<AffectedRows> {
430        let _timer = MITO_OPERATION_ELAPSED
431            .with_label_values(&["put"])
432            .start_timer();
433
434        let (physical_region_id, data_region_id, primary_key_encoding) =
435            self.find_data_region_meta(logical_region_id)?;
436
437        self.verify_rows(
438            logical_region_id,
439            physical_region_id,
440            &mut request.rows,
441            true,
442        )
443        .await?;
444
445        // write to data region
446        // TODO: retrieve table name
447        self.modify_rows(
448            physical_region_id,
449            logical_region_id.table_id(),
450            &mut request.rows,
451            primary_key_encoding,
452        )?;
453        if primary_key_encoding == PrimaryKeyEncoding::Sparse {
454            request.hint = Some(WriteHint {
455                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
456            });
457        }
458        self.data_region
459            .write_data(data_region_id, RegionRequest::Put(request))
460            .await
461    }
462
463    async fn delete_logical_region(
464        &self,
465        logical_region_id: RegionId,
466        mut request: RegionDeleteRequest,
467    ) -> Result<AffectedRows> {
468        let _timer = MITO_OPERATION_ELAPSED
469            .with_label_values(&["delete"])
470            .start_timer();
471
472        let (physical_region_id, data_region_id, primary_key_encoding) =
473            self.find_data_region_meta(logical_region_id)?;
474
475        self.verify_rows(
476            logical_region_id,
477            physical_region_id,
478            &mut request.rows,
479            false,
480        )
481        .await?;
482
483        // write to data region
484        // TODO: retrieve table name
485        self.modify_rows(
486            physical_region_id,
487            logical_region_id.table_id(),
488            &mut request.rows,
489            primary_key_encoding,
490        )?;
491        if primary_key_encoding == PrimaryKeyEncoding::Sparse {
492            request.hint = Some(WriteHint {
493                primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
494            });
495        }
496        self.data_region
497            .write_data(data_region_id, RegionRequest::Delete(request))
498            .await
499    }
500
501    pub(crate) fn find_data_region_meta(
502        &self,
503        logical_region_id: RegionId,
504    ) -> Result<(RegionId, RegionId, PrimaryKeyEncoding)> {
505        let state = self.state.read().unwrap();
506        let physical_region_id = *state
507            .logical_regions()
508            .get(&logical_region_id)
509            .with_context(|| LogicalRegionNotFoundSnafu {
510                region_id: logical_region_id,
511            })?;
512        let data_region_id = to_data_region_id(physical_region_id);
513        let primary_key_encoding = state.get_primary_key_encoding(data_region_id).context(
514            PhysicalRegionNotFoundSnafu {
515                region_id: data_region_id,
516            },
517        )?;
518        Ok((physical_region_id, data_region_id, primary_key_encoding))
519    }
520
521    /// Verifies a request for a logical region against its corresponding metadata region.
522    ///
523    /// Includes:
524    /// - Check if the logical region exists
525    /// - Check if every column in the request exists in the physical region
526    /// - Check each column's datatype and semantic type match the physical region's schema
527    /// - Check the time index column is present
528    /// - When `check_fields` is true, check every logical field column is present.
529    ///   Set this to `false` for delete requests, which legitimately carry only
530    ///   the primary key + timestamp.
531    async fn verify_rows(
532        &self,
533        logical_region_id: RegionId,
534        physical_region_id: RegionId,
535        rows: &mut Rows,
536        check_fields: bool,
537    ) -> Result<()> {
538        // Check if the region exists
539        let data_region_id = to_data_region_id(physical_region_id);
540        let (physical_columns, ts_name) = {
541            let state = self.state.read().unwrap();
542            if !state.is_logical_region_exist(logical_region_id) {
543                error!("Trying to write to an nonexistent region {logical_region_id}");
544                return LogicalRegionNotFoundSnafu {
545                    region_id: logical_region_id,
546                }
547                .fail();
548            }
549
550            let physical_state = state
551                .physical_region_states()
552                .get(&data_region_id)
553                .context(PhysicalRegionNotFoundSnafu {
554                    region_id: data_region_id,
555                })?;
556            (
557                physical_state.physical_columns().clone(),
558                physical_state.time_index_column_name().to_string(),
559            )
560        };
561
562        // Type + semantic check on every column in the request schema.
563        for col in &rows.schema {
564            let info = physical_columns
565                .get(&col.column_name)
566                .context(ColumnNotFoundSnafu {
567                    name: &col.column_name,
568                    region_id: logical_region_id,
569                })?;
570
571            ensure!(
572                api::helper::is_column_type_value_eq(
573                    col.datatype,
574                    col.datatype_extension.clone(),
575                    &info.column_schema.data_type
576                ),
577                InvalidRequestSnafu {
578                    region_id: logical_region_id,
579                    reason: format!(
580                        "column {} expect type {:?}, given: {}({})",
581                        col.column_name,
582                        info.column_schema.data_type,
583                        api::v1::ColumnDataType::try_from(col.datatype)
584                            .map(|v| v.as_str_name())
585                            .unwrap_or("Unknown"),
586                        col.datatype,
587                    ),
588                }
589            );
590
591            ensure!(
592                api::helper::is_semantic_type_eq(col.semantic_type, info.semantic_type),
593                InvalidRequestSnafu {
594                    region_id: logical_region_id,
595                    reason: format!(
596                        "column {} expect semantic type {:?}, given: {}({})",
597                        col.column_name,
598                        info.semantic_type,
599                        api::v1::SemanticType::try_from(col.semantic_type)
600                            .map(|v| v.as_str_name())
601                            .unwrap_or("Unknown"),
602                        col.semantic_type,
603                    ),
604                }
605            );
606        }
607
608        ensure!(
609            rows.schema.iter().any(|col| col.column_name == ts_name),
610            InvalidRequestSnafu {
611                region_id: logical_region_id,
612                reason: format!("missing required time index column {ts_name}"),
613            }
614        );
615
616        let logical_columns = self
617            .load_logical_columns(physical_region_id, logical_region_id)
618            .await?;
619        let logical_fields = logical_columns
620            .iter()
621            .filter(|col| col.semantic_type == SemanticType::Field)
622            .map(|col| (col.column_schema.name.as_str(), col))
623            .collect::<HashMap<_, _>>();
624
625        for col in &rows.schema {
626            if api::helper::is_semantic_type_eq(col.semantic_type, SemanticType::Field) {
627                ensure!(
628                    logical_fields.contains_key(col.column_name.as_str()),
629                    InvalidRequestSnafu {
630                        region_id: logical_region_id,
631                        reason: format!(
632                            "field column {} does not belong to logical region {logical_region_id}",
633                            col.column_name,
634                        ),
635                    }
636                );
637            }
638        }
639
640        if check_fields {
641            // Sparse logical writes may omit nullable field columns. Fill them
642            // before the rows are rewritten for the shared physical table.
643            for (field_name, field_meta) in logical_fields {
644                if !rows.schema.iter().any(|col| col.column_name == field_name) {
645                    Self::fill_missing_field_column(
646                        logical_region_id,
647                        field_name,
648                        field_meta,
649                        rows,
650                    )?;
651                }
652            }
653
654            for (field_name, field_meta) in physical_columns
655                .iter()
656                .filter(|(_, col)| col.semantic_type == SemanticType::Field)
657            {
658                if !rows.schema.iter().any(|col| col.column_name == *field_name) {
659                    Self::fill_missing_field_column(
660                        logical_region_id,
661                        field_name,
662                        field_meta,
663                        rows,
664                    )?;
665                }
666            }
667        }
668
669        Ok(())
670    }
671
672    fn fill_missing_field_column(
673        logical_region_id: RegionId,
674        field_name: &str,
675        field_meta: &ColumnMetadata,
676        rows: &mut Rows,
677    ) -> Result<()> {
678        // This is only for schema columns with a concrete default, usually NULL
679        // for field columns from other logical tables sharing this physical table.
680        ensure!(
681            !field_meta.column_schema.is_default_impure(),
682            UnexpectedRequestSnafu {
683                reason: format!(
684                    "unexpected impure default value with region_id: {logical_region_id}, column: {field_name}, default_value: {:?}",
685                    field_meta.column_schema.default_constraint(),
686                ),
687            }
688        );
689
690        let default_value = field_meta
691            .column_schema
692            .create_default()
693            .context(CreateDefaultSnafu {
694                region_id: logical_region_id,
695                column: field_name,
696            })?
697            .with_context(|| InvalidRequestSnafu {
698                region_id: logical_region_id,
699                reason: format!("missing required field column {field_name}"),
700            })?;
701        let default_value = api::helper::to_grpc_value(default_value);
702        let (datatype, datatype_extension) =
703            ColumnDataTypeWrapper::try_from(field_meta.column_schema.data_type.clone())
704                .map_err(|e| {
705                    InvalidRequestSnafu {
706                        region_id: logical_region_id,
707                        reason: format!(
708                            "no protobuf type for field column {field_name} ({:?}): {e}",
709                            field_meta.column_schema.data_type
710                        ),
711                    }
712                    .build()
713                })?
714                .to_parts();
715
716        rows.schema.push(ColumnSchema {
717            column_name: field_name.to_string(),
718            datatype: datatype as i32,
719            semantic_type: SemanticType::Field as i32,
720            datatype_extension,
721            options: None,
722        });
723
724        for row in &mut rows.rows {
725            row.values.push(default_value.clone());
726        }
727
728        Ok(())
729    }
730
731    /// Perform metric engine specific logic to incoming rows.
732    /// - Add table_id column
733    /// - Generate tsid
734    fn modify_rows(
735        &self,
736        physical_region_id: RegionId,
737        table_id: TableId,
738        rows: &mut Rows,
739        encoding: PrimaryKeyEncoding,
740    ) -> Result<()> {
741        let input = std::mem::take(rows);
742        let iter = {
743            let state = self.state.read().unwrap();
744            let physical_columns = state
745                .physical_region_states()
746                .get(&physical_region_id)
747                .with_context(|| PhysicalRegionNotFoundSnafu {
748                    region_id: physical_region_id,
749                })?
750                .physical_columns();
751            RowsIter::new(input, physical_columns)
752        };
753        let output =
754            self.row_modifier
755                .modify_rows(iter, TableIdInput::Single(table_id), encoding)?;
756        *rows = output;
757        Ok(())
758    }
759}
760
761#[cfg(test)]
762mod tests {
763    use std::collections::HashSet;
764
765    use api::v1::value::ValueData;
766    use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema};
767    use common_error::ext::ErrorExt;
768    use common_error::status_code::StatusCode;
769    use common_function::utils::partition_expr_version;
770    use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
771    use common_recordbatch::RecordBatches;
772    use datatypes::prelude::ConcreteDataType;
773    use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
774    use datatypes::value::Value as PartitionValue;
775    use partition::expr::col;
776    use store_api::metadata::ColumnMetadata;
777    use store_api::metric_engine_consts::{
778        DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, PRIMARY_KEY_ENCODING,
779    };
780    use store_api::path_utils::table_dir;
781    use store_api::region_engine::RegionEngine;
782    use store_api::region_request::{
783        EnterStagingRequest, RegionRequest, StagingPartitionDirective,
784    };
785    use store_api::storage::ScanRequest;
786    use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
787
788    use super::*;
789    use crate::test_util::{self, TestEnv};
790
791    fn assert_merged_schema(rows: &Rows, expect_sparse: bool) {
792        let column_names: HashSet<String> = rows
793            .schema
794            .iter()
795            .map(|col| col.column_name.clone())
796            .collect();
797
798        if expect_sparse {
799            assert!(
800                column_names.contains(PRIMARY_KEY_COLUMN_NAME),
801                "sparse encoding should include primary key column"
802            );
803            assert!(
804                !column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME),
805                "sparse encoding should not include table id column"
806            );
807            assert!(
808                !column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME),
809                "sparse encoding should not include tsid column"
810            );
811            assert!(
812                !column_names.contains("job"),
813                "sparse encoding should not include tag columns"
814            );
815            assert!(
816                !column_names.contains("instance"),
817                "sparse encoding should not include tag columns"
818            );
819        } else {
820            assert!(
821                !column_names.contains(PRIMARY_KEY_COLUMN_NAME),
822                "dense encoding should not include primary key column"
823            );
824            assert!(
825                column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME),
826                "dense encoding should include table id column"
827            );
828            assert!(
829                column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME),
830                "dense encoding should include tsid column"
831            );
832            assert!(
833                column_names.contains("job"),
834                "dense encoding should keep tag columns"
835            );
836            assert!(
837                column_names.contains("instance"),
838                "dense encoding should keep tag columns"
839            );
840        }
841    }
842
843    fn job_partition_expr_json() -> String {
844        let expr = col("job")
845            .gt_eq(PartitionValue::String("job-0".into()))
846            .and(col("job").lt(PartitionValue::String("job-9".into())));
847        expr.as_json_str().unwrap()
848    }
849
850    async fn create_logical_region_with_tags(
851        env: &TestEnv,
852        physical_region_id: RegionId,
853        logical_region_id: RegionId,
854        tags: &[&str],
855    ) {
856        let region_create_request = test_util::create_logical_region_request(
857            tags,
858            physical_region_id,
859            &table_dir("test", logical_region_id.table_id()),
860        );
861        env.metric()
862            .handle_request(
863                logical_region_id,
864                RegionRequest::Create(region_create_request),
865            )
866            .await
867            .unwrap();
868    }
869
870    fn column_index(rows: &Rows, name: &str) -> usize {
871        rows.schema
872            .iter()
873            .position(|col| col.column_name == name)
874            .unwrap()
875    }
876
877    async fn run_batch_write_with_schema_variants(
878        env: &TestEnv,
879        physical_region_id: RegionId,
880        options: Vec<(String, String)>,
881        expect_sparse: bool,
882    ) {
883        env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), options)
884            .await;
885
886        let logical_region_1 = env.default_logical_region_id();
887        let logical_region_2 = RegionId::new(1024, 1);
888
889        create_logical_region_with_tags(env, physical_region_id, logical_region_1, &["job"]).await;
890        create_logical_region_with_tags(
891            env,
892            physical_region_id,
893            logical_region_2,
894            &["job", "instance"],
895        )
896        .await;
897
898        let schema_1 = test_util::row_schema_with_tags(&["job"]);
899        let schema_2 = test_util::row_schema_with_tags(&["job", "instance"]);
900
901        let data_region_id = RegionId::new(physical_region_id.table_id(), 2);
902        let primary_key_encoding = env
903            .metric()
904            .inner
905            .get_primary_key_encoding(data_region_id)
906            .unwrap();
907        assert_eq!(
908            primary_key_encoding,
909            if expect_sparse {
910                PrimaryKeyEncoding::Sparse
911            } else {
912                PrimaryKeyEncoding::Dense
913            }
914        );
915
916        let build_requests = || {
917            let rows_1 = test_util::build_rows(1, 3);
918            let rows_2 = test_util::build_rows(2, 2);
919
920            vec![
921                (
922                    logical_region_1,
923                    RegionPutRequest {
924                        rows: Rows {
925                            schema: schema_1.clone(),
926                            rows: rows_1,
927                        },
928                        hint: None,
929                        partition_expr_version: None,
930                    },
931                ),
932                (
933                    logical_region_2,
934                    RegionPutRequest {
935                        rows: Rows {
936                            schema: schema_2.clone(),
937                            rows: rows_2,
938                        },
939                        hint: None,
940                        partition_expr_version: None,
941                    },
942                ),
943            ]
944        };
945
946        let merged_request = if expect_sparse {
947            let (merged_request, _) = env
948                .metric()
949                .inner
950                .merge_sparse_batch(physical_region_id, build_requests())
951                .unwrap();
952            let hint = merged_request
953                .hint
954                .as_ref()
955                .expect("missing sparse write hint");
956            assert_eq!(
957                hint.primary_key_encoding,
958                PrimaryKeyEncodingProto::Sparse as i32
959            );
960            merged_request
961        } else {
962            let (merged_request, _) = env
963                .metric()
964                .inner
965                .merge_dense_batch(data_region_id, build_requests())
966                .unwrap();
967            assert!(merged_request.hint.is_none());
968            merged_request
969        };
970
971        assert_merged_schema(&merged_request.rows, expect_sparse);
972
973        let affected_rows = env
974            .metric()
975            .inner
976            .put_regions_batch(build_requests().into_iter())
977            .await
978            .unwrap();
979        assert_eq!(affected_rows, 5);
980
981        let request = ScanRequest::default();
982        let stream = env
983            .mito()
984            .scan_to_stream(data_region_id, request)
985            .await
986            .unwrap();
987        let batches = RecordBatches::try_collect(stream).await.unwrap();
988
989        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 5);
990    }
991
992    #[test]
993    fn test_sparse_batch_aligns_mixed_field_order() {
994        let primary_key = PbColumnSchema {
995            column_name: PRIMARY_KEY_COLUMN_NAME.to_string(),
996            datatype: ColumnDataType::Binary as i32,
997            semantic_type: SemanticType::Tag as _,
998            datatype_extension: None,
999            options: None,
1000        };
1001        let timestamp = PbColumnSchema {
1002            column_name: greptime_timestamp().to_string(),
1003            datatype: ColumnDataType::TimestampMillisecond as i32,
1004            semantic_type: SemanticType::Timestamp as _,
1005            datatype_extension: None,
1006            options: None,
1007        };
1008        let value = PbColumnSchema {
1009            column_name: greptime_value().to_string(),
1010            datatype: ColumnDataType::Float64 as i32,
1011            semantic_type: SemanticType::Field as _,
1012            datatype_extension: None,
1013            options: None,
1014        };
1015        let histogram = PbColumnSchema {
1016            column_name: greptime_native_histogram().to_string(),
1017            datatype: ColumnDataType::Struct as i32,
1018            semantic_type: SemanticType::Field as _,
1019            datatype_extension: None,
1020            options: None,
1021        };
1022
1023        let sample_rows = Rows {
1024            schema: vec![
1025                primary_key.clone(),
1026                timestamp.clone(),
1027                value.clone(),
1028                histogram.clone(),
1029            ],
1030            rows: vec![Row {
1031                values: vec![
1032                    ValueData::BinaryValue(vec![1]).into(),
1033                    ValueData::TimestampMillisecondValue(0).into(),
1034                    ValueData::F64Value(1.0).into(),
1035                    Value { value_data: None },
1036                ],
1037            }],
1038        };
1039        let histogram_rows = Rows {
1040            schema: vec![primary_key, timestamp, histogram, value],
1041            rows: vec![Row {
1042                values: vec![
1043                    ValueData::BinaryValue(vec![2]).into(),
1044                    ValueData::TimestampMillisecondValue(0).into(),
1045                    ValueData::StructValue(api::v1::StructValue { items: vec![] }).into(),
1046                    Value { value_data: None },
1047                ],
1048            }],
1049        };
1050
1051        let schema = MetricEngineInner::build_union_schema([
1052            sample_rows.schema.as_slice(),
1053            histogram_rows.schema.as_slice(),
1054        ]);
1055        let merged_rows = MetricEngineInner::align_rows_to_schema(sample_rows, &schema)
1056            .into_iter()
1057            .chain(MetricEngineInner::align_rows_to_schema(
1058                histogram_rows,
1059                &schema,
1060            ))
1061            .collect();
1062        let merged_request = Rows {
1063            schema,
1064            rows: merged_rows,
1065        };
1066
1067        let value_idx = column_index(&merged_request, greptime_value());
1068        let histogram_idx = column_index(&merged_request, greptime_native_histogram());
1069        assert!(matches!(
1070            merged_request.rows[0].values[value_idx].value_data,
1071            Some(ValueData::F64Value(_))
1072        ));
1073        assert!(
1074            merged_request.rows[0].values[histogram_idx]
1075                .value_data
1076                .is_none()
1077        );
1078        assert!(
1079            merged_request.rows[1].values[value_idx]
1080                .value_data
1081                .is_none()
1082        );
1083        assert!(matches!(
1084            merged_request.rows[1].values[histogram_idx].value_data,
1085            Some(ValueData::StructValue(_))
1086        ));
1087    }
1088
1089    #[tokio::test]
1090    async fn test_write_logical_region() {
1091        let env = TestEnv::new().await;
1092        env.init_metric_region().await;
1093
1094        // prepare data
1095        let schema = test_util::row_schema_with_tags(&["job"]);
1096        let rows = test_util::build_rows(1, 5);
1097        let request = RegionRequest::Put(RegionPutRequest {
1098            rows: Rows { schema, rows },
1099            hint: None,
1100            partition_expr_version: None,
1101        });
1102
1103        // write data
1104        let logical_region_id = env.default_logical_region_id();
1105        let result = env
1106            .metric()
1107            .handle_request(logical_region_id, request)
1108            .await
1109            .unwrap();
1110        assert_eq!(result.affected_rows, 5);
1111
1112        // read data from physical region
1113        let physical_region_id = env.default_physical_region_id();
1114        let request = ScanRequest::default();
1115        let stream = env
1116            .metric()
1117            .scan_to_stream(physical_region_id, request)
1118            .await
1119            .unwrap();
1120        let batches = RecordBatches::try_collect(stream).await.unwrap();
1121        let expected = "\
1122+-------------------------+----------------+------------+---------------------+-------+
1123| greptime_timestamp      | greptime_value | __table_id | __tsid              | job   |
1124+-------------------------+----------------+------------+---------------------+-------+
1125| 1970-01-01T00:00:00     | 0.0            | 3          | 2955007454552897459 | tag_0 |
1126| 1970-01-01T00:00:00.001 | 1.0            | 3          | 2955007454552897459 | tag_0 |
1127| 1970-01-01T00:00:00.002 | 2.0            | 3          | 2955007454552897459 | tag_0 |
1128| 1970-01-01T00:00:00.003 | 3.0            | 3          | 2955007454552897459 | tag_0 |
1129| 1970-01-01T00:00:00.004 | 4.0            | 3          | 2955007454552897459 | tag_0 |
1130+-------------------------+----------------+------------+---------------------+-------+";
1131        assert_eq!(expected, batches.pretty_print().unwrap(), "physical region");
1132
1133        // read data from logical region
1134        let request = ScanRequest::default();
1135        let stream = env
1136            .metric()
1137            .scan_to_stream(logical_region_id, request)
1138            .await
1139            .unwrap();
1140        let batches = RecordBatches::try_collect(stream).await.unwrap();
1141        let expected = "\
1142+-------------------------+----------------+-------+
1143| greptime_timestamp      | greptime_value | job   |
1144+-------------------------+----------------+-------+
1145| 1970-01-01T00:00:00     | 0.0            | tag_0 |
1146| 1970-01-01T00:00:00.001 | 1.0            | tag_0 |
1147| 1970-01-01T00:00:00.002 | 2.0            | tag_0 |
1148| 1970-01-01T00:00:00.003 | 3.0            | tag_0 |
1149| 1970-01-01T00:00:00.004 | 4.0            | tag_0 |
1150+-------------------------+----------------+-------+";
1151        assert_eq!(expected, batches.pretty_print().unwrap(), "logical region");
1152    }
1153
1154    #[tokio::test]
1155    async fn test_write_logical_region_row_count() {
1156        let env = TestEnv::new().await;
1157        env.init_metric_region().await;
1158        let engine = env.metric();
1159
1160        // add columns
1161        let logical_region_id = env.default_logical_region_id();
1162        let columns = &["odd", "even", "Ev_En"];
1163        let alter_request = test_util::alter_logical_region_add_tag_columns(123456, columns);
1164        engine
1165            .handle_request(logical_region_id, RegionRequest::Alter(alter_request))
1166            .await
1167            .unwrap();
1168
1169        // prepare data
1170        let schema = test_util::row_schema_with_tags(columns);
1171        let rows = test_util::build_rows(3, 100);
1172        let request = RegionRequest::Put(RegionPutRequest {
1173            rows: Rows { schema, rows },
1174            hint: None,
1175            partition_expr_version: None,
1176        });
1177
1178        // write data
1179        let result = engine
1180            .handle_request(logical_region_id, request)
1181            .await
1182            .unwrap();
1183        assert_eq!(100, result.affected_rows);
1184    }
1185
1186    #[tokio::test]
1187    async fn test_write_physical_region() {
1188        let env = TestEnv::new().await;
1189        env.init_metric_region().await;
1190        let engine = env.metric();
1191
1192        let physical_region_id = env.default_physical_region_id();
1193        let schema = test_util::row_schema_with_tags(&["abc"]);
1194        let rows = test_util::build_rows(1, 100);
1195        let request = RegionRequest::Put(RegionPutRequest {
1196            rows: Rows { schema, rows },
1197            hint: None,
1198            partition_expr_version: None,
1199        });
1200
1201        engine
1202            .handle_request(physical_region_id, request)
1203            .await
1204            .unwrap_err();
1205    }
1206
1207    #[tokio::test]
1208    async fn test_write_nonexist_logical_region() {
1209        let env = TestEnv::new().await;
1210        env.init_metric_region().await;
1211        let engine = env.metric();
1212
1213        let logical_region_id = RegionId::new(175, 8345);
1214        let schema = test_util::row_schema_with_tags(&["def"]);
1215        let rows = test_util::build_rows(1, 100);
1216        let request = RegionRequest::Put(RegionPutRequest {
1217            rows: Rows { schema, rows },
1218            hint: None,
1219            partition_expr_version: None,
1220        });
1221
1222        engine
1223            .handle_request(logical_region_id, request)
1224            .await
1225            .unwrap_err();
1226    }
1227
1228    #[tokio::test]
1229    async fn test_batch_write_multiple_logical_regions() {
1230        let env = TestEnv::new().await;
1231        env.init_metric_region().await;
1232        let engine = env.metric();
1233
1234        // Create two additional logical regions
1235        let physical_region_id = env.default_physical_region_id();
1236        let logical_region_1 = env.default_logical_region_id();
1237        let logical_region_2 = RegionId::new(1024, 1);
1238        let logical_region_3 = RegionId::new(1024, 2);
1239
1240        env.create_logical_region(physical_region_id, logical_region_2)
1241            .await;
1242        env.create_logical_region(physical_region_id, logical_region_3)
1243            .await;
1244
1245        // Prepare batch requests with non-overlapping timestamps
1246        let schema = test_util::row_schema_with_tags(&["job"]);
1247
1248        // Use build_rows_with_ts to create non-overlapping timestamps
1249        // logical_region_1: ts 0, 1, 2
1250        // logical_region_2: ts 10, 11  (offset to avoid overlap)
1251        // logical_region_3: ts 20, 21, 22, 23, 24  (offset to avoid overlap)
1252        let rows1 = test_util::build_rows(1, 3);
1253        let mut rows2 = test_util::build_rows(1, 2);
1254        let mut rows3 = test_util::build_rows(1, 5);
1255
1256        // Adjust timestamps to avoid conflicts
1257        use api::v1::value::ValueData;
1258        for (i, row) in rows2.iter_mut().enumerate() {
1259            if let Some(ValueData::TimestampMillisecondValue(ts)) =
1260                row.values.get_mut(0).and_then(|v| v.value_data.as_mut())
1261            {
1262                *ts = (10 + i) as i64;
1263            }
1264        }
1265        for (i, row) in rows3.iter_mut().enumerate() {
1266            if let Some(ValueData::TimestampMillisecondValue(ts)) =
1267                row.values.get_mut(0).and_then(|v| v.value_data.as_mut())
1268            {
1269                *ts = (20 + i) as i64;
1270            }
1271        }
1272
1273        let requests = vec![
1274            (
1275                logical_region_1,
1276                RegionPutRequest {
1277                    rows: Rows {
1278                        schema: schema.clone(),
1279                        rows: rows1,
1280                    },
1281                    hint: None,
1282                    partition_expr_version: None,
1283                },
1284            ),
1285            (
1286                logical_region_2,
1287                RegionPutRequest {
1288                    rows: Rows {
1289                        schema: schema.clone(),
1290                        rows: rows2,
1291                    },
1292                    hint: None,
1293                    partition_expr_version: None,
1294                },
1295            ),
1296            (
1297                logical_region_3,
1298                RegionPutRequest {
1299                    rows: Rows {
1300                        schema: schema.clone(),
1301                        rows: rows3,
1302                    },
1303                    hint: None,
1304                    partition_expr_version: None,
1305                },
1306            ),
1307        ];
1308
1309        // Batch write
1310        let affected_rows = engine
1311            .inner
1312            .put_regions_batch(requests.into_iter())
1313            .await
1314            .unwrap();
1315        assert_eq!(affected_rows, 10);
1316
1317        // Verify physical region contains data from all logical regions
1318        let request = ScanRequest::default();
1319        let stream = env
1320            .metric()
1321            .scan_to_stream(physical_region_id, request)
1322            .await
1323            .unwrap();
1324        let batches = RecordBatches::try_collect(stream).await.unwrap();
1325
1326        // Should have 3 + 2 + 5 = 10 rows total
1327        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 10);
1328    }
1329
1330    #[tokio::test]
1331    async fn test_batch_write_with_partial_failure() {
1332        let env = TestEnv::new().await;
1333        env.init_metric_region().await;
1334        let engine = env.metric();
1335
1336        let physical_region_id = env.default_physical_region_id();
1337        let logical_region_1 = env.default_logical_region_id();
1338        let logical_region_2 = RegionId::new(1024, 1);
1339        let nonexistent_region = RegionId::new(9999, 9999);
1340
1341        env.create_logical_region(physical_region_id, logical_region_2)
1342            .await;
1343
1344        // Prepare batch with one invalid region
1345        let schema = test_util::row_schema_with_tags(&["job"]);
1346        let requests = vec![
1347            (
1348                logical_region_1,
1349                RegionPutRequest {
1350                    rows: Rows {
1351                        schema: schema.clone(),
1352                        rows: test_util::build_rows(1, 3),
1353                    },
1354                    hint: None,
1355                    partition_expr_version: None,
1356                },
1357            ),
1358            (
1359                nonexistent_region,
1360                RegionPutRequest {
1361                    rows: Rows {
1362                        schema: schema.clone(),
1363                        rows: test_util::build_rows(1, 2),
1364                    },
1365                    hint: None,
1366                    partition_expr_version: None,
1367                },
1368            ),
1369            (
1370                logical_region_2,
1371                RegionPutRequest {
1372                    rows: Rows {
1373                        schema: schema.clone(),
1374                        rows: test_util::build_rows(1, 5),
1375                    },
1376                    hint: None,
1377                    partition_expr_version: None,
1378                },
1379            ),
1380        ];
1381
1382        // Batch write
1383        let result = engine.inner.put_regions_batch(requests.into_iter()).await;
1384        assert!(result.is_err());
1385
1386        // Invalid region is detected before any write, so the physical region remains empty.
1387        // Fail-fast is per physical-region group; cross-group partial success is possible.
1388        let request = ScanRequest::default();
1389        let stream = env
1390            .metric()
1391            .scan_to_stream(physical_region_id, request)
1392            .await
1393            .unwrap();
1394        let batches = RecordBatches::try_collect(stream).await.unwrap();
1395
1396        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 0);
1397    }
1398
1399    #[tokio::test]
1400    async fn test_batch_write_single_physical_region_forbidden() {
1401        let env = TestEnv::new().await;
1402        env.init_metric_region().await;
1403        let engine = env.metric();
1404
1405        let physical_region_id = env.default_physical_region_id();
1406        let schema = test_util::row_schema_with_tags(&["job"]);
1407        let requests = vec![(
1408            physical_region_id,
1409            RegionPutRequest {
1410                rows: Rows {
1411                    schema,
1412                    rows: test_util::build_rows(1, 1),
1413                },
1414                hint: None,
1415                partition_expr_version: None,
1416            },
1417        )];
1418
1419        let err = engine
1420            .inner
1421            .put_regions_batch(requests.into_iter())
1422            .await
1423            .unwrap_err();
1424
1425        assert!(matches!(
1426            err,
1427            crate::error::Error::ForbiddenPhysicalWrite { .. }
1428        ));
1429    }
1430
1431    #[tokio::test]
1432    async fn test_batch_write_physical_region_forbidden() {
1433        let env = TestEnv::new().await;
1434        env.init_metric_region().await;
1435        let engine = env.metric();
1436
1437        let physical_region_id = env.default_physical_region_id();
1438        let logical_region_id = env.default_logical_region_id();
1439        let schema = test_util::row_schema_with_tags(&["job"]);
1440        let requests = vec![
1441            (
1442                logical_region_id,
1443                RegionPutRequest {
1444                    rows: Rows {
1445                        schema: schema.clone(),
1446                        rows: test_util::build_rows(1, 1),
1447                    },
1448                    hint: None,
1449                    partition_expr_version: None,
1450                },
1451            ),
1452            (
1453                physical_region_id,
1454                RegionPutRequest {
1455                    rows: Rows {
1456                        schema,
1457                        rows: test_util::build_rows(1, 1),
1458                    },
1459                    hint: None,
1460                    partition_expr_version: None,
1461                },
1462            ),
1463        ];
1464
1465        let err = engine
1466            .inner
1467            .put_regions_batch(requests.into_iter())
1468            .await
1469            .unwrap_err();
1470
1471        assert!(matches!(
1472            err,
1473            crate::error::Error::ForbiddenPhysicalWrite { .. }
1474        ));
1475    }
1476
1477    #[tokio::test]
1478    async fn test_batch_write_single_request_fast_path() {
1479        let env = TestEnv::new().await;
1480        env.init_metric_region().await;
1481        let engine = env.metric();
1482
1483        let logical_region_id = env.default_logical_region_id();
1484        let schema = test_util::row_schema_with_tags(&["job"]);
1485
1486        // Single request should use fast path
1487        let requests = vec![(
1488            logical_region_id,
1489            RegionPutRequest {
1490                rows: Rows {
1491                    schema,
1492                    rows: test_util::build_rows(1, 5),
1493                },
1494                hint: None,
1495                partition_expr_version: None,
1496            },
1497        )];
1498
1499        let affected_rows = engine
1500            .inner
1501            .put_regions_batch(requests.into_iter())
1502            .await
1503            .unwrap();
1504        assert_eq!(affected_rows, 5);
1505    }
1506
1507    #[tokio::test]
1508    async fn test_batch_write_empty_requests() {
1509        let env = TestEnv::new().await;
1510        env.init_metric_region().await;
1511        let engine = env.metric();
1512
1513        // Empty batch should return zero affected rows
1514        let requests = vec![];
1515        let affected_rows = engine
1516            .inner
1517            .put_regions_batch(requests.into_iter())
1518            .await
1519            .unwrap();
1520
1521        assert_eq!(affected_rows, 0);
1522    }
1523
1524    #[tokio::test]
1525    async fn test_batch_write_sparse_encoding() {
1526        let env = TestEnv::new().await;
1527        let physical_region_id = env.default_physical_region_id();
1528
1529        run_batch_write_with_schema_variants(
1530            &env,
1531            physical_region_id,
1532            vec![(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string())],
1533            true,
1534        )
1535        .await;
1536    }
1537
1538    #[tokio::test]
1539    async fn test_batch_write_dense_encoding() {
1540        let env = TestEnv::new().await;
1541        let physical_region_id = env.default_physical_region_id();
1542
1543        run_batch_write_with_schema_variants(
1544            &env,
1545            physical_region_id,
1546            vec![(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string())],
1547            false,
1548        )
1549        .await;
1550    }
1551
1552    #[tokio::test]
1553    async fn test_metric_put_rejects_bad_partition_expr_version() {
1554        let env = TestEnv::new().await;
1555        env.init_metric_region().await;
1556
1557        let logical_region_id = env.default_logical_region_id();
1558        let rows = Rows {
1559            schema: test_util::row_schema_with_tags(&["job"]),
1560            rows: test_util::build_rows(1, 3),
1561        };
1562
1563        let err = env
1564            .metric()
1565            .handle_request(
1566                logical_region_id,
1567                RegionRequest::Put(RegionPutRequest {
1568                    rows,
1569                    hint: None,
1570                    partition_expr_version: Some(1),
1571                }),
1572            )
1573            .await
1574            .unwrap_err();
1575
1576        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1577    }
1578
1579    #[tokio::test]
1580    async fn test_metric_put_respects_staging_partition_expr_version() {
1581        let env = TestEnv::new().await;
1582        env.init_metric_region().await;
1583
1584        let logical_region_id = env.default_logical_region_id();
1585        let physical_region_id = env.default_physical_region_id();
1586        let partition_expr = job_partition_expr_json();
1587        env.metric()
1588            .handle_request(
1589                physical_region_id,
1590                RegionRequest::EnterStaging(EnterStagingRequest {
1591                    partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
1592                        partition_expr.clone(),
1593                    ),
1594                }),
1595            )
1596            .await
1597            .unwrap();
1598
1599        let expected_version = partition_expr_version(Some(&partition_expr));
1600        let rows = Rows {
1601            schema: test_util::row_schema_with_tags(&["job"]),
1602            rows: test_util::build_rows(1, 3),
1603        };
1604
1605        let err = env
1606            .metric()
1607            .handle_request(
1608                logical_region_id,
1609                RegionRequest::Put(RegionPutRequest {
1610                    rows: rows.clone(),
1611                    hint: None,
1612                    partition_expr_version: Some(expected_version.wrapping_add(1)),
1613                }),
1614            )
1615            .await
1616            .unwrap_err();
1617        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1618
1619        let response = env
1620            .metric()
1621            .handle_request(
1622                logical_region_id,
1623                RegionRequest::Put(RegionPutRequest {
1624                    rows: rows.clone(),
1625                    hint: None,
1626                    partition_expr_version: None,
1627                }),
1628            )
1629            .await
1630            .unwrap();
1631        assert_eq!(response.affected_rows, 3);
1632
1633        let response = env
1634            .metric()
1635            .handle_request(
1636                logical_region_id,
1637                RegionRequest::Put(RegionPutRequest {
1638                    rows,
1639                    hint: None,
1640                    partition_expr_version: Some(expected_version),
1641                }),
1642            )
1643            .await
1644            .unwrap();
1645        assert_eq!(response.affected_rows, 3);
1646    }
1647
1648    /// Regression test for issue #7990: the metric engine must reject a row
1649    /// whose timestamp column carries a non-timestamp datatype, rather than
1650    /// letting it panic inside mito's `ValueBuilder::push`.
1651    #[tokio::test]
1652    async fn test_verify_rows_rejects_wrong_type() {
1653        use api::v1::value::ValueData;
1654        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1655        use common_query::prelude::{greptime_timestamp, greptime_value};
1656
1657        let env = TestEnv::new().await;
1658        env.init_metric_region().await;
1659
1660        let logical_region_id = env.default_logical_region_id();
1661
1662        // Timestamp column is declared as String — the very payload that
1663        // caused #7990. It should surface a typed error rather than panic.
1664        let schema = vec![
1665            PbColumnSchema {
1666                column_name: greptime_timestamp().to_string(),
1667                datatype: ColumnDataType::String as i32,
1668                semantic_type: SemanticType::Timestamp as _,
1669                datatype_extension: None,
1670                options: None,
1671            },
1672            PbColumnSchema {
1673                column_name: greptime_value().to_string(),
1674                datatype: ColumnDataType::Float64 as i32,
1675                semantic_type: SemanticType::Field as _,
1676                datatype_extension: None,
1677                options: None,
1678            },
1679            PbColumnSchema {
1680                column_name: "job".to_string(),
1681                datatype: ColumnDataType::String as i32,
1682                semantic_type: SemanticType::Tag as _,
1683                datatype_extension: None,
1684                options: None,
1685            },
1686        ];
1687        let rows = vec![Row {
1688            values: vec![
1689                Value {
1690                    value_data: Some(ValueData::StringValue("not-a-timestamp".to_string())),
1691                },
1692                Value {
1693                    value_data: Some(ValueData::F64Value(1.0)),
1694                },
1695                Value {
1696                    value_data: Some(ValueData::StringValue("tag_0".to_string())),
1697                },
1698            ],
1699        }];
1700
1701        let err = env
1702            .metric()
1703            .handle_request(
1704                logical_region_id,
1705                RegionRequest::Put(RegionPutRequest {
1706                    rows: Rows { schema, rows },
1707                    hint: None,
1708                    partition_expr_version: None,
1709                }),
1710            )
1711            .await
1712            .unwrap_err();
1713        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1714    }
1715
1716    /// The completeness check must reject requests that omit the time index
1717    /// column, since mito cannot default-fill a `TimeIndex` column and would
1718    /// previously panic on the empty builder.
1719    #[tokio::test]
1720    async fn test_verify_rows_rejects_missing_time_index() {
1721        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1722        use common_query::prelude::greptime_value;
1723
1724        let env = TestEnv::new().await;
1725        env.init_metric_region().await;
1726
1727        let logical_region_id = env.default_logical_region_id();
1728
1729        // Payload only carries the field and a tag — no timestamp column.
1730        let schema = vec![
1731            PbColumnSchema {
1732                column_name: greptime_value().to_string(),
1733                datatype: ColumnDataType::Float64 as i32,
1734                semantic_type: SemanticType::Field as _,
1735                datatype_extension: None,
1736                options: None,
1737            },
1738            PbColumnSchema {
1739                column_name: "job".to_string(),
1740                datatype: ColumnDataType::String as i32,
1741                semantic_type: SemanticType::Tag as _,
1742                datatype_extension: None,
1743                options: None,
1744            },
1745        ];
1746        let rows = vec![Row {
1747            values: vec![
1748                Value {
1749                    value_data: Some(api::v1::value::ValueData::F64Value(1.0)),
1750                },
1751                Value {
1752                    value_data: Some(api::v1::value::ValueData::StringValue("tag_0".to_string())),
1753                },
1754            ],
1755        }];
1756
1757        let err = env
1758            .metric()
1759            .handle_request(
1760                logical_region_id,
1761                RegionRequest::Put(RegionPutRequest {
1762                    rows: Rows { schema, rows },
1763                    hint: None,
1764                    partition_expr_version: None,
1765                }),
1766            )
1767            .await
1768            .unwrap_err();
1769        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1770    }
1771
1772    #[tokio::test]
1773    async fn test_verify_rows_rejects_missing_field() {
1774        use api::v1::value::ValueData;
1775        use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1776        use common_query::prelude::greptime_timestamp;
1777
1778        let env = TestEnv::new().await;
1779        env.init_metric_region().await;
1780
1781        let logical_region_id = env.default_logical_region_id();
1782
1783        // Schema has timestamp + tag but no field column.
1784        let schema = vec![
1785            PbColumnSchema {
1786                column_name: greptime_timestamp().to_string(),
1787                datatype: ColumnDataType::TimestampMillisecond as i32,
1788                semantic_type: SemanticType::Timestamp as _,
1789                datatype_extension: None,
1790                options: None,
1791            },
1792            PbColumnSchema {
1793                column_name: "job".to_string(),
1794                datatype: ColumnDataType::String as i32,
1795                semantic_type: SemanticType::Tag as _,
1796                datatype_extension: None,
1797                options: None,
1798            },
1799        ];
1800        let rows = vec![Row {
1801            values: vec![
1802                Value {
1803                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
1804                },
1805                Value {
1806                    value_data: Some(ValueData::StringValue("tag_0".to_string())),
1807                },
1808            ],
1809        }];
1810
1811        let err = env
1812            .metric()
1813            .handle_request(
1814                logical_region_id,
1815                RegionRequest::Put(RegionPutRequest {
1816                    rows: Rows { schema, rows },
1817                    hint: None,
1818                    partition_expr_version: None,
1819                }),
1820            )
1821            .await
1822            .unwrap_err();
1823        let message = err.to_string();
1824        assert!(
1825            message.contains("missing required field column"),
1826            "expected field-completeness rejection, got: {message}"
1827        );
1828        assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1829    }
1830
1831    #[test]
1832    fn test_fill_missing_field_column_nullable_no_default() {
1833        let field_meta = ColumnMetadata {
1834            column_id: 1,
1835            semantic_type: SemanticType::Field,
1836            column_schema: ColumnSchema::new(
1837                "greptime_value".to_string(),
1838                ConcreteDataType::float64_datatype(),
1839                true, // nullable, no default
1840            ),
1841        };
1842        let mut rows = Rows {
1843            schema: vec![PbColumnSchema {
1844                column_name: "ts".to_string(),
1845                datatype: ColumnDataType::TimestampMillisecond as i32,
1846                semantic_type: SemanticType::Timestamp as _,
1847                datatype_extension: None,
1848                options: None,
1849            }],
1850            rows: vec![Row {
1851                values: vec![Value {
1852                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
1853                }],
1854            }],
1855        };
1856
1857        MetricEngineInner::fill_missing_field_column(
1858            RegionId::new(1, 1),
1859            "greptime_value",
1860            &field_meta,
1861            &mut rows,
1862        )
1863        .unwrap();
1864
1865        assert_eq!(rows.schema.len(), 2);
1866        assert_eq!(rows.schema[1].column_name, "greptime_value");
1867        assert_eq!(rows.rows[0].values.len(), 2);
1868        assert!(
1869            rows.rows[0].values[1].value_data.is_none(),
1870            "missing nullable field should be filled with null"
1871        );
1872    }
1873
1874    #[test]
1875    fn test_fill_missing_field_column_rejects_impure_default() {
1876        let field_meta = ColumnMetadata {
1877            column_id: 1,
1878            semantic_type: SemanticType::Field,
1879            column_schema: ColumnSchema::new(
1880                "greptime_value".to_string(),
1881                ConcreteDataType::timestamp_millisecond_datatype(),
1882                false,
1883            )
1884            .with_default_constraint(Some(ColumnDefaultConstraint::Function("now()".to_string())))
1885            .unwrap(),
1886        };
1887        let mut rows = Rows {
1888            schema: vec![PbColumnSchema {
1889                column_name: "ts".to_string(),
1890                datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
1891                semantic_type: SemanticType::Timestamp as _,
1892                datatype_extension: None,
1893                options: None,
1894            }],
1895            rows: vec![Row {
1896                values: vec![Value {
1897                    value_data: Some(ValueData::TimestampMillisecondValue(0)),
1898                }],
1899            }],
1900        };
1901
1902        let err = MetricEngineInner::fill_missing_field_column(
1903            RegionId::new(1, 1),
1904            "greptime_value",
1905            &field_meta,
1906            &mut rows,
1907        )
1908        .unwrap_err();
1909        assert!(
1910            err.to_string().contains("impure default value"),
1911            "expected impure-default rejection, got: {err}"
1912        );
1913    }
1914}