linkrs_query/optimizer/heuristic/join_optimization/
join_to_expand.rs

1//! Rules for converting Vertex-Edge JOIN to ExpandAll
2//!
3//! This rule converts a JOIN between vertices and edges into an ExpandAll operation,
4//! which is more efficient for graph traversal patterns.
5//!
6//! # Conversion example
7//!
8//! Before:
9//! ```text
10//!   ScanVertices(v) → InnerJoin(ON v.id = e._src) → ScanEdges(e)
11//! ```
12//!
13//! After:
14//! ```text
15//!   ScanVertices(v) → ExpandAll(edge_types, direction=OUT)
16//! ```
17//!
18//! # Applicable Conditions
19//!
20//! One side is ScanVertices, the other is ScanEdges
21//! JOIN condition connects vertex ID to edge source/destination
22//! The edge types can be determined from the ScanEdges
23
24use crate::optimizer::heuristic::context::RewriteContext;
25use crate::optimizer::heuristic::pattern::Pattern;
26use crate::optimizer::heuristic::result::{RewriteResult, TransformResult};
27use crate::optimizer::heuristic::rule::RewriteRule;
28use crate::planning::plan::core::nodes::base::plan_node_traits::MultipleInputNode;
29use crate::planning::plan::core::nodes::join::join_node::InnerJoinNode;
30use crate::planning::plan::core::nodes::traversal::traversal_node::ExpandAllNode;
31use crate::planning::plan::PlanNodeEnum;
32use linkrs_core::types::expr::contextual::ContextualExpression;
33use linkrs_core::types::expr::visitor::ExpressionVisitor;
34use linkrs_core::types::expr::visitor_collectors::VariableCollector;
35
36/// Rules for converting Vertex-Edge JOIN to ExpandAll
37#[derive(Debug)]
38pub struct JoinToExpandRule;
39
40impl JoinToExpandRule {
41    pub fn new() -> Self {
42        Self
43    }
44
45    fn extract_join_key_variable(&self, expr: &ContextualExpression) -> Option<String> {
46        if let Some(expr_meta) = expr.expression() {
47            let mut collector = VariableCollector::new();
48            collector.visit(expr_meta.inner());
49            collector.variables.into_iter().next()
50        } else {
51            None
52        }
53    }
54
55    fn analyze_join_condition(&self, probe_keys: &[ContextualExpression]) -> Option<String> {
56        if probe_keys.len() != 1 {
57            return None;
58        }
59
60        self.extract_join_key_variable(&probe_keys[0])
61    }
62
63    fn determine_direction(&self, _edge_var: &str, join_key: &str) -> Option<&'static str> {
64        if join_key.ends_with("._src") || join_key.contains("src") {
65            Some("OUT")
66        } else if join_key.ends_with("._dst") || join_key.contains("dst") {
67            Some("IN")
68        } else {
69            None
70        }
71    }
72
73    fn apply_to_inner_join(&self, join: &InnerJoinNode) -> RewriteResult<Option<TransformResult>> {
74        let left = join.left_input();
75        let right = join.right_input();
76
77        let (scan_vertices, scan_edges, vertex_on_left) = match (left, right) {
78            (PlanNodeEnum::ScanVertices(v), PlanNodeEnum::ScanEdges(e)) => (v, e, true),
79            (PlanNodeEnum::ScanEdges(e), PlanNodeEnum::ScanVertices(v)) => (v, e, false),
80            _ => return Ok(None),
81        };
82
83        let probe_keys = if vertex_on_left {
84            join.probe_keys()
85        } else {
86            join.hash_keys()
87        };
88
89        let probe_var = match self.analyze_join_condition(probe_keys) {
90            Some(var) => var,
91            None => return Ok(None),
92        };
93
94        let direction = self.determine_direction(&probe_var, &probe_var);
95        let direction = match direction {
96            Some(d) => d,
97            None => return Ok(None),
98        };
99
100        let edge_types = scan_edges
101            .edge_type()
102            .map(|et| vec![et])
103            .unwrap_or_default();
104
105        let mut expand_all = ExpandAllNode::new(scan_vertices.space_id(), edge_types, direction);
106
107        expand_all.add_input(left.clone());
108
109        let mut result = TransformResult::new();
110        result.erase_curr = true;
111        result.add_new_node(PlanNodeEnum::ExpandAll(expand_all));
112
113        Ok(Some(result))
114    }
115}
116
117impl Default for JoinToExpandRule {
118    fn default() -> Self {
119        Self::new()
120    }
121}
122
123impl RewriteRule for JoinToExpandRule {
124    fn name(&self) -> &'static str {
125        "JoinToExpandRule"
126    }
127
128    fn pattern(&self) -> Pattern {
129        Pattern::new_with_name("InnerJoin")
130    }
131
132    fn apply(
133        &self,
134        _ctx: &mut RewriteContext,
135        node: &PlanNodeEnum,
136    ) -> RewriteResult<Option<TransformResult>> {
137        match node {
138            PlanNodeEnum::InnerJoin(join) => self.apply_to_inner_join(join),
139            _ => Ok(None),
140        }
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn test_rule_name() {
150        let rule = JoinToExpandRule::new();
151        assert_eq!(rule.name(), "JoinToExpandRule");
152    }
153
154    #[test]
155    fn test_rule_pattern() {
156        let rule = JoinToExpandRule::new();
157        let pattern = rule.pattern();
158        assert!(pattern.node.is_some());
159    }
160}
linkrs_query/optimizer/heuristic/join_optimization/join_to_expand.rs

linkrs_query/optimizer/heuristic/join_optimization/
join_to_expand.rs
