metric_engine/engine/
bulk_insert.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::HashSet;
16
17use api::v1::{ArrowIpc, SemanticType};
18use bytes::Bytes;
19use common_grpc::flight::{FlightEncoder, FlightMessage};
20use datatypes::arrow::record_batch::RecordBatch;
21use snafu::{OptionExt, ensure};
22use store_api::codec::PrimaryKeyEncoding;
23use store_api::metadata::RegionMetadataRef;
24use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest};
25use store_api::storage::RegionId;
26
27use crate::batch_modifier::{TagColumnInfo, modify_batch_sparse};
28use crate::engine::MetricEngineInner;
29use crate::error;
30use crate::error::Result;
31use crate::metrics::MITO_OPERATION_ELAPSED;
32
33impl MetricEngineInner {
34    /// Bulk-inserts rows into a metric region.
35    ///
36    /// **Logical region path:** The request payload is a logical `RecordBatch`
37    /// (timestamp, value and tag columns). It is transformed to physical format
38    /// via `modify_batch_sparse`, encoded to Arrow IPC, and forwarded as a
39    /// `BulkInserts` request to the data region.
40    ///
41    /// **Physical region path:** The request payload is already in physical format
42    /// (produced by the batcher's `flush_batch_physical`). It is forwarded directly
43    /// to the data region with no transformation.
44    ///
45    /// Returns the number of affected rows, or `0` if the input batch is empty.
46    pub async fn bulk_insert_region(
47        &self,
48        region_id: RegionId,
49        request: RegionBulkInsertsRequest,
50    ) -> Result<AffectedRows> {
51        if request.payload.num_rows() == 0 {
52            return Ok(0);
53        }
54        if self.is_physical_region(region_id) {
55            let _timer = MITO_OPERATION_ELAPSED
56                .with_label_values(&["bulk_insert_physical"])
57                .start_timer();
58            return self.bulk_insert_physical_region(region_id, request).await;
59        }
60
61        let _timer = MITO_OPERATION_ELAPSED
62            .with_label_values(&["bulk_insert_logical"])
63            .start_timer();
64        self.bulk_insert_logical_region(region_id, request).await
65    }
66
67    /// Passthrough for bulk inserts targeting a physical data region.
68    ///
69    /// The batch is already in physical format (with `__primary_key`, timestamp,
70    /// value columns), so no logical-to-physical transformation is needed.
71    async fn bulk_insert_physical_region(
72        &self,
73        region_id: RegionId,
74        mut request: RegionBulkInsertsRequest,
75    ) -> Result<AffectedRows> {
76        // Other logical tables can add fields to the shared physical schema.
77        // Let Mito fill fields absent from this batch before writing it.
78        request.aligned_schema_version = None;
79        self.data_region
80            .write_data(region_id, RegionRequest::BulkInserts(request))
81            .await
82    }
83
84    /// Bulk-inserts logical rows, transforming them to physical format first.
85    async fn bulk_insert_logical_region(
86        &self,
87        region_id: RegionId,
88        request: RegionBulkInsertsRequest,
89    ) -> Result<AffectedRows> {
90        let (physical_region_id, data_region_id, primary_key_encoding) =
91            self.find_data_region_meta(region_id)?;
92
93        if primary_key_encoding != PrimaryKeyEncoding::Sparse {
94            return error::UnsupportedRegionRequestSnafu {
95                request: RegionRequest::BulkInserts(request),
96            }
97            .fail();
98        }
99
100        let batch = request.payload;
101        if batch.num_rows() == 0 {
102            return Ok(0);
103        }
104
105        let logical_metadata = self
106            .logical_region_metadata(physical_region_id, region_id)
107            .await?;
108        let (tag_columns, non_tag_indices) = self.resolve_tag_columns_from_metadata(
109            region_id,
110            data_region_id,
111            &batch,
112            &logical_metadata,
113        )?;
114        let modified_batch = modify_batch_sparse(
115            batch.clone(),
116            region_id.table_id(),
117            &tag_columns,
118            &non_tag_indices,
119        )?;
120        let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?;
121
122        let partition_expr_version = request.partition_expr_version;
123
124        let request = RegionBulkInsertsRequest {
125            skip_wal: request.skip_wal,
126            region_id: data_region_id,
127            payload: modified_batch,
128            raw_data: ArrowIpc {
129                schema,
130                data_header,
131                payload,
132            },
133            partition_expr_version,
134            aligned_schema_version: None,
135        };
136        self.data_region
137            .write_data(data_region_id, RegionRequest::BulkInserts(request))
138            .await
139    }
140
141    fn resolve_tag_columns_from_metadata(
142        &self,
143        logical_region_id: RegionId,
144        data_region_id: RegionId,
145        batch: &RecordBatch,
146        logical_metadata: &RegionMetadataRef,
147    ) -> Result<(Vec<TagColumnInfo>, Vec<usize>)> {
148        let tag_names: HashSet<&str> = logical_metadata
149            .column_metadatas
150            .iter()
151            .filter_map(|column| {
152                if column.semantic_type == SemanticType::Tag {
153                    Some(column.column_schema.name.as_str())
154                } else {
155                    None
156                }
157            })
158            .collect();
159
160        let mut tag_columns = Vec::new();
161        let mut non_tag_indices = Vec::new();
162        {
163            let state = self.state.read().unwrap();
164            let physical_columns = state
165                .physical_region_states()
166                .get(&data_region_id)
167                .context(error::PhysicalRegionNotFoundSnafu {
168                    region_id: data_region_id,
169                })?
170                .physical_columns();
171
172            for (index, field) in batch.schema().fields().iter().enumerate() {
173                let name = field.name();
174                let column_id = physical_columns
175                    .get(name)
176                    .map(|info| info.column_id)
177                    .with_context(|| error::ColumnNotFoundSnafu {
178                        name: name.clone(),
179                        region_id: logical_region_id,
180                    })?;
181                if tag_names.contains(name.as_str()) {
182                    tag_columns.push(TagColumnInfo {
183                        name: name.clone(),
184                        index,
185                        column_id,
186                    });
187                } else {
188                    non_tag_indices.push(index);
189                }
190            }
191        }
192
193        tag_columns.sort_by(|a, b| a.name.cmp(&b.name));
194        Ok((tag_columns, non_tag_indices))
195    }
196}
197
198fn record_batch_to_ipc(record_batch: &RecordBatch) -> Result<(Bytes, Bytes, Bytes)> {
199    let mut encoder = FlightEncoder::default();
200    let schema = encoder.encode_schema(record_batch.schema().as_ref());
201    let mut iter = encoder
202        .encode(FlightMessage::RecordBatch(record_batch.clone()))
203        .into_iter();
204
205    let Some(flight_data) = iter.next() else {
206        return error::UnexpectedRequestSnafu {
207            reason: "Failed to encode empty flight data",
208        }
209        .fail();
210    };
211    ensure!(
212        iter.next().is_none(),
213        error::UnexpectedRequestSnafu {
214            reason: "Bulk insert RecordBatch with dictionary arrays is unsupported".to_string(),
215        }
216    );
217
218    Ok((
219        schema.data_header,
220        flight_data.data_header,
221        flight_data.data_body,
222    ))
223}
224
225#[cfg(test)]
226mod tests {
227    use std::assert_matches;
228    use std::sync::Arc;
229
230    use api::v1::ArrowIpc;
231    use common_error::ext::ErrorExt;
232    use common_query::prelude::{greptime_timestamp, greptime_value};
233    use common_recordbatch::RecordBatches;
234    use datafusion::parquet::basic::Encoding;
235    use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray};
236    use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit};
237    use datatypes::arrow::record_batch::RecordBatch;
238    use mito2::config::MitoConfig;
239    use mito2::sst::parquet::metadata::MetadataLoader;
240    use mito2::sst::parquet::reader::MetadataCacheMetrics;
241    use store_api::metric_engine_consts::{
242        METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING,
243    };
244    use store_api::mito_engine_options::EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING;
245    use store_api::path_utils::table_dir;
246    use store_api::region_engine::RegionEngine;
247    use store_api::region_request::{
248        PathType, RegionBulkInsertsRequest, RegionCloseRequest, RegionFlushRequest,
249        RegionOpenRequest, RegionPutRequest, RegionRequest,
250    };
251    use store_api::storage::{RegionId, ScanRequest};
252
253    use super::record_batch_to_ipc;
254    use crate::batch_modifier::{TagColumnInfo, modify_batch_sparse};
255    use crate::engine::MetricEngine;
256    use crate::error::Error;
257    use crate::test_util::{self, TestEnv};
258
259    fn build_logical_batch(start: usize, rows: usize) -> RecordBatch {
260        let schema = Arc::new(ArrowSchema::new(vec![
261            Field::new(
262                greptime_timestamp(),
263                DataType::Timestamp(TimeUnit::Millisecond, None),
264                false,
265            ),
266            Field::new(greptime_value(), DataType::Float64, true),
267            Field::new("job", DataType::Utf8, true),
268        ]));
269
270        let mut ts = Vec::with_capacity(rows);
271        let mut values = Vec::with_capacity(rows);
272        let mut tags = Vec::with_capacity(rows);
273        for i in start..start + rows {
274            ts.push(i as i64);
275            values.push(i as f64);
276            tags.push("tag_0".to_string());
277        }
278
279        RecordBatch::try_new(
280            schema,
281            vec![
282                Arc::new(TimestampMillisecondArray::from(ts)),
283                Arc::new(Float64Array::from(values)),
284                Arc::new(StringArray::from(tags)),
285            ],
286        )
287        .unwrap()
288    }
289
290    fn build_bulk_request(
291        logical_region_id: RegionId,
292        batch: RecordBatch,
293        skip_wal: bool,
294    ) -> RegionRequest {
295        let (schema, data_header, payload) = record_batch_to_ipc(&batch).unwrap();
296        RegionRequest::BulkInserts(RegionBulkInsertsRequest {
297            skip_wal,
298            region_id: logical_region_id,
299            payload: batch,
300            raw_data: ArrowIpc {
301                schema,
302                data_header,
303                payload,
304            },
305            partition_expr_version: None,
306            aligned_schema_version: None,
307        })
308    }
309
310    async fn init_dense_metric_region(env: &TestEnv) -> RegionId {
311        let physical_region_id = env.default_physical_region_id();
312        env.create_physical_region(
313            physical_region_id,
314            &TestEnv::default_table_dir(),
315            vec![(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string())],
316        )
317        .await;
318
319        let logical_region_id = env.default_logical_region_id();
320        let request = test_util::create_logical_region_request(
321            &["job"],
322            physical_region_id,
323            &table_dir("test", logical_region_id.table_id()),
324        );
325        env.metric()
326            .handle_request(logical_region_id, RegionRequest::Create(request))
327            .await
328            .unwrap();
329        logical_region_id
330    }
331
332    #[tokio::test]
333    async fn test_bulk_insert_empty_batch_returns_zero() {
334        let env = TestEnv::new().await;
335        env.init_metric_region().await;
336        let logical_region_id = env.default_logical_region_id();
337
338        let batch = build_logical_batch(0, 0);
339        let request = RegionRequest::BulkInserts(RegionBulkInsertsRequest {
340            skip_wal: false,
341            region_id: logical_region_id,
342            payload: batch,
343            raw_data: ArrowIpc::default(),
344            partition_expr_version: None,
345            aligned_schema_version: None,
346        });
347        let response = env
348            .metric()
349            .handle_request(logical_region_id, request)
350            .await
351            .unwrap();
352        assert_eq!(response.affected_rows, 0);
353    }
354
355    #[tokio::test]
356    async fn test_bulk_insert_physical_region_passthrough() {
357        check_bulk_insert_physical_region_passthrough(false, false).await;
358        check_bulk_insert_physical_region_passthrough(true, false).await;
359    }
360
361    #[tokio::test]
362    async fn test_bulk_insert_physical_region_close_recovery() {
363        check_bulk_insert_physical_region_passthrough(false, true).await;
364        check_bulk_insert_physical_region_passthrough(true, true).await;
365    }
366
367    async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) {
368        // Use flat format so that BulkMemtable is used (supports write_bulk).
369        let mito_config = MitoConfig {
370            default_flat_format: true,
371            ..Default::default()
372        };
373        let mut env = TestEnv::with_mito_config("", mito_config, Default::default()).await;
374        env.init_metric_region().await;
375        env.metric().inner.flush_task.stop().await.unwrap();
376        let physical_region_id = env.default_physical_region_id();
377        let logical_region_id = env.default_logical_region_id();
378
379        // First, do a normal logical bulk insert so we can compare results.
380        let logical_batch = build_logical_batch(0, 3);
381        let logical_request = build_bulk_request(logical_region_id, logical_batch, skip_wal);
382        let response = env
383            .metric()
384            .handle_request(logical_region_id, logical_request)
385            .await
386            .unwrap();
387        assert_eq!(response.affected_rows, 3);
388
389        // Now build a physical-format batch using modify_batch_sparse (simulating
390        // what the batcher's flush_batch_physical does) and send it directly to
391        // the physical region.
392        let tag_columns = vec![TagColumnInfo {
393            name: "job".to_string(),
394            index: 2,
395            column_id: 2, // column_id for "job" in the physical table
396        }];
397        let non_tag_indices = vec![0, 1]; // timestamp, value
398        let second_batch = build_logical_batch(3, 3);
399        let physical_batch = modify_batch_sparse(
400            second_batch,
401            logical_region_id.table_id(),
402            &tag_columns,
403            &non_tag_indices,
404        )
405        .unwrap();
406        let request = build_bulk_request(physical_region_id, physical_batch, skip_wal);
407        let response = env
408            .metric()
409            .handle_request(physical_region_id, request)
410            .await
411            .unwrap();
412        assert_eq!(response.affected_rows, 3);
413
414        // Verify all 6 rows are readable from the logical region.
415        let stream = env
416            .metric()
417            .scan_to_stream(logical_region_id, ScanRequest::default())
418            .await
419            .unwrap();
420        let batches = RecordBatches::try_collect(stream).await.unwrap();
421        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 6);
422
423        let stat = env.mito().region_statistic(physical_region_id).unwrap();
424        assert_eq!(stat.sst_num, 0);
425        let reopened = if close {
426            env.metric()
427                .handle_request(
428                    physical_region_id,
429                    RegionRequest::Close(RegionCloseRequest::default()),
430                )
431                .await
432                .unwrap();
433            // Recreating the wrapper prevents its metadata cache from hiding data loss.
434            MetricEngine::try_new(env.mito(), Default::default()).unwrap()
435        } else {
436            // Stopping workers without Close exercises recovery from WAL alone.
437            env.reopen_engine(Default::default()).await;
438            env.metric()
439        };
440        reopened.inner.flush_task.stop().await.unwrap();
441        reopened
442            .handle_request(
443                physical_region_id,
444                RegionRequest::Open(RegionOpenRequest {
445                    engine: METRIC_ENGINE_NAME.to_string(),
446                    table_dir: TestEnv::default_table_dir(),
447                    path_type: PathType::Bare,
448                    options: [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
449                        .into_iter()
450                        .collect(),
451                    skip_wal_replay: false,
452                    checkpoint: None,
453                    requirements: Default::default(),
454                }),
455            )
456            .await
457            .unwrap();
458        let stat = env
459            .mito()
460            .region_statistic(crate::utils::to_data_region_id(physical_region_id))
461            .unwrap();
462        assert_eq!(stat.sst_num > 0, skip_wal && close);
463        let stream = reopened
464            .scan_to_stream(logical_region_id, ScanRequest::default())
465            .await
466            .unwrap();
467        let batches = RecordBatches::try_collect(stream).await.unwrap();
468        assert_eq!(
469            batches.iter().map(|b| b.num_rows()).sum::<usize>(),
470            if skip_wal && !close { 0 } else { 6 },
471        );
472    }
473
474    #[tokio::test]
475    async fn test_bulk_insert_physical_region_empty_batch() {
476        // Use flat format so that BulkMemtable is used (supports write_bulk).
477        let mito_config = MitoConfig {
478            default_flat_format: true,
479            ..Default::default()
480        };
481        let env = TestEnv::with_mito_config("", mito_config, Default::default()).await;
482        env.init_metric_region().await;
483        let physical_region_id = env.default_physical_region_id();
484
485        let batch = build_logical_batch(0, 0);
486        let request = build_bulk_request(physical_region_id, batch, false);
487        let response = env
488            .metric()
489            .handle_request(physical_region_id, request)
490            .await
491            .unwrap();
492        assert_eq!(response.affected_rows, 0);
493    }
494
495    #[tokio::test]
496    async fn test_bulk_insert_unknown_column_errors() {
497        let env = TestEnv::new().await;
498        env.init_metric_region().await;
499        let logical_region_id = env.default_logical_region_id();
500
501        let schema = Arc::new(ArrowSchema::new(vec![
502            Field::new(
503                greptime_timestamp(),
504                DataType::Timestamp(TimeUnit::Millisecond, None),
505                false,
506            ),
507            Field::new(greptime_value(), DataType::Float64, true),
508            Field::new("nonexistent_column", DataType::Utf8, true),
509        ]));
510        let batch = RecordBatch::try_new(
511            schema,
512            vec![
513                Arc::new(TimestampMillisecondArray::from(vec![0i64])),
514                Arc::new(Float64Array::from(vec![1.0])),
515                Arc::new(StringArray::from(vec!["val"])),
516            ],
517        )
518        .unwrap();
519
520        let request = build_bulk_request(logical_region_id, batch, false);
521        let err = env
522            .metric()
523            .handle_request(logical_region_id, request)
524            .await
525            .unwrap_err();
526        let Some(err) = err.as_any().downcast_ref::<Error>() else {
527            panic!("unexpected error type");
528        };
529        assert_matches!(err, Error::ColumnNotFound { .. });
530    }
531
532    #[tokio::test]
533    async fn test_bulk_insert_multiple_tag_columns() {
534        let env = TestEnv::new().await;
535        let physical_region_id = env.default_physical_region_id();
536        env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), vec![])
537            .await;
538        let logical_region_id = env.default_logical_region_id();
539        let request = test_util::create_logical_region_request(
540            &["host", "region"],
541            physical_region_id,
542            &table_dir("test", logical_region_id.table_id()),
543        );
544        env.metric()
545            .handle_request(logical_region_id, RegionRequest::Create(request))
546            .await
547            .unwrap();
548
549        let schema = Arc::new(ArrowSchema::new(vec![
550            Field::new(
551                greptime_timestamp(),
552                DataType::Timestamp(TimeUnit::Millisecond, None),
553                false,
554            ),
555            Field::new(greptime_value(), DataType::Float64, true),
556            Field::new("host", DataType::Utf8, true),
557            Field::new("region", DataType::Utf8, true),
558        ]));
559        let batch = RecordBatch::try_new(
560            schema,
561            vec![
562                Arc::new(TimestampMillisecondArray::from(vec![0i64, 1, 2])),
563                Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])),
564                Arc::new(StringArray::from(vec!["h1", "h2", "h1"])),
565                Arc::new(StringArray::from(vec!["us-east", "us-west", "eu-west"])),
566            ],
567        )
568        .unwrap();
569
570        let request = build_bulk_request(logical_region_id, batch, false);
571        let response = env
572            .metric()
573            .handle_request(logical_region_id, request)
574            .await
575            .unwrap();
576        assert_eq!(response.affected_rows, 3);
577
578        let stream = env
579            .metric()
580            .scan_to_stream(logical_region_id, ScanRequest::default())
581            .await
582            .unwrap();
583        let batches = RecordBatches::try_collect(stream).await.unwrap();
584        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
585    }
586
587    #[tokio::test]
588    async fn test_bulk_insert_accumulates_rows() {
589        let env = TestEnv::new().await;
590        env.init_metric_region().await;
591        let logical_region_id = env.default_logical_region_id();
592
593        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 3), false);
594        let response = env
595            .metric()
596            .handle_request(logical_region_id, request)
597            .await
598            .unwrap();
599        assert_eq!(response.affected_rows, 3);
600
601        let request = build_bulk_request(logical_region_id, build_logical_batch(3, 5), false);
602        let response = env
603            .metric()
604            .handle_request(logical_region_id, request)
605            .await
606            .unwrap();
607        assert_eq!(response.affected_rows, 5);
608
609        let stream = env
610            .metric()
611            .scan_to_stream(logical_region_id, ScanRequest::default())
612            .await
613            .unwrap();
614        let batches = RecordBatches::try_collect(stream).await.unwrap();
615        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 8);
616    }
617
618    #[tokio::test]
619    async fn test_bulk_insert_sparse_encoding() {
620        let env = TestEnv::new().await;
621        let physical_region_id = env.default_physical_region_id();
622        env.create_physical_region(
623            physical_region_id,
624            &TestEnv::default_table_dir(),
625            vec![(
626                EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING.to_string(),
627                "byte_stream_split".to_string(),
628            )],
629        )
630        .await;
631        let logical_region_id = env.default_logical_region_id();
632        env.create_logical_region(physical_region_id, logical_region_id)
633            .await;
634
635        let rows = 4;
636        let request = build_bulk_request(logical_region_id, build_logical_batch(0, rows), false);
637        let response = env
638            .metric()
639            .handle_request(logical_region_id, request)
640            .await
641            .unwrap();
642        assert_eq!(response.affected_rows, rows);
643
644        let data_region_id = crate::utils::to_data_region_id(physical_region_id);
645        env.mito()
646            .handle_request(
647                data_region_id,
648                RegionRequest::Flush(RegionFlushRequest::default()),
649            )
650            .await
651            .unwrap();
652
653        let region = env.mito().find_region(data_region_id).unwrap();
654        let entry = region
655            .manifest_sst_entries()
656            .await
657            .into_iter()
658            .find(|entry| entry.visible && entry.file_path.ends_with(".parquet"))
659            .unwrap();
660        let mut cache_metrics = MetadataCacheMetrics::default();
661        let footer = MetadataLoader::new(
662            region.access_layer().object_store().clone(),
663            &entry.file_path,
664            entry.file_size,
665        )
666        .load(&mut cache_metrics)
667        .await
668        .unwrap();
669        let field_column_index = footer
670            .file_metadata()
671            .schema_descr()
672            .columns()
673            .iter()
674            .position(|column| column.name() == greptime_value())
675            .unwrap();
676        assert!(!footer.row_groups().is_empty());
677        assert!(footer.row_groups().iter().all(|row_group| {
678            row_group
679                .column(field_column_index)
680                .encodings()
681                .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT)
682        }));
683
684        let stream = env
685            .metric()
686            .scan_to_stream(logical_region_id, ScanRequest::default())
687            .await
688            .unwrap();
689        let batches = RecordBatches::try_collect(stream).await.unwrap();
690        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), rows);
691    }
692
693    #[tokio::test]
694    async fn test_bulk_insert_dense_encoding_rejected() {
695        let env = TestEnv::new().await;
696        let logical_region_id = init_dense_metric_region(&env).await;
697
698        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 2), false);
699        let err = env
700            .metric()
701            .handle_request(logical_region_id, request)
702            .await
703            .unwrap_err();
704        let Some(err) = err.as_any().downcast_ref::<Error>() else {
705            panic!("unexpected error type");
706        };
707        assert_matches!(err, Error::UnsupportedRegionRequest { .. });
708    }
709
710    #[tokio::test]
711    async fn test_bulk_insert_matches_put() {
712        let env_put = TestEnv::new().await;
713        env_put.init_metric_region().await;
714        let logical_region_id = env_put.default_logical_region_id();
715        let schema = test_util::row_schema_with_tags(&["job"]);
716        let rows = test_util::build_rows(1, 5);
717        env_put
718            .metric()
719            .handle_request(
720                logical_region_id,
721                RegionRequest::Put(RegionPutRequest {
722                    skip_wal: false,
723                    rows: api::v1::Rows { schema, rows },
724                    hint: None,
725                    partition_expr_version: None,
726                }),
727            )
728            .await
729            .unwrap();
730        let put_stream = env_put
731            .metric()
732            .scan_to_stream(logical_region_id, ScanRequest::default())
733            .await
734            .unwrap();
735        let put_batches = RecordBatches::try_collect(put_stream).await.unwrap();
736        let put_output = put_batches.pretty_print().unwrap();
737
738        let env_bulk = TestEnv::new().await;
739        env_bulk.init_metric_region().await;
740        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 5), false);
741        env_bulk
742            .metric()
743            .handle_request(logical_region_id, request)
744            .await
745            .unwrap();
746        let bulk_stream = env_bulk
747            .metric()
748            .scan_to_stream(logical_region_id, ScanRequest::default())
749            .await
750            .unwrap();
751        let bulk_batches = RecordBatches::try_collect(bulk_stream).await.unwrap();
752        let bulk_output = bulk_batches.pretty_print().unwrap();
753
754        assert_eq!(put_output, bulk_output);
755    }
756}
metric_engine/engine/bulk_insert.rs

metric_engine/engine/
bulk_insert.rs
