linkrs-query 0.1.0

//! ORDER BY Clause Planner
//!
//! Responsible for planning the execution of the ORDER BY clause and sorting the results.
//! Supports both simple column references and complex expressions (e.g., function calls).

use crate::binder::validation::CypherClauseKind;
use crate::parser::ast::Stmt;
use crate::parser::OrderByItem;
use crate::planning::plan::core::next_node_id;
use crate::planning::plan::core::nodes::base::plan_node_traits::PlanNode;
use crate::planning::plan::core::nodes::operation::sort_node::{SortItem, SortNode};
use crate::planning::plan::logical::logical_nodes::operation::LogicalSortNode;
use crate::planning::plan::logical::LogicalNodeEnum;
use crate::planning::plan::SubPlan;
use crate::planning::planner::PlannerError;
use crate::planning::statements::clauses::exists_planner;
use crate::planning::statements::plan_combiner::wrap_logical;
use crate::planning::statements::statement_planner::ClausePlanner;
use crate::QueryContext;
use linkrs_core::types::ContextualExpression;
use std::sync::Arc;

/// The ORDER BY clause planner
///
/// Responsible for planning the execution of the ORDER BY clause and sorting the results.
#[derive(Debug, Clone)]
pub struct OrderByClausePlanner {}

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

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

fn extract_order_by_items(stmt: &Stmt) -> Vec<OrderByItem> {
    if let Stmt::Match(match_stmt) = stmt {
        if let Some(order_by_clause) = &match_stmt.order_by {
            return order_by_clause.items.clone();
        }
    }
    Vec::new()
}

/// Extract the expression from a ContextualExpression.
///
/// Returns a clone of the inner Expression, or a Variable expression as fallback.
fn extract_expression(expr: &ContextualExpression) -> linkrs_core::Expression {
    if let Some(expr_meta) = expr.expression() {
        expr_meta.inner().clone()
    } else {
        // Fallback: create a variable expression from the string representation
        let expr_string = expr
            .expression()
            .map(|e| e.inner().to_expression_string())
            .unwrap_or_default();
        linkrs_core::Expression::Variable(expr_string)
    }
}

impl ClausePlanner for OrderByClausePlanner {
    fn clause_kind(&self) -> CypherClauseKind {
        CypherClauseKind::OrderBy
    }

    fn transform_clause(
        &self,
        qctx: Arc<QueryContext>,
        stmt: &Stmt,
        input_plan: SubPlan,
    ) -> Result<SubPlan, PlannerError> {
        let order_by_items = extract_order_by_items(stmt);

        if order_by_items.is_empty() {
            return Ok(input_plan);
        }

        let input_node = input_plan.root().as_ref().ok_or_else(|| {
            PlannerError::PlanGenerationFailed(
                "The ORDER BY clause requires an input plan".to_string(),
            )
        })?;

        // Unified entry for expression-level EXISTS / IN: subqueries in ORDER
        // BY expressions are rejected at planning time with a precise error.
        let space_id = qctx.space_id().unwrap_or(1);
        let space_name = qctx.space_name().unwrap_or_else(|| "default".to_string());
        let outer_col_names = input_plan
            .root()
            .as_ref()
            .map(|root| root.col_names().to_vec())
            .unwrap_or_default();
        for item in &order_by_items {
            if let Some(expr_meta) = item.expression.expression() {
                exists_planner::check_expression_subqueries(
                    expr_meta.inner(),
                    &qctx,
                    space_id,
                    &space_name,
                    &outer_col_names,
                )?;
            }
        }

        let sort_items: Vec<SortItem> = order_by_items
            .into_iter()
            .map(|item| {
                let expression = extract_expression(&item.expression);
                SortItem::new(expression, item.direction)
            })
            .collect();

        let sort_node = SortNode::new(input_node.clone(), sort_items.clone())?;
        let logical_root = wrap_logical(&input_plan, |input| {
            LogicalNodeEnum::Sort(LogicalSortNode {
                id: next_node_id(),
                input: Some(Box::new(input)),
                sort_items,
                limit: None,
                output_var: None,
                col_names: vec![],
                column_types: vec![],
            })
        });
        Ok(SubPlan {
            root: Some(sort_node.into_enum()),
            tail: input_plan.tail,
            logical_root,
        })
    }
}

#[cfg(test)]
#[allow(clippy::arc_with_non_send_sync)]
mod tests {
    use super::*;
    use crate::parser::ast::{OrderByItem, 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::expr::FunctionArg;
    use linkrs_core::types::OrderDirection;
    use linkrs_core::Expression;
    use std::sync::Arc;

    #[test]
    fn test_order_by_clause_planner_creation() {
        let planner = OrderByClausePlanner::new();
        assert_eq!(planner.clause_kind(), CypherClauseKind::OrderBy);
    }

    #[test]
    fn test_extract_order_by_items() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("age".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 match_stmt = Stmt::Match(crate::parser::ast::stmt::MatchStmt {
            span: Span::default(),
            patterns: vec![],
            where_clause: None,
            return_clause: None,
            order_by: Some(crate::parser::ast::stmt::OrderByClause {
                span: Span::default(),
                items: vec![OrderByItem {
                    expression: ctx_expr.clone(),
                    direction: OrderDirection::Asc,
                }],
            }),
            limit: None,
            skip: None,
            join_hint: None,
            optional: false,
            delete_clause: None,
        });

        let items = extract_order_by_items(&match_stmt);
        assert_eq!(items.len(), 1);
        assert_eq!(items[0].direction, OrderDirection::Asc);
    }

    #[test]
    fn test_extract_order_by_items_empty() {
        let match_stmt = Stmt::Match(crate::parser::ast::stmt::MatchStmt {
            span: Span::default(),
            patterns: vec![],
            where_clause: None,
            return_clause: None,
            order_by: None,
            limit: None,
            skip: None,
            join_hint: None,
            optional: false,
            delete_clause: None,
        });

        let items = extract_order_by_items(&match_stmt);
        assert_eq!(items.len(), 0);
    }

    #[test]
    fn test_extract_expression() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("age".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 result = extract_expression(&ctx_expr);
        assert_eq!(result, Expression::Variable("age".to_string()));
    }

    #[test]
    fn test_extract_expression_complex() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Property {
            object: Box::new(Expression::Variable("n".to_string())),
            property: "name".to_string(),
        };
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr.clone());
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = linkrs_core::types::ContextualExpression::new(id, ctx);

        let result = extract_expression(&ctx_expr);
        assert_eq!(result, expr);
    }

    #[test]
    fn test_extract_expression_function_call() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Function {
            name: "cosine_similarity".to_string(),
            args: vec![
                FunctionArg::Positional(Expression::Variable("a".to_string())),
                FunctionArg::Positional(Expression::Variable("b".to_string())),
            ],
        };
        let expr_meta = linkrs_core::types::expr::ExpressionMeta::new(expr.clone());
        let id = ctx.register_expression(expr_meta);
        let ctx_expr = linkrs_core::types::ContextualExpression::new(id, ctx);

        let result = extract_expression(&ctx_expr);
        assert_eq!(result, expr);
    }

    #[test]
    fn test_transform_clause() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("age".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 match_stmt = Stmt::Match(crate::parser::ast::stmt::MatchStmt {
            span: Span::default(),
            patterns: vec![],
            where_clause: None,
            return_clause: None,
            order_by: Some(crate::parser::ast::stmt::OrderByClause {
                span: Span::default(),
                items: vec![OrderByItem {
                    expression: ctx_expr.clone(),
                    direction: OrderDirection::Asc,
                }],
            }),
            limit: None,
            skip: None,
            join_hint: None,
            optional: false,
            delete_clause: 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 = OrderByClausePlanner::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, &match_stmt, input_plan);
        assert!(result.is_ok());

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

        if let Some(PlanNodeEnum::Sort(_)) = sub_plan.root {
        } else {
            panic!("Expected SortNode");
        }
    }

    #[test]
    fn test_transform_clause_empty_order_by() {
        let match_stmt = Stmt::Match(crate::parser::ast::stmt::MatchStmt {
            span: Span::default(),
            patterns: vec![],
            where_clause: None,
            return_clause: None,
            order_by: None,
            limit: None,
            skip: None,
            join_hint: None,
            optional: false,
            delete_clause: 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 = OrderByClausePlanner::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, &match_stmt, input_plan);
        assert!(result.is_ok());

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

    #[test]
    fn test_transform_clause_empty_input_plan() {
        let ctx = Arc::new(ExpressionAnalysisContext::new());
        let expr = Expression::Variable("age".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 match_stmt = Stmt::Match(crate::parser::ast::stmt::MatchStmt {
            span: Span::default(),
            patterns: vec![],
            where_clause: None,
            return_clause: None,
            order_by: Some(crate::parser::ast::stmt::OrderByClause {
                span: Span::default(),
                items: vec![OrderByItem {
                    expression: ctx_expr.clone(),
                    direction: OrderDirection::Asc,
                }],
            }),
            limit: None,
            skip: None,
            join_hint: None,
            optional: false,
            delete_clause: None,
        });

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

        let planner = OrderByClausePlanner::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, &match_stmt, input_plan);
        assert!(result.is_err());
    }
}