featherdb_query/optimizer/

mod.rs

//! Query optimizer

pub mod cost;
mod helpers;
mod join_order;
mod rule;
mod rules;

use crate::planner::LogicalPlan;
use featherdb_catalog::Catalog;
use rule::OptimizationRule;
use rules::{
    ConstantFolding, IndexSelection, PkSeekRule, PredicatePushdown, ProjectionPushdown,
    SubqueryToJoinConversion,
};
use std::sync::Arc;

// Re-export public types
pub use cost::{constants as cost_constants, CostEstimator};

/// Query optimizer that applies transformation rules
pub struct Optimizer {
    rules: Vec<Box<dyn OptimizationRule>>,
    /// Catalog reference for index selection
    #[allow(dead_code)]
    catalog: Option<Arc<Catalog>>,
}

impl Optimizer {
    /// Create a new optimizer with default rules
    pub fn new() -> Self {
        Optimizer {
            rules: vec![
                Box::new(PredicatePushdown),
                Box::new(PkSeekRule::new()),
                Box::new(ProjectionPushdown),
                Box::new(ConstantFolding),
            ],
            catalog: None,
        }
    }

    /// Create an optimizer with catalog access for index selection and join ordering
    pub fn with_catalog(catalog: Arc<Catalog>) -> Self {
        Optimizer {
            rules: vec![
                // SubqueryToJoinConversion should run early to convert subqueries to joins
                Box::new(SubqueryToJoinConversion),
                // PredicatePushdown before JoinOrder so that single-table filters
                // are already on scans when we estimate relation cardinalities
                Box::new(PredicatePushdown),
                // JoinOrderOptimizer reorders joins using cardinality estimates
                Box::new(join_order::JoinOrderOptimizer::new(catalog.clone())),
                // Run PredicatePushdown again to push any new predicates from reordered joins
                Box::new(PredicatePushdown),
                // PkSeekRule after PredicatePushdown, before IndexSelection
                Box::new(PkSeekRule::new()),
                // IndexSelection converts Scan + Filter into IndexScan when beneficial
                Box::new(IndexSelection::new(catalog.clone())),
                Box::new(ProjectionPushdown),
                Box::new(ConstantFolding),
            ],
            catalog: Some(catalog),
        }
    }

    /// Optimize a logical plan
    pub fn optimize(&self, plan: LogicalPlan) -> LogicalPlan {
        let mut current = plan;

        // Apply each rule
        for rule in &self.rules {
            current = rule.apply(current);
        }

        current
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::expr::{BinaryOp, Expr};
    use featherdb_catalog::{Catalog, ColumnConstraint, Index, TableBuilder};
    use featherdb_core::{ColumnType, PageId, Value};

    fn create_test_catalog() -> Catalog {
        let catalog = Catalog::new();
        let users = TableBuilder::new("users")
            .column_with(
                "id",
                ColumnType::Integer,
                vec![ColumnConstraint::PrimaryKey],
            )
            .column("name", ColumnType::Text { max_len: None })
            .column("age", ColumnType::Integer)
            .build(1, PageId(10));
        catalog.create_table(users).unwrap();

        // Add an index on 'id'
        let idx = Index::new("idx_users_id", "users", vec![0], true, PageId(20));
        catalog.add_index("users", idx).unwrap();

        // Add an index on 'age'
        let idx_age = Index::new("idx_users_age", "users", vec![2], false, PageId(30));
        catalog.add_index("users", idx_age).unwrap();

        catalog
    }

    #[test]
    fn test_constant_folding() {
        let expr = Expr::binary(Expr::literal(1i64), BinaryOp::Add, Expr::literal(2i64));

        let folded = helpers::fold_constants(expr);
        assert!(matches!(folded, Expr::Literal(Value::Integer(3))));
    }

    #[test]
    fn test_split_conjunctions() {
        let expr = Expr::and(
            Expr::column("a"),
            Expr::and(Expr::column("b"), Expr::column("c")),
        );

        let parts = helpers::split_conjunctions(expr);
        assert_eq!(parts.len(), 3);
    }

    #[test]
    fn test_index_selection_equality() {
        let catalog = Arc::new(create_test_catalog());
        let table = catalog.get_table("users").unwrap();

        // Create a Filter(Scan) plan with an equality predicate: WHERE id = 5
        // Since id is the primary key, this should use PkSeek, not IndexScan
        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };
        let filter = LogicalPlan::Filter {
            input: Box::new(scan),
            predicate: Expr::binary(Expr::column("id"), BinaryOp::Eq, Expr::literal(5i64)),
        };

        // Apply optimization
        let optimizer = Optimizer::with_catalog(catalog);
        let optimized = optimizer.optimize(filter);

        // Should be converted to PkSeek (primary key equality)
        match optimized {
            LogicalPlan::PkSeek {
                key_values,
                residual_filter,
                ..
            } => {
                assert_eq!(key_values.len(), 1);
                assert!(matches!(key_values[0], Expr::Literal(Value::Integer(5))));
                assert!(residual_filter.is_none());
            }
            other => panic!("Expected PkSeek, got {:?}", other),
        }
    }

    #[test]
    fn test_index_selection_range() {
        let catalog = Arc::new(create_test_catalog());
        let table = catalog.get_table("users").unwrap();

        // Create a plan: WHERE age > 18 AND age < 65
        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };
        let filter = LogicalPlan::Filter {
            input: Box::new(scan),
            predicate: Expr::and(
                Expr::binary(Expr::column("age"), BinaryOp::Gt, Expr::literal(18i64)),
                Expr::binary(Expr::column("age"), BinaryOp::Lt, Expr::literal(65i64)),
            ),
        };

        let optimizer = Optimizer::with_catalog(catalog);
        let optimized = optimizer.optimize(filter);

        // Should be converted to IndexScan with range
        match optimized {
            LogicalPlan::IndexScan {
                index,
                range,
                residual_filter,
                ..
            } => {
                assert_eq!(index.name, "idx_users_age");
                assert!(!range.is_point_lookup());
                assert!(residual_filter.is_none()); // Both predicates consumed by range
            }
            other => panic!("Expected IndexScan, got {:?}", other),
        }
    }

    #[test]
    fn test_index_selection_with_residual_filter() {
        let catalog = Arc::new(create_test_catalog());
        let table = catalog.get_table("users").unwrap();

        // Create: WHERE age > 18 AND name = 'Alice'
        // Only age predicate can use index, name becomes residual
        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };
        let filter = LogicalPlan::Filter {
            input: Box::new(scan),
            predicate: Expr::and(
                Expr::binary(Expr::column("age"), BinaryOp::Gt, Expr::literal(18i64)),
                Expr::binary(
                    Expr::column("name"),
                    BinaryOp::Eq,
                    Expr::literal(Value::Text("Alice".to_string())),
                ),
            ),
        };

        let optimizer = Optimizer::with_catalog(catalog);
        let optimized = optimizer.optimize(filter);

        match optimized {
            LogicalPlan::IndexScan {
                index,
                range,
                residual_filter,
                ..
            } => {
                assert_eq!(index.name, "idx_users_age");
                assert!(!range.is_point_lookup());
                // name = 'Alice' should be in residual filter
                assert!(residual_filter.is_some());
            }
            other => panic!("Expected IndexScan, got {:?}", other),
        }
    }

    #[test]
    fn test_index_selection_no_index() {
        let catalog = Arc::new(create_test_catalog());
        let table = catalog.get_table("users").unwrap();

        // Create: WHERE name = 'Alice' (no index on name)
        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };
        let filter = LogicalPlan::Filter {
            input: Box::new(scan),
            predicate: Expr::binary(
                Expr::column("name"),
                BinaryOp::Eq,
                Expr::literal(Value::Text("Alice".to_string())),
            ),
        };

        let optimizer = Optimizer::with_catalog(catalog);
        let optimized = optimizer.optimize(filter);

        // Should stay as Scan with filter (predicate pushed down)
        match optimized {
            LogicalPlan::Scan { filter, .. } => {
                assert!(filter.is_some());
            }
            other => panic!("Expected Scan with filter, got {:?}", other),
        }
    }

    #[test]
    fn test_cost_estimator_basic() {
        let catalog = create_test_catalog();
        let table = catalog.get_table("users").unwrap();

        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let estimator = CostEstimator::new(&catalog);
        let cardinality = estimator.estimate_cardinality(&scan);
        let cost = estimator.estimate_cost(&scan);

        // Default row count is 1000 for new tables
        assert_eq!(cardinality, 1000.0);
        assert!(cost > 0.0);
    }

    #[test]
    fn test_cost_estimator_filter() {
        let catalog = create_test_catalog();
        let table = catalog.get_table("users").unwrap();

        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let filter = LogicalPlan::Filter {
            input: Box::new(scan),
            predicate: Expr::binary(Expr::column("age"), BinaryOp::Gt, Expr::literal(18i64)),
        };

        let estimator = CostEstimator::new(&catalog);
        let cardinality = estimator.estimate_cardinality(&filter);

        // Filter should reduce cardinality
        assert!(cardinality < 1000.0);
        assert!(cardinality > 0.0);
    }

    #[test]
    fn test_cost_estimator_join() {
        let catalog = create_test_catalog();

        // Create another table for joining
        let orders = TableBuilder::new("orders")
            .column("id", ColumnType::Integer)
            .column("user_id", ColumnType::Integer)
            .build(2, PageId(100));
        catalog.create_table(orders).unwrap();

        let users = catalog.get_table("users").unwrap();
        let orders_table = catalog.get_table("orders").unwrap();

        let left_scan = LogicalPlan::Scan {
            table: users.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let right_scan = LogicalPlan::Scan {
            table: orders_table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let join = LogicalPlan::Join {
            left: Box::new(left_scan),
            right: Box::new(right_scan),
            join_type: crate::planner::JoinType::Inner,
            condition: Some(Expr::binary(
                Expr::column("users.id"),
                BinaryOp::Eq,
                Expr::column("orders.user_id"),
            )),
        };

        let estimator = CostEstimator::new(&catalog);
        let cardinality = estimator.estimate_cardinality(&join);
        let cost = estimator.estimate_cost(&join);

        // Join cardinality should be > 0
        assert!(cardinality > 0.0);
        assert!(cost > 0.0);
    }

    #[test]
    fn test_cost_estimator_format_plan() {
        let catalog = create_test_catalog();
        let table = catalog.get_table("users").unwrap();

        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let estimator = CostEstimator::new(&catalog);
        let formatted = estimator.format_plan_with_costs(&scan, 0);

        assert!(formatted.contains("Scan: users"));
        assert!(formatted.contains("cost="));
        assert!(formatted.contains("rows="));
    }

    #[test]
    fn test_join_order_optimizer() {
        let catalog = Arc::new(create_test_catalog());

        // Create additional tables for multi-way join
        let orders = TableBuilder::new("orders")
            .column("id", ColumnType::Integer)
            .column("user_id", ColumnType::Integer)
            .build(2, PageId(100));
        catalog.create_table(orders).unwrap();

        let products = TableBuilder::new("products")
            .column("id", ColumnType::Integer)
            .column("order_id", ColumnType::Integer)
            .build(3, PageId(200));
        catalog.create_table(products).unwrap();

        let users = catalog.get_table("users").unwrap();
        let orders_table = catalog.get_table("orders").unwrap();
        let products_table = catalog.get_table("products").unwrap();

        // Create a 3-way join: users JOIN orders JOIN products
        let users_scan = LogicalPlan::Scan {
            table: users.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let orders_scan = LogicalPlan::Scan {
            table: orders_table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        let products_scan = LogicalPlan::Scan {
            table: products_table.clone(),
            alias: None,
            projection: None,
            filter: None,
        };

        // Join users with orders
        let join1 = LogicalPlan::Join {
            left: Box::new(users_scan),
            right: Box::new(orders_scan),
            join_type: crate::planner::JoinType::Inner,
            condition: Some(Expr::binary(
                Expr::column("users.id"),
                BinaryOp::Eq,
                Expr::column("orders.user_id"),
            )),
        };

        // Join result with products
        let join2 = LogicalPlan::Join {
            left: Box::new(join1),
            right: Box::new(products_scan),
            join_type: crate::planner::JoinType::Inner,
            condition: Some(Expr::binary(
                Expr::column("orders.id"),
                BinaryOp::Eq,
                Expr::column("products.order_id"),
            )),
        };

        let optimizer = Optimizer::with_catalog(catalog);
        let optimized = optimizer.optimize(join2);

        // Verify it's still a join (may be reordered)
        match optimized {
            LogicalPlan::Join { .. } => {
                // Success - join order may have been optimized
            }
            other => panic!("Expected Join, got {:?}", other),
        }
    }

    #[test]
    fn test_selectivity_estimation() {
        let catalog = create_test_catalog();
        let table = catalog.get_table("users").unwrap();

        let estimator = CostEstimator::new(&catalog);

        // Test AND selectivity (should multiply)
        let and_pred = Expr::and(
            Expr::binary(Expr::column("age"), BinaryOp::Gt, Expr::literal(18i64)),
            Expr::binary(Expr::column("age"), BinaryOp::Lt, Expr::literal(65i64)),
        );

        let scan = LogicalPlan::Scan {
            table: table.clone(),
            alias: None,
            projection: None,
            filter: Some(and_pred),
        };

        let cardinality = estimator.estimate_cardinality(&scan);

        // Should be significantly less than base rows due to AND selectivity
        assert!(cardinality < 1000.0);
        assert!(cardinality > 0.0);
    }
}