akar_processor/physical/write_ops/
merge.rs1use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
4use crate::physical::write_ops::set::{PhysicalSet, append_pipeline_columns};
5use akar_common::types::{PhysicalTypeID, Value};
6use akar_common::vector::{DataChunk, ValueVector};
7use akar_storage::table::TableCatalog;
8use akar_storage::wal::{WalSink, log_insert_record};
9use akar_transaction::UndoRecord;
10use std::sync::{Arc, Mutex};
11
12pub struct PhysicalMerge {
15 pub table_name: String,
16 pub table_id: u64,
17 pub properties: Vec<(String, akar_parser::ast::Expression)>,
18 pub on_match: Vec<PhysicalSet>,
19 pub on_create: Vec<PhysicalSet>,
20 pub table_catalog: Arc<TableCatalog>,
21 pub txn_id: Option<u64>,
23 pub undo_sink: Option<Arc<Mutex<Vec<UndoRecord>>>>,
25 pub wal_sink: Option<WalSink>,
27}
28
29fn row_id_chunk(row_ids: &[u64]) -> DataChunk {
33 let mut v = ValueVector::new(PhysicalTypeID::Int64, row_ids.len());
34 v.resize(row_ids.len());
35 for (i, r) in row_ids.iter().enumerate() {
36 v.set_i64(i, *r as i64);
37 }
38 let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
39 DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64]).with_names(vec!["_id".to_string()])
40}
41
42impl PhysicalOperatorExec for PhysicalMerge {
43 fn operator_type(&self) -> &str {
44 "merge"
45 }
46
47 fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
48 if !input.is_empty() && input.iter().all(|c| c.size == 0) {
50 return Ok(vec![DataChunk::new(vec![], vec![])]);
51 }
52
53 let eval_const = |expr: &akar_parser::ast::Expression, chunk: Option<&DataChunk>, row: usize| -> Value {
55 if let Some(c) = chunk {
56 crate::physical::write_ops::set::evaluate_expression_for_row(expr, c, row)
57 } else {
58 match expr {
59 akar_parser::ast::Expression::Constant(c) => match c {
60 akar_parser::ast::Constant::Null => Value::Null,
61 akar_parser::ast::Constant::Bool(b) => Value::Bool(*b),
62 akar_parser::ast::Constant::Integer(i) => Value::Int64(*i),
63 akar_parser::ast::Constant::Float(f) => Value::Double(*f),
64 akar_parser::ast::Constant::String(s) => Value::String(s.clone()),
65 },
66 _ => Value::Null,
67 }
68 }
69 };
70
71 let num_cols = {
73 let tbl = self
74 .table_catalog
75 .get_node_table_by_name(&self.table_name)
76 .ok_or_else(|| format!("Table '{}' not found for MERGE", self.table_name))?;
77 tbl.columns.len()
78 };
79
80 let chunks = if input.is_empty() {
82 let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
83 v.resize(1);
84 v.set_i64(0, 0);
85 let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
86 vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]
87 } else {
88 input
89 };
90
91 let mut merged_row_ids: Vec<u64> = Vec::new();
95 let mut source_rows: Vec<(usize, usize)> = Vec::new();
96 let mut matched_ids: Vec<u64> = Vec::new();
97 let mut created_ids: Vec<u64> = Vec::new();
98
99 for (ci, chunk) in chunks.iter().enumerate() {
100 for row in 0..chunk.size {
101 let table_info = self
102 .table_catalog
103 .get_node_table_by_name(&self.table_name)
104 .ok_or_else(|| format!("Table '{}' not found", self.table_name))?;
105
106 let mut pk_val = None;
108 for col_idx in 0..num_cols {
109 let col = &table_info.columns[col_idx];
110 if col.is_primary_key {
111 if let Some((_, expr)) = self.properties.iter().find(|(n, _)| n == &col.name) {
112 let val = eval_const(expr, Some(chunk), row);
113 pk_val = Some(val);
114 }
115 }
116 }
117
118 let mut matched: Option<u64> = None;
123 if let Some(val) = &pk_val {
124 let row_ids = table_info.lookup_by_pk_range(Some(val), true, Some(val), true, 1);
125 if !row_ids.is_empty() {
126 matched = Some(row_ids[0]);
127 }
128 } else if let Some((prop_name, expr)) = self.properties.first() {
129 let first_val = eval_const(expr, Some(chunk), row);
130 if let Some(prop_col) = table_info.columns.iter().position(|c| &c.name == prop_name) {
131 for row_idx in 0..table_info.num_rows as usize {
132 if let Some(val) = table_info.get_value(row_idx, prop_col)
133 && val == &first_val
134 {
135 matched = Some(row_idx as u64);
136 break;
137 }
138 }
139 }
140 }
141 drop(table_info);
142
143 if let Some(row_id) = matched {
144 matched_ids.push(row_id);
145 merged_row_ids.push(row_id);
146 source_rows.push((ci, row));
147 } else {
148 let table_info = self
149 .table_catalog
150 .get_node_table_by_name(&self.table_name)
151 .ok_or_else(|| format!("Table '{}' not found", self.table_name))?;
152 let mut new_values: Vec<Value> = Vec::new();
153 for col_idx in 0..num_cols {
154 let col_name = &table_info.columns[col_idx].name;
155 if let Some((_, expr)) = self.properties.iter().find(|(n, _)| n == col_name) {
156 new_values.push(eval_const(expr, Some(chunk), row));
157 } else if table_info.columns[col_idx].is_primary_key {
158 return Err(format!("MERGE CREATE requires primary key '{}'", col_name).into());
159 } else {
160 new_values.push(Value::Null);
161 }
162 }
163 drop(table_info);
164
165 if let Some(mut tbl) = self.table_catalog.get_node_table_by_name_mut(&self.table_name) {
166 let logged_row = self.wal_sink.is_some().then(|| new_values.clone());
167 let row_id = tbl
168 .insert_row_with_txn(new_values, self.txn_id)
169 .map_err(|e| format!("MERGE CREATE failed: {e}"))?;
170 log_insert_record(&self.wal_sink, self.table_id, logged_row.as_deref().unwrap_or(&[]));
171 if let Some(sink) = self.undo_sink.as_ref()
172 && let Ok(mut u) = sink.lock()
173 {
174 u.push(UndoRecord::insert(self.table_id, row_id));
175 }
176 created_ids.push(row_id);
177 merged_row_ids.push(row_id);
178 source_rows.push((ci, row));
179 }
180 }
181 }
182 }
183
184 for set_op in &self.on_match {
187 if !matched_ids.is_empty() {
188 let chunk = row_id_chunk(&matched_ids);
189 set_op.execute(vec![chunk])?;
190 }
191 }
192 for set_op in &self.on_create {
193 if !created_ids.is_empty() {
194 let chunk = row_id_chunk(&created_ids);
195 set_op.execute(vec![chunk])?;
196 }
197 }
198
199 tracing::info!("MERGE: processed merges in '{}'", self.table_name);
200
201 if merged_row_ids.is_empty() {
202 return Ok(vec![DataChunk::new(vec![], vec![])]);
203 }
204
205 let mut output = {
211 let table = self
212 .table_catalog
213 .get_node_table_by_name(&self.table_name)
214 .ok_or_else(|| format!("Node table '{}' not found for MERGE", self.table_name))?;
215 crate::physical::write_ops::set::build_old_row_chunk(&table.columns, &merged_row_ids, &|row_id, col| {
216 table.get_value(row_id as usize, col).cloned()
217 })?
218 };
219 append_pipeline_columns(&mut output, &chunks, &source_rows)?;
220
221 let mut v = ValueVector::new(PhysicalTypeID::Int64, merged_row_ids.len());
222 v.resize(merged_row_ids.len());
223 for (i, r) in merged_row_ids.iter().enumerate() {
224 v.set_i64(i, *r as i64);
225 }
226 output
227 .fields
228 .push(akar_common::arrow_vector::ArrowVector::from_legacy(&v).array);
229 output.field_types.push(PhysicalTypeID::Int64);
230 output.field_names.push("_id".to_string());
231
232 Ok(vec![output])
233 }
234}