Skip to main content

akar_processor/physical/
batch_insert.rs

1//! PhysicalBatchInsert — dedicated batch insert operator.
2//!
3//! Wraps `NodeTable::insert_rows_batch()` / `RelTable::insert_rels_batch()`
4//! for use in query plans where multiple CREATE statements can be fused
5//! into a single efficient batch operation.
6
7use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
8use akar_common::types::{PhysicalTypeID, Value};
9use akar_common::vector::{DataChunk, ValueVector};
10use akar_storage::table::TableCatalog;
11use std::sync::Arc;
12
13/// Physical operator for BATCH INSERT — inserts pre-collected rows/rels
14/// into a table using batch APIs for maximum throughput.
15pub struct PhysicalBatchInsert {
16    pub table_name: String,
17    pub table_id: u64,
18    /// Rows to insert: each row is a Vec<Value> matching column order.
19    pub rows: Vec<Vec<Value>>,
20    pub table_catalog: Arc<TableCatalog>,
21}
22
23impl PhysicalOperatorExec for PhysicalBatchInsert {
24    fn operator_type(&self) -> &str {
25        "batch_insert"
26    }
27
28    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
29        let num_rows = self.rows.len();
30        if num_rows == 0 {
31            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
32            v.resize(1);
33            v.set_i64(0, 0);
34            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
35            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
36        }
37
38        // Try node table first, then rel table
39        if let Some(mut table) = self.table_catalog.get_node_table_by_name_mut(&self.table_name) {
40            let count = table
41                .insert_rows_batch(&self.rows)
42                .map_err(|e| format!("BatchInsert node error: {e}"))?;
43            tracing::info!(
44                "BATCH INSERT: inserted {count} rows into node table '{}'",
45                self.table_name
46            );
47            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
48            v.resize(1);
49            v.set_i64(0, count as i64);
50            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
51            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
52        }
53
54        if let Some(mut table) = self.table_catalog.get_rel_table_by_name_mut(&self.table_name) {
55            let rels: Vec<(u64, u64, Vec<Value>)> = self
56                .rows
57                .iter()
58                .map(|row| {
59                    let from = match &row[0] {
60                        Value::Int64(v) => *v as u64,
61                        _ => 0,
62                    };
63                    let to = match &row[1] {
64                        Value::Int64(v) => *v as u64,
65                        _ => 0,
66                    };
67                    let props = row[2..].to_vec();
68                    (from, to, props)
69                })
70                .collect();
71            let count = table
72                .insert_rels_batch(&rels)
73                .map_err(|e| format!("BatchInsert rel error: {e}"))?;
74            tracing::info!(
75                "BATCH INSERT: inserted {count} rels into rel table '{}'",
76                self.table_name
77            );
78            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
79            v.resize(1);
80            v.set_i64(0, count as i64);
81            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
82            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
83        }
84
85        Err(format!(
86            "Table '{}' not found in storage catalog for BatchInsert",
87            self.table_name
88        )
89        .into())
90    }
91}