featherdb_query/optimizer/cost/

estimator.rs

//! Cost estimator for query plans

use super::constants as cost_constants;
use crate::expr::{BinaryOp, Expr};
use crate::planner::{IndexBound, IndexRange, JoinType, LogicalPlan};
use featherdb_catalog::{Catalog, Table};
use std::sync::Arc;

/// Cost estimator for query plans
///
/// Provides cardinality and cost estimates for logical plan nodes.
/// Used by the query optimizer to choose optimal join orders and
/// access methods.
pub struct CostEstimator<'a> {
    catalog: &'a Catalog,
}

impl<'a> CostEstimator<'a> {
    /// Create a new cost estimator with catalog access
    pub fn new(catalog: &'a Catalog) -> Self {
        CostEstimator { catalog }
    }

    /// Estimate the number of rows output by a plan node
    pub fn estimate_cardinality(&self, plan: &LogicalPlan) -> f64 {
        match plan {
            LogicalPlan::Scan { table, filter, .. } => {
                let base_rows = self.catalog.estimated_row_count(&table.name) as f64;
                match filter {
                    Some(pred) => base_rows * self.estimate_selectivity(pred, table),
                    None => base_rows,
                }
            }

            LogicalPlan::IndexScan {
                table,
                range,
                residual_filter,
                index_column,
                ..
            } => {
                let base_rows = self.catalog.estimated_row_count(&table.name) as f64;

                // Estimate selectivity based on range
                let range_selectivity = if range.is_point_lookup() {
                    self.catalog.selectivity_eq(&table.name, *index_column)
                } else {
                    self.estimate_range_selectivity(range, &table.name, *index_column)
                };

                let rows_after_index = base_rows * range_selectivity;

                // Apply residual filter selectivity
                match residual_filter {
                    Some(pred) => rows_after_index * self.estimate_selectivity(pred, table),
                    None => rows_after_index,
                }
            }

            LogicalPlan::PkSeek {
                table,
                residual_filter,
                ..
            } => {
                // PK seek returns at most 1 row (point lookup)
                let base_cardinality = 1.0;

                // Apply residual filter selectivity if present
                match residual_filter {
                    Some(pred) => base_cardinality * self.estimate_selectivity(pred, table),
                    None => base_cardinality,
                }
            }

            LogicalPlan::PkRangeScan {
                table,
                range,
                residual_filter,
                ..
            } => {
                let base_rows = self.catalog.estimated_row_count(&table.name) as f64;
                // Use range selectivity estimate - for PK ranges, treat like the first PK column
                let pk_col = if !table.primary_key.is_empty() {
                    table.primary_key[0]
                } else {
                    0
                };
                let range_selectivity = self.estimate_range_selectivity(range, &table.name, pk_col);
                let rows_after_range = base_rows * range_selectivity;
                match residual_filter {
                    Some(pred) => rows_after_range * self.estimate_selectivity(pred, table),
                    None => rows_after_range,
                }
            }

            LogicalPlan::Filter { input, predicate } => {
                let input_rows = self.estimate_cardinality(input);
                let table = self.extract_table(input);
                let selectivity = table
                    .map(|t| self.estimate_selectivity(predicate, &t))
                    .unwrap_or(cost_constants::DEFAULT_SELECTIVITY);
                input_rows * selectivity
            }

            LogicalPlan::Project { input, .. } => self.estimate_cardinality(input),

            LogicalPlan::Join {
                left,
                right,
                condition,
                join_type,
            } => {
                let left_rows = self.estimate_cardinality(left);
                let right_rows = self.estimate_cardinality(right);

                let join_selectivity = condition
                    .as_ref()
                    .map(|c| self.estimate_join_selectivity(c, left, right))
                    .unwrap_or(1.0); // Cross join

                let base_result = left_rows * right_rows * join_selectivity;

                // Adjust for join type
                match join_type {
                    JoinType::Inner => base_result,
                    JoinType::Left => base_result.max(left_rows), // At least left rows
                    JoinType::Right => base_result.max(right_rows), // At least right rows
                    JoinType::Full => base_result.max(left_rows).max(right_rows),
                }
            }

            LogicalPlan::Aggregate {
                input, group_by, ..
            } => {
                if group_by.is_empty() {
                    1.0 // Single aggregate result
                } else {
                    // Estimate distinct groups
                    let input_rows = self.estimate_cardinality(input);
                    // Assume group by reduces rows significantly
                    (input_rows * 0.1).max(1.0)
                }
            }

            LogicalPlan::Sort { input, .. } => self.estimate_cardinality(input),

            LogicalPlan::Limit {
                input,
                limit,
                offset,
            } => {
                let input_rows = self.estimate_cardinality(input);
                let remaining = input_rows - *offset as f64;
                match limit {
                    Some(l) => remaining.min(*l as f64).max(0.0),
                    None => remaining.max(0.0),
                }
            }

            LogicalPlan::Distinct { input } => {
                let input_rows = self.estimate_cardinality(input);
                // Assume distinct reduces by some factor
                (input_rows * 0.5).max(1.0)
            }

            LogicalPlan::EmptyRelation => 1.0,

            _ => 1000.0, // Default for other plan types
        }
    }

    /// Estimate the total cost of executing a plan
    pub fn estimate_cost(&self, plan: &LogicalPlan) -> f64 {
        match plan {
            LogicalPlan::Scan { table, .. } => {
                let rows = self.catalog.estimated_row_count(&table.name) as f64;
                rows * cost_constants::SEQ_SCAN_COST_PER_ROW
            }

            LogicalPlan::IndexScan {
                table,
                range,
                index_column,
                ..
            } => {
                let base_rows = self.catalog.estimated_row_count(&table.name) as f64;
                let selectivity = if range.is_point_lookup() {
                    self.catalog.selectivity_eq(&table.name, *index_column)
                } else {
                    self.estimate_range_selectivity(range, &table.name, *index_column)
                };
                let rows_scanned = base_rows * selectivity;
                rows_scanned * cost_constants::INDEX_SCAN_COST_PER_ROW
            }

            LogicalPlan::PkSeek { .. } => {
                // PK seek: single B-tree lookup (O(log N))
                // Very low cost - similar to index point lookup
                cost_constants::INDEX_SCAN_COST_PER_ROW
            }

            LogicalPlan::PkRangeScan { table, range, .. } => {
                let base_rows = self.catalog.estimated_row_count(&table.name) as f64;
                let pk_col = if !table.primary_key.is_empty() {
                    table.primary_key[0]
                } else {
                    0
                };
                let selectivity = self.estimate_range_selectivity(range, &table.name, pk_col);
                let rows_scanned = base_rows * selectivity;
                // Similar cost to index scan - direct B-tree range access
                rows_scanned * cost_constants::INDEX_SCAN_COST_PER_ROW
            }

            LogicalPlan::Filter { input, .. } => {
                let input_cost = self.estimate_cost(input);
                let rows = self.estimate_cardinality(input);
                input_cost + rows * cost_constants::CPU_COST_MULTIPLIER
            }

            LogicalPlan::Project { input, .. } => {
                let input_cost = self.estimate_cost(input);
                let rows = self.estimate_cardinality(input);
                input_cost + rows * cost_constants::CPU_COST_MULTIPLIER
            }

            LogicalPlan::Join { left, right, .. } => {
                let left_cost = self.estimate_cost(left);
                let right_cost = self.estimate_cost(right);
                let left_rows = self.estimate_cardinality(left);
                let right_rows = self.estimate_cardinality(right);

                // Hash join cost: build hash table + probe
                // Smaller table goes on build side
                let (build_rows, probe_rows) = if left_rows <= right_rows {
                    (left_rows, right_rows)
                } else {
                    (right_rows, left_rows)
                };

                let build_cost = build_rows * cost_constants::HASH_BUILD_COST_PER_ROW;
                let probe_cost = probe_rows * cost_constants::HASH_JOIN_COST_PER_ROW;

                left_cost + right_cost + build_cost + probe_cost
            }

            LogicalPlan::Aggregate { input, .. } => {
                let input_cost = self.estimate_cost(input);
                let rows = self.estimate_cardinality(input);
                input_cost + rows * cost_constants::CPU_COST_MULTIPLIER
            }

            LogicalPlan::Sort { input, .. } => {
                let input_cost = self.estimate_cost(input);
                let rows = self.estimate_cardinality(input);
                // n log n cost
                let sort_cost = if rows > 1.0 {
                    rows * rows.log2() * cost_constants::SORT_COST_PER_ROW
                } else {
                    0.0
                };
                input_cost + sort_cost
            }

            LogicalPlan::Limit { input, .. } => self.estimate_cost(input),

            LogicalPlan::Distinct { input } => {
                let input_cost = self.estimate_cost(input);
                let rows = self.estimate_cardinality(input);
                input_cost + rows * cost_constants::CPU_COST_MULTIPLIER
            }

            LogicalPlan::EmptyRelation => 0.0,

            _ => 1000.0, // Default cost for other plan types
        }
    }

    /// Estimate selectivity of a predicate
    pub fn estimate_selectivity(&self, predicate: &Expr, table: &Table) -> f64 {
        match predicate {
            Expr::BinaryOp { left, op, right } => {
                match op {
                    BinaryOp::And => {
                        // AND: multiply selectivities (independence assumption)
                        let left_sel = self.estimate_selectivity(left, table);
                        let right_sel = self.estimate_selectivity(right, table);
                        left_sel * right_sel
                    }
                    BinaryOp::Or => {
                        // OR: p(A) + p(B) - p(A)*p(B)
                        let left_sel = self.estimate_selectivity(left, table);
                        let right_sel = self.estimate_selectivity(right, table);
                        left_sel + right_sel - (left_sel * right_sel)
                    }
                    BinaryOp::Eq => self.estimate_equality_selectivity(left, right, table),
                    BinaryOp::Ne => 1.0 - self.estimate_equality_selectivity(left, right, table),
                    BinaryOp::Lt | BinaryOp::Le | BinaryOp::Gt | BinaryOp::Ge => {
                        cost_constants::DEFAULT_RANGE_SELECTIVITY
                    }
                    _ => cost_constants::DEFAULT_SELECTIVITY,
                }
            }
            Expr::UnaryOp { op, .. } => {
                match op {
                    crate::expr::UnaryOp::Not => 1.0 - cost_constants::DEFAULT_SELECTIVITY,
                    crate::expr::UnaryOp::IsNull => 0.01, // Assume 1% nulls
                    crate::expr::UnaryOp::IsNotNull => 0.99,
                    _ => cost_constants::DEFAULT_SELECTIVITY,
                }
            }
            Expr::Between { .. } => cost_constants::DEFAULT_RANGE_SELECTIVITY,
            Expr::InList { list, .. } => {
                // IN list: each value has equality selectivity
                let eq_sel = cost_constants::DEFAULT_SELECTIVITY;
                (eq_sel * list.len() as f64).min(1.0)
            }
            Expr::Like { .. } => 0.2, // LIKE usually quite selective
            _ => cost_constants::DEFAULT_SELECTIVITY,
        }
    }

    /// Estimate selectivity for equality predicate
    fn estimate_equality_selectivity(&self, left: &Expr, right: &Expr, table: &Table) -> f64 {
        // Check if this is a column = literal pattern
        if let Expr::Column { name, .. } = left {
            if matches!(right, Expr::Literal(_)) {
                if let Some(col_idx) = table.get_column_index(name) {
                    return self.catalog.selectivity_eq(&table.name, col_idx);
                }
            }
        }
        if let Expr::Column { name, .. } = right {
            if matches!(left, Expr::Literal(_)) {
                if let Some(col_idx) = table.get_column_index(name) {
                    return self.catalog.selectivity_eq(&table.name, col_idx);
                }
            }
        }
        cost_constants::DEFAULT_SELECTIVITY
    }

    /// Estimate selectivity for join conditions
    fn estimate_join_selectivity(
        &self,
        condition: &Expr,
        _left: &LogicalPlan,
        _right: &LogicalPlan,
    ) -> f64 {
        match condition {
            Expr::BinaryOp {
                op: BinaryOp::Eq, ..
            } => cost_constants::DEFAULT_JOIN_SELECTIVITY,
            Expr::BinaryOp {
                op: BinaryOp::And,
                left,
                right,
            } => {
                let left_sel = self.estimate_join_selectivity(left, _left, _right);
                let right_sel = self.estimate_join_selectivity(right, _left, _right);
                left_sel * right_sel
            }
            _ => cost_constants::DEFAULT_JOIN_SELECTIVITY,
        }
    }

    /// Estimate selectivity for index range
    fn estimate_range_selectivity(
        &self,
        range: &IndexRange,
        table_name: &str,
        col_index: usize,
    ) -> f64 {
        let low = match &range.start {
            IndexBound::Inclusive(v) | IndexBound::Exclusive(v) => Some(v),
            IndexBound::Unbounded => None,
        };
        let high = match &range.end {
            IndexBound::Inclusive(v) | IndexBound::Exclusive(v) => Some(v),
            IndexBound::Unbounded => None,
        };
        self.catalog
            .selectivity_range(table_name, col_index, low, high)
    }

    /// Extract the primary table from a plan (for selectivity estimation)
    fn extract_table(&self, plan: &LogicalPlan) -> Option<Arc<Table>> {
        match plan {
            LogicalPlan::Scan { table, .. } => Some(table.clone()),
            LogicalPlan::IndexScan { table, .. } => Some(table.clone()),
            LogicalPlan::PkSeek { table, .. } => Some(table.clone()),
            LogicalPlan::PkRangeScan { table, .. } => Some(table.clone()),
            LogicalPlan::Filter { input, .. } => self.extract_table(input),
            LogicalPlan::Project { input, .. } => self.extract_table(input),
            _ => None,
        }
    }

    /// Format a plan with cost annotations for EXPLAIN output
    pub fn format_plan_with_costs(&self, plan: &LogicalPlan, indent: usize) -> String {
        let prefix = "  ".repeat(indent);
        let cost = self.estimate_cost(plan);
        let rows = self.estimate_cardinality(plan);

        match plan {
            LogicalPlan::Scan {
                table,
                alias,
                filter,
                ..
            } => {
                let alias_str = alias
                    .as_ref()
                    .map(|a| format!(" AS {}", a))
                    .unwrap_or_default();
                let filter_str = filter
                    .as_ref()
                    .map(|f| format!(" (filter: {:?})", f))
                    .unwrap_or_default();
                format!(
                    "{}Scan: {}{}{} (cost={:.2}, rows={:.0})",
                    prefix, table.name, alias_str, filter_str, cost, rows
                )
            }

            LogicalPlan::IndexScan {
                table,
                index,
                range,
                alias,
                ..
            } => {
                let alias_str = alias
                    .as_ref()
                    .map(|a| format!(" AS {}", a))
                    .unwrap_or_default();
                let range_str = if range.is_point_lookup() {
                    "point lookup"
                } else {
                    "range scan"
                };
                format!(
                    "{}IndexScan: {}{} using {} ({}) (cost={:.2}, rows={:.0})",
                    prefix, table.name, alias_str, index.name, range_str, cost, rows
                )
            }

            LogicalPlan::PkSeek {
                table,
                alias,
                key_values,
                residual_filter,
                ..
            } => {
                let alias_str = alias
                    .as_ref()
                    .map(|a| format!(" AS {}", a))
                    .unwrap_or_default();
                let keys_str = key_values
                    .iter()
                    .map(|e| format!("{:?}", e))
                    .collect::<Vec<_>>()
                    .join(", ");
                let filter_str = residual_filter
                    .as_ref()
                    .map(|f| format!(" (residual: {:?})", f))
                    .unwrap_or_default();
                format!(
                    "{}PkSeek: {}{} [{}]{} (cost={:.2}, rows={:.0})",
                    prefix, table.name, alias_str, keys_str, filter_str, cost, rows
                )
            }

            LogicalPlan::PkRangeScan {
                table,
                alias,
                range,
                residual_filter,
                ..
            } => {
                let alias_str = alias
                    .as_ref()
                    .map(|a| format!(" AS {}", a))
                    .unwrap_or_default();
                let range_str = format!("{:?}..{:?}", range.start, range.end);
                let filter_str = residual_filter
                    .as_ref()
                    .map(|f| format!(" (residual: {:?})", f))
                    .unwrap_or_default();
                format!(
                    "{}PkRangeScan: {}{} [{}]{} (cost={:.2}, rows={:.0})",
                    prefix, table.name, alias_str, range_str, filter_str, cost, rows
                )
            }

            LogicalPlan::Filter { input, predicate } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                format!(
                    "{}Filter: {:?} (cost={:.2}, rows={:.0})\n{}",
                    prefix, predicate, cost, rows, child
                )
            }

            LogicalPlan::Project { input, exprs } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                let cols: Vec<_> = exprs.iter().map(|(_, name)| name.as_str()).collect();
                format!(
                    "{}Project: [{}] (cost={:.2}, rows={:.0})\n{}",
                    prefix,
                    cols.join(", "),
                    cost,
                    rows,
                    child
                )
            }

            LogicalPlan::Join {
                left,
                right,
                join_type,
                condition,
            } => {
                let left_child = self.format_plan_with_costs(left, indent + 1);
                let right_child = self.format_plan_with_costs(right, indent + 1);
                let join_str = match join_type {
                    JoinType::Inner => "Inner",
                    JoinType::Left => "Left",
                    JoinType::Right => "Right",
                    JoinType::Full => "Full",
                };
                let cond_str = condition
                    .as_ref()
                    .map(|c| format!(" ON {:?}", c))
                    .unwrap_or_default();
                format!(
                    "{}{} Join{} (cost={:.2}, rows={:.0})\n{}\n{}",
                    prefix, join_str, cond_str, cost, rows, left_child, right_child
                )
            }

            LogicalPlan::Aggregate {
                input,
                group_by,
                aggregates,
            } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                let aggs: Vec<_> = aggregates.iter().map(|(_, name)| name.as_str()).collect();
                format!(
                    "{}Aggregate: group_by={}, aggs=[{}] (cost={:.2}, rows={:.0})\n{}",
                    prefix,
                    group_by.len(),
                    aggs.join(", "),
                    cost,
                    rows,
                    child
                )
            }

            LogicalPlan::Sort { input, order_by } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                format!(
                    "{}Sort: {} columns (cost={:.2}, rows={:.0})\n{}",
                    prefix,
                    order_by.len(),
                    cost,
                    rows,
                    child
                )
            }

            LogicalPlan::Limit {
                input,
                limit,
                offset,
            } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                format!(
                    "{}Limit: {:?} offset {} (cost={:.2}, rows={:.0})\n{}",
                    prefix, limit, offset, cost, rows, child
                )
            }

            LogicalPlan::Distinct { input } => {
                let child = self.format_plan_with_costs(input, indent + 1);
                format!(
                    "{}Distinct (cost={:.2}, rows={:.0})\n{}",
                    prefix, cost, rows, child
                )
            }

            _ => format!("{}Plan node (cost={:.2}, rows={:.0})", prefix, cost, rows),
        }
    }
}