linkrs-query 0.1.0

//! The YIELD clause planner
//!
//! Responsible for converting the YIELD clause into an execution plan node.
//! Support for the YIELD ... WHERE ... syntax

use crate::binder::validation::CypherClauseKind;
use crate::parser::ast::Stmt;
use crate::planning::plan::core::next_node_id;
use crate::planning::plan::core::nodes::graph_operations::graph_operations_node::DedupNode;
use crate::planning::plan::core::nodes::operation::sample_node::SampleNode;
use crate::planning::plan::core::nodes::{FilterNode, LimitNode, PlanNodeEnum, ProjectNode};
use crate::planning::plan::logical::logical_nodes::operation::{
    LogicalDedupNode, LogicalFilterNode, LogicalLimitNode, LogicalProjectNode, LogicalSampleNode,
};
use crate::planning::plan::logical::LogicalNodeEnum;
use crate::planning::plan::SubPlan;
use crate::planning::planner::PlannerError;
use crate::planning::statements::plan_combiner::wrap_logical;
use crate::planning::statements::statement_planner::ClausePlanner;
use crate::QueryContext;
use linkrs_core::YieldColumn;
use std::sync::Arc;

/// Yield info result type alias
type YieldInfoResult = Result<
    (
        Vec<YieldColumn>,
        Option<linkrs_core::types::ContextualExpression>,
        Option<usize>,
        Option<usize>,
        Option<usize>,
    ),
    PlannerError,
>;

/// The YIELD clause planner
#[derive(Debug)]
pub struct YieldClausePlanner {}

impl YieldClausePlanner {
    pub fn new() -> Self {
        Self {}
    }

    /// Planning the YIELD clause
    ///
    /// Processing procedure:
    /// 1. Constructing the projection node (YIELD column)
    /// 2. If there are WHERE conditions, add a Filter node.
    /// 3. If there are LIMIT/SKIP settings, add pagination nodes.
    pub fn plan_yield_clause(
        &self,
        yield_columns: &[YieldColumn],
        filter_condition: Option<linkrs_core::types::ContextualExpression>,
        skip: Option<usize>,
        limit: Option<usize>,
        sample: Option<usize>,
        distinct: bool,
        input_plan: &SubPlan,
    ) -> Result<SubPlan, PlannerError> {
        let mut current_plan = input_plan.clone();

        // 1. Build the projection node (if there is a specific YIELD column).
        if !yield_columns.is_empty() {
            let project_node = self.create_project_node(&current_plan, yield_columns)?;
            let logical_root = wrap_logical(&current_plan, |input| {
                LogicalNodeEnum::Project(LogicalProjectNode {
                    id: next_node_id(),
                    input: Some(Box::new(input)),
                    columns: yield_columns.to_vec(),
                    subqueries: Vec::new(),
                    has_folded_expressions: false,
                    output_var: None,
                    col_names: project_node.col_names().to_vec(),
                    column_types: vec![],
                })
            });
            current_plan = SubPlan {
                root: Some(project_node),
                tail: current_plan.tail.clone(),
                logical_root,
            };
        }

        // 2. If there are WHERE conditions, add a Filter node.
        if let Some(ref filter_condition) = filter_condition {
            let filter_node = self.create_filter_node(&current_plan, filter_condition.clone())?;
            let logical_root = wrap_logical(&current_plan, |input| {
                LogicalNodeEnum::Filter(LogicalFilterNode {
                    id: next_node_id(),
                    input: Some(Box::new(input)),
                    condition: filter_condition.clone(),
                    output_var: None,
                    col_names: filter_node.col_names().to_vec(),
                    column_types: vec![],
                })
            });
            current_plan = SubPlan {
                root: Some(filter_node),
                tail: current_plan.tail.clone(),
                logical_root,
            };
        }

        // 2.5 Apply DISTINCT if requested.
        if distinct {
            let input_node = current_plan.root().as_ref().ok_or_else(|| {
                PlannerError::PlanGenerationFailed("The input plan has no root node".to_string())
            })?;
            if let Ok(dedup_node) = DedupNode::new(input_node.clone()) {
                let logical_root = wrap_logical(&current_plan, |input| {
                    LogicalNodeEnum::Dedup(LogicalDedupNode {
                        id: next_node_id(),
                        input: Some(Box::new(input)),
                        output_var: None,
                        col_names: vec![],
                        column_types: vec![],
                    })
                });
                current_plan = SubPlan {
                    root: Some(PlanNodeEnum::Dedup(dedup_node)),
                    tail: current_plan.tail.clone(),
                    logical_root,
                };
            }
        }

        // 3. Apply SAMPLE clause if present (before pagination)
        if let Some(sample_count) = sample {
            let input_node = current_plan.root().as_ref().ok_or_else(|| {
                PlannerError::PlanGenerationFailed("The input plan has no root node".to_string())
            })?;
            let sample_node =
                SampleNode::new(input_node.clone(), sample_count as i64).map_err(|e| {
                    PlannerError::PlanGenerationFailed(format!(
                        "Failed to create SampleNode: {}",
                        e
                    ))
                })?;
            let logical_root = wrap_logical(&current_plan, |input| {
                LogicalNodeEnum::Sample(LogicalSampleNode {
                    id: next_node_id(),
                    input: Some(Box::new(input)),
                    count: sample_count as i64,
                    output_var: None,
                    col_names: vec![],
                    column_types: vec![],
                })
            });
            current_plan = SubPlan {
                root: Some(PlanNodeEnum::Sample(sample_node)),
                tail: current_plan.tail.clone(),
                logical_root,
            };
        }

        // 4. Handling pagination (LIMIT/SKIP)
        if limit.is_some() || skip.is_some() {
            current_plan = self.apply_pagination(current_plan, skip, limit)?;
        }

        Ok(current_plan)
    }

    /// Create a projection node.
    fn create_project_node(
        &self,
        input_plan: &SubPlan,
        columns: &[YieldColumn],
    ) -> Result<PlanNodeEnum, PlannerError> {
        let input_node = input_plan.root().as_ref().ok_or_else(|| {
            PlannerError::PlanGenerationFailed("The input plan has no root node".to_string())
        })?;

        ProjectNode::new(input_node.clone(), columns.to_vec())
            .map_err(|e| {
                PlannerError::PlanGenerationFailed(format!(
                    "Failed to create projection node: {}",
                    e
                ))
            })
            .map(PlanNodeEnum::Project)
    }

    /// Create a filter node.
    fn create_filter_node(
        &self,
        input_plan: &SubPlan,
        condition: linkrs_core::types::ContextualExpression,
    ) -> Result<PlanNodeEnum, PlannerError> {
        let input_node = input_plan.root().as_ref().ok_or_else(|| {
            PlannerError::PlanGenerationFailed("The input plan has no root node".to_string())
        })?;

        FilterNode::new(input_node.clone(), condition)
            .map_err(|e| {
                PlannerError::PlanGenerationFailed(format!("Failed to create filter node: {}", e))
            })
            .map(PlanNodeEnum::Filter)
    }

    /// Application pagination
    fn apply_pagination(
        &self,
        input_plan: SubPlan,
        skip: Option<usize>,
        limit: Option<usize>,
    ) -> Result<SubPlan, PlannerError> {
        let input_node = input_plan.root().as_ref().ok_or_else(|| {
            PlannerError::PlanGenerationFailed("The input plan has no root node".to_string())
        })?;

        let offset = skip.unwrap_or(0) as i64;
        let count = limit.map(|l| l as i64).unwrap_or(i64::MAX);

        let limit_node = LimitNode::new(input_node.clone(), offset, count).map_err(|e| {
            PlannerError::PlanGenerationFailed(format!("Failed to create paging node: {}", e))
        })?;
        let logical_root = wrap_logical(&input_plan, |input| {
            LogicalNodeEnum::Limit(LogicalLimitNode {
                id: next_node_id(),
                input: Some(Box::new(input)),
                offset,
                count,
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            })
        });

        Ok(SubPlan {
            root: Some(PlanNodeEnum::Limit(limit_node)),
            tail: input_plan.tail.clone(),
            logical_root,
        })
    }
}

impl ClausePlanner for YieldClausePlanner {
    fn clause_kind(&self) -> CypherClauseKind {
        CypherClauseKind::Yield
    }

    fn transform_clause(
        &self,
        _qctx: Arc<QueryContext>,
        stmt: &Stmt,
        input_plan: SubPlan,
    ) -> Result<SubPlan, PlannerError> {
        // Extract the information from the YIELD clause in the sentence.
        let (yield_columns, filter_condition, skip, limit, sample) =
            Self::extract_yield_info(stmt)?;

        let distinct = match stmt {
            Stmt::Yield(y) => y.distinct,
            _ => false,
        };

        // Ensure standalone YIELD skip/limit/sample are extracted from the
        // statement itself (not only from an inner YieldClause).
        let (skip, limit, sample) = match stmt {
            Stmt::Yield(y) => (
                y.skip.as_ref().map(|s| s.count).or(skip),
                y.limit.as_ref().map(|l| l.count).or(limit),
                y.sample.as_ref().map(|s| s.count).or(sample),
            ),
            _ => (skip, limit, sample),
        };

        self.plan_yield_clause(
            &yield_columns,
            filter_condition,
            skip,
            limit,
            sample,
            distinct,
            &input_plan,
        )
    }
}

impl YieldClausePlanner {
    /// Extract the information from the YIELD clause in the sentence.
    ///
    /// The improved implementation includes:
    /// Support for the YIELD clause in various statement types
    /// The complete conversion from YieldItem to YieldColumn
    /// Aggregation expression detection
    /// Alias handling
    fn extract_yield_info(stmt: &Stmt) -> YieldInfoResult {
        use crate::parser::ast::Stmt;

        // The “YIELD” keyword can appear either as an independent statement or as a clause within other statements.
        match stmt {
            Stmt::Yield(yield_stmt) => {
                let yield_columns = Self::convert_yield_items(&yield_stmt.items)?;
                let skip = yield_stmt.skip.as_ref().map(|s| s.count);
                let limit = yield_stmt.limit.as_ref().map(|l| l.count);
                let sample = yield_stmt.sample.as_ref().map(|s| s.count);
                Ok((
                    yield_columns,
                    yield_stmt.where_clause.clone(),
                    skip,
                    limit,
                    sample,
                ))
            }
            Stmt::Go(go_stmt) => {
                // Extract the YIELD clause from the GO statement.
                if let Some(ref yield_clause) = go_stmt.yield_clause {
                    let yield_columns = Self::convert_yield_items(&yield_clause.items)?;
                    let skip = yield_clause.skip.as_ref().map(|s| s.count);
                    let limit = yield_clause.limit.as_ref().map(|l| l.count);
                    let sample = yield_clause.sample.as_ref().map(|s| s.count);
                    Ok((
                        yield_columns,
                        yield_clause.where_clause.clone(),
                        skip,
                        limit,
                        sample,
                    ))
                } else {
                    Ok((vec![], None, None, None, None))
                }
            }
            Stmt::Fetch(_fetch_stmt) => {
                // The FETCH statement may have an implicit YIELD.
                Ok((vec![], None, None, None, None))
            }
            _ => {
                // Other statement types are not currently supported for extracting the YIELD value.
                Ok((vec![], None, None, None, None))
            }
        }
    }

    /// Convert the YieldItem list to the YieldColumn list.
    fn convert_yield_items(
        items: &[crate::parser::ast::stmt::YieldItem],
    ) -> Result<Vec<YieldColumn>, PlannerError> {
        crate::planning::statements::projection_util::yield_items_to_columns(items)
    }
}

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

#[cfg(test)]
#[allow(clippy::arc_with_non_send_sync)]
mod tests {
    use super::*;
    use crate::parser::ast::Span;
    use crate::planning::plan::core::nodes::StartNode;
    use crate::planning::plan::core::PlanNodeEnum;
    use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
    use linkrs_core::types::ContextualExpression;
    use linkrs_core::Expression;
    use std::sync::Arc;

    #[test]
    fn test_yield_clause_planner_creation() {
        let planner = YieldClausePlanner::new();
        assert_eq!(planner.clause_kind(), CypherClauseKind::Yield);
    }

    #[test]
    fn test_extract_yield_info_from_yield_stmt() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("n".to_string());
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = ContextualExpression::new(id, ctx);

        let yield_stmt = Stmt::Yield(crate::parser::ast::stmt::YieldStmt {
            span: Span::default(),
            items: vec![crate::parser::ast::stmt::YieldItem {
                expression: ctx_expr.clone(),
                alias: None,
            }],
            where_clause: None,
            distinct: false,
            order_by: None,
            skip: None,
            limit: None,
            sample: None,
        });

        let (columns, filter, skip, limit, _sample) =
            YieldClausePlanner::extract_yield_info(&yield_stmt).expect("extract failed");
        assert_eq!(columns.len(), 1);
        assert_eq!(columns[0].alias, "n");
        assert!(filter.is_none());
        assert!(skip.is_none());
        assert!(limit.is_none());
    }

    #[test]
    fn test_extract_yield_info_from_go_stmt() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("n".to_string());
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = ContextualExpression::new(id, ctx);

        let go_stmt = Stmt::Go(crate::parser::ast::stmt::GoStmt {
            span: Span::default(),
            steps: crate::parser::ast::Steps::Fixed(1),
            from: crate::parser::ast::stmt::FromClause {
                span: Span::default(),
                vertices: vec![],
            },
            over: None,
            where_clause: None,
            yield_clause: Some(crate::parser::ast::stmt::YieldClause {
                span: Span::default(),
                distinct: false,
                items: vec![crate::parser::ast::stmt::YieldItem {
                    expression: ctx_expr.clone(),
                    alias: None,
                }],
                where_clause: None,
                order_by: None,
                limit: Some(crate::parser::ast::types::LimitClause {
                    span: Span::default(),
                    count: 10,
                }),
                skip: Some(crate::parser::ast::types::SkipClause {
                    span: Span::default(),
                    count: 5,
                }),
                sample: None,
            }),
        });

        let (columns, filter, skip, limit, _sample) =
            YieldClausePlanner::extract_yield_info(&go_stmt).expect("extract failed");
        assert_eq!(columns.len(), 1);
        assert!(filter.is_none());
        assert_eq!(skip, Some(5));
        assert_eq!(limit, Some(10));
    }

    #[test]
    fn test_convert_yield_items() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("n".to_string());
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = ContextualExpression::new(id, ctx);

        let items = vec![crate::parser::ast::stmt::YieldItem {
            expression: ctx_expr.clone(),
            alias: Some("node".to_string()),
        }];

        let yield_columns =
            YieldClausePlanner::convert_yield_items(&items).expect("convert failed");
        assert_eq!(yield_columns.len(), 1);
        assert_eq!(yield_columns[0].alias, "node");
    }

    #[test]
    fn test_generate_default_alias() {
        use linkrs_core::types::expr::ExpressionMeta;
        fn ctx_of(expr: Expression) -> ContextualExpression {
            let ctx = Arc::new(ExpressionAnalysisContext::new());
            let id = ctx.register_expression(ExpressionMeta::new(expr));
            ContextualExpression::new(id, ctx)
        }
        let alias = crate::planning::statements::projection_util::default_projection_alias(
            &ctx_of(Expression::Variable("n".to_string())),
        );
        assert_eq!(alias, "n");

        let alias = crate::planning::statements::projection_util::default_projection_alias(
            &ctx_of(Expression::Property {
                object: Box::new(Expression::Variable("n".to_string())),
                property: "name".to_string(),
            }),
        );
        assert_eq!(alias, "n.name");

        let alias = crate::planning::statements::projection_util::default_projection_alias(
            &ctx_of(Expression::Function {
                name: "count".to_string(),
                args: vec![],
            }),
        );
        assert_eq!(alias, "count");
    }

    #[test]
    fn test_transform_clause() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("n".to_string());
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = ContextualExpression::new(id, ctx);

        let yield_stmt = Stmt::Yield(crate::parser::ast::stmt::YieldStmt {
            span: Span::default(),
            items: vec![crate::parser::ast::stmt::YieldItem {
                expression: ctx_expr.clone(),
                alias: None,
            }],
            where_clause: None,
            distinct: false,
            order_by: None,
            skip: None,
            limit: None,
            sample: None,
        });

        let start_node = StartNode::new();
        let start_node_enum = PlanNodeEnum::Start(start_node.clone());
        let input_plan = SubPlan {
            root: Some(start_node_enum.clone()),
            tail: Some(start_node_enum),
            logical_root: None,
        };

        let planner = YieldClausePlanner::new();
        let qctx = Arc::new(crate::QueryContext::new(Arc::new(
            crate::QueryRequestContext {
                session_id: None,
                user_name: None,
                space_name: None,
                query: String::new(),
                parameters: std::collections::HashMap::new(),
                ..Default::default()
            },
        )));

        let result = planner.transform_clause(qctx, &yield_stmt, input_plan);
        assert!(result.is_ok());

        let sub_plan = result.expect("transform_clause should succeed");
        assert!(sub_plan.root.is_some());

        match sub_plan.root {
            Some(PlanNodeEnum::Project(_)) => {}
            Some(PlanNodeEnum::Filter(_)) => {}
            Some(PlanNodeEnum::Limit(_)) => {}
            _ => panic!("Expected ProjectNode, FilterNode, or LimitNode"),
        }
    }

    #[test]
    fn test_transform_clause_empty_input_plan() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("n".to_string());
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = linkrs_core::types::ContextualExpression::new(id, ctx);

        let yield_stmt = Stmt::Yield(crate::parser::ast::stmt::YieldStmt {
            span: Span::default(),
            items: vec![crate::parser::ast::stmt::YieldItem {
                expression: ctx_expr.clone(),
                alias: None,
            }],
            where_clause: None,
            distinct: false,
            order_by: None,
            skip: None,
            limit: None,
            sample: None,
        });

        let input_plan = SubPlan {
            root: None,
            tail: None,
            logical_root: None,
        };

        let planner = YieldClausePlanner::new();
        let qctx = Arc::new(crate::QueryContext::new(Arc::new(
            crate::QueryRequestContext {
                session_id: None,
                user_name: None,
                space_name: None,
                query: String::new(),
                parameters: std::collections::HashMap::new(),
                ..Default::default()
            },
        )));

        let result = planner.transform_clause(qctx, &yield_stmt, input_plan);
        assert!(result.is_err());
    }
}