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.clone()).collect()),
20        PhysicalOp::HashAggregation(p) => Some(
21            p.group_by
22                .iter()
23                .chain(p.aggregates.iter())
24                .map(|i| i.name.clone())
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        PhysicalOp::Filter(p) => plan_columns_at(plan, p.input),
32        PhysicalOp::Unwind(p) => plan_columns_at(plan, p.input),
33        PhysicalOp::Create(p) => plan_columns_at(plan, p.input),
34        PhysicalOp::Merge(p) => plan_columns_at(plan, p.input),
35        PhysicalOp::Delete(p) => plan_columns_at(plan, p.input),
36        PhysicalOp::Set(p) => plan_columns_at(plan, p.input),
37        PhysicalOp::Remove(p) => plan_columns_at(plan, p.input),
38        PhysicalOp::Argument(_)
39        | PhysicalOp::NodeScan(_)
40        | PhysicalOp::NodeByLabelScan(_)
41        | PhysicalOp::NodeByPropertyScan(_)
42        | PhysicalOp::Expand(_) => None,
43    }
44}
45
46/// Result column names for a compiled query (head plan; UNION branches
47/// must produce the same shape so the head's columns are authoritative).
48pub fn compiled_result_columns(compiled: &CompiledQuery) -> Vec<String> {
49    plan_result_columns(&compiled.physical)
50}