Skip to main content

lora_compiler/
plan_tree.rs

1//! Public-API-friendly mirror of a [`CompiledQuery`]'s operator tree.
2//!
3//! `PlanTree` walks the physical plan and produces a flat, serializable
4//! description suitable for surfacing through `Database::explain` /
5//! `Database::profile` and onwards through the language bindings.
6//!
7//! The internal `PhysicalOp` nodes reference analyzer-internal types
8//! (`ResolvedExpr`, `VarId`); we deliberately do not expose those. Each
9//! operator becomes a `PlanTreeNode` with an opaque, human-readable
10//! `details` map keyed on stable strings. Future cost-modelling can fill
11//! in `estimated_rows` without breaking the type.
12use std::collections::BTreeMap;
13use std::fmt::Write as _;
14
15use lora_analyzer::{ResolvedExpr, ResolvedProjection};
16use lora_ast::Direction;
17
18use crate::physical::{PhysicalNodeId, PhysicalOp, PhysicalPlan};
19use crate::{CompiledQuery, CompiledUnionBranch};
20
21/// One node in the rendered plan tree.
22#[derive(Debug, Clone)]
23pub struct PlanTreeNode {
24    /// Stable `PhysicalNodeId` within the owning plan. Synthetic
25    /// nodes (e.g. the Union root, branch wrappers) reuse a sentinel
26    /// id of `usize::MAX`.
27    pub id: usize,
28    /// Operator label, e.g. `NodeByLabelScan`, `Expand`, `Projection`.
29    pub operator: String,
30    /// Human-readable operator details. Values are stringified so the
31    /// public API never leaks internal expression / `VarId` types.
32    pub details: BTreeMap<String, String>,
33    /// Reserved for a future cost model. Always `None` today.
34    pub estimated_rows: Option<u64>,
35    /// Children in physical execution order (leaf-most first).
36    pub children: Vec<PlanTreeNode>,
37}
38
39/// Top-level plan tree.
40#[derive(Debug, Clone)]
41pub struct PlanTree {
42    pub root: PlanTreeNode,
43}
44
45const SYNTHETIC_ID: usize = usize::MAX;
46
47/// Build a `PlanTree` from a compiled query, including UNION branches.
48pub fn plan_tree_from_compiled(compiled: &CompiledQuery) -> PlanTree {
49    let head = build_node(&compiled.physical, compiled.physical.root);
50    if compiled.unions.is_empty() {
51        return PlanTree { root: head };
52    }
53
54    let mut children = Vec::with_capacity(compiled.unions.len() + 1);
55    children.push(head);
56    for branch in &compiled.unions {
57        children.push(build_union_branch(branch));
58    }
59    let mut details = BTreeMap::new();
60    details.insert("kind".to_string(), union_kind(&compiled.unions).to_string());
61    PlanTree {
62        root: PlanTreeNode {
63            id: SYNTHETIC_ID,
64            operator: "Union".to_string(),
65            details,
66            estimated_rows: None,
67            children,
68        },
69    }
70}
71
72fn build_union_branch(branch: &CompiledUnionBranch) -> PlanTreeNode {
73    let mut details = BTreeMap::new();
74    details.insert(
75        "kind".to_string(),
76        if branch.all { "ALL" } else { "DISTINCT" }.to_string(),
77    );
78    PlanTreeNode {
79        id: SYNTHETIC_ID,
80        operator: "UnionBranch".to_string(),
81        details,
82        estimated_rows: None,
83        children: vec![build_node(&branch.physical, branch.physical.root)],
84    }
85}
86
87fn build_node(plan: &PhysicalPlan, id: PhysicalNodeId) -> PlanTreeNode {
88    let op = &plan.nodes[id];
89    let description = describe(op);
90    let children = description
91        .child_ids
92        .into_iter()
93        .map(|cid| build_node(plan, cid))
94        .collect();
95    PlanTreeNode {
96        id,
97        operator: description.operator,
98        details: description.details,
99        estimated_rows: None,
100        children,
101    }
102}
103
104struct PlanDescription {
105    operator: String,
106    details: BTreeMap<String, String>,
107    child_ids: Vec<PhysicalNodeId>,
108}
109
110impl PlanDescription {
111    fn leaf(operator: &str) -> Self {
112        Self::new(operator, BTreeMap::new(), Vec::new())
113    }
114
115    fn with_children(
116        operator: &str,
117        details: BTreeMap<String, String>,
118        child_ids: Vec<PhysicalNodeId>,
119    ) -> Self {
120        Self::new(operator, details, child_ids)
121    }
122
123    fn new(
124        operator: &str,
125        details: BTreeMap<String, String>,
126        child_ids: Vec<PhysicalNodeId>,
127    ) -> Self {
128        Self {
129            operator: operator.to_string(),
130            details,
131            child_ids,
132        }
133    }
134}
135
136fn describe(op: &PhysicalOp) -> PlanDescription {
137    let mut d = BTreeMap::new();
138    match op {
139        PhysicalOp::Argument(_) => PlanDescription::leaf("Argument"),
140        PhysicalOp::NodeScan(n) => {
141            d.insert("var".to_string(), var_str(n.var));
142            PlanDescription::with_children("NodeScan", d, opt_input(n.input))
143        }
144        PhysicalOp::NodeByLabelScan(n) => {
145            d.insert("var".to_string(), var_str(n.var));
146            d.insert("labels".to_string(), label_groups_str(&n.labels));
147            PlanDescription::with_children("NodeByLabelScan", d, opt_input(n.input))
148        }
149        PhysicalOp::NodeByIdSeek(n) => {
150            d.insert("var".to_string(), var_str(n.var));
151            if !n.labels.is_empty() {
152                d.insert("labels".to_string(), label_groups_str(&n.labels));
153            }
154            d.insert("ids".to_string(), expr_str(&n.ids));
155            if n.in_list {
156                d.insert("mode".to_string(), "in".to_string());
157            }
158            PlanDescription::with_children("NodeByIdSeek", d, opt_input(n.input))
159        }
160        PhysicalOp::RelByIdSeek(n) => {
161            d.insert("rel".to_string(), var_str(n.rel));
162            d.insert("src".to_string(), var_str(n.src));
163            d.insert("dst".to_string(), var_str(n.dst));
164            if !n.src_labels.is_empty() {
165                d.insert("src_labels".to_string(), label_groups_str(&n.src_labels));
166            }
167            if !n.types.is_empty() {
168                d.insert("types".to_string(), n.types.join("|"));
169            }
170            d.insert(
171                "direction".to_string(),
172                direction_str(n.direction).to_string(),
173            );
174            d.insert("ids".to_string(), expr_str(&n.ids));
175            if n.in_list {
176                d.insert("mode".to_string(), "in".to_string());
177            }
178            PlanDescription::with_children("RelByIdSeek", d, opt_input(n.input))
179        }
180        PhysicalOp::NodeByPropertyScan(n) => {
181            d.insert("var".to_string(), var_str(n.var));
182            if !n.labels.is_empty() {
183                d.insert("labels".to_string(), label_groups_str(&n.labels));
184            }
185            d.insert("key".to_string(), n.key.clone());
186            d.insert("value".to_string(), expr_str(&n.value));
187            if n.in_list {
188                d.insert("mode".to_string(), "in".to_string());
189            }
190            PlanDescription::with_children("NodeByPropertyScan", d, opt_input(n.input))
191        }
192        PhysicalOp::NodeByPropertyRangeScan(n) => {
193            d.insert("var".to_string(), var_str(n.var));
194            if !n.labels.is_empty() {
195                d.insert("labels".to_string(), label_groups_str(&n.labels));
196            }
197            d.insert("key".to_string(), n.key.clone());
198            if let Some(lo) = &n.lo {
199                d.insert(
200                    "lo".to_string(),
201                    format!(
202                        "{} {}",
203                        if n.lo_inclusive { ">=" } else { ">" },
204                        expr_str(lo)
205                    ),
206                );
207            }
208            if let Some(hi) = &n.hi {
209                d.insert(
210                    "hi".to_string(),
211                    format!(
212                        "{} {}",
213                        if n.hi_inclusive { "<=" } else { "<" },
214                        expr_str(hi)
215                    ),
216                );
217            }
218            PlanDescription::with_children("NodeByPropertyRangeScan", d, opt_input(n.input))
219        }
220        PhysicalOp::NodeByPointScan(n) => {
221            d.insert("var".to_string(), var_str(n.var));
222            if !n.labels.is_empty() {
223                d.insert("labels".to_string(), label_groups_str(&n.labels));
224            }
225            d.insert("key".to_string(), n.key.clone());
226            match &n.predicate {
227                crate::PointPredicate::WithinBBox {
228                    lower_left,
229                    upper_right,
230                } => {
231                    d.insert("predicate".to_string(), "withinBBox".to_string());
232                    d.insert("lowerLeft".to_string(), expr_str(lower_left));
233                    d.insert("upperRight".to_string(), expr_str(upper_right));
234                }
235                crate::PointPredicate::WithinDistance {
236                    center,
237                    max_distance,
238                    inclusive,
239                } => {
240                    d.insert(
241                        "predicate".to_string(),
242                        if *inclusive {
243                            "distance<="
244                        } else {
245                            "distance<"
246                        }
247                        .to_string(),
248                    );
249                    d.insert("center".to_string(), expr_str(center));
250                    d.insert("maxDistance".to_string(), expr_str(max_distance));
251                }
252            }
253            PlanDescription::with_children("NodeByPointScan", d, opt_input(n.input))
254        }
255        PhysicalOp::NodeByTextScan(n) => {
256            d.insert("var".to_string(), var_str(n.var));
257            if !n.labels.is_empty() {
258                d.insert("labels".to_string(), label_groups_str(&n.labels));
259            }
260            d.insert("key".to_string(), n.key.clone());
261            d.insert(
262                "predicate".to_string(),
263                match n.predicate {
264                    crate::TextPredicate::StartsWith => "STARTS WITH",
265                    crate::TextPredicate::EndsWith => "ENDS WITH",
266                    crate::TextPredicate::Contains => "CONTAINS",
267                }
268                .to_string(),
269            );
270            d.insert("query".to_string(), expr_str(&n.query));
271            PlanDescription::with_children("NodeByTextScan", d, opt_input(n.input))
272        }
273        PhysicalOp::RelByPropertyRangeScan(n) => {
274            d.insert("rel".to_string(), var_str(n.rel));
275            d.insert("src".to_string(), var_str(n.src));
276            d.insert("dst".to_string(), var_str(n.dst));
277            if !n.types.is_empty() {
278                d.insert("types".to_string(), n.types.join("|"));
279            }
280            d.insert(
281                "direction".to_string(),
282                direction_str(n.direction).to_string(),
283            );
284            d.insert("key".to_string(), n.key.clone());
285            if let Some(lo) = &n.lo {
286                d.insert(
287                    "lo".to_string(),
288                    format!(
289                        "{} {}",
290                        if n.lo_inclusive { ">=" } else { ">" },
291                        expr_str(lo)
292                    ),
293                );
294            }
295            if let Some(hi) = &n.hi {
296                d.insert(
297                    "hi".to_string(),
298                    format!(
299                        "{} {}",
300                        if n.hi_inclusive { "<=" } else { "<" },
301                        expr_str(hi)
302                    ),
303                );
304            }
305            PlanDescription::with_children("RelByPropertyRangeScan", d, opt_input(n.input))
306        }
307        PhysicalOp::RelByTextScan(n) => {
308            d.insert("rel".to_string(), var_str(n.rel));
309            d.insert("src".to_string(), var_str(n.src));
310            d.insert("dst".to_string(), var_str(n.dst));
311            if !n.types.is_empty() {
312                d.insert("types".to_string(), n.types.join("|"));
313            }
314            d.insert(
315                "direction".to_string(),
316                direction_str(n.direction).to_string(),
317            );
318            d.insert("key".to_string(), n.key.clone());
319            d.insert(
320                "predicate".to_string(),
321                match n.predicate {
322                    crate::TextPredicate::StartsWith => "STARTS WITH",
323                    crate::TextPredicate::EndsWith => "ENDS WITH",
324                    crate::TextPredicate::Contains => "CONTAINS",
325                }
326                .to_string(),
327            );
328            d.insert("query".to_string(), expr_str(&n.query));
329            PlanDescription::with_children("RelByTextScan", d, opt_input(n.input))
330        }
331        PhysicalOp::RelByPointScan(n) => {
332            d.insert("rel".to_string(), var_str(n.rel));
333            d.insert("src".to_string(), var_str(n.src));
334            d.insert("dst".to_string(), var_str(n.dst));
335            if !n.types.is_empty() {
336                d.insert("types".to_string(), n.types.join("|"));
337            }
338            d.insert(
339                "direction".to_string(),
340                direction_str(n.direction).to_string(),
341            );
342            d.insert("key".to_string(), n.key.clone());
343            match &n.predicate {
344                crate::PointPredicate::WithinBBox {
345                    lower_left,
346                    upper_right,
347                } => {
348                    d.insert("predicate".to_string(), "withinBBox".to_string());
349                    d.insert("lowerLeft".to_string(), expr_str(lower_left));
350                    d.insert("upperRight".to_string(), expr_str(upper_right));
351                }
352                crate::PointPredicate::WithinDistance {
353                    center,
354                    max_distance,
355                    inclusive,
356                } => {
357                    d.insert(
358                        "predicate".to_string(),
359                        if *inclusive {
360                            "distance<="
361                        } else {
362                            "distance<"
363                        }
364                        .to_string(),
365                    );
366                    d.insert("center".to_string(), expr_str(center));
367                    d.insert("maxDistance".to_string(), expr_str(max_distance));
368                }
369            }
370            PlanDescription::with_children("RelByPointScan", d, opt_input(n.input))
371        }
372        PhysicalOp::Expand(n) => describe_expand(n),
373        PhysicalOp::Filter(n) => {
374            d.insert("predicate".to_string(), expr_str(&n.predicate));
375            PlanDescription::with_children("Filter", d, vec![n.input])
376        }
377        PhysicalOp::Projection(n) => describe_projection(n),
378        PhysicalOp::Unwind(n) => {
379            d.insert("alias".to_string(), var_str(n.alias));
380            d.insert("expr".to_string(), expr_str(&n.expr));
381            PlanDescription::with_children("Unwind", d, vec![n.input])
382        }
383        PhysicalOp::HashAggregation(n) => describe_hash_aggregation(n),
384        PhysicalOp::Sort(n) => {
385            d.insert(
386                "items".to_string(),
387                format!("{} sort key(s)", n.items.len()),
388            );
389            if let Some(top_k) = n.top_k {
390                d.insert("top_k".to_string(), top_k.to_string());
391            } else if let Some(bound) = &n.limit {
392                let limit = expr_str(&bound.limit);
393                let top_k = match &bound.skip {
394                    Some(skip) => format!("{} + {limit}", expr_str(skip)),
395                    None => limit,
396                };
397                d.insert("top_k".to_string(), top_k);
398            }
399            PlanDescription::with_children("Sort", d, vec![n.input])
400        }
401        PhysicalOp::Limit(n) => {
402            if let Some(skip) = &n.skip {
403                d.insert("skip".to_string(), expr_str(skip));
404            }
405            if let Some(limit) = &n.limit {
406                d.insert("limit".to_string(), expr_str(limit));
407            }
408            PlanDescription::with_children("Limit", d, vec![n.input])
409        }
410        PhysicalOp::Create(n) => {
411            d.insert(
412                "elements".to_string(),
413                pattern_summary(n.pattern.parts.len()),
414            );
415            PlanDescription::with_children("Create", d, vec![n.input])
416        }
417        PhysicalOp::Merge(n) => describe_merge(n),
418        PhysicalOp::Delete(n) => {
419            d.insert("detach".to_string(), n.detach.to_string());
420            d.insert("targets".to_string(), n.expressions.len().to_string());
421            PlanDescription::with_children("Delete", d, vec![n.input])
422        }
423        PhysicalOp::Set(n) => {
424            d.insert("items".to_string(), n.items.len().to_string());
425            PlanDescription::with_children("Set", d, vec![n.input])
426        }
427        PhysicalOp::Remove(n) => {
428            d.insert("items".to_string(), n.items.len().to_string());
429            PlanDescription::with_children("Remove", d, vec![n.input])
430        }
431        PhysicalOp::Foreach(n) => {
432            d.insert("variable".to_string(), var_str(n.variable));
433            d.insert("list".to_string(), expr_str(&n.list));
434            d.insert("body".to_string(), n.body.len().to_string());
435            PlanDescription::with_children("Foreach", d, vec![n.input])
436        }
437        PhysicalOp::OptionalMatch(n) => describe_optional_match(n),
438        PhysicalOp::CallSubquery(n) => {
439            d.insert(
440                "new_vars".to_string(),
441                n.new_vars
442                    .iter()
443                    .copied()
444                    .map(var_str)
445                    .collect::<Vec<_>>()
446                    .join(", "),
447            );
448            PlanDescription::with_children("CallSubquery", d, vec![n.input, n.inner])
449        }
450        PhysicalOp::PathBuild(n) => {
451            d.insert("output".to_string(), var_str(n.output));
452            d.insert("nodes".to_string(), n.node_vars.len().to_string());
453            d.insert("rels".to_string(), n.rel_vars.len().to_string());
454            if let Some(all) = n.shortest_path_all {
455                d.insert("shortest_path_all".to_string(), all.to_string());
456            }
457            PlanDescription::with_children("PathBuild", d, vec![n.input])
458        }
459    }
460}
461
462fn union_kind(branches: &[CompiledUnionBranch]) -> &'static str {
463    let all = branches.iter().all(|b| b.all);
464    let all_distinct = branches.iter().all(|b| !b.all);
465    if all {
466        "ALL"
467    } else if all_distinct {
468        "DISTINCT"
469    } else {
470        "MIXED"
471    }
472}
473
474fn describe_expand(n: &crate::physical::ExpandExec) -> PlanDescription {
475    let mut d = BTreeMap::new();
476    d.insert("src".to_string(), var_str(n.src));
477    d.insert("dst".to_string(), var_str(n.dst));
478    if let Some(rel) = n.rel {
479        d.insert("rel".to_string(), var_str(rel));
480    }
481    if !n.types.is_empty() {
482        d.insert("types".to_string(), n.types.join("|"));
483    }
484    d.insert(
485        "direction".to_string(),
486        direction_str(n.direction).to_string(),
487    );
488    if let Some(props) = &n.rel_properties {
489        d.insert("rel_properties".to_string(), expr_str(props));
490    }
491    if let Some(range) = &n.range {
492        d.insert("range".to_string(), format!("{:?}", range));
493    }
494    PlanDescription::with_children("Expand", d, vec![n.input])
495}
496
497fn describe_projection(n: &crate::physical::ProjectionExec) -> PlanDescription {
498    let mut d = BTreeMap::new();
499    d.insert("distinct".to_string(), n.distinct.to_string());
500    d.insert(
501        "include_existing".to_string(),
502        n.include_existing.to_string(),
503    );
504    d.insert("items".to_string(), projection_names(&n.items));
505    PlanDescription::with_children("Projection", d, vec![n.input])
506}
507
508fn describe_hash_aggregation(n: &crate::physical::HashAggregationExec) -> PlanDescription {
509    let mut d = BTreeMap::new();
510    d.insert("group_by".to_string(), projection_names(&n.group_by));
511    d.insert("aggregates".to_string(), projection_names(&n.aggregates));
512    PlanDescription::with_children("HashAggregation", d, vec![n.input])
513}
514
515fn describe_merge(n: &crate::physical::MergeExec) -> PlanDescription {
516    let mut d = BTreeMap::new();
517    d.insert("actions".to_string(), n.actions.len().to_string());
518    PlanDescription::with_children("Merge", d, vec![n.input])
519}
520
521fn describe_optional_match(n: &crate::physical::OptionalMatchExec) -> PlanDescription {
522    let mut d = BTreeMap::new();
523    d.insert(
524        "new_vars".to_string(),
525        n.new_vars
526            .iter()
527            .copied()
528            .map(var_str)
529            .collect::<Vec<_>>()
530            .join(", "),
531    );
532    PlanDescription::with_children("OptionalMatch", d, vec![n.input, n.inner])
533}
534
535fn opt_input(input: Option<PhysicalNodeId>) -> Vec<PhysicalNodeId> {
536    input.map(|i| vec![i]).unwrap_or_default()
537}
538
539fn var_str(v: lora_analyzer::symbols::VarId) -> String {
540    format!("v{}", v.0)
541}
542
543fn label_groups_str(groups: &[Vec<String>]) -> String {
544    groups
545        .iter()
546        .map(|or_group| or_group.join("|"))
547        .collect::<Vec<_>>()
548        .join("&")
549}
550
551fn projection_names(items: &[ResolvedProjection]) -> String {
552    items
553        .iter()
554        .map(|p| p.name.clone())
555        .collect::<Vec<_>>()
556        .join(", ")
557}
558
559fn direction_str(d: Direction) -> &'static str {
560    match d {
561        Direction::Right => "->",
562        Direction::Left => "<-",
563        Direction::Undirected => "-",
564    }
565}
566
567fn expr_str(e: &ResolvedExpr) -> String {
568    let mut out = String::new();
569    let _ = write!(&mut out, "{:?}", e);
570    out
571}
572
573fn pattern_summary(part_count: usize) -> String {
574    format!("{} pattern part(s)", part_count)
575}