Skip to main content

lora_executor/pull/
columns.rs

1use lora_compiler::physical::{PhysicalNodeId, PhysicalOp, PhysicalPlan};
2use lora_compiler::CompiledQuery;
3
4/// Result column names derived from the compiled plan.
5///
6/// Walks the plan from `root` looking for the topmost projection-shaped
7/// node (Projection, HashAggregation). Other operators that wrap a
8/// projection (Limit, Sort, PathBuild, OptionalMatch, Filter, Unwind,
9/// Create/Merge/Set/Delete/Remove) defer to their input. Returns an
10/// empty `Vec` for plans that have no named output (e.g. a bare
11/// scan-only plan), preserving the previous "infer from first row"
12/// behaviour for those cases.
13pub fn plan_result_columns(plan: &PhysicalPlan) -> Vec<String> {
14    plan_columns_at(plan, plan.root).unwrap_or_default()
15}
16
17fn plan_columns_at(plan: &PhysicalPlan, node: PhysicalNodeId) -> Option<Vec<String>> {
18    match &plan.nodes[node] {
19        PhysicalOp::Projection(p) => Some(p.items.iter().map(|i| i.name.to_string()).collect()),
20        PhysicalOp::HashAggregation(p) => Some(
21            p.group_by
22                .iter()
23                .chain(p.aggregates.iter())
24                .map(|i| i.name.to_string())
25                .collect(),
26        ),
27        PhysicalOp::Limit(p) => plan_columns_at(plan, p.input),
28        PhysicalOp::Sort(p) => plan_columns_at(plan, p.input),
29        PhysicalOp::PathBuild(p) => plan_columns_at(plan, p.input),
30        PhysicalOp::OptionalMatch(p) => plan_columns_at(plan, p.input),
31        // CALL { ... } emits an outerĂ—inner row pair; the visible result
32        // columns are still whatever the surrounding pipeline named.
33        PhysicalOp::CallSubquery(p) => plan_columns_at(plan, p.input),
34        PhysicalOp::Filter(p) => plan_columns_at(plan, p.input),
35        PhysicalOp::Unwind(p) => plan_columns_at(plan, p.input),
36        PhysicalOp::Create(p) => plan_columns_at(plan, p.input),
37        PhysicalOp::Merge(p) => plan_columns_at(plan, p.input),
38        PhysicalOp::Delete(p) => plan_columns_at(plan, p.input),
39        PhysicalOp::Set(p) => plan_columns_at(plan, p.input),
40        PhysicalOp::Remove(p) => plan_columns_at(plan, p.input),
41        PhysicalOp::Foreach(p) => plan_columns_at(plan, p.input),
42        PhysicalOp::Argument(_)
43        | PhysicalOp::NodeScan(_)
44        | PhysicalOp::NodeByLabelScan(_)
45        | PhysicalOp::NodeByPropertyScan(_)
46        | PhysicalOp::NodeByPropertyRangeScan(_)
47        | PhysicalOp::NodeByTextScan(_)
48        | PhysicalOp::NodeByPointScan(_)
49        | PhysicalOp::RelByPropertyRangeScan(_)
50        | PhysicalOp::RelByTextScan(_)
51        | PhysicalOp::RelByPointScan(_)
52        | PhysicalOp::Expand(_) => None,
53    }
54}
55
56/// Result column names for a compiled query (head plan; UNION branches
57/// must produce the same shape so the head's columns are authoritative).
58pub fn compiled_result_columns(compiled: &CompiledQuery) -> Vec<String> {
59    plan_result_columns(&compiled.physical)
60}