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 akar_transaction::UndoRecord;
12use std::sync::{Arc, Mutex};
13
14/// Physical operator for BATCH INSERT — inserts pre-collected rows/rels
15/// into a table using batch APIs for maximum throughput.
16pub struct PhysicalBatchInsert {
17    pub table_name: String,
18    pub table_id: u64,
19    /// Rows to insert: each row is a Vec<Value> matching column order.
20    pub rows: Vec<Vec<Value>>,
21    pub table_catalog: Arc<TableCatalog>,
22    /// Active transaction id (P52.18).
23    pub txn_id: Option<u64>,
24    /// Undo sink for rollback records (P52.18).
25    pub undo_sink: Option<Arc<Mutex<Vec<UndoRecord>>>>,
26}
27
28impl PhysicalOperatorExec for PhysicalBatchInsert {
29    fn operator_type(&self) -> &str {
30        "batch_insert"
31    }
32
33    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
34        let num_rows = self.rows.len();
35        if num_rows == 0 {
36            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
37            v.resize(1);
38            v.set_i64(0, 0);
39            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
40            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
41        }
42
43        // Try node table first, then rel table
44        if let Some(mut table) = self.table_catalog.get_node_table_by_name_mut(&self.table_name) {
45            let start = table.num_rows;
46            let count = table
47                .insert_rows_batch_with_txn(&self.rows, self.txn_id)
48                .map_err(|e| format!("BatchInsert node error: {e}"))?;
49            if let Some(sink) = self.undo_sink.as_ref()
50                && let Ok(mut u) = sink.lock()
51            {
52                for row in start..start + count {
53                    u.push(UndoRecord::insert(self.table_id, row));
54                }
55            }
56            tracing::info!(
57                "BATCH INSERT: inserted {count} rows into node table '{}'",
58                self.table_name
59            );
60            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
61            v.resize(1);
62            v.set_i64(0, count as i64);
63            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
64            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
65        }
66
67        if let Some(mut table) = self.table_catalog.get_rel_table_by_name_mut(&self.table_name) {
68            let rels: Vec<(u64, u64, Vec<Value>)> = self
69                .rows
70                .iter()
71                .map(|row| {
72                    let from = match &row[0] {
73                        Value::Int64(v) => *v as u64,
74                        _ => 0,
75                    };
76                    let to = match &row[1] {
77                        Value::Int64(v) => *v as u64,
78                        _ => 0,
79                    };
80                    let props = row[2..].to_vec();
81                    (from, to, props)
82                })
83                .collect();
84            let start = table.edges.len();
85            let count = table
86                .insert_rels_batch(&rels)
87                .map_err(|e| format!("BatchInsert rel error: {e}"))?;
88            if let Some(sink) = self.undo_sink.as_ref()
89                && let Ok(mut u) = sink.lock()
90            {
91                for idx in start..start + count as usize {
92                    u.push(UndoRecord::insert(self.table_id, idx as u64));
93                }
94            }
95            tracing::info!(
96                "BATCH INSERT: inserted {count} rels into rel table '{}'",
97                self.table_name
98            );
99            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
100            v.resize(1);
101            v.set_i64(0, count as i64);
102            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
103            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
104        }
105
106        Err(format!(
107            "Table '{}' not found in storage catalog for BatchInsert",
108            self.table_name
109        )
110        .into())
111    }
112}