operator/
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 ahash::{HashMap, HashMapExt};
18use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests};
19use api::v1::region::{
20    BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
21};
22use api::v1::{ArrowIpc, PartitionExprVersion};
23use arrow::array::Array;
24use arrow::record_batch::RecordBatch;
25use bytes::Bytes;
26use common_base::AffectedRows;
27use common_grpc::FlightData;
28use common_grpc::flight::{FlightEncoder, FlightMessage};
29use common_telemetry::error;
30use common_telemetry::tracing_context::TracingContext;
31use snafu::{OptionExt, ResultExt, ensure};
32use store_api::storage::RegionId;
33use table::TableRef;
34use table::metadata::TableInfoRef;
35
36use crate::insert::Inserter;
37use crate::{error, metrics};
38
39impl Inserter {
40    /// Handle bulk insert request.
41    pub async fn handle_bulk_insert(
42        &self,
43        table: TableRef,
44        raw_flight_data: FlightData,
45        record_batch: RecordBatch,
46        schema_bytes: Bytes,
47        skip_wal: bool,
48    ) -> error::Result<AffectedRows> {
49        let table_info = table.table_info();
50        let table_id = table_info.table_id();
51        let db_name = table_info.get_db_string();
52
53        if record_batch.num_rows() == 0 {
54            return Ok(0);
55        }
56
57        let body_size = raw_flight_data.data_body.len();
58        // TODO(yingwen): Fill record batch impure default values. Note that we should override `raw_flight_data` if we have to fill defaults.
59        // notify flownode to update dirty timestamps if flow is configured.
60        self.maybe_update_flow_dirty_window(table_info.clone(), record_batch.clone());
61
62        metrics::BULK_REQUEST_MESSAGE_SIZE.observe(body_size as f64);
63        metrics::BULK_REQUEST_ROWS
64            .with_label_values(&["raw"])
65            .observe(record_batch.num_rows() as f64);
66
67        let partition_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
68            .with_label_values(&["partition"])
69            .start_timer();
70        let (partition_rule, partition_versions) = self
71            .partition_manager
72            .find_table_partition_rule(&table_info)
73            .await
74            .context(error::InvalidPartitionSnafu)?;
75
76        // find partitions for each row in the record batch
77        let region_masks = partition_rule
78            .split_record_batch(&record_batch)
79            .context(error::SplitInsertSnafu)?;
80        partition_timer.observe_duration();
81
82        // fast path: only one region.
83        if region_masks.len() == 1 {
84            metrics::BULK_REQUEST_ROWS
85                .with_label_values(&["rows_per_region"])
86                .observe(record_batch.num_rows() as f64);
87
88            // SAFETY: region masks length checked
89            let (region_number, _) = region_masks.into_iter().next().unwrap();
90            let region_id = RegionId::new(table_id, region_number);
91            let partition_expr_version = partition_versions
92                .get(&region_number)
93                .copied()
94                .unwrap_or_default();
95            let datanode = self
96                .partition_manager
97                .find_region_leader(region_id)
98                .await
99                .context(error::FindRegionLeaderSnafu)?;
100
101            let request = RegionRequest {
102                header: Some(RegionRequestHeader {
103                    tracing_context: TracingContext::from_current_span().to_w3c(),
104                    ..Default::default()
105                }),
106                body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
107                    skip_wal,
108                    region_id: region_id.as_u64(),
109                    partition_expr_version: partition_expr_version
110                        .map(|value| PartitionExprVersion { value }),
111                    aligned_schema_version: None,
112                    body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
113                        schema: schema_bytes.clone(),
114                        data_header: raw_flight_data.data_header,
115                        payload: raw_flight_data.data_body,
116                    })),
117                })),
118            };
119
120            let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
121                .with_label_values(&["datanode_handle"])
122                .start_timer();
123            let datanode = self.node_manager.datanode(&datanode).await;
124            let result = datanode
125                .handle(request)
126                .await
127                .context(error::RequestRegionSnafu)
128                .map(|r| r.affected_rows);
129            if let Ok(rows) = result {
130                crate::metrics::DIST_INGEST_ROW_COUNT
131                    .with_label_values(&[db_name.as_str()])
132                    .inc_by(rows as u64);
133            }
134            return result;
135        }
136
137        let mut mask_per_datanode = HashMap::with_capacity(region_masks.len());
138        for (region_number, mask) in region_masks {
139            let region_id = RegionId::new(table_id, region_number);
140            let datanode = self
141                .partition_manager
142                .find_region_leader(region_id)
143                .await
144                .context(error::FindRegionLeaderSnafu)?;
145            mask_per_datanode
146                .entry(datanode)
147                .or_insert_with(Vec::new)
148                .push((region_id, mask));
149        }
150
151        let wait_all_datanode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
152            .with_label_values(&["wait_all_datanode"])
153            .start_timer();
154
155        let mut handles = Vec::with_capacity(mask_per_datanode.len());
156
157        for (peer, masks) in mask_per_datanode {
158            for (region_id, mask) in masks {
159                if mask.select_none() {
160                    continue;
161                }
162                let partition_expr_version = partition_versions
163                    .get(&region_id.region_number())
164                    .copied()
165                    .unwrap_or_default();
166                let rb = record_batch.clone();
167                let schema_bytes = schema_bytes.clone();
168                let node_manager = self.node_manager.clone();
169                let peer = peer.clone();
170                let raw_header_and_data = if mask.select_all() {
171                    Some((
172                        raw_flight_data.data_header.clone(),
173                        raw_flight_data.data_body.clone(),
174                    ))
175                } else {
176                    None
177                };
178                let handle: common_runtime::JoinHandle<error::Result<api::region::RegionResponse>> =
179                    common_runtime::spawn_global(async move {
180                        let (header, payload) = if mask.select_all() {
181                            // SAFETY: raw data must be present, we can avoid re-encoding.
182                            raw_header_and_data.unwrap()
183                        } else {
184                            let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
185                                .with_label_values(&["filter"])
186                                .start_timer();
187                            let batch = arrow::compute::filter_record_batch(&rb, mask.array())
188                                .context(error::ComputeArrowSnafu)?;
189                            filter_timer.observe_duration();
190                            metrics::BULK_REQUEST_ROWS
191                                .with_label_values(&["rows_per_region"])
192                                .observe(batch.num_rows() as f64);
193
194                            let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
195                                .with_label_values(&["encode"])
196                                .start_timer();
197                            let mut iter = FlightEncoder::default()
198                                .encode(FlightMessage::RecordBatch(batch))
199                                .into_iter();
200                            let Some(flight_data) = iter.next() else {
201                                // Safety: `iter` on a type of `Vec1`, which is guaranteed to have
202                                // at least one element.
203                                unreachable!()
204                            };
205                            ensure!(
206                                iter.next().is_none(),
207                                error::NotSupportedSnafu {
208                                    feat: "bulk insert RecordBatch with dictionary arrays",
209                                }
210                            );
211                            encode_timer.observe_duration();
212                            (flight_data.data_header, flight_data.data_body)
213                        };
214                        let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
215                            .with_label_values(&["datanode_handle"])
216                            .start_timer();
217                        let request = RegionRequest {
218                            header: Some(RegionRequestHeader {
219                                tracing_context: TracingContext::from_current_span().to_w3c(),
220                                ..Default::default()
221                            }),
222                            body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
223                                skip_wal,
224                                region_id: region_id.as_u64(),
225                                partition_expr_version: partition_expr_version
226                                    .map(|value| PartitionExprVersion { value }),
227                                aligned_schema_version: None,
228                                body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
229                                    schema: schema_bytes,
230                                    data_header: header,
231                                    payload,
232                                })),
233                            })),
234                        };
235
236                        let datanode = node_manager.datanode(&peer).await;
237                        datanode
238                            .handle(request)
239                            .await
240                            .context(error::RequestRegionSnafu)
241                    });
242                handles.push(handle);
243            }
244        }
245
246        let region_responses = futures::future::try_join_all(handles)
247            .await
248            .context(error::JoinTaskSnafu)?;
249        wait_all_datanode_timer.observe_duration();
250        let mut rows_inserted: usize = 0;
251        for res in region_responses {
252            rows_inserted += res?.affected_rows;
253        }
254        crate::metrics::DIST_INGEST_ROW_COUNT
255            .with_label_values(&[db_name.as_str()])
256            .inc_by(rows_inserted as u64);
257        Ok(rows_inserted)
258    }
259
260    fn maybe_update_flow_dirty_window(&self, table_info: TableInfoRef, record_batch: RecordBatch) {
261        let table_id = table_info.table_id();
262        let table_flownode_set_cache = self.table_flownode_set_cache.clone();
263        let node_manager = self.node_manager.clone();
264        common_runtime::spawn_global(async move {
265            let result = table_flownode_set_cache
266                .get(table_id)
267                .await
268                .context(error::RequestInsertsSnafu);
269            let flownodes = match result {
270                Ok(flownodes) => flownodes.unwrap_or_default(),
271                Err(e) => {
272                    error!(e; "Failed to get flownodes for table id: {}", table_id);
273                    return;
274                }
275            };
276
277            let peers: HashSet<_> = flownodes.values().cloned().collect();
278            if peers.is_empty() {
279                return;
280            }
281
282            let Ok(timestamps) = extract_timestamps(
283                &record_batch,
284                &table_info
285                    .meta
286                    .schema
287                    .timestamp_column()
288                    .as_ref()
289                    .unwrap()
290                    .name,
291            )
292            .inspect_err(|e| {
293                error!(e; "Failed to extract timestamps from record batch");
294            }) else {
295                return;
296            };
297
298            for peer in peers {
299                let node_manager = node_manager.clone();
300                let timestamps = timestamps.clone();
301                common_runtime::spawn_global(async move {
302                    if let Err(e) = node_manager
303                        .flownode(&peer)
304                        .await
305                        .handle_mark_window_dirty(DirtyWindowRequests {
306                            requests: vec![DirtyWindowRequest {
307                                table_id,
308                                timestamps,
309                                time_ranges: vec![],
310                            }],
311                        })
312                        .await
313                        .context(error::RequestInsertsSnafu)
314                    {
315                        error!(e; "Failed to mark timestamps as dirty, table: {}", table_id);
316                    }
317                });
318            }
319        });
320    }
321}
322
323/// Calculate the timestamp range of record batch. Return `None` if record batch is empty.
324fn extract_timestamps(rb: &RecordBatch, timestamp_index_name: &str) -> error::Result<Vec<i64>> {
325    let ts_col = rb
326        .column_by_name(timestamp_index_name)
327        .context(error::ColumnNotFoundSnafu {
328            msg: timestamp_index_name,
329        })?;
330    if rb.num_rows() == 0 {
331        return Ok(vec![]);
332    }
333    let (primitive, _) =
334        datatypes::timestamp::timestamp_array_to_primitive(ts_col).with_context(|| {
335            error::InvalidTimeIndexTypeSnafu {
336                ty: ts_col.data_type().clone(),
337            }
338        })?;
339    Ok(primitive.iter().flatten().collect())
340}
operator/bulk_insert.rs

operator/
bulk_insert.rs
