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

uqa_planner/
query_optimizer.rs
