linkrs-query 0.1.0

//! The FETCH VERTICES query planner
//! Planning for the FETCH VERTICES query

use crate::binder::BoundStatement;
use crate::parser::ast::{FetchTarget, Stmt};
use crate::planning::plan::core::common::TagProp;
use crate::planning::plan::core::node_id_generator::next_node_id;
use crate::planning::plan::core::nodes::{
    ArgumentNode, GetVerticesNode, PlanNodeEnum, ProjectNode,
};
use crate::planning::plan::logical::logical_nodes::access::LogicalGetVerticesNode;
use crate::planning::plan::logical::logical_nodes::operation::LogicalProjectNode;
use crate::planning::plan::logical::LogicalNodeEnum;
use crate::planning::plan::SubPlan;
use crate::planning::planner::{Planner, PlannerError, ValidatedStatement};
use crate::planning::statements::plan_combiner::logical_argument_root;
use crate::planning::statements::projection_util;
use crate::QueryContext;
use std::sync::Arc;

/// The FETCH VERTICES query planner
/// Responsible for converting the FETCH VERTICES query into an execution plan.
#[derive(Debug, Clone)]
pub struct FetchVerticesPlanner;

impl FetchVerticesPlanner {
    /// Create a new FETCH VERTICES planner.
    pub fn new() -> Self {
        Self
    }
}

impl Planner for FetchVerticesPlanner {
    fn transform(
        &mut self,
        validated: &ValidatedStatement,
        qctx: Arc<QueryContext>,
    ) -> Result<SubPlan, PlannerError> {
        let fetch_stmt = match validated.stmt() {
            Stmt::Fetch(fetch_stmt) => fetch_stmt,
            _ => {
                return Err(PlannerError::InvalidOperation(
                    "FetchVerticesPlanner requires the Fetch statement.".to_string(),
                ));
            }
        };

        // Check whether it is a FETCH VERTICES operation.
        let (tag_name, ids, properties) = match &fetch_stmt.target {
            FetchTarget::Vertices {
                tag_name,
                ids,
                properties,
                ..
            } => (tag_name, ids, properties),
            _ => {
                return Err(PlannerError::InvalidOperation(
                    "FetchVerticesPlanner requires the FETCH VERTICES statement.".to_string(),
                ));
            }
        };

        // Get space_id and space_name from query context
        let space_id = qctx.space_id().unwrap_or(0);
        let space_name = qctx.space_name().unwrap_or_else(|| "default".to_string());

        // Convert ids to comma-separated string for GetVerticesNode
        let ids_str = ids
            .iter()
            .map(|ctx_expr| {
                if let Some(expr_meta) = ctx_expr.expression() {
                    let expr = expr_meta.inner();
                    match expr {
                        linkrs_core::Expression::Literal(linkrs_core::Value::Int(i)) => {
                            i.to_string()
                        }
                        linkrs_core::Expression::Literal(linkrs_core::Value::BigInt(i)) => {
                            i.to_string()
                        }
                        linkrs_core::Expression::Literal(linkrs_core::Value::String(s)) => {
                            s.to_string()
                        }
                        _ => expr.to_string(),
                    }
                } else {
                    String::new()
                }
            })
            .collect::<Vec<_>>()
            .join(",");

        let var_name = "v";
        let entity_var = tag_name.clone().unwrap_or_else(|| "vertex".to_string());

        // 1. Create a parameter node to provide the vertex IDs.
        let mut arg_node = ArgumentNode::new(next_node_id(), var_name);
        arg_node.set_col_names(vec!["vid".to_string()]);
        arg_node.set_output_var("vertex_ids".to_string());

        let arg_node_enum = PlanNodeEnum::Argument(arg_node.clone());

        // 2. Create GetVerticesNode with proper space_id and space_name
        // Use actual ids_str instead of var_name to pass vertex IDs directly
        let mut get_vertices_node = GetVerticesNode::new(space_id, &space_name, &ids_str);
        get_vertices_node.add_dependency(arg_node_enum.clone());
        get_vertices_node.set_output_var("fetched_vertices".to_string());
        get_vertices_node.set_col_names(vec![entity_var.clone()]);
        let fetch_tag = tag_name.clone().ok_or_else(|| {
            PlannerError::InvalidOperation("FETCH VERTICES requires a tag qualifier".to_string())
        })?;
        get_vertices_node.set_tag(fetch_tag);

        // Set the tag attributes (obtained from the properties field)
        let tag_props = if let Some(props) = properties {
            vec![TagProp::new("default", props.clone())]
        } else {
            vec![]
        };
        get_vertices_node.set_tag_props(tag_props);

        let mut current_logical = get_vertices_mirror(&get_vertices_node, logical_argument());

        let get_vertices_node_enum = PlanNodeEnum::GetVertices(get_vertices_node);

        // 3. Apply the YIELD clause as a projection when present; without one
        // the fetched vertices are returned directly.
        let root = if let Some(ref yield_clause) = fetch_stmt.yield_clause {
            let columns: Vec<linkrs_core::YieldColumn> = yield_clause
                .items
                .iter()
                .map(projection_util::yield_item_to_yield_column)
                .collect();
            let project_node = ProjectNode::new(get_vertices_node_enum.clone(), columns.clone())?;
            let root = PlanNodeEnum::Project(project_node);
            let column_types = projection_util::project_column_types(&columns);
            current_logical = LogicalNodeEnum::Project(LogicalProjectNode {
                id: next_node_id(),
                input: Some(Box::new(current_logical)),
                columns,
                subqueries: Vec::new(),
                has_folded_expressions: false,
                output_var: None,
                col_names: root.col_names().to_vec(),
                column_types,
            });
            root
        } else {
            get_vertices_node_enum.clone()
        };

        // 4. Create a SubPlan
        let sub_plan = SubPlan {
            root: Some(root),
            tail: Some(arg_node_enum),
            logical_root: Some(current_logical),
        };

        Ok(sub_plan)
    }

    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 fetch = match bound {
            BoundStatement::FetchVertices(f) => f,
            _ => {
                return Err(PlannerError::InvalidOperation(
                    "FetchVerticesPlanner requires the FetchVertices statement.".to_string(),
                ));
            }
        };

        let space_id = qctx.space_id().unwrap_or(0);
        let space_name = qctx.space_name().unwrap_or_else(|| "default".to_string());

        let expr_ctx = std::sync::Arc::new(
            linkrs_core::types::expr::expression_context::ExpressionAnalysisContext::new(),
        );

        let ids_str = fetch
            .ids
            .iter()
            .filter_map(|id| {
                let ctx_expr =
                    crate::binder::expr_converter::bound_expr_to_contextual(id, &expr_ctx).ok()?;
                let meta = ctx_expr.expression()?;
                let expr = meta.inner();
                Some(match expr {
                    linkrs_core::Expression::Literal(linkrs_core::Value::Int(i)) => i.to_string(),
                    linkrs_core::Expression::Literal(linkrs_core::Value::BigInt(i)) => {
                        i.to_string()
                    }
                    linkrs_core::Expression::Literal(linkrs_core::Value::String(s)) => {
                        s.to_string()
                    }
                    _ => expr.to_string(),
                })
            })
            .collect::<Vec<_>>()
            .join(",");

        let var_name = "v";
        let entity_var = fetch
            .tag_name
            .clone()
            .unwrap_or_else(|| "vertex".to_string());

        let mut arg_node = ArgumentNode::new(next_node_id(), var_name);
        arg_node.set_col_names(vec!["vid".to_string()]);
        arg_node.set_output_var("vertex_ids".to_string());
        let arg_node_enum = PlanNodeEnum::Argument(arg_node.clone());

        let mut get_vertices_node = GetVerticesNode::new(space_id, &space_name, &ids_str);
        get_vertices_node.add_dependency(arg_node_enum.clone());
        get_vertices_node.set_output_var("fetched_vertices".to_string());
        get_vertices_node.set_col_names(vec![entity_var.clone()]);
        let fetch_tag = fetch.tag_name.clone().ok_or_else(|| {
            PlannerError::InvalidOperation("FETCH VERTICES requires a tag qualifier".to_string())
        })?;
        get_vertices_node.set_tag(fetch_tag);

        let tag_props = if let Some(ref props) = fetch.properties {
            vec![TagProp::new("default", props.clone())]
        } else {
            vec![]
        };
        get_vertices_node.set_tag_props(tag_props);

        let current_logical = get_vertices_mirror(&get_vertices_node, logical_argument());

        let root = PlanNodeEnum::GetVertices(get_vertices_node);

        let sub_plan = SubPlan {
            root: Some(root),
            tail: Some(arg_node_enum),
            logical_root: Some(current_logical),
        };
        Ok(sub_plan)
    }

    fn match_planner(&self, stmt: &Stmt) -> bool {
        match stmt {
            Stmt::Fetch(fetch_stmt) => {
                matches!(&fetch_stmt.target, FetchTarget::Vertices { .. })
            }
            _ => false,
        }
    }
}

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

/// Seed logical argument matching the physical `v` argument node.
fn logical_argument() -> LogicalNodeEnum {
    logical_argument_root("v", vec!["vid".to_string()], Some("vertex_ids".to_string()))
}

/// Mirror a physical fetch-vertices lookup over the logical argument seed.
fn get_vertices_mirror(
    get_vertices_node: &GetVerticesNode,
    arg_logical: LogicalNodeEnum,
) -> LogicalNodeEnum {
    LogicalNodeEnum::GetVertices(LogicalGetVerticesNode {
        id: next_node_id(),
        deps: vec![arg_logical],
        space_id: get_vertices_node.space_id(),
        space_name: get_vertices_node.space_name().to_string(),
        tag: get_vertices_node.tag().map(|s| s.to_string()),
        src_ref: get_vertices_node.src_ref().clone(),
        src_vids: get_vertices_node.src_vids().to_string(),
        tag_props: get_vertices_node.tag_props().to_vec(),
        expression: get_vertices_node.filter().cloned(),
        dedup: get_vertices_node.dedup(),
        limit: get_vertices_node.limit(),
        projected_properties: get_vertices_node.projected_properties().to_vec(),
        output_var: get_vertices_node.output_var().map(|s| s.to_string()),
        col_names: get_vertices_node.col_names().to_vec(),
        column_types: projection_util::unknown_column_types(get_vertices_node.col_names().len()),
    })
}