Skip to main content

akar_processor/physical/write_ops/
merge.rs

1//! Physical operator for MERGE.
2
3use 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
12/// Physical operator for MERGE.
13/// Represents a combination of MATCH and INSERT (Upsert).
14pub 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    /// Active transaction id (P52.18).
22    pub txn_id: Option<u64>,
23    /// Undo sink for rollback records (P52.18).
24    pub undo_sink: Option<Arc<Mutex<Vec<UndoRecord>>>>,
25    /// Typed WAL sink so MERGE-created rows survive restarts via replay (P60.2).
26    pub wal_sink: Option<WalSink>,
27}
28
29/// Build a single-column chunk carrying physical row indices under the `_id`
30/// pseudo-column name, so a `PhysicalSet` passed as ON MATCH / ON CREATE can
31/// re-target exactly those rows.
32fn 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        // A filtered-out pipeline stays empty (P53.25): nothing was merged.
49        if !input.is_empty() && input.iter().all(|c| c.size == 0) {
50            return Ok(vec![DataChunk::new(vec![], vec![])]);
51        }
52
53        // Evaluate constant helper (fallback if input is empty)
54        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        // Get table info to build the row
72        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        // Handle input chunks for pipeline (or just 1 iteration if empty)
81        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        // One output row per processed input row, carrying the physical row id
92        // (matched or newly created) plus its source (chunk, row) so pipeline
93        // columns (e.g. UNWIND variables) survive into the output (P53.31).
94        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                // Determine PK property
107                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                // Match against the existing table. Without a PK property, fall
119                // back to a full scan over the first pattern property (Cypher
120                // MERGE on non-PK fields has no O(1) lookup without a secondary
121                // index).
122                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        // Apply ON MATCH / ON CREATE SET once per group, targeting the affected
185        // row ids via the `_id` pseudo-column (mirrors the edge MERGE path).
186        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        // Output one row per processed input row: the post-update table columns
206        // (named, so a following RETURN resolves `<alias>.<prop>`), the input
207        // pipeline columns, and the `_id` pseudo-column. Previously MERGE
208        // emitted a single count chunk, so `MERGE ... SET ... RETURN n.prop`
209        // resolved the projection against the count (P53.31).
210        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}