linkrs-query 0.1.0

//! GO Statement Planner
//! Planning for handling Nebula GO queries
//!
//! ## Improvement Notes
//!
//! Implement the complete logic for filtering expressions.
//! Improving the handling of JOIN operations
//! - Add support for attribute projection.

use crate::binder::BoundStatement;
use crate::parser::ast::{GoStmt, Stmt};
use crate::planning::plan::core::node_id_generator::next_node_id;
use crate::planning::plan::logical::logical_nodes::access::LogicalStartNode;
use crate::planning::plan::logical::logical_nodes::control_flow::LogicalArgumentNode;
use crate::planning::plan::logical::logical_nodes::operation::{
    LogicalDedupNode, LogicalFilterNode, LogicalProjectNode,
};
use crate::planning::plan::logical::logical_nodes::traversal::LogicalExpandAllNode;
use crate::planning::plan::logical::LogicalNodeEnum;
use crate::planning::plan::SubPlan;
use crate::planning::planner::{Planner, PlannerError, ValidatedStatement};
use crate::planning::statements::projection_util;
use crate::QueryContext;
use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
use linkrs_core::types::{ContextualExpression, EdgeDirection};
use std::sync::Arc;

/// GO Query Planner
/// Responsible for converting GO statements into execution plans.
#[derive(Debug, Clone)]
pub struct GoPlanner {}

impl GoPlanner {
    /// Create a new GO planner.
    pub fn new() -> Self {
        Self {}
    }
}

impl Planner for GoPlanner {
    fn transform(
        &mut self,
        validated: &ValidatedStatement,
        qctx: Arc<QueryContext>,
    ) -> Result<SubPlan, PlannerError> {
        let space_id = qctx.space_id().unwrap_or(1);

        let go_stmt = match validated.stmt() {
            Stmt::Go(go_stmt) => go_stmt,
            _ => {
                return Err(PlannerError::InvalidOperation(
                    "GoPlanner requires Go statements".to_string(),
                ));
            }
        };

        // GO syntax carries no vertex label: neighbor vertices are
        // materialized at execution from the edge type schema, so the plan
        // leaves the neighbor label empty instead of guessing one.
        // Use the verification information to optimize the planning process.
        let validation_info = &validated.validation_info;

        // 1. Check the optimization suggestions.
        for hint in &validation_info.optimization_hints {
            log::debug!("GO Optimization Tip: {:?}", hint);
        }

        // 2. Use the path analysis information
        for path_analysis in &validation_info.path_analysis {
            if path_analysis.edge_count > 5 {
                log::warn!(
                    "The GO path contains {} edges, which may affect performance.",
                    path_analysis.edge_count
                );
            }
        }

        // 3. Use semantic information
        let referenced_edges = &validation_info.semantic_info.referenced_edges;
        if !referenced_edges.is_empty() {
            log::debug!("GO referenced edge type: {:?}", referenced_edges);
        }

        // Handle FROM clause - extract source vertex IDs
        let from_vertices = &go_stmt.from.vertices;
        if from_vertices.is_empty() {
            return Err(PlannerError::PlanGenerationFailed(
                "GO statement must have FROM clause".to_string(),
            ));
        }

        // Check if the first from expression is a literal (vertex ID)
        let first_from = &from_vertices[0];
        let (use_start_node, from_var) = if first_from.is_literal() {
            // If it's a literal like "1", we need to create a variable in context
            // Use StartNode as the tail and set the variable in execution context
            (true, "v".to_string())
        } else if let Some(var_name) = first_from.as_variable() {
            // If it's already a variable, use ArgumentNode
            (false, var_name.clone())
        } else {
            // For other expressions, use ArgumentNode with a default variable name
            (false, "v".to_string())
        };

        // Create the tail node — this becomes the input to ExpandAllNode.
        let tail_logical = if use_start_node {
            LogicalNodeEnum::Start(LogicalStartNode::new())
        } else {
            LogicalNodeEnum::Argument(LogicalArgumentNode {
                id: next_node_id(),
                var: from_var.clone(),
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            })
        };

        let (direction_str, edge_types) = if let Some(over_clause) = &go_stmt.over {
            let direction_str = match over_clause.direction {
                EdgeDirection::Out => "out",
                EdgeDirection::In => "in",
                EdgeDirection::Both => "both",
            };
            (direction_str, over_clause.edge_types.clone())
        } else {
            ("both", vec![])
        };

        // Set step_limit based on GO statement steps
        let step_limit = match go_stmt.steps {
            crate::parser::ast::Steps::Fixed(n) => n as u32,
            crate::parser::ast::Steps::Range { min: _, max } => max as u32,
            crate::parser::ast::Steps::Variable(_) => 1,
        };

        // Set column names to match ExpandAll's output format: [src, edge, dst]
        // Also add edge type name as variable for accessing edge properties
        let mut col_names = vec!["src".to_string(), "edge".to_string(), "dst".to_string()];
        if edge_types.len() == 1 {
            col_names.push(edge_types[0].clone());
        }

        // Set src_vids from FROM clause if they are literals
        let src_vids: Vec<linkrs_core::Value> = if use_start_node {
            from_vertices
                .iter()
                .filter_map(|expr| expr.as_literal())
                .collect()
        } else {
            vec![]
        };

        // Build a pure logical tree natively (GO statements have no
        // planner-level physical choices) and convert it to the physical plan
        // exactly once at the plan exit via `SubPlan::from_logical_root`.
        let expand_width = col_names.len();
        let mut logical_root = LogicalNodeEnum::ExpandAll(LogicalExpandAllNode {
            id: next_node_id(),
            deps: vec![tail_logical],
            space_id,
            edge_types: edge_types.clone(),
            direction: direction_str.to_string(),
            any_edge_type: false,
            step_limit: Some(step_limit),
            step_limits: None,
            join_input: false,
            sample: false,
            edge_props: vec![],
            vertex_props: vec![],
            filter: None,
            src_vids,
            include_empty_paths: false,
            input_var: None,
            path_semantic: None,
            dst_tag: None,
            output_var: None,
            col_names,
            column_types: projection_util::unknown_column_types(expand_width),
        });

        if let Some(ref condition) = go_stmt.where_clause {
            logical_root = LogicalNodeEnum::Filter(LogicalFilterNode {
                id: next_node_id(),
                input: Some(Box::new(logical_root)),
                condition: condition.clone(),
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            });
        }

        let project_columns = Self::build_yield_columns(go_stmt, validated.expr_context())?;
        let column_types = projection_util::project_column_types(&project_columns);
        logical_root = LogicalNodeEnum::Project(LogicalProjectNode {
            id: next_node_id(),
            input: Some(Box::new(logical_root)),
            columns: project_columns.clone(),
            subqueries: Vec::new(),
            has_folded_expressions: false,
            output_var: None,
            col_names: project_columns
                .iter()
                .map(|col| col.alias.clone())
                .collect(),
            column_types,
        });

        // `YIELD DISTINCT` and multi-step traversal both require collapsing
        // duplicate rows; a single dedup node covers either case.
        let yield_distinct = go_stmt
            .yield_clause
            .as_ref()
            .is_some_and(|clause| clause.distinct);
        if yield_distinct || step_limit > 1 {
            logical_root = LogicalNodeEnum::Dedup(LogicalDedupNode {
                id: next_node_id(),
                input: Some(Box::new(logical_root)),
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            });
        }

        Ok(SubPlan::from_logical_root(logical_root))
    }

    fn plan_bound(
        &mut self,
        ctx: &crate::planning::context::PlanContext<'_>,
    ) -> Result<SubPlan, PlannerError> {
        let bound = ctx.bound;
        let qctx = ctx.qctx.clone();
        let metadata = ctx.metadata;
        let validated = ctx.validated;
        let _ = (&bound, &qctx, &metadata, &validated);
        let go_stmt = match bound {
            BoundStatement::Go(go) => go,
            _ => {
                return Err(PlannerError::InvalidOperation(
                    "GoPlanner requires Go statements".to_string(),
                ));
            }
        };

        // GO syntax carries no vertex label: neighbor vertices are
        // materialized at execution from the edge type schema, so the plan
        // leaves the neighbor label empty instead of guessing one.
        let space_id = qctx.space_id().unwrap_or(1);

        // Handle FROM clause - extract source vertex IDs from bound expressions
        if go_stmt.from.is_empty() {
            return Err(PlannerError::PlanGenerationFailed(
                "GO statement must have FROM clause".to_string(),
            ));
        }

        let first_from = &go_stmt.from[0];
        let (use_start_node, from_var) = match first_from {
            crate::binder::bound::BoundExpression::Literal(_, _) => (true, "v".to_string()),
            crate::binder::bound::BoundExpression::Variable(name, _) => (false, name.clone()),
            _ => (false, "v".to_string()),
        };

        let tail_logical = if use_start_node {
            LogicalNodeEnum::Start(LogicalStartNode::new())
        } else {
            LogicalNodeEnum::Argument(LogicalArgumentNode {
                id: next_node_id(),
                var: from_var.clone(),
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            })
        };

        let (direction_str, edge_types) = match &go_stmt.over {
            Some(edge_types) => {
                let direction_str = match go_stmt.direction {
                    EdgeDirection::Out => "out",
                    EdgeDirection::In => "in",
                    EdgeDirection::Both => "both",
                };
                (direction_str, edge_types.clone())
            }
            None => ("both", vec![]),
        };

        let step_limit = match go_stmt.steps {
            crate::parser::ast::Steps::Fixed(n) => n as u32,
            crate::parser::ast::Steps::Range { min: _, max } => max as u32,
            crate::parser::ast::Steps::Variable(_) => 1,
        };

        let mut col_names = vec!["src".to_string(), "edge".to_string(), "dst".to_string()];
        if edge_types.len() == 1 {
            col_names.push(edge_types[0].clone());
        }

        let src_vids: Vec<linkrs_core::Value> = if use_start_node {
            go_stmt
                .from
                .iter()
                .filter_map(|expr| match expr {
                    crate::binder::bound::BoundExpression::Literal(v, _) => Some(v.clone()),
                    _ => None,
                })
                .collect()
        } else {
            vec![]
        };

        let mut logical_root = LogicalNodeEnum::ExpandAll(LogicalExpandAllNode {
            id: next_node_id(),
            deps: vec![tail_logical],
            space_id,
            edge_types: edge_types.clone(),
            direction: direction_str.to_string(),
            any_edge_type: false,
            step_limit: Some(step_limit),
            step_limits: None,
            join_input: false,
            sample: false,
            edge_props: vec![],
            vertex_props: vec![],
            filter: None,
            src_vids,
            include_empty_paths: false,
            input_var: None,
            path_semantic: None,
            dst_tag: None,
            output_var: None,
            col_names,
            column_types: vec![],
        });

        if let Some(ref where_clause) = go_stmt.where_clause {
            let condition = crate::binder::expr_converter::bound_expr_to_contextual(
                &where_clause.condition,
                validated.expr_context(),
            )
            .map_err(PlannerError::PlanGenerationFailed)?;
            logical_root = LogicalNodeEnum::Filter(LogicalFilterNode {
                id: next_node_id(),
                input: Some(Box::new(logical_root)),
                condition,
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            });
        }

        let project_columns =
            self.build_yield_columns_from_bound(go_stmt, validated.expr_context())?;
        let column_types = projection_util::project_column_types(&project_columns);
        logical_root = LogicalNodeEnum::Project(LogicalProjectNode {
            id: next_node_id(),
            input: Some(Box::new(logical_root)),
            columns: project_columns.clone(),
            subqueries: Vec::new(),
            has_folded_expressions: false,
            output_var: None,
            col_names: project_columns
                .iter()
                .map(|col| col.alias.clone())
                .collect(),
            column_types,
        });

        // `YIELD DISTINCT` and multi-step traversal both require collapsing
        // duplicate rows; a single dedup node covers either case.
        let yield_distinct = go_stmt
            .yield_clause
            .as_ref()
            .is_some_and(|clause| clause.distinct);
        if yield_distinct || step_limit > 1 {
            logical_root = LogicalNodeEnum::Dedup(LogicalDedupNode {
                id: next_node_id(),
                input: Some(Box::new(logical_root)),
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            });
        }

        Ok(SubPlan::from_logical_root(logical_root))
    }

    fn match_planner(&self, stmt: &Stmt) -> bool {
        matches!(stmt, Stmt::Go(_))
    }
}

impl GoPlanner {
    /// Create the YIELD column from AST statement
    fn build_yield_columns(
        go_stmt: &GoStmt,
        expr_context: &Arc<ExpressionAnalysisContext>,
    ) -> Result<Vec<linkrs_core::YieldColumn>, PlannerError> {
        let mut columns = Vec::new();

        if let Some(ref yield_clause) = go_stmt.yield_clause {
            for item in &yield_clause.items {
                columns.push(projection_util::yield_item_to_yield_column(item));
            }
        } else {
            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("dst".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "dst".to_string(),
            });

            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("edge".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "edge".to_string(),
            });
        }

        if columns.is_empty() {
            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("*".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "result".to_string(),
            });
        }

        Ok(columns)
    }

    /// Create YIELD columns from bound statement (plan_bound path)
    fn build_yield_columns_from_bound(
        &self,
        go_stmt: &crate::binder::bound::BoundGoStatement,
        expr_context: &Arc<ExpressionAnalysisContext>,
    ) -> Result<Vec<linkrs_core::YieldColumn>, PlannerError> {
        let mut columns = Vec::new();

        if let Some(ref yield_clause) = go_stmt.yield_clause {
            for item in &yield_clause.items {
                let column = crate::binder::expr_converter::bound_projection_to_yield_column(
                    item,
                    expr_context,
                )
                .map_err(PlannerError::PlanGenerationFailed)?;
                columns.push(column);
            }
        } else {
            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("dst".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "dst".to_string(),
            });

            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("edge".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "edge".to_string(),
            });
        }

        if columns.is_empty() {
            let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(
                linkrs_core::Expression::Variable("*".to_string()),
            );
            let id = expr_context.register_expression(expr_meta);
            let ctx_expr = ContextualExpression::new(id, expr_context.clone());
            columns.push(linkrs_core::YieldColumn {
                expression: ctx_expr,
                alias: "result".to_string(),
            });
        }

        Ok(columns)
    }
}

impl Default for GoPlanner {
    fn default() -> Self {
        Self::new()
    }
}