Skip to main content

marsdb_query/
executor.rs

1use std::collections::{BTreeMap, HashMap, HashSet};
2
3use marsdb_graph::{EdgeId, GraphStore, NodeId, PropertyValue, WriteTransaction};
4
5use crate::ast::{CompareOp, Expr, Literal, Pattern, PropAccess, RelDirection, ReturnExpr, ReturnItem, Statement, Tail};
6use crate::error::QueryError;
7use crate::ir::LogicalPlan;
8use crate::planner::build_match_plan;
9use crate::result::QueryResult;
10use crate::value::Value;
11
12#[derive(Debug, Clone, Copy)]
13enum Binding {
14    Node(NodeId),
15    Edge(EdgeId),
16}
17
18type BindingRow = HashMap<String, Binding>;
19
20pub struct Executor<'a> {
21    store: &'a GraphStore,
22}
23
24impl<'a> Executor<'a> {
25    pub fn new(store: &'a GraphStore) -> Self {
26        Self { store }
27    }
28
29    /// Runs the whole statement inside a single write transaction — the
30    /// crash-safety boundary from the plan (one statement = one commit).
31    /// Every graph access below this point must go through `write_txn` and
32    /// the `*_in_txn` GraphStore methods, never the standalone
33    /// `self.store.*` methods, which open (and would deadlock trying to
34    /// re-open) their own transaction.
35    pub fn execute(&self, stmt: &Statement) -> Result<QueryResult, QueryError> {
36        let write_txn = self.store.begin_write()?;
37        let outcome = match stmt {
38            Statement::Create(patterns) => self.execute_create(&write_txn, patterns),
39            Statement::Match {
40                pattern,
41                where_clause,
42                tail,
43                limit,
44            } => self.execute_match(&write_txn, pattern, where_clause, tail, *limit),
45        };
46        match outcome {
47            Ok(result) => {
48                GraphStore::commit(write_txn)?;
49                Ok(result)
50            }
51            Err(e) => {
52                // Best-effort rollback; the original error is what matters.
53                let _ = GraphStore::abort(write_txn);
54                Err(e)
55            }
56        }
57    }
58
59    fn execute_create(&self, write_txn: &WriteTransaction, patterns: &[Pattern]) -> Result<QueryResult, QueryError> {
60        for pattern in patterns {
61            let start_label = pattern.start.label.clone().unwrap_or_else(|| "Node".to_string());
62            let start_props = literal_props_to_values(&pattern.start.props);
63            let mut prev_id = GraphStore::create_node_in_txn(write_txn, &start_label, start_props)?;
64
65            for (rel, node) in &pattern.hops {
66                let label = node.label.clone().unwrap_or_else(|| "Node".to_string());
67                let props = literal_props_to_values(&node.props);
68                let node_id = GraphStore::create_node_in_txn(write_txn, &label, props)?;
69
70                let rel_label = rel.rel_type.clone().unwrap_or_else(|| "REL".to_string());
71                let rel_props = literal_props_to_values(&rel.props);
72                let (src, dst) = match rel.direction {
73                    RelDirection::Right => (prev_id, node_id),
74                    RelDirection::Left => (node_id, prev_id),
75                };
76                GraphStore::create_edge_in_txn(write_txn, &rel_label, src, dst, rel_props)?;
77                prev_id = node_id;
78            }
79        }
80        Ok(QueryResult {
81            columns: vec![],
82            rows: vec![],
83        })
84    }
85
86    fn execute_match(
87        &self,
88        write_txn: &WriteTransaction,
89        pattern: &Pattern,
90        where_clause: &Option<Expr>,
91        tail: &Tail,
92        limit: Option<i64>,
93    ) -> Result<QueryResult, QueryError> {
94        let mut plan = build_match_plan(pattern, where_clause);
95        if let Some(count) = limit {
96            plan = LogicalPlan::Limit {
97                input: Box::new(plan),
98                count,
99            };
100        }
101        let rows = self.eval_plan(write_txn, &plan)?;
102        match tail {
103            Tail::Return(items) => self.materialize_return(write_txn, items, &rows),
104            Tail::Delete(vars) => self.materialize_delete(write_txn, vars, &rows, false),
105            Tail::DetachDelete(vars) => self.materialize_delete(write_txn, vars, &rows, true),
106            Tail::Set(items) => self.materialize_set(write_txn, items, &rows),
107        }
108    }
109
110    fn eval_plan(&self, write_txn: &WriteTransaction, plan: &LogicalPlan) -> Result<Vec<BindingRow>, QueryError> {
111        match plan {
112            LogicalPlan::AllNodesScan { var } => self.scan(write_txn, var, None),
113            LogicalPlan::NodeByLabelScan { var, label } => self.scan(write_txn, var, Some(label)),
114            LogicalPlan::Expand {
115                input,
116                from_var,
117                to_var,
118                rel_var,
119                rel_label,
120                direction,
121            } => {
122                let base_rows = self.eval_plan(write_txn, input)?;
123                let mut out = Vec::new();
124                for row in base_rows {
125                    let Some(Binding::Node(from_id)) = row.get(from_var).copied() else {
126                        return Err(QueryError::UnboundVariable(from_var.clone()));
127                    };
128                    let entries =
129                        GraphStore::neighbors_in_txn(write_txn, from_id, *direction, rel_label.as_deref())?;
130                    for entry in entries {
131                        let mut new_row = row.clone();
132                        new_row.insert(to_var.clone(), Binding::Node(entry.other));
133                        if let Some(rv) = rel_var {
134                            new_row.insert(rv.clone(), Binding::Edge(entry.edge_id));
135                        }
136                        out.push(new_row);
137                    }
138                }
139                Ok(out)
140            }
141            LogicalPlan::Filter { input, predicate } => {
142                let rows = self.eval_plan(write_txn, input)?;
143                let mut out = Vec::with_capacity(rows.len());
144                for row in rows {
145                    if self.eval_expr(write_txn, predicate, &row)? {
146                        out.push(row);
147                    }
148                }
149                Ok(out)
150            }
151            LogicalPlan::Limit { input, count } => {
152                let mut rows = self.eval_plan(write_txn, input)?;
153                rows.truncate((*count).max(0) as usize);
154                Ok(rows)
155            }
156        }
157    }
158
159    fn scan(&self, write_txn: &WriteTransaction, var: &str, label: Option<&str>) -> Result<Vec<BindingRow>, QueryError> {
160        let nodes = GraphStore::all_nodes_in_txn(write_txn, label)?;
161        Ok(nodes
162            .into_iter()
163            .map(|n| {
164                let mut row = BindingRow::new();
165                row.insert(var.to_string(), Binding::Node(n.id));
166                row
167            })
168            .collect())
169    }
170
171    fn eval_expr(&self, write_txn: &WriteTransaction, expr: &Expr, row: &BindingRow) -> Result<bool, QueryError> {
172        Ok(match expr {
173            Expr::And(l, r) => self.eval_expr(write_txn, l, row)? && self.eval_expr(write_txn, r, row)?,
174            Expr::Or(l, r) => self.eval_expr(write_txn, l, row)? || self.eval_expr(write_txn, r, row)?,
175            Expr::Not(e) => !self.eval_expr(write_txn, e, row)?,
176            Expr::Compare(pa, op, lit) => {
177                let prop_value = self.lookup_prop(write_txn, pa, row)?;
178                compare(&prop_value, *op, lit)
179            }
180        })
181    }
182
183    fn lookup_prop(
184        &self,
185        write_txn: &WriteTransaction,
186        pa: &PropAccess,
187        row: &BindingRow,
188    ) -> Result<Option<PropertyValue>, QueryError> {
189        let binding = row
190            .get(&pa.var)
191            .ok_or_else(|| QueryError::UnboundVariable(pa.var.clone()))?;
192        match binding {
193            Binding::Node(id) => {
194                let node = GraphStore::get_node_in_txn(write_txn, *id)?;
195                Ok(node.and_then(|n| n.props.get(&pa.prop).cloned()))
196            }
197            Binding::Edge(id) => {
198                let edge = GraphStore::get_edge_in_txn(write_txn, *id)?;
199                Ok(edge.and_then(|e| e.props.get(&pa.prop).cloned()))
200            }
201        }
202    }
203
204    fn materialize_return(
205        &self,
206        write_txn: &WriteTransaction,
207        items: &[ReturnItem],
208        rows: &[BindingRow],
209    ) -> Result<QueryResult, QueryError> {
210        let columns = items
211            .iter()
212            .enumerate()
213            .map(|(i, item)| item.alias.clone().unwrap_or_else(|| default_column_name(&item.expr, i)))
214            .collect();
215        let mut out_rows = Vec::with_capacity(rows.len());
216        for row in rows {
217            let mut out_row = Vec::with_capacity(items.len());
218            for item in items {
219                out_row.push(self.eval_return_expr(write_txn, &item.expr, row)?);
220            }
221            out_rows.push(out_row);
222        }
223        Ok(QueryResult {
224            columns,
225            rows: out_rows,
226        })
227    }
228
229    fn eval_return_expr(
230        &self,
231        write_txn: &WriteTransaction,
232        expr: &ReturnExpr,
233        row: &BindingRow,
234    ) -> Result<Value, QueryError> {
235        match expr {
236            ReturnExpr::Var(var) => {
237                let binding = row.get(var).ok_or_else(|| QueryError::UnboundVariable(var.clone()))?;
238                match binding {
239                    Binding::Node(id) => {
240                        let node = GraphStore::get_node_in_txn(write_txn, *id)?
241                            .expect("bound node exists within this statement's transaction");
242                        Ok(Value::Node(node))
243                    }
244                    Binding::Edge(id) => {
245                        let edge = GraphStore::get_edge_in_txn(write_txn, *id)?
246                            .expect("bound edge exists within this statement's transaction");
247                        Ok(Value::Edge(edge))
248                    }
249                }
250            }
251            ReturnExpr::Prop(pa) => {
252                let value = self.lookup_prop(write_txn, pa, row)?;
253                Ok(match value {
254                    Some(pv) => Value::Property(pv),
255                    None => Value::Null,
256                })
257            }
258            ReturnExpr::Lit(lit) => Ok(Value::Literal(lit.clone())),
259        }
260    }
261
262    fn materialize_delete(
263        &self,
264        write_txn: &WriteTransaction,
265        vars: &[String],
266        rows: &[BindingRow],
267        detach: bool,
268    ) -> Result<QueryResult, QueryError> {
269        let mut deleted_nodes = HashSet::new();
270        let mut deleted_edges = HashSet::new();
271        for row in rows {
272            for var in vars {
273                let binding = row.get(var).ok_or_else(|| QueryError::UnboundVariable(var.clone()))?;
274                match binding {
275                    Binding::Node(id) => {
276                        if deleted_nodes.insert(*id) {
277                            GraphStore::delete_node_in_txn(write_txn, *id, detach)?;
278                        }
279                    }
280                    Binding::Edge(id) => {
281                        if deleted_edges.insert(*id) {
282                            GraphStore::delete_edge_in_txn(write_txn, *id)?;
283                        }
284                    }
285                }
286            }
287        }
288        Ok(QueryResult {
289            columns: vec![],
290            rows: vec![],
291        })
292    }
293
294    fn materialize_set(
295        &self,
296        write_txn: &WriteTransaction,
297        items: &[(PropAccess, Literal)],
298        rows: &[BindingRow],
299    ) -> Result<QueryResult, QueryError> {
300        for row in rows {
301            for (pa, lit) in items {
302                let binding = row.get(&pa.var).ok_or_else(|| QueryError::UnboundVariable(pa.var.clone()))?;
303                let value = literal_to_value(lit);
304                match binding {
305                    Binding::Node(id) => {
306                        GraphStore::set_node_prop_in_txn(write_txn, *id, &pa.prop, value)?;
307                    }
308                    Binding::Edge(id) => {
309                        GraphStore::set_edge_prop_in_txn(write_txn, *id, &pa.prop, value)?;
310                    }
311                }
312            }
313        }
314        Ok(QueryResult {
315            columns: vec![],
316            rows: vec![],
317        })
318    }
319}
320
321fn default_column_name(expr: &ReturnExpr, idx: usize) -> String {
322    match expr {
323        ReturnExpr::Var(v) => v.clone(),
324        ReturnExpr::Prop(pa) => format!("{}.{}", pa.var, pa.prop),
325        ReturnExpr::Lit(_) => format!("col{idx}"),
326    }
327}
328
329fn literal_to_value(lit: &Literal) -> PropertyValue {
330    match lit {
331        Literal::Int(i) => PropertyValue::Int(*i),
332        Literal::Float(f) => PropertyValue::Float(*f),
333        Literal::String(s) => PropertyValue::String(s.clone()),
334        Literal::Bool(b) => PropertyValue::Bool(*b),
335        Literal::Null => PropertyValue::Null,
336    }
337}
338
339fn literal_props_to_values(props: &[(String, Literal)]) -> BTreeMap<String, PropertyValue> {
340    props.iter().map(|(k, v)| (k.clone(), literal_to_value(v))).collect()
341}
342
343fn compare(prop: &Option<PropertyValue>, op: CompareOp, lit: &Literal) -> bool {
344    let Some(prop) = prop else { return false };
345    match (prop, lit) {
346        (PropertyValue::Int(a), Literal::Int(b)) => cmp_f64(op, *a as f64, *b as f64),
347        (PropertyValue::Int(a), Literal::Float(b)) => cmp_f64(op, *a as f64, *b),
348        (PropertyValue::Float(a), Literal::Float(b)) => cmp_f64(op, *a, *b),
349        (PropertyValue::Float(a), Literal::Int(b)) => cmp_f64(op, *a, *b as f64),
350        (PropertyValue::String(a), Literal::String(b)) => cmp_ord(op, a.as_str(), b.as_str()),
351        (PropertyValue::Bool(a), Literal::Bool(b)) => match op {
352            CompareOp::Eq => a == b,
353            CompareOp::Ne => a != b,
354            _ => false,
355        },
356        (PropertyValue::Null, Literal::Null) => matches!(op, CompareOp::Eq),
357        _ => false,
358    }
359}
360
361fn cmp_f64(op: CompareOp, a: f64, b: f64) -> bool {
362    match op {
363        CompareOp::Eq => a == b,
364        CompareOp::Ne => a != b,
365        CompareOp::Lt => a < b,
366        CompareOp::Le => a <= b,
367        CompareOp::Gt => a > b,
368        CompareOp::Ge => a >= b,
369    }
370}
371
372fn cmp_ord<T: PartialOrd>(op: CompareOp, a: T, b: T) -> bool {
373    match op {
374        CompareOp::Eq => a == b,
375        CompareOp::Ne => a != b,
376        CompareOp::Lt => a < b,
377        CompareOp::Le => a <= b,
378        CompareOp::Gt => a > b,
379        CompareOp::Ge => a >= b,
380    }
381}