servers/batcher/table/
batch.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::sync::Arc;
16use std::time::Instant;
17
18use arrow::compute::concat_batches;
19use arrow::record_batch::RecordBatch;
20use common_query::WriteProgress;
21use common_telemetry::error;
22use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu};
23use operator::insert::Inserter;
24use operator::metrics::DIST_INGEST_ROW_COUNT;
25use snafu::ResultExt;
26
27use crate::batcher::table::flow_notifier::FlowNotifier;
28use crate::batcher::table::metrics::{
29    FLUSH_DROPPED_ROWS, FLUSH_ELAPSED, FLUSH_FAILURES, FLUSH_ROWS, FLUSH_TOTAL,
30};
31use crate::batcher::table::pending_batch::{PendingBatch, notify_batches};
32
33/// A detached batch owns every submission until its completion is reported.
34pub(in crate::batcher::table) struct Batch {
35    pub submissions: Vec<PendingBatch>,
36    pub total_rows: usize,
37}
38
39pub(in crate::batcher::table) async fn flush_batch(
40    batch: Batch,
41    inserter: Arc<Inserter>,
42    notifier: FlowNotifier,
43) {
44    let started = Instant::now();
45    let result = send_batch(&batch, &inserter).await;
46    FLUSH_ELAPSED.observe(started.elapsed().as_secs_f64());
47    match result {
48        Ok(progress) => {
49            if progress.iter().any(|result| result.error.is_some()) {
50                FLUSH_FAILURES.inc();
51            } else {
52                FLUSH_TOTAL.inc();
53                FLUSH_ROWS.observe(batch.total_rows as f64);
54            }
55            if let Some(first) = batch.submissions.first() {
56                notifier.notify(
57                    first.table_info.clone(),
58                    batch
59                        .submissions
60                        .iter()
61                        .zip(&progress)
62                        .filter(|(_, result)| result.error.is_none())
63                        .map(|(pending, _)| &pending.batch),
64                );
65            }
66            for (pending, result) in batch.submissions.into_iter().zip(progress) {
67                if let Some(counter) = pending.ctx.write_progress() {
68                    counter.add_written(result.affected_rows);
69                }
70                DIST_INGEST_ROW_COUNT
71                    .with_label_values(&[pending.ctx.get_db_string().as_str()])
72                    .inc_by(result.affected_rows as u64);
73                if result.error.is_some() {
74                    FLUSH_DROPPED_ROWS
75                        .inc_by((pending.batch.num_rows() - result.affected_rows) as u64);
76                }
77                let _ = pending.response_tx.send(result);
78            }
79        }
80        Err(error) => {
81            error!(error; "Failed to flush table batch, rows: {}", batch.total_rows);
82            FLUSH_FAILURES.inc();
83            FLUSH_DROPPED_ROWS.inc_by(batch.total_rows as u64);
84            notify_batches(batch.submissions, Err(Arc::new(error)));
85        }
86    }
87}
88
89async fn send_batch(batch: &Batch, inserter: &Inserter) -> Result<Vec<WriteProgress>> {
90    let first = batch.submissions.first().ok_or_else(|| {
91        UnexpectedSnafu {
92            violated: "cannot execute an empty flush".to_string(),
93        }
94        .build()
95    })?;
96    let combined = combine_batches(batch.submissions.iter().map(|pending| &pending.batch))?;
97    let mut owners = Vec::with_capacity(batch.total_rows);
98    for (owner, pending) in batch.submissions.iter().enumerate() {
99        owners.extend(std::iter::repeat_n(owner, pending.batch.num_rows()));
100    }
101    inserter
102        .flush_bulk_batch_tracked(
103            first.table_info.clone(),
104            combined,
105            first.ctx.clone(),
106            owners,
107            batch.submissions.len(),
108        )
109        .await
110}
111
112fn combine_batches<'a>(batches: impl IntoIterator<Item = &'a RecordBatch>) -> Result<RecordBatch> {
113    let batches = batches.into_iter().collect::<Vec<_>>();
114    let Some(first) = batches.first() else {
115        return UnexpectedSnafu {
116            violated: "cannot combine an empty flush".to_string(),
117        }
118        .fail();
119    };
120    let schema = first.schema();
121    if batches.iter().any(|batch| batch.schema() != schema) {
122        return UnexpectedSnafu {
123            violated: "cannot combine different Arrow schemas".to_string(),
124        }
125        .fail();
126    }
127    concat_batches(&schema, batches).context(ComputeArrowSnafu)
128}
129
130#[cfg(test)]
131mod tests {
132    use arrow::array::Int32Array;
133    use arrow::datatypes::{DataType, Field, Schema};
134
135    use crate::batcher::table::batch::*;
136
137    fn batch(name: &str, values: Vec<i32>) -> RecordBatch {
138        RecordBatch::try_new(
139            Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)])),
140            vec![Arc::new(Int32Array::from(values))],
141        )
142        .unwrap()
143    }
144
145    #[test]
146    fn test_combine_preserves_rows_and_rejects_schema_changes() {
147        let first = batch("a", vec![1, 2]);
148        let second = batch("a", vec![3]);
149        assert_eq!(
150            batch("a", vec![1, 2, 3]),
151            combine_batches([&first, &second]).unwrap()
152        );
153        let different = batch("b", vec![4]);
154        assert!(combine_batches([&first, &different]).is_err());
155        let schema = first.schema();
156        let metadata = Schema::new_with_metadata(
157            schema.fields().clone(),
158            [("version".to_string(), "2".to_string())]
159                .into_iter()
160                .collect(),
161        );
162        let changed = RecordBatch::try_new(Arc::new(metadata), first.columns().to_vec()).unwrap();
163        assert!(combine_batches([&first, &changed]).is_err());
164    }
165}
servers/batcher/table/batch.rs

servers/batcher/table/
batch.rs
