operator/

bulk_insert.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use ahash::{HashMap, HashMapExt};
use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests};
use api::v1::region::{
    BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
};
use api::v1::{ArrowIpc, PartitionExprVersion};
use arrow::compute::filter_record_batch;
use arrow::record_batch::RecordBatch;
use bytes::Bytes;
use common_base::AffectedRows;
use common_error::ext::BoxedError;
use common_grpc::FlightData;
use common_grpc::flight::{FlightEncoder, FlightMessage, record_batch_to_ipc};
use common_telemetry::error;
use common_telemetry::tracing_context::TracingContext;
use futures::future::{join_all, try_join_all};
use meter_core::data::MeterRecord;
use meter_macros::write_meter;
use session::context::{Channel, QueryContextRef};
use snafu::{ResultExt, ensure};
use store_api::storage::RegionId;
use table::TableRef;
use table::metadata::TableInfoRef;

use crate::error::Result;
use crate::insert::Inserter;
use crate::req_convert::insert::extract_timestamps;
use crate::{error, metrics};

impl Inserter {
    /// Routes and writes a prepared table batch, returning the affected row count.
    ///
    /// Callers must exclude instant-TTL tables and handle Flow notifications after
    /// successful writes. This execution helper does not perform either step.
    pub async fn flush_bulk_batch(
        &self,
        table_info: TableInfoRef,
        batch: RecordBatch,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows> {
        let (rule, versions) = self
            .partition_manager
            .find_table_partition_rule(&table_info)
            .await
            .context(error::InvalidPartitionSnafu)?;
        let masks = rule
            .split_record_batch(&batch)
            .context(error::SplitInsertSnafu)?;
        let mut writes = Vec::with_capacity(masks.len());
        for (region_number, mask) in masks {
            if mask.select_none() {
                continue;
            }
            let region_id = RegionId::new(table_info.table_id(), region_number);
            let selected = if mask.select_all() {
                batch.clone()
            } else {
                filter_record_batch(&batch, mask.array()).context(error::ComputeArrowSnafu)?
            };
            let (schema, data_header, payload) = record_batch_to_ipc(selected)
                .map_err(BoxedError::new)
                .context(error::ExternalSnafu)?;
            let peer = self
                .partition_manager
                .find_region_leader(region_id)
                .await
                .context(error::FindRegionLeaderSnafu)?;
            let request = RegionRequest {
                header: Some(RegionRequestHeader {
                    dbname: ctx.get_db_string(),
                    tracing_context: TracingContext::from_current_span().to_w3c(),
                    ..Default::default()
                }),
                body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
                    skip_wal: ctx.skip_wal(),
                    region_id: region_id.as_u64(),
                    partition_expr_version: versions
                        .get(&region_number)
                        .copied()
                        .flatten()
                        .map(|value| PartitionExprVersion { value }),
                    // Let the datanode revalidate against its current schema.
                    aligned_schema_version: None,
                    body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
                        schema,
                        data_header,
                        payload,
                    })),
                })),
            };
            writes.push((peer, request));
        }
        let results = join_all(writes.into_iter().map(|(peer, request)| async move {
            self.node_manager
                .datanode(&peer)
                .await
                .handle(request)
                .await
                .context(error::RequestInsertsSnafu)
        }))
        .await;
        let affected_rows = results
            .into_iter()
            .map(|result| result.map(|response| response.affected_rows))
            .sum::<Result<usize>>()?;
        Ok(affected_rows)
    }

    /// Handle bulk insert request.
    pub async fn handle_bulk_insert(
        &self,
        table: TableRef,
        raw_flight_data: FlightData,
        record_batch: RecordBatch,
        schema_bytes: Bytes,
        skip_wal: bool,
        channel: Channel,
    ) -> Result<AffectedRows> {
        let table_info = table.table_info();
        let table_id = table_info.table_id();
        let db_name = table_info.get_db_string();

        if record_batch.num_rows() == 0 {
            return Ok(0);
        }

        // The zero value is WCU, not bytes. Bulk writes have no WCU accounting;
        // preserve that behavior while admitting their rows before dispatch.
        write_meter!(MeterRecord::new(
            table_info.catalog_name.clone(),
            table_info.schema_name.clone(),
            0,
            record_batch.num_rows() as u64,
            channel as u8,
        ))
        .await
        .context(error::WriteRejectedSnafu)?;

        let body_size = raw_flight_data.data_body.len();
        // TODO(yingwen): Fill record batch impure default values. Note that we should override `raw_flight_data` if we have to fill defaults.
        // notify flownode to update dirty timestamps if flow is configured.
        self.maybe_update_flow_dirty_window(table_info.clone(), record_batch.clone());

        metrics::BULK_REQUEST_MESSAGE_SIZE.observe(body_size as f64);
        metrics::BULK_REQUEST_ROWS
            .with_label_values(&["raw"])
            .observe(record_batch.num_rows() as f64);

        let partition_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
            .with_label_values(&["partition"])
            .start_timer();
        let (partition_rule, partition_versions) = self
            .partition_manager
            .find_table_partition_rule(&table_info)
            .await
            .context(error::InvalidPartitionSnafu)?;

        // find partitions for each row in the record batch
        let region_masks = partition_rule
            .split_record_batch(&record_batch)
            .context(error::SplitInsertSnafu)?;
        partition_timer.observe_duration();

        // fast path: only one region.
        if region_masks.len() == 1 {
            metrics::BULK_REQUEST_ROWS
                .with_label_values(&["rows_per_region"])
                .observe(record_batch.num_rows() as f64);

            // SAFETY: region masks length checked
            let (region_number, _) = region_masks.into_iter().next().unwrap();
            let region_id = RegionId::new(table_id, region_number);
            let partition_expr_version = partition_versions
                .get(&region_number)
                .copied()
                .unwrap_or_default();
            let datanode = self
                .partition_manager
                .find_region_leader(region_id)
                .await
                .context(error::FindRegionLeaderSnafu)?;

            let request = RegionRequest {
                header: Some(RegionRequestHeader {
                    tracing_context: TracingContext::from_current_span().to_w3c(),
                    ..Default::default()
                }),
                body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
                    skip_wal,
                    region_id: region_id.as_u64(),
                    partition_expr_version: partition_expr_version
                        .map(|value| PartitionExprVersion { value }),
                    aligned_schema_version: None,
                    body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
                        schema: schema_bytes.clone(),
                        data_header: raw_flight_data.data_header,
                        payload: raw_flight_data.data_body,
                    })),
                })),
            };

            let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
                .with_label_values(&["datanode_handle"])
                .start_timer();
            let datanode = self.node_manager.datanode(&datanode).await;
            let result = datanode
                .handle(request)
                .await
                .context(error::RequestRegionSnafu)
                .map(|r| r.affected_rows);
            if let Ok(rows) = result {
                crate::metrics::DIST_INGEST_ROW_COUNT
                    .with_label_values(&[db_name.as_str()])
                    .inc_by(rows as u64);
            }
            return result;
        }

        let mut mask_per_datanode = HashMap::with_capacity(region_masks.len());
        for (region_number, mask) in region_masks {
            let region_id = RegionId::new(table_id, region_number);
            let datanode = self
                .partition_manager
                .find_region_leader(region_id)
                .await
                .context(error::FindRegionLeaderSnafu)?;
            mask_per_datanode
                .entry(datanode)
                .or_insert_with(Vec::new)
                .push((region_id, mask));
        }

        let wait_all_datanode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
            .with_label_values(&["wait_all_datanode"])
            .start_timer();

        let mut handles = Vec::with_capacity(mask_per_datanode.len());

        for (peer, masks) in mask_per_datanode {
            for (region_id, mask) in masks {
                if mask.select_none() {
                    continue;
                }
                let partition_expr_version = partition_versions
                    .get(&region_id.region_number())
                    .copied()
                    .unwrap_or_default();
                let rb = record_batch.clone();
                let schema_bytes = schema_bytes.clone();
                let node_manager = self.node_manager.clone();
                let peer = peer.clone();
                let raw_header_and_data = if mask.select_all() {
                    Some((
                        raw_flight_data.data_header.clone(),
                        raw_flight_data.data_body.clone(),
                    ))
                } else {
                    None
                };
                let handle: common_runtime::JoinHandle<Result<api::region::RegionResponse>> =
                    common_runtime::spawn_global(async move {
                        let (header, payload) = if mask.select_all() {
                            // SAFETY: raw data must be present, we can avoid re-encoding.
                            raw_header_and_data.unwrap()
                        } else {
                            let filter_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
                                .with_label_values(&["filter"])
                                .start_timer();
                            let batch = filter_record_batch(&rb, mask.array())
                                .context(error::ComputeArrowSnafu)?;
                            filter_timer.observe_duration();
                            metrics::BULK_REQUEST_ROWS
                                .with_label_values(&["rows_per_region"])
                                .observe(batch.num_rows() as f64);

                            let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
                                .with_label_values(&["encode"])
                                .start_timer();
                            let mut iter = FlightEncoder::default()
                                .encode(FlightMessage::RecordBatch(batch))
                                .into_iter();
                            let Some(flight_data) = iter.next() else {
                                // Safety: `iter` on a type of `Vec1`, which is guaranteed to have
                                // at least one element.
                                unreachable!()
                            };
                            ensure!(
                                iter.next().is_none(),
                                error::NotSupportedSnafu {
                                    feat: "bulk insert RecordBatch with dictionary arrays",
                                }
                            );
                            encode_timer.observe_duration();
                            (flight_data.data_header, flight_data.data_body)
                        };
                        let _datanode_handle_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
                            .with_label_values(&["datanode_handle"])
                            .start_timer();
                        let request = RegionRequest {
                            header: Some(RegionRequestHeader {
                                tracing_context: TracingContext::from_current_span().to_w3c(),
                                ..Default::default()
                            }),
                            body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
                                skip_wal,
                                region_id: region_id.as_u64(),
                                partition_expr_version: partition_expr_version
                                    .map(|value| PartitionExprVersion { value }),
                                aligned_schema_version: None,
                                body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
                                    schema: schema_bytes,
                                    data_header: header,
                                    payload,
                                })),
                            })),
                        };

                        let datanode = node_manager.datanode(&peer).await;
                        datanode
                            .handle(request)
                            .await
                            .context(error::RequestRegionSnafu)
                    });
                handles.push(handle);
            }
        }

        let region_responses = try_join_all(handles).await.context(error::JoinTaskSnafu)?;
        wait_all_datanode_timer.observe_duration();
        let mut rows_inserted: usize = 0;
        for res in region_responses {
            rows_inserted += res?.affected_rows;
        }
        crate::metrics::DIST_INGEST_ROW_COUNT
            .with_label_values(&[db_name.as_str()])
            .inc_by(rows_inserted as u64);
        Ok(rows_inserted)
    }

    fn maybe_update_flow_dirty_window(&self, table_info: TableInfoRef, record_batch: RecordBatch) {
        let table_id = table_info.table_id();
        let table_flownode_set_cache = self.table_flownode_set_cache.clone();
        let node_manager = self.node_manager.clone();
        common_runtime::spawn_global(async move {
            let result = table_flownode_set_cache
                .get(table_id)
                .await
                .context(error::RequestInsertsSnafu);
            let flownodes = match result {
                Ok(flownodes) => flownodes.unwrap_or_default(),
                Err(e) => {
                    error!(e; "Failed to get flownodes for table id: {}", table_id);
                    return;
                }
            };

            let peers: HashSet<_> = flownodes.values().cloned().collect();
            if peers.is_empty() {
                return;
            }

            let Ok(timestamps) = extract_timestamps(
                &record_batch,
                &table_info
                    .meta
                    .schema
                    .timestamp_column()
                    .as_ref()
                    .unwrap()
                    .name,
            )
            .inspect_err(|e| {
                error!(e; "Failed to extract timestamps from record batch");
            }) else {
                return;
            };

            for peer in peers {
                let node_manager = node_manager.clone();
                let timestamps = timestamps.clone();
                common_runtime::spawn_global(async move {
                    if let Err(e) = node_manager
                        .flownode(&peer)
                        .await
                        .handle_mark_window_dirty(DirtyWindowRequests {
                            requests: vec![DirtyWindowRequest {
                                table_id,
                                timestamps,
                                time_ranges: vec![],
                            }],
                        })
                        .await
                        .context(error::RequestInsertsSnafu)
                    {
                        error!(e; "Failed to mark timestamps as dirty, table: {}", table_id);
                    }
                });
            }
        });
    }
}