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

operator/
bulk_insert.rs
