uqa_planner/
query_optimizer.rs1use 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#[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 #[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#[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
102pub 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 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 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;