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

metric_engine/engine/
bulk_insert.rs
