mito2/worker/
handle_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
15//! Handles bulk insert requests.
16
17use datatypes::arrow;
18use store_api::logstore::LogStore;
19use store_api::metadata::RegionMetadataRef;
20use store_api::region_request::RegionBulkInsertsRequest;
21
22use crate::error::InconsistentTimestampLengthSnafu;
23use crate::memtable::bulk::part::BulkPart;
24use crate::request::{OptionOutputTx, SenderBulkRequest};
25use crate::worker::RegionWorkerLoop;
26use crate::{error, metrics};
27
28impl<S: LogStore> RegionWorkerLoop<S> {
29    pub(crate) async fn handle_bulk_insert_batch(
30        &mut self,
31        region_metadata: RegionMetadataRef,
32        request: RegionBulkInsertsRequest,
33        pending_bulk_request: &mut Vec<SenderBulkRequest>,
34        sender: OptionOutputTx,
35    ) {
36        let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED
37            .with_label_values(&["process_bulk_req"])
38            .start_timer();
39        let batch = request.payload;
40        if batch.num_rows() == 0 {
41            sender.send(Ok(0));
42            return;
43        }
44
45        let Some((ts_index, ts)) = batch
46            .schema()
47            .column_with_name(&region_metadata.time_index_column().column_schema.name)
48            .map(|(index, _)| (index, batch.column(index)))
49        else {
50            sender.send(
51                error::InvalidRequestSnafu {
52                    region_id: region_metadata.region_id,
53                    reason: format!(
54                        "timestamp column `{}` not found",
55                        region_metadata.time_index_column().column_schema.name
56                    ),
57                }
58                .fail(),
59            );
60            return;
61        };
62
63        if batch.num_rows() != ts.len() {
64            sender.send(
65                InconsistentTimestampLengthSnafu {
66                    expected: batch.num_rows(),
67                    actual: ts.len(),
68                }
69                .fail(),
70            );
71            return;
72        }
73
74        // safety: `WorkerRequest::try_from_region_request` checked the batch against
75        // `region_metadata`, so the time index is a timestamp array without nulls.
76        let (ts_primitive, _) = datatypes::timestamp::timestamp_array_to_primitive(ts).unwrap();
77
78        // safety: we've checked ts.len() == batch.num_rows(), batch is not empty and
79        // the time index has no nulls.
80        let min_ts = arrow::compute::min(&ts_primitive).unwrap();
81        let max_ts = arrow::compute::max(&ts_primitive).unwrap();
82
83        let part = BulkPart {
84            batch,
85            max_timestamp: max_ts,
86            min_timestamp: min_ts,
87            sequence: 0,
88            min_sequence: 0,
89            timestamp_index: ts_index,
90            raw_data: Some(request.raw_data),
91        };
92
93        let aligned_region_metadata = if request
94            .aligned_schema_version
95            .is_some_and(|aligned| aligned == region_metadata.schema_version)
96        {
97            Some(region_metadata)
98        } else {
99            None
100        };
101        pending_bulk_request.push(SenderBulkRequest {
102            skip_wal: request.skip_wal,
103            sender,
104            request: part,
105            region_id: request.region_id,
106            region_metadata: aligned_region_metadata,
107            partition_expr_version: request.partition_expr_version,
108        });
109    }
110}
mito2/worker/handle_bulk_insert.rs

mito2/worker/
handle_bulk_insert.rs
