linkrs_query/optimizer/heuristic/join_optimization/
index_join_selection.rs

1//! Rules for index-based JOIN selection
2//!
3//! This rule analyzes JOIN conditions and determines whether to use
4//! index-based JOIN instead of hash JOIN when appropriate indexes exist.
5//!
6//! # Conversion examples
7//!
8//! ## Case 1: Index JOIN for vertex ID lookup
9//! Before:
10//! ```text
11//!   InnerJoin(ON v.id = e._src) → ScanVertices(v) → ScanEdges(e)
12//! ```
13//! After (if index exists on e._src):
14//! ```text
15//!   IndexJoin(index=e._src_idx) → ScanVertices(v) → IndexScanEdge(e)
16//! ```
17//!
18//! ## Case 2: Index JOIN for edge properties
19//! Before:
20//! ```text
21//!   InnerJoin(ON v.prop = e.prop) → ScanVertices(v) → ScanEdges(e)
22//! ```
23//! After (if indexes exist):
24//! ```text
25//!   IndexJoin using property indexes
26//! ```
27
28use crate::optimizer::heuristic::context::RewriteContext;
29use crate::optimizer::heuristic::pattern::Pattern;
30use crate::optimizer::heuristic::result::{RewriteResult, TransformResult};
31use crate::optimizer::heuristic::rule::RewriteRule;
32use crate::planning::plan::core::nodes::join::join_node::InnerJoinNode;
33use crate::planning::plan::PlanNodeEnum;
34use linkrs_core::types::expr::contextual::ContextualExpression;
35use linkrs_core::types::expr::visitor::ExpressionVisitor;
36use linkrs_core::types::expr::visitor_collectors::VariableCollector;
37
38/// Rules for index-based JOIN selection
39#[derive(Debug)]
40pub struct IndexJoinSelectionRule;
41
42impl IndexJoinSelectionRule {
43    pub fn new() -> Self {
44        Self
45    }
46
47    fn extract_join_key_variable(&self, expr: &ContextualExpression) -> Option<String> {
48        if let Some(expr_meta) = expr.expression() {
49            let mut collector = VariableCollector::new();
50            collector.visit(expr_meta.inner());
51            collector.variables.into_iter().next()
52        } else {
53            None
54        }
55    }
56
57    fn is_indexable_join_key(&self, key: &str) -> bool {
58        key.ends_with(".id")
59            || key.ends_with("._src")
60            || key.ends_with("._dst")
61            || key.contains("id")
62    }
63
64    fn apply_to_inner_join(&self, join: &InnerJoinNode) -> RewriteResult<Option<TransformResult>> {
65        let hash_keys = join.hash_keys();
66        let probe_keys = join.probe_keys();
67
68        if hash_keys.len() != 1 || probe_keys.len() != 1 {
69            return Ok(None);
70        }
71
72        let hash_var = match self.extract_join_key_variable(&hash_keys[0]) {
73            Some(v) => v,
74            None => return Ok(None),
75        };
76
77        let probe_var = match self.extract_join_key_variable(&probe_keys[0]) {
78            Some(v) => v,
79            None => return Ok(None),
80        };
81
82        if !self.is_indexable_join_key(&hash_var) && !self.is_indexable_join_key(&probe_var) {
83            return Ok(None);
84        }
85
86        Ok(None)
87    }
88}
89
90impl Default for IndexJoinSelectionRule {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96impl RewriteRule for IndexJoinSelectionRule {
97    fn name(&self) -> &'static str {
98        "IndexJoinSelectionRule"
99    }
100
101    fn pattern(&self) -> Pattern {
102        Pattern::new_with_name("InnerJoin")
103    }
104
105    fn apply(
106        &self,
107        _ctx: &mut RewriteContext,
108        node: &PlanNodeEnum,
109    ) -> RewriteResult<Option<TransformResult>> {
110        match node {
111            PlanNodeEnum::InnerJoin(join) => self.apply_to_inner_join(join),
112            _ => Ok(None),
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_rule_name() {
123        let rule = IndexJoinSelectionRule::new();
124        assert_eq!(rule.name(), "IndexJoinSelectionRule");
125    }
126
127    #[test]
128    fn test_rule_pattern() {
129        let rule = IndexJoinSelectionRule::new();
130        let pattern = rule.pattern();
131        assert!(pattern.node.is_some());
132    }
133
134    #[test]
135    fn test_is_indexable_join_key() {
136        let rule = IndexJoinSelectionRule::new();
137        assert!(rule.is_indexable_join_key("v.id"));
138        assert!(rule.is_indexable_join_key("e._src"));
139        assert!(rule.is_indexable_join_key("e._dst"));
140    }
141}
linkrs_query/optimizer/heuristic/join_optimization/index_join_selection.rs

linkrs_query/optimizer/heuristic/join_optimization/
index_join_selection.rs
