linkrs_query/planning/statements/seeks/
seek_strategy_base.rs

1//! Basic Module of Search Strategy Development
2//!
3//! Define the basic types of search strategies and the selectors used for their selection.
4//!
5//! ## Cost-Based Optimization
6//!
7//! This module uses cost-based optimization (CBO) to select the optimal seek strategy.
8//! It considers I/O costs, CPU costs, and selectivity estimates to choose between:
9//! - VertexSeek: Direct vertex ID lookup
10//! - IndexSeek: Tag/label index scan
11//! - PropIndexSeek: Property index scan
12//! - ScanSeek: Full table scan
13
14use std::sync::Arc;
15
16use crate::optimizer::cost::{CostModelConfig, SelectivityEstimator};
17use linkrs_core::types::expr::visitor_checkers::PropertyContainsChecker;
18use linkrs_core::types::Expression;
19use linkrs_core::Value;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum SeekStrategyType {
23    VertexSeek,
24    IndexSeek,
25    PropIndexSeek,
26    VariablePropIndexSeek,
27    EdgeSeek,
28    ScanSeek,
29}
30
31#[derive(Debug)]
32pub struct SeekStrategyContext {
33    pub space_id: u64,
34    pub node_pattern: NodePattern,
35    pub predicates: Vec<Expression>,
36    pub estimated_rows: usize,
37    pub available_indexes: Vec<IndexInfo>,
38    pub total_vertices: usize,
39}
40
41#[derive(Debug, Clone, PartialEq)]
42pub struct NodePattern {
43    pub vid: Option<Value>,
44    pub labels: Vec<String>,
45    pub properties: Vec<(Arc<str>, Value)>,
46}
47
48#[derive(Debug, Clone)]
49pub struct IndexInfo {
50    pub name: String,
51    pub target_type: String,
52    pub target_name: String,
53    pub properties: Vec<String>,
54    pub selectivity: f32,
55    pub field_count: usize,
56    pub is_composite: bool,
57    pub composite_targets: Vec<String>,
58}
59
60impl IndexInfo {
61    pub fn new(
62        name: String,
63        target_type: String,
64        target_name: String,
65        properties: Vec<String>,
66    ) -> Self {
67        let field_count = properties.len();
68        Self {
69            name,
70            target_type,
71            target_name,
72            properties,
73            selectivity: 0.5,
74            field_count,
75            is_composite: false,
76            composite_targets: Vec::new(),
77        }
78    }
79
80    pub fn with_selectivity(mut self, selectivity: f32) -> Self {
81        self.selectivity = selectivity;
82        self
83    }
84
85    pub fn with_composite(mut self, targets: Vec<String>) -> Self {
86        self.is_composite = true;
87        self.composite_targets = targets;
88        self
89    }
90
91    pub fn covers_labels(&self, labels: &[String]) -> bool {
92        if self.is_composite {
93            self.composite_targets.iter().all(|t| labels.contains(t))
94        } else {
95            labels.contains(&self.target_name)
96        }
97    }
98
99    pub fn coverage_score(&self, labels: &[String]) -> f32 {
100        if labels.is_empty() {
101            return 0.0;
102        }
103
104        let covered = if self.is_composite {
105            self.composite_targets
106                .iter()
107                .filter(|t| labels.contains(t))
108                .count()
109        } else {
110            if labels.contains(&self.target_name) {
111                1
112            } else {
113                0
114            }
115        };
116
117        covered as f32 / labels.len() as f32
118    }
119}
120
121#[derive(Debug)]
122pub struct SeekResult {
123    pub vertex_ids: Vec<Value>,
124    pub strategy_used: SeekStrategyType,
125    pub rows_scanned: usize,
126}
127
128impl SeekStrategyContext {
129    pub fn new(space_id: u64, node_pattern: NodePattern, predicates: Vec<Expression>) -> Self {
130        Self {
131            space_id,
132            node_pattern,
133            predicates,
134            estimated_rows: 0,
135            available_indexes: Vec::new(),
136            total_vertices: 10000,
137        }
138    }
139
140    pub fn with_estimated_rows(mut self, rows: usize) -> Self {
141        self.estimated_rows = rows;
142        self
143    }
144
145    pub fn with_indexes(mut self, indexes: Vec<IndexInfo>) -> Self {
146        self.available_indexes = indexes;
147        self
148    }
149
150    pub fn with_total_vertices(mut self, total: usize) -> Self {
151        self.total_vertices = total;
152        self
153    }
154
155    pub fn has_explicit_vid(&self) -> bool {
156        self.node_pattern.vid.is_some()
157    }
158
159    pub fn has_labels(&self) -> bool {
160        !self.node_pattern.labels.is_empty()
161    }
162
163    pub fn has_predicates(&self) -> bool {
164        !self.predicates.is_empty()
165    }
166
167    pub fn get_index_for_labels(&self, labels: &[String]) -> Option<&IndexInfo> {
168        self.available_indexes
169            .iter()
170            .find(|idx| idx.target_type == "tag" && labels.contains(&idx.target_name))
171    }
172
173    pub fn get_index_for_property(&self, property: &str) -> Option<&IndexInfo> {
174        self.available_indexes
175            .iter()
176            .find(|idx| idx.properties.contains(&property.to_string()))
177    }
178
179    pub fn has_property_predicates(&self) -> bool {
180        self.predicates
181            .iter()
182            .any(|pred| matches!(pred, Expression::Binary { .. }))
183    }
184
185    pub fn has_index_for_properties(&self) -> bool {
186        !self.available_indexes.is_empty() && !self.predicates.is_empty()
187    }
188}
189
190#[derive(Debug)]
191pub struct SeekStrategySelector {
192    use_index_threshold: usize,
193    scan_threshold: usize,
194    cost_config: CostModelConfig,
195    selectivity_estimator: Option<Arc<SelectivityEstimator>>,
196}
197
198impl SeekStrategySelector {
199    pub fn new(cost_config: CostModelConfig) -> Self {
200        Self {
201            use_index_threshold: 1000,
202            scan_threshold: 10000,
203            cost_config,
204            selectivity_estimator: None,
205        }
206    }
207
208    pub fn with_thresholds(mut self, use_index: usize, scan: usize) -> Self {
209        self.use_index_threshold = use_index;
210        self.scan_threshold = scan;
211        self
212    }
213
214    pub fn with_selectivity_estimator(mut self, estimator: Arc<SelectivityEstimator>) -> Self {
215        self.selectivity_estimator = Some(estimator);
216        self
217    }
218
219    pub fn select_best_index<'a>(
220        &self,
221        indexes: &'a [IndexInfo],
222        predicates: &[Expression],
223    ) -> Option<&'a IndexInfo> {
224        if indexes.is_empty() {
225            return None;
226        }
227
228        let candidate_indexes: Vec<&IndexInfo> = indexes
229            .iter()
230            .filter(|idx| {
231                idx.properties.iter().any(|prop| {
232                    predicates.iter().any(|pred| {
233                        PropertyContainsChecker::check(pred, std::slice::from_ref(prop))
234                    })
235                })
236            })
237            .collect();
238
239        if candidate_indexes.is_empty() {
240            return indexes.iter().min_by_key(|idx| idx.field_count);
241        }
242
243        candidate_indexes.into_iter().min_by(|a, b| {
244            let field_cmp = a.field_count.cmp(&b.field_count);
245            if field_cmp != std::cmp::Ordering::Equal {
246                return field_cmp;
247            }
248            b.selectivity
249                .partial_cmp(&a.selectivity)
250                .unwrap_or(std::cmp::Ordering::Equal)
251        })
252    }
253
254    pub fn select_strategy(&self, context: &SeekStrategyContext) -> SeekStrategyType {
255        self.select_strategy_with_estimation(context).strategy_type
256    }
257
258    pub fn select_strategy_with_estimation(
259        &self,
260        context: &SeekStrategyContext,
261    ) -> StrategySelection {
262        let candidates = self.evaluate_all_strategies(context);
263
264        candidates
265            .into_iter()
266            .flatten()
267            .min_by(|a, b| {
268                a.cost
269                    .partial_cmp(&b.cost)
270                    .unwrap_or(std::cmp::Ordering::Equal)
271            })
272            .map(|e| StrategySelection {
273                strategy_type: e.strategy_type,
274                estimated_rows: e.estimated_rows,
275            })
276            .unwrap_or(StrategySelection {
277                strategy_type: SeekStrategyType::ScanSeek,
278                estimated_rows: context.total_vertices,
279            })
280    }
281
282    fn evaluate_all_strategies(
283        &self,
284        context: &SeekStrategyContext,
285    ) -> Vec<Option<StrategyEvaluation>> {
286        vec![
287            self.evaluate_vertex_seek(context),
288            self.evaluate_index_seek(context),
289            self.evaluate_prop_index_seek(context),
290            self.evaluate_scan_seek(context),
291        ]
292    }
293
294    fn evaluate_vertex_seek(&self, context: &SeekStrategyContext) -> Option<StrategyEvaluation> {
295        if !context.has_explicit_vid() {
296            return None;
297        }
298
299        let cost = self.cost_config.random_page_cost;
300        Some(StrategyEvaluation {
301            strategy_type: SeekStrategyType::VertexSeek,
302            cost,
303            estimated_rows: 1,
304        })
305    }
306
307    fn evaluate_index_seek(&self, context: &SeekStrategyContext) -> Option<StrategyEvaluation> {
308        let index = context.get_index_for_labels(&context.node_pattern.labels)?;
309        let total_rows = context.total_vertices;
310        let selectivity = index.selectivity as f64;
311        let matching = ((total_rows as f64 * selectivity).max(1.0)) as u64;
312        let cost = Self::shared_index_seek_cost(&self.cost_config, matching);
313        let estimated_rows = matching as usize;
314
315        Some(StrategyEvaluation {
316            strategy_type: SeekStrategyType::IndexSeek,
317            cost,
318            estimated_rows,
319        })
320    }
321
322    /// Shared index-seek cost matching the optimizer calculator:
323    /// sequential index pages plus per-tuple index and table costs.
324    fn shared_index_seek_cost(cost_config: &CostModelConfig, matching_rows: u64) -> f64 {
325        use crate::optimizer::cost::config::ROWS_PER_PAGE;
326        let pages = (matching_rows / ROWS_PER_PAGE).max(1);
327        pages as f64 * cost_config.seq_page_cost
328            + matching_rows as f64 * cost_config.cpu_index_tuple_cost
329            + matching_rows as f64 * cost_config.random_page_cost
330            + matching_rows as f64 * cost_config.cpu_tuple_cost
331    }
332
333    fn evaluate_prop_index_seek(
334        &self,
335        context: &SeekStrategyContext,
336    ) -> Option<StrategyEvaluation> {
337        if !context.has_property_predicates() || !context.has_index_for_properties() {
338            return None;
339        }
340
341        let _index = self.select_best_index(&context.available_indexes, &context.predicates)?;
342        let selectivity = self.estimate_predicate_selectivity(context);
343        let total_rows = context.total_vertices;
344        let matching = ((total_rows as f64 * selectivity).max(1.0)) as u64;
345
346        let cost = Self::shared_index_seek_cost(&self.cost_config, matching);
347        let estimated_rows = matching as usize;
348
349        Some(StrategyEvaluation {
350            strategy_type: SeekStrategyType::PropIndexSeek,
351            cost,
352            estimated_rows,
353        })
354    }
355
356    fn evaluate_scan_seek(&self, context: &SeekStrategyContext) -> Option<StrategyEvaluation> {
357        let total_rows = context.total_vertices;
358        let cost = self.cost_config.seq_page_cost * total_rows as f64
359            + self.cost_config.cpu_tuple_cost * total_rows as f64;
360
361        Some(StrategyEvaluation {
362            strategy_type: SeekStrategyType::ScanSeek,
363            cost,
364            estimated_rows: total_rows,
365        })
366    }
367
368    fn estimate_predicate_selectivity(&self, context: &SeekStrategyContext) -> f64 {
369        if let Some(ref estimator) = self.selectivity_estimator {
370            let mut total_selectivity = 1.0;
371            for pred in &context.predicates {
372                let sel = estimator.estimate_from_expression(None, pred, None);
373                total_selectivity *= sel;
374            }
375            return total_selectivity;
376        }
377
378        if context.predicates.len() == 1 {
379            return 0.1;
380        }
381
382        0.1_f64.powi(context.predicates.len() as i32)
383    }
384}
385
386#[derive(Debug, Clone)]
387struct StrategyEvaluation {
388    strategy_type: SeekStrategyType,
389    cost: f64,
390    estimated_rows: usize,
391}
392
393#[derive(Debug, Clone)]
394pub struct StrategySelection {
395    pub strategy_type: SeekStrategyType,
396    pub estimated_rows: usize,
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402
403    #[test]
404    fn test_index_info_creation() {
405        let index = IndexInfo::new(
406            "person_name_idx".to_string(),
407            "tag".to_string(),
408            "person".to_string(),
409            vec!["name".to_string()],
410        );
411        assert_eq!(index.name, "person_name_idx");
412        assert_eq!(index.field_count, 1);
413        assert!(!index.is_composite);
414    }
415
416    #[test]
417    fn test_index_info_with_composite() {
418        let index = IndexInfo::new(
419            "composite_idx".to_string(),
420            "tag".to_string(),
421            "person".to_string(),
422            vec!["name".to_string(), "age".to_string()],
423        )
424        .with_composite(vec!["person".to_string(), "employee".to_string()]);
425
426        assert!(index.is_composite);
427        assert!(index.covers_labels(&["person".to_string(), "employee".to_string()]));
428        assert!(!index.covers_labels(&["person".to_string(), "other".to_string()]));
429    }
430
431    #[test]
432    fn test_seek_strategy_context() {
433        let node_pattern = NodePattern {
434            vid: Some(Value::Int(1)),
435            labels: vec!["person".to_string()],
436            properties: vec![],
437        };
438        let context = SeekStrategyContext::new(1, node_pattern, vec![]);
439
440        assert!(context.has_explicit_vid());
441        assert!(context.has_labels());
442        assert!(!context.has_predicates());
443    }
444
445    #[test]
446    fn test_seek_strategy_selector_vertex_seek() {
447        let selector = SeekStrategySelector::new(CostModelConfig::default());
448
449        let node_pattern = NodePattern {
450            vid: Some(Value::Int(1)),
451            labels: vec![],
452            properties: vec![],
453        };
454        let context = SeekStrategyContext::new(1, node_pattern, vec![]);
455
456        let strategy = selector.select_strategy(&context);
457        assert_eq!(strategy, SeekStrategyType::VertexSeek);
458    }
459
460    #[test]
461    fn test_seek_strategy_selector_scan_seek() {
462        let selector = SeekStrategySelector::new(CostModelConfig::default());
463
464        let node_pattern = NodePattern {
465            vid: None,
466            labels: vec![],
467            properties: vec![],
468        };
469        let context = SeekStrategyContext::new(1, node_pattern, vec![]);
470
471        let strategy = selector.select_strategy(&context);
472        assert_eq!(strategy, SeekStrategyType::ScanSeek);
473    }
474}
linkrs_query/planning/statements/seeks/seek_strategy_base.rs

linkrs_query/planning/statements/seeks/
seek_strategy_base.rs
