Skip to main content

uqa_planner/
query_optimizer.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Rule-based and cost-aware operator-tree optimizer.
8//!
9//! Walks an [`OperatorTree`] and applies algebraic, graph, and physical rewrites:
10//!
11//! 1. `simplify_algebra` -- address-independent idempotence /
12//!    absorption / empty elimination on membership-only Intersect /
13//!    Union operands. Score-bearing operands remain distinct because
14//!    posting-list merges add their scores.
15//! 2. `push_filters_down` -- sink Filter into Intersect children when
16//!    the field applies.
17//! 3. `push_graph_pattern_filters` -- fold vertex / edge property
18//!    filters into PatternMatch constraints.
19//! 4. `push_filter_into_traverse` -- absorb vertex predicates into
20//!    Traverse so BFS prunes during expansion.
21//! 5. `push_filter_below_graph_join` -- move filters past graph joins
22//!    when the field belongs to the left side.
23//! 6. `fuse_join_pattern` -- merge intersected PatternMatch operators
24//!    that share a vertex variable.
25//! 7. `reorder_intersect` -- sort Intersect children by estimated operator cost (cheapest first).
26//! 8. `reorder_fusion_signals` -- sort fusion signals by cost; graph
27//!    operators receive a 0.5x discount when graph stats are
28//!    available.
29//! 9. `apply_index_scan` -- substitute leaf Filter with IndexScan when a covering index is registered and cheaper.
30//!
31//! 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.
32
33use std::{collections::BTreeMap, sync::Arc};
34
35mod algebra;
36mod graph_rewrites;
37mod index_selection;
38mod reorder;
39mod tree_map;
40
41use uqa_core::{IndexStats, Predicate};
42use uqa_operators::OperatorTree;
43
44use crate::cardinality::{CardinalityEstimator, ColumnStats, GraphStats, GraphStoreSampler};
45use crate::cost_model::CostModel;
46
47/// Fluent configuration for the optimizer pipeline. Lets callers
48/// disable individual stages for testing without poking at private
49/// fields.
50#[derive(Debug, Clone)]
51pub struct OptimizerConfig {
52    pub enable_simplify_algebra: bool,
53    pub enable_push_filters_down: bool,
54    pub enable_push_graph_pattern_filters: bool,
55    pub enable_push_filter_into_traverse: bool,
56    pub enable_push_filter_below_graph_join: bool,
57    pub enable_fuse_join_pattern: bool,
58    /// Retained for source compatibility; ignored because merging vector thresholds loses scores and can suppress validation errors.
59    #[deprecated(note = "vector threshold merging was removed; this field has no effect")]
60    pub enable_merge_vector_thresholds: bool,
61    pub enable_reorder_intersect: bool,
62    pub enable_reorder_fusion_signals: bool,
63    pub enable_apply_index_scan: bool,
64}
65
66/// Query-local physical index candidate supplied by an engine catalog.
67///
68/// The candidate already contains the scan cost for this predicate. This
69/// keeps the planner independent of an engine's index implementation while
70/// allowing the final optimizer pass to emit a concrete
71/// [`OperatorTree::IndexScan`].
72#[derive(Debug, Clone, PartialEq)]
73pub struct IndexScanCandidate {
74    pub index_name: String,
75    pub table_name: String,
76    pub field: String,
77    pub predicate: Predicate,
78    pub scan_cost: f64,
79}
80
81impl Default for OptimizerConfig {
82    #[allow(
83        deprecated,
84        reason = "initialize the retained source-compatibility field"
85    )]
86    fn default() -> Self {
87        Self {
88            enable_simplify_algebra: true,
89            enable_push_filters_down: true,
90            enable_push_graph_pattern_filters: true,
91            enable_push_filter_into_traverse: true,
92            enable_push_filter_below_graph_join: true,
93            enable_fuse_join_pattern: true,
94            enable_merge_vector_thresholds: false,
95            enable_reorder_intersect: true,
96            enable_reorder_fusion_signals: true,
97            enable_apply_index_scan: true,
98        }
99    }
100}
101
102/// Operator-tree query optimizer.
103pub struct QueryOptimizer {
104    pub estimator: CardinalityEstimator,
105    pub cost_model: CostModel,
106    pub graph_stats: Option<GraphStats>,
107    pub index_candidates: Vec<IndexScanCandidate>,
108    pub table_name: Option<String>,
109    pub row_count: Option<u64>,
110    pub index_stats: IndexStats,
111    pub config: OptimizerConfig,
112}
113
114impl QueryOptimizer {
115    pub fn new() -> Self {
116        Self {
117            estimator: CardinalityEstimator::new(),
118            cost_model: CostModel::new(),
119            graph_stats: None,
120            index_candidates: Vec::new(),
121            table_name: None,
122            row_count: None,
123            index_stats: IndexStats::new(1_000),
124            config: OptimizerConfig::default(),
125        }
126    }
127
128    /// Attach immutable scan candidates discovered from the caller's catalog snapshot. The optimizer selects among these candidates without retaining a storage handle.
129    pub fn with_index_candidates(
130        mut self,
131        candidates: impl IntoIterator<Item = IndexScanCandidate>,
132        table: impl Into<String>,
133    ) -> Self {
134        self.index_candidates = candidates.into_iter().collect();
135        self.table_name = Some(table.into());
136        self
137    }
138
139    pub fn with_graph_stats(mut self, gs: GraphStats) -> Self {
140        self.cost_model = std::mem::take(&mut self.cost_model).with_graph_stats(gs.clone());
141        self.estimator = std::mem::take(&mut self.estimator).with_graph_stats(gs.clone());
142        self.graph_stats = Some(gs);
143        self
144    }
145
146    pub fn with_graph_store(mut self, store: Arc<dyn GraphStoreSampler>) -> Self {
147        self.estimator = std::mem::take(&mut self.estimator).with_graph_store(store);
148        self
149    }
150
151    pub fn with_row_count(mut self, n: u64) -> Self {
152        self.row_count = Some(n);
153        self.index_stats.total_docs = n;
154        self
155    }
156
157    pub fn with_index_stats(mut self, stats: IndexStats) -> Self {
158        self.row_count = Some(stats.total_docs);
159        self.index_stats = stats;
160        self
161    }
162
163    pub fn with_column_stats(mut self, stats: BTreeMap<String, ColumnStats>) -> Self {
164        self.cost_model = std::mem::take(&mut self.cost_model).with_column_stats(stats.clone());
165        self.estimator = std::mem::take(&mut self.estimator).with_column_stats(stats);
166        self
167    }
168
169    /// Optimize a query through the complete rewrite pipeline.
170    pub fn optimize(&self, op: OperatorTree) -> OperatorTree {
171        let mut op = op;
172        if self.config.enable_simplify_algebra {
173            op = self.simplify_algebra(op);
174        }
175        if self.config.enable_push_filters_down {
176            op = self.push_filters_down(op);
177        }
178        if self.config.enable_push_graph_pattern_filters {
179            op = self.push_graph_pattern_filters(op);
180        }
181        if self.config.enable_push_filter_into_traverse {
182            op = self.push_filter_into_traverse(op);
183        }
184        if self.config.enable_push_filter_below_graph_join {
185            op = self.push_filter_below_graph_join(op);
186        }
187        if self.config.enable_fuse_join_pattern {
188            op = Self::fuse_join_pattern(op);
189        }
190        if self.config.enable_reorder_intersect {
191            op = self.reorder_intersect(op);
192        }
193        if self.config.enable_reorder_fusion_signals {
194            op = self.reorder_fusion_signals(op);
195        }
196        if self.config.enable_apply_index_scan {
197            op = self.apply_index_scan(op);
198        }
199        op
200    }
201}
202
203impl Default for QueryOptimizer {
204    fn default() -> Self {
205        Self::new()
206    }
207}
208
209#[cfg(test)]
210mod tests;