use lora_compiler::physical::{PhysicalNodeId, PhysicalOp, PhysicalPlan};
use lora_compiler::CompiledQuery;
pub fn plan_result_columns(plan: &PhysicalPlan) -> Vec<String> {
plan_columns_at(plan, plan.root).unwrap_or_default()
}
fn plan_columns_at(plan: &PhysicalPlan, node: PhysicalNodeId) -> Option<Vec<String>> {
match &plan.nodes[node] {
PhysicalOp::Projection(p) => Some(p.items.iter().map(|i| i.name.to_string()).collect()),
PhysicalOp::HashAggregation(p) => Some(
p.group_by
.iter()
.chain(p.aggregates.iter())
.map(|i| i.name.to_string())
.collect(),
),
PhysicalOp::Limit(p) => plan_columns_at(plan, p.input),
PhysicalOp::Sort(p) => plan_columns_at(plan, p.input),
PhysicalOp::PathBuild(p) => plan_columns_at(plan, p.input),
PhysicalOp::OptionalMatch(p) => plan_columns_at(plan, p.input),
PhysicalOp::CallSubquery(p) => plan_columns_at(plan, p.input),
PhysicalOp::Filter(p) => plan_columns_at(plan, p.input),
PhysicalOp::Unwind(p) => plan_columns_at(plan, p.input),
PhysicalOp::Create(p) => plan_columns_at(plan, p.input),
PhysicalOp::Merge(p) => plan_columns_at(plan, p.input),
PhysicalOp::Delete(p) => plan_columns_at(plan, p.input),
PhysicalOp::Set(p) => plan_columns_at(plan, p.input),
PhysicalOp::Remove(p) => plan_columns_at(plan, p.input),
PhysicalOp::Foreach(p) => plan_columns_at(plan, p.input),
PhysicalOp::Argument(_)
| PhysicalOp::NodeScan(_)
| PhysicalOp::NodeByLabelScan(_)
| PhysicalOp::NodeByPropertyScan(_)
| PhysicalOp::NodeByPropertyRangeScan(_)
| PhysicalOp::NodeByTextScan(_)
| PhysicalOp::NodeByPointScan(_)
| PhysicalOp::RelByPropertyRangeScan(_)
| PhysicalOp::RelByTextScan(_)
| PhysicalOp::RelByPointScan(_)
| PhysicalOp::Expand(_) => None,
}
}
pub fn compiled_result_columns(compiled: &CompiledQuery) -> Vec<String> {
plan_result_columns(&compiled.physical)
}