uqa-planner 0.3.8

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Rule-based and cost-aware operator-tree optimizer.
//!
//! Walks an [`OperatorTree`] and applies algebraic, graph, and physical rewrites:
//!
//! 1. `simplify_algebra` -- address-independent idempotence /
//!    absorption / empty elimination on membership-only Intersect /
//!    Union operands. Score-bearing operands remain distinct because
//!    posting-list merges add their scores.
//! 2. `push_filters_down` -- sink Filter into Intersect children when
//!    the field applies.
//! 3. `push_graph_pattern_filters` -- fold vertex / edge property
//!    filters into PatternMatch constraints.
//! 4. `push_filter_into_traverse` -- absorb vertex predicates into
//!    Traverse so BFS prunes during expansion.
//! 5. `push_filter_below_graph_join` -- move filters past graph joins
//!    when the field belongs to the left side.
//! 6. `fuse_join_pattern` -- merge intersected PatternMatch operators
//!    that share a vertex variable.
//! 7. `reorder_intersect` -- sort Intersect children by estimated operator cost (cheapest first).
//! 8. `reorder_fusion_signals` -- sort fusion signals by cost; graph
//!    operators receive a 0.5x discount when graph stats are
//!    available.
//! 9. `apply_index_scan` -- substitute leaf Filter with IndexScan when a covering index is registered and cheaper.
//!
//! Vector-threshold operands remain distinct: intersection adds each raw cosine score, and approximate query-vector equality cannot preserve threshold support. Rewrites must also retain invalid-threshold errors.

use std::{collections::BTreeMap, sync::Arc};

mod algebra;
mod graph_rewrites;
mod index_selection;
mod reorder;
mod tree_map;

use uqa_core::{IndexStats, Predicate};
use uqa_operators::OperatorTree;

use crate::cardinality::{CardinalityEstimator, ColumnStats, GraphStats, GraphStoreSampler};
use crate::cost_model::CostModel;

/// Fluent configuration for the optimizer pipeline. Lets callers
/// disable individual stages for testing without poking at private
/// fields.
#[derive(Debug, Clone)]
pub struct OptimizerConfig {
    pub enable_simplify_algebra: bool,
    pub enable_push_filters_down: bool,
    pub enable_push_graph_pattern_filters: bool,
    pub enable_push_filter_into_traverse: bool,
    pub enable_push_filter_below_graph_join: bool,
    pub enable_fuse_join_pattern: bool,
    /// Retained for source compatibility; ignored because merging vector thresholds loses scores and can suppress validation errors.
    #[deprecated(note = "vector threshold merging was removed; this field has no effect")]
    pub enable_merge_vector_thresholds: bool,
    pub enable_reorder_intersect: bool,
    pub enable_reorder_fusion_signals: bool,
    pub enable_apply_index_scan: bool,
}

/// Query-local physical index candidate supplied by an engine catalog.
///
/// The candidate already contains the scan cost for this predicate. This
/// keeps the planner independent of an engine's index implementation while
/// allowing the final optimizer pass to emit a concrete
/// [`OperatorTree::IndexScan`].
#[derive(Debug, Clone, PartialEq)]
pub struct IndexScanCandidate {
    pub index_name: String,
    pub table_name: String,
    pub field: String,
    pub predicate: Predicate,
    pub scan_cost: f64,
}

impl Default for OptimizerConfig {
    #[allow(
        deprecated,
        reason = "initialize the retained source-compatibility field"
    )]
    fn default() -> Self {
        Self {
            enable_simplify_algebra: true,
            enable_push_filters_down: true,
            enable_push_graph_pattern_filters: true,
            enable_push_filter_into_traverse: true,
            enable_push_filter_below_graph_join: true,
            enable_fuse_join_pattern: true,
            enable_merge_vector_thresholds: false,
            enable_reorder_intersect: true,
            enable_reorder_fusion_signals: true,
            enable_apply_index_scan: true,
        }
    }
}

/// Operator-tree query optimizer.
pub struct QueryOptimizer {
    pub estimator: CardinalityEstimator,
    pub cost_model: CostModel,
    pub graph_stats: Option<GraphStats>,
    pub index_candidates: Vec<IndexScanCandidate>,
    pub table_name: Option<String>,
    pub row_count: Option<u64>,
    pub index_stats: IndexStats,
    pub config: OptimizerConfig,
}

impl QueryOptimizer {
    pub fn new() -> Self {
        Self {
            estimator: CardinalityEstimator::new(),
            cost_model: CostModel::new(),
            graph_stats: None,
            index_candidates: Vec::new(),
            table_name: None,
            row_count: None,
            index_stats: IndexStats::new(1_000),
            config: OptimizerConfig::default(),
        }
    }

    /// Attach immutable scan candidates discovered from the caller's catalog snapshot. The optimizer selects among these candidates without retaining a storage handle.
    pub fn with_index_candidates(
        mut self,
        candidates: impl IntoIterator<Item = IndexScanCandidate>,
        table: impl Into<String>,
    ) -> Self {
        self.index_candidates = candidates.into_iter().collect();
        self.table_name = Some(table.into());
        self
    }

    pub fn with_graph_stats(mut self, gs: GraphStats) -> Self {
        self.cost_model = std::mem::take(&mut self.cost_model).with_graph_stats(gs.clone());
        self.estimator = std::mem::take(&mut self.estimator).with_graph_stats(gs.clone());
        self.graph_stats = Some(gs);
        self
    }

    pub fn with_graph_store(mut self, store: Arc<dyn GraphStoreSampler>) -> Self {
        self.estimator = std::mem::take(&mut self.estimator).with_graph_store(store);
        self
    }

    pub fn with_row_count(mut self, n: u64) -> Self {
        self.row_count = Some(n);
        self.index_stats.total_docs = n;
        self
    }

    pub fn with_index_stats(mut self, stats: IndexStats) -> Self {
        self.row_count = Some(stats.total_docs);
        self.index_stats = stats;
        self
    }

    pub fn with_column_stats(mut self, stats: BTreeMap<String, ColumnStats>) -> Self {
        self.cost_model = std::mem::take(&mut self.cost_model).with_column_stats(stats.clone());
        self.estimator = std::mem::take(&mut self.estimator).with_column_stats(stats);
        self
    }

    /// Optimize a query through the complete rewrite pipeline.
    pub fn optimize(&self, op: OperatorTree) -> OperatorTree {
        let mut op = op;
        if self.config.enable_simplify_algebra {
            op = self.simplify_algebra(op);
        }
        if self.config.enable_push_filters_down {
            op = self.push_filters_down(op);
        }
        if self.config.enable_push_graph_pattern_filters {
            op = self.push_graph_pattern_filters(op);
        }
        if self.config.enable_push_filter_into_traverse {
            op = self.push_filter_into_traverse(op);
        }
        if self.config.enable_push_filter_below_graph_join {
            op = self.push_filter_below_graph_join(op);
        }
        if self.config.enable_fuse_join_pattern {
            op = Self::fuse_join_pattern(op);
        }
        if self.config.enable_reorder_intersect {
            op = self.reorder_intersect(op);
        }
        if self.config.enable_reorder_fusion_signals {
            op = self.reorder_fusion_signals(op);
        }
        if self.config.enable_apply_index_scan {
            op = self.apply_index_scan(op);
        }
        op
    }
}

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

#[cfg(test)]
mod tests;