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, false).await;
358        check_bulk_insert_physical_region_passthrough(true, false, 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, false).await;
364        check_bulk_insert_physical_region_passthrough(true, true, false).await;
365    }
366
367    #[tokio::test]
368    async fn test_physical_bulk_partition_columns_preserve_stored_rows() {
369        for skip_wal in [false, true] {
370            let stripped =
371                check_bulk_insert_physical_region_passthrough(skip_wal, true, false).await;
372            let retained =
373                check_bulk_insert_physical_region_passthrough(skip_wal, true, true).await;
374            assert_eq!(retained, stripped, "skip_wal={skip_wal}");
375        }
376    }
377
378    async fn check_bulk_insert_physical_region_passthrough(
379        skip_wal: bool,
380        close: bool,
381        retain_partition_columns: bool,
382    ) -> String {
383        // Use flat format so that BulkMemtable is used (supports write_bulk).
384        let mito_config = MitoConfig {
385            default_flat_format: true,
386            ..Default::default()
387        };
388        let mut env = TestEnv::with_mito_config("", mito_config, Default::default()).await;
389        env.init_metric_region().await;
390        env.metric().inner.flush_task.stop().await.unwrap();
391        let physical_region_id = env.default_physical_region_id();
392        let logical_region_id = env.default_logical_region_id();
393
394        // First, do a normal logical bulk insert so we can compare results.
395        let logical_batch = build_logical_batch(0, 3);
396        let logical_request = build_bulk_request(logical_region_id, logical_batch, skip_wal);
397        let response = env
398            .metric()
399            .handle_request(logical_region_id, logical_request)
400            .await
401            .unwrap();
402        assert_eq!(response.affected_rows, 3);
403
404        // Now build a physical-format batch using modify_batch_sparse (simulating
405        // what the batcher's flush_batch_physical does) and send it directly to
406        // the physical region.
407        let tag_columns = vec![TagColumnInfo {
408            name: "job".to_string(),
409            index: 2,
410            column_id: 2, // column_id for "job" in the physical table
411        }];
412        let non_tag_indices = if retain_partition_columns {
413            vec![0, 1, 2] // timestamp, value, and the job partition column
414        } else {
415            vec![0, 1] // timestamp, value
416        };
417        let second_batch = build_logical_batch(3, 3);
418        let physical_batch = modify_batch_sparse(
419            second_batch,
420            logical_region_id.table_id(),
421            &tag_columns,
422            &non_tag_indices,
423        )
424        .unwrap();
425        assert_eq!(
426            physical_batch.column_by_name("job").is_some(),
427            retain_partition_columns
428        );
429        let request = build_bulk_request(physical_region_id, physical_batch, skip_wal);
430        let response = env
431            .metric()
432            .handle_request(physical_region_id, request)
433            .await
434            .unwrap();
435        assert_eq!(response.affected_rows, 3);
436
437        // Verify all 6 rows are readable from the logical region.
438        let stream = env
439            .metric()
440            .scan_to_stream(logical_region_id, ScanRequest::default())
441            .await
442            .unwrap();
443        let batches = RecordBatches::try_collect(stream).await.unwrap();
444        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 6);
445
446        let stat = env.mito().region_statistic(physical_region_id).unwrap();
447        assert_eq!(stat.sst_num, 0);
448        let reopened = if close {
449            env.metric()
450                .handle_request(
451                    physical_region_id,
452                    RegionRequest::Close(RegionCloseRequest::default()),
453                )
454                .await
455                .unwrap();
456            // Recreating the wrapper prevents its metadata cache from hiding data loss.
457            MetricEngine::try_new(env.mito(), Default::default()).unwrap()
458        } else {
459            // Stopping workers without Close exercises recovery from WAL alone.
460            env.reopen_engine(Default::default()).await;
461            env.metric()
462        };
463        reopened.inner.flush_task.stop().await.unwrap();
464        reopened
465            .handle_request(
466                physical_region_id,
467                RegionRequest::Open(RegionOpenRequest {
468                    engine: METRIC_ENGINE_NAME.to_string(),
469                    table_dir: TestEnv::default_table_dir(),
470                    path_type: PathType::Bare,
471                    options: [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
472                        .into_iter()
473                        .collect(),
474                    skip_wal_replay: false,
475                    checkpoint: None,
476                    requirements: Default::default(),
477                }),
478            )
479            .await
480            .unwrap();
481        let stat = env
482            .mito()
483            .region_statistic(crate::utils::to_data_region_id(physical_region_id))
484            .unwrap();
485        assert_eq!(stat.sst_num > 0, skip_wal && close);
486        let stream = reopened
487            .scan_to_stream(logical_region_id, ScanRequest::default())
488            .await
489            .unwrap();
490        let batches = RecordBatches::try_collect(stream).await.unwrap();
491        assert_eq!(
492            batches.iter().map(|b| b.num_rows()).sum::<usize>(),
493            if skip_wal && !close { 0 } else { 6 },
494        );
495        batches.pretty_print().unwrap()
496    }
497
498    #[tokio::test]
499    async fn test_bulk_insert_physical_region_empty_batch() {
500        // Use flat format so that BulkMemtable is used (supports write_bulk).
501        let mito_config = MitoConfig {
502            default_flat_format: true,
503            ..Default::default()
504        };
505        let env = TestEnv::with_mito_config("", mito_config, Default::default()).await;
506        env.init_metric_region().await;
507        let physical_region_id = env.default_physical_region_id();
508
509        let batch = build_logical_batch(0, 0);
510        let request = build_bulk_request(physical_region_id, batch, false);
511        let response = env
512            .metric()
513            .handle_request(physical_region_id, request)
514            .await
515            .unwrap();
516        assert_eq!(response.affected_rows, 0);
517    }
518
519    #[tokio::test]
520    async fn test_bulk_insert_unknown_column_errors() {
521        let env = TestEnv::new().await;
522        env.init_metric_region().await;
523        let logical_region_id = env.default_logical_region_id();
524
525        let schema = Arc::new(ArrowSchema::new(vec![
526            Field::new(
527                greptime_timestamp(),
528                DataType::Timestamp(TimeUnit::Millisecond, None),
529                false,
530            ),
531            Field::new(greptime_value(), DataType::Float64, true),
532            Field::new("nonexistent_column", DataType::Utf8, true),
533        ]));
534        let batch = RecordBatch::try_new(
535            schema,
536            vec![
537                Arc::new(TimestampMillisecondArray::from(vec![0i64])),
538                Arc::new(Float64Array::from(vec![1.0])),
539                Arc::new(StringArray::from(vec!["val"])),
540            ],
541        )
542        .unwrap();
543
544        let request = build_bulk_request(logical_region_id, batch, false);
545        let err = env
546            .metric()
547            .handle_request(logical_region_id, request)
548            .await
549            .unwrap_err();
550        let Some(err) = err.as_any().downcast_ref::<Error>() else {
551            panic!("unexpected error type");
552        };
553        assert_matches!(err, Error::ColumnNotFound { .. });
554    }
555
556    #[tokio::test]
557    async fn test_bulk_insert_multiple_tag_columns() {
558        let env = TestEnv::new().await;
559        let physical_region_id = env.default_physical_region_id();
560        env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), vec![])
561            .await;
562        let logical_region_id = env.default_logical_region_id();
563        let request = test_util::create_logical_region_request(
564            &["host", "region"],
565            physical_region_id,
566            &table_dir("test", logical_region_id.table_id()),
567        );
568        env.metric()
569            .handle_request(logical_region_id, RegionRequest::Create(request))
570            .await
571            .unwrap();
572
573        let schema = Arc::new(ArrowSchema::new(vec![
574            Field::new(
575                greptime_timestamp(),
576                DataType::Timestamp(TimeUnit::Millisecond, None),
577                false,
578            ),
579            Field::new(greptime_value(), DataType::Float64, true),
580            Field::new("host", DataType::Utf8, true),
581            Field::new("region", DataType::Utf8, true),
582        ]));
583        let batch = RecordBatch::try_new(
584            schema,
585            vec![
586                Arc::new(TimestampMillisecondArray::from(vec![0i64, 1, 2])),
587                Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])),
588                Arc::new(StringArray::from(vec!["h1", "h2", "h1"])),
589                Arc::new(StringArray::from(vec!["us-east", "us-west", "eu-west"])),
590            ],
591        )
592        .unwrap();
593
594        let request = build_bulk_request(logical_region_id, batch, false);
595        let response = env
596            .metric()
597            .handle_request(logical_region_id, request)
598            .await
599            .unwrap();
600        assert_eq!(response.affected_rows, 3);
601
602        let stream = env
603            .metric()
604            .scan_to_stream(logical_region_id, ScanRequest::default())
605            .await
606            .unwrap();
607        let batches = RecordBatches::try_collect(stream).await.unwrap();
608        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
609    }
610
611    #[tokio::test]
612    async fn test_bulk_insert_accumulates_rows() {
613        let env = TestEnv::new().await;
614        env.init_metric_region().await;
615        let logical_region_id = env.default_logical_region_id();
616
617        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 3), false);
618        let response = env
619            .metric()
620            .handle_request(logical_region_id, request)
621            .await
622            .unwrap();
623        assert_eq!(response.affected_rows, 3);
624
625        let request = build_bulk_request(logical_region_id, build_logical_batch(3, 5), false);
626        let response = env
627            .metric()
628            .handle_request(logical_region_id, request)
629            .await
630            .unwrap();
631        assert_eq!(response.affected_rows, 5);
632
633        let stream = env
634            .metric()
635            .scan_to_stream(logical_region_id, ScanRequest::default())
636            .await
637            .unwrap();
638        let batches = RecordBatches::try_collect(stream).await.unwrap();
639        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 8);
640    }
641
642    #[tokio::test]
643    async fn test_bulk_insert_sparse_encoding() {
644        let env = TestEnv::new().await;
645        let physical_region_id = env.default_physical_region_id();
646        env.create_physical_region(
647            physical_region_id,
648            &TestEnv::default_table_dir(),
649            vec![(
650                EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING.to_string(),
651                "byte_stream_split".to_string(),
652            )],
653        )
654        .await;
655        let logical_region_id = env.default_logical_region_id();
656        env.create_logical_region(physical_region_id, logical_region_id)
657            .await;
658
659        let rows = 4;
660        let request = build_bulk_request(logical_region_id, build_logical_batch(0, rows), false);
661        let response = env
662            .metric()
663            .handle_request(logical_region_id, request)
664            .await
665            .unwrap();
666        assert_eq!(response.affected_rows, rows);
667
668        let data_region_id = crate::utils::to_data_region_id(physical_region_id);
669        env.mito()
670            .handle_request(
671                data_region_id,
672                RegionRequest::Flush(RegionFlushRequest::default()),
673            )
674            .await
675            .unwrap();
676
677        let region = env.mito().find_region(data_region_id).unwrap();
678        let entry = region
679            .manifest_sst_entries()
680            .await
681            .into_iter()
682            .find(|entry| entry.visible && entry.file_path.ends_with(".parquet"))
683            .unwrap();
684        let mut cache_metrics = MetadataCacheMetrics::default();
685        let footer = MetadataLoader::new(
686            region.access_layer().object_store().clone(),
687            &entry.file_path,
688            entry.file_size,
689        )
690        .load(&mut cache_metrics)
691        .await
692        .unwrap();
693        let field_column_index = footer
694            .file_metadata()
695            .schema_descr()
696            .columns()
697            .iter()
698            .position(|column| column.name() == greptime_value())
699            .unwrap();
700        assert!(!footer.row_groups().is_empty());
701        assert!(footer.row_groups().iter().all(|row_group| {
702            row_group
703                .column(field_column_index)
704                .encodings()
705                .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT)
706        }));
707
708        let stream = env
709            .metric()
710            .scan_to_stream(logical_region_id, ScanRequest::default())
711            .await
712            .unwrap();
713        let batches = RecordBatches::try_collect(stream).await.unwrap();
714        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), rows);
715    }
716
717    #[tokio::test]
718    async fn test_bulk_insert_dense_encoding_rejected() {
719        let env = TestEnv::new().await;
720        let logical_region_id = init_dense_metric_region(&env).await;
721
722        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 2), false);
723        let err = env
724            .metric()
725            .handle_request(logical_region_id, request)
726            .await
727            .unwrap_err();
728        let Some(err) = err.as_any().downcast_ref::<Error>() else {
729            panic!("unexpected error type");
730        };
731        assert_matches!(err, Error::UnsupportedRegionRequest { .. });
732    }
733
734    #[tokio::test]
735    async fn test_bulk_insert_matches_put() {
736        let env_put = TestEnv::new().await;
737        env_put.init_metric_region().await;
738        let logical_region_id = env_put.default_logical_region_id();
739        let schema = test_util::row_schema_with_tags(&["job"]);
740        let rows = test_util::build_rows(1, 5);
741        env_put
742            .metric()
743            .handle_request(
744                logical_region_id,
745                RegionRequest::Put(RegionPutRequest {
746                    skip_wal: false,
747                    rows: api::v1::Rows { schema, rows },
748                    hint: None,
749                    partition_expr_version: None,
750                }),
751            )
752            .await
753            .unwrap();
754        let put_stream = env_put
755            .metric()
756            .scan_to_stream(logical_region_id, ScanRequest::default())
757            .await
758            .unwrap();
759        let put_batches = RecordBatches::try_collect(put_stream).await.unwrap();
760        let put_output = put_batches.pretty_print().unwrap();
761
762        let env_bulk = TestEnv::new().await;
763        env_bulk.init_metric_region().await;
764        let request = build_bulk_request(logical_region_id, build_logical_batch(0, 5), false);
765        env_bulk
766            .metric()
767            .handle_request(logical_region_id, request)
768            .await
769            .unwrap();
770        let bulk_stream = env_bulk
771            .metric()
772            .scan_to_stream(logical_region_id, ScanRequest::default())
773            .await
774            .unwrap();
775        let bulk_batches = RecordBatches::try_collect(bulk_stream).await.unwrap();
776        let bulk_output = bulk_batches.pretty_print().unwrap();
777
778        assert_eq!(put_output, bulk_output);
779    }
780}
metric_engine/engine/bulk_insert.rs

metric_engine/engine/
bulk_insert.rs
