1use std::collections::HashMap;
24
25#[derive(Debug, Default)]
27pub struct QueryOptimizer {
28 cost_model: CostModel,
29 cardinality_hints: HashMap<String, CardinalityHint>,
30 total_edges: usize,
31}
32
33#[derive(Debug, Clone)]
35pub struct CostModel {
36 pub row_read_cost: f64,
38 pub index_lookup_cost: f64,
40 pub seq_scan_cost: f64,
42 pub vector_search_cost: f64,
44}
45
46impl Default for CostModel {
47 fn default() -> Self {
48 Self {
49 row_read_cost: 1.0,
50 index_lookup_cost: 0.5,
51 seq_scan_cost: 0.1,
52 vector_search_cost: 10.0,
53 }
54 }
55}
56
57#[derive(Debug, Clone)]
59pub struct CardinalityHint {
60 pub cardinality: usize,
62 pub confidence: f64,
64 pub source: CardinalitySource,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum CardinalitySource {
71 Exact,
73 HyperLogLog,
75 Histogram,
77 Default,
79}
80
81#[derive(Debug, Clone)]
83pub enum QueryPredicate {
84 Eq(String, String),
86 Range {
88 column: String,
89 min: Option<String>,
90 max: Option<String>,
91 },
92 In(String, Vec<String>),
94 Prefix(String, String),
96 TimeRange(u64, u64),
98 Project(u16),
100 Tenant(u32),
102 SpanType(String),
104}
105
106#[derive(Debug, Clone, PartialEq)]
108pub enum QueryOperation {
109 PointLookup,
111 RangeScan,
113 FullScan,
115 VectorSearch { k: usize },
117 LsmRangeScan { start_us: u64, end_us: u64 },
119 GraphTraversal {
121 direction: TraversalDirection,
122 max_depth: usize,
123 },
124}
125
126#[derive(Debug, Clone)]
128pub enum IndexSelection {
129 PrimaryKey,
131 Secondary(String),
133 TimeIndex,
135 VectorIndex,
137 FullScan,
139 LsmScan,
141 CausalIndex,
143 ProjectIndex,
145 MultiIndex(Vec<IndexSelection>),
147}
148
149#[derive(Debug, Clone)]
151pub struct QueryCost {
152 pub estimated_cost: f64,
154 pub total_cost: f64,
156 pub estimated_rows: usize,
158 pub records_returned: usize,
160 pub index: IndexSelection,
162 pub operation: QueryOperation,
164 pub breakdown: Vec<(QueryOperation, f64)>,
166}
167
168#[derive(Debug, Clone, Copy, PartialEq, Eq)]
170pub enum TraversalDirection {
171 Forward,
172 Backward,
173}
174
175#[derive(Debug, Clone)]
177pub struct QueryPlan {
178 pub cost: QueryCost,
180 pub predicates: Vec<QueryPredicate>,
182 pub direction: TraversalDirection,
184 pub limit: Option<usize>,
186 pub index_selection: IndexSelection,
188 pub operations: Vec<QueryOperation>,
190}
191
192impl QueryOptimizer {
193 pub fn new() -> Self {
195 Self::default()
196 }
197
198 pub fn with_cost_model(cost_model: CostModel) -> Self {
200 Self {
201 cost_model,
202 ..Default::default()
203 }
204 }
205
206 pub fn update_total_edges(&mut self, count: usize, _source: CardinalitySource) {
208 self.total_edges = count;
209 }
210
211 pub fn set_cardinality_hint(&mut self, column: &str, hint: CardinalityHint) {
213 self.cardinality_hints.insert(column.to_string(), hint);
214 }
215
216 pub fn plan_query(&self, predicates: &[QueryPredicate], limit: Option<usize>) -> QueryPlan {
218 let (index, operation) = if predicates.is_empty() {
220 (IndexSelection::FullScan, QueryOperation::FullScan)
221 } else {
222 let has_eq = predicates
224 .iter()
225 .any(|p| matches!(p, QueryPredicate::Eq(_, _)));
226 if has_eq {
227 (IndexSelection::PrimaryKey, QueryOperation::PointLookup)
228 } else {
229 (IndexSelection::FullScan, QueryOperation::RangeScan)
230 }
231 };
232
233 let estimated_rows = match &index {
234 IndexSelection::PrimaryKey => 1,
235 IndexSelection::FullScan => self.total_edges.max(1000),
236 _ => self.total_edges / 10,
237 };
238
239 let estimated_cost = match &operation {
240 QueryOperation::PointLookup => self.cost_model.index_lookup_cost,
241 QueryOperation::RangeScan => estimated_rows as f64 * self.cost_model.row_read_cost,
242 QueryOperation::FullScan => self.total_edges as f64 * self.cost_model.seq_scan_cost,
243 QueryOperation::VectorSearch { .. } => self.cost_model.vector_search_cost,
244 _ => estimated_rows as f64 * self.cost_model.row_read_cost,
245 };
246
247 QueryPlan {
248 cost: QueryCost {
249 estimated_cost,
250 total_cost: estimated_cost,
251 estimated_rows,
252 records_returned: estimated_rows,
253 index: index.clone(),
254 operation: operation.clone(),
255 breakdown: vec![(operation.clone(), estimated_cost)],
256 },
257 predicates: predicates.to_vec(),
258 direction: TraversalDirection::Forward,
259 limit,
260 index_selection: index,
261 operations: vec![operation],
262 }
263 }
264
265 pub fn estimate_selectivity(&self, predicate: &QueryPredicate) -> f64 {
267 match predicate {
268 QueryPredicate::Eq(col, _) => {
269 if let Some(hint) = self.cardinality_hints.get(col) {
270 1.0 / hint.cardinality.max(1) as f64
271 } else {
272 0.1 }
274 }
275 QueryPredicate::Range { .. } => 0.25, QueryPredicate::In(_, values) => (values.len() as f64 * 0.1).min(0.5),
277 QueryPredicate::Prefix(_, _) => 0.15, QueryPredicate::TimeRange(_, _) => 0.2,
279 QueryPredicate::Project(_) => 0.1,
280 QueryPredicate::Tenant(_) => 0.05,
281 QueryPredicate::SpanType(_) => 0.15,
282 }
283 }
284}