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 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;
use arrow::record_batch::RecordBatch;
use bytes::Bytes;
use common_base::AffectedRows;
use common_grpc::FlightData;
use common_grpc::flight::{FlightEncoder, FlightMessage};
use common_telemetry::error;
use common_telemetry::tracing_context::TracingContext;
use futures::future;
use meter_core::data::MeterRecord;
use meter_macros::write_meter;
use session::context::{Channel, QueryContextRef};
use snafu::{OptionExt, ResultExt, ensure};
use store_api::storage::RegionId;
use table::TableRef;
use table::metadata::TableInfoRef;
use tokio_util::sync::CancellationToken;

use crate::error::Result;
use crate::insert::Inserter;
use crate::req_convert::TablePartitioner;
use crate::req_convert::insert;
use crate::write_dispatcher::{WriteAttempt, WriteAttribution};
use crate::{error, metrics};

impl Inserter {
    /// Writes a coalesced batch while retaining each original input's acknowledgements.
    pub async fn flush_bulk_batch_tracked(
        &self,
        table_info: TableInfoRef,
        batch: RecordBatch,
        ctx: QueryContextRef,
        owners: Vec<usize>,
        inputs: usize,
    ) -> Result<Vec<common_query::WriteProgress>> {
        let requests = TablePartitioner::try_new(&self.partition_manager, &table_info)
            .await
            .context(error::InvalidPartitionSnafu)?
            .partition_tracked_bulk_insert_requests(batch, owners, ctx.skip_wal())?;
        let requests = requests
            .into_iter()
            .map(|(request, owners)| crate::write_dispatcher::TrackedWrite {
                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(request)),
                },
                owners,
            })
            .collect();
        Ok(self
            .write_dispatcher
            .dispatch_tracked(requests, inputs, None)
            .await)
    }

    /// 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,
        attribution: WriteAttribution,
    ) -> Result<AffectedRows> {
        let requests = TablePartitioner::try_new(&self.partition_manager, &table_info)
            .await
            .context(error::InvalidPartitionSnafu)?
            .partition_attributed_bulk_requests(batch, ctx.skip_wal(), attribution)?;
        let writes = requests
            .into_iter()
            .map(|(request, attribution)| {
                (
                    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(request)),
                    },
                    attribution,
                )
            })
            .collect();
        Ok(self
            .write_dispatcher
            .dispatch_attributed(writes, None)
            .await
            .context(error::RequestInsertsSnafu)?
            .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);
        }

        // Bulk storage paths may discard unknown columns instead of rejecting them.
        for field in record_batch.schema_ref().fields() {
            ensure!(
                table_info
                    .meta
                    .schema
                    .column_schema_by_name(field.name())
                    .is_some(),
                error::InvalidInsertRequestSnafu {
                    reason: format!(
                        "Column '{}' not found in table '{}'",
                        field.name(),
                        table_info.full_table_name()
                    ),
                }
            );
        }

        // 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 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 writes = self.write_dispatcher.prepare(vec![request]).await?;
            let result = self
                .write_dispatcher
                .execute(
                    writes,
                    Some(
                        metrics::HANDLE_BULK_INSERT_ELAPSED.with_label_values(&["datanode_handle"]),
                    ),
                    None,
                )
                .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 regions = region_masks
            .keys()
            .map(|region| RegionId::new(table_id, *region))
            .collect::<Vec<_>>();
        let mut peers = self
            .write_dispatcher
            .resolve_peers(&regions)
            .await
            .context(error::FindRegionLeaderSnafu)?;

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

        // Dropping the caller stops retries without abandoning in-flight writes.
        let cancellation = CancellationToken::new();
        let _guard = cancellation.clone().drop_guard();
        let mut handles = Vec::with_capacity(region_masks.len());

        for (region_number, mask) in region_masks {
            let region_id = RegionId::new(table_id, region_number);
            if mask.select_none() {
                continue;
            }
            let partition_expr_version = partition_versions
                .get(&region_id.region_number())
                .copied()
                .unwrap_or_default();
            let peer = peers
                .remove(&region_id)
                .context(partition::error::FindLeaderSnafu {
                    region_id,
                    table_id,
                })
                .context(error::FindRegionLeaderSnafu)?;
            let dispatcher = self.write_dispatcher.clone();
            let rb = record_batch.clone();
            let schema_bytes = schema_bytes.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<WriteAttempt>> =
                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 = compute::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 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,
                            })),
                        })),
                    };

                    Ok(dispatcher
                        .send(
                            peer,
                            request,
                            Some(
                                &metrics::HANDLE_BULK_INSERT_ELAPSED
                                    .with_label_values(&["datanode_handle"]),
                            ),
                        )
                        .await)
                });
            handles.push(handle);
        }

        let dispatcher = self.write_dispatcher.clone();
        let rows_inserted = common_runtime::spawn_global(async move {
            let results = future::try_join_all(handles)
                .await
                .context(error::JoinTaskSnafu)?;
            let attempts = results.into_iter().collect::<Result<Vec<_>>>()?;
            dispatcher
                .finish(
                    attempts,
                    Some(
                        &metrics::HANDLE_BULK_INSERT_ELAPSED
                            .with_label_values(&["datanode_handle"]),
                    ),
                    Some(&cancellation),
                )
                .await
                .context(error::RequestRegionSnafu)
                .map(|response| response.affected_rows)
        })
        .await
        .context(error::JoinTaskSnafu)??;
        wait_all_datanode_timer.observe_duration();
        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) = insert::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);
                    }
                });
            }
        });
    }
}