linkrs_query/planning/join_order/
query_graph.rs

1//! Query graph representation for join ordering.
2//!
3//! A `QueryGraph` captures the pattern part of a MATCH query as a set of
4//! query nodes and query rels with positional indexes. The join order
5//! enumerator works on this graph instead of the AST so that DP enumeration
6//! and WCO detection share one structural view.
7
8use std::collections::HashMap;
9use std::str::FromStr;
10use std::sync::Arc;
11
12/// Traversal direction of a query rel.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
14pub enum ExtendDirection {
15    /// Outgoing: src -> dst.
16    #[default]
17    Out,
18    /// Incoming: src <- dst.
19    In,
20    /// Bidirectional.
21    Both,
22}
23
24impl ExtendDirection {
25    pub fn as_str(&self) -> &'static str {
26        match self {
27            ExtendDirection::Out => "OUT",
28            ExtendDirection::In => "IN",
29            ExtendDirection::Both => "BOTH",
30        }
31    }
32}
33
34impl FromStr for ExtendDirection {
35    type Err = String;
36
37    fn from_str(s: &str) -> Result<Self, Self::Err> {
38        match s.to_ascii_uppercase().as_str() {
39            "OUT" => Ok(ExtendDirection::Out),
40            "IN" => Ok(ExtendDirection::In),
41            "BOTH" | "BIDIRECT" | "UNDIRECTED" => Ok(ExtendDirection::Both),
42            other => Err(format!("Unknown direction: {}", other)),
43        }
44    }
45}
46
47impl From<linkrs_core::types::graph_schema::EdgeDirection> for ExtendDirection {
48    fn from(d: linkrs_core::types::graph_schema::EdgeDirection) -> Self {
49        use linkrs_core::types::graph_schema::EdgeDirection as Core;
50        match d {
51            Core::Out => ExtendDirection::Out,
52            Core::In => ExtendDirection::In,
53            Core::Both => ExtendDirection::Both,
54        }
55    }
56}
57
58/// A single query node (pattern variable bound to vertices).
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct QueryNode {
61    /// Positional name used for graph wiring (usually the variable).
62    pub name: String,
63    /// Query variable, e.g. `a` in `(a)`.
64    pub variable: String,
65    /// Label predicates on the node.
66    pub labels: Vec<String>,
67}
68
69impl QueryNode {
70    pub fn new(name: impl Into<String>, variable: impl Into<String>) -> Self {
71        let variable = variable.into();
72        let name = name.into();
73        Self {
74            name,
75            variable,
76            labels: Vec::new(),
77        }
78    }
79
80    pub fn with_labels(mut self, labels: Vec<String>) -> Self {
81        self.labels = labels;
82        self
83    }
84}
85
86/// A single query rel (pattern variable bound to edges).
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct QueryRel {
89    /// Positional name used for graph wiring (usually the variable).
90    pub name: String,
91    /// Query variable, e.g. `e` in `-[e]->`.
92    pub variable: String,
93    /// Edge type predicates.
94    pub edge_types: Vec<String>,
95    /// Source node name.
96    pub src_node_name: String,
97    /// Destination node name.
98    pub dst_node_name: String,
99    /// Traversal direction.
100    pub direction: ExtendDirection,
101}
102
103impl QueryRel {
104    pub fn new(
105        name: impl Into<String>,
106        variable: impl Into<String>,
107        src_node_name: impl Into<String>,
108        dst_node_name: impl Into<String>,
109    ) -> Self {
110        let variable = variable.into();
111        Self {
112            name: name.into(),
113            variable,
114            edge_types: Vec::new(),
115            src_node_name: src_node_name.into(),
116            dst_node_name: dst_node_name.into(),
117            direction: ExtendDirection::Out,
118        }
119    }
120
121    pub fn with_direction(mut self, direction: ExtendDirection) -> Self {
122        self.direction = direction;
123        self
124    }
125
126    pub fn with_edge_types(mut self, edge_types: Vec<String>) -> Self {
127        self.edge_types = edge_types;
128        self
129    }
130
131    /// The other endpoint of this rel given one endpoint name.
132    pub fn other_node(&self, node_name: &str) -> Option<&str> {
133        if self.src_node_name == node_name {
134            Some(self.dst_node_name.as_str())
135        } else if self.dst_node_name == node_name {
136            Some(self.src_node_name.as_str())
137        } else {
138            None
139        }
140    }
141
142    /// Whether this rel touches the given node name.
143    pub fn touches(&self, node_name: &str) -> bool {
144        self.src_node_name == node_name || self.dst_node_name == node_name
145    }
146}
147
148/// Pattern graph used by the join order enumerator.
149#[derive(Debug, Clone, Default)]
150pub struct QueryGraph {
151    /// All query nodes in positional order.
152    pub query_nodes: Vec<Arc<QueryNode>>,
153    /// All query rels in positional order.
154    pub query_rels: Vec<Arc<QueryRel>>,
155    /// Node name -> position index.
156    pub node_name_to_pos: HashMap<String, usize>,
157    /// Rel name -> position index.
158    pub rel_name_to_pos: HashMap<String, usize>,
159}
160
161impl QueryGraph {
162    pub fn new() -> Self {
163        Self::default()
164    }
165
166    /// Add a query node and return its position.
167    pub fn add_node(&mut self, node: QueryNode) -> usize {
168        let pos = self.query_nodes.len();
169        self.node_name_to_pos.insert(node.name.clone(), pos);
170        self.query_nodes.push(Arc::new(node));
171        pos
172    }
173
174    /// Add a query rel and return its position.
175    pub fn add_rel(&mut self, rel: QueryRel) -> usize {
176        let pos = self.query_rels.len();
177        self.rel_name_to_pos.insert(rel.name.clone(), pos);
178        self.query_rels.push(Arc::new(rel));
179        pos
180    }
181
182    pub fn num_nodes(&self) -> usize {
183        self.query_nodes.len()
184    }
185
186    pub fn num_rels(&self) -> usize {
187        self.query_rels.len()
188    }
189
190    /// Total pattern variables (nodes + rels), capped by the bitset width.
191    pub fn num_variables(&self) -> usize {
192        self.num_nodes() + self.num_rels()
193    }
194
195    pub fn node_pos(&self, name: &str) -> Option<usize> {
196        self.node_name_to_pos.get(name).copied()
197    }
198
199    pub fn rel_pos(&self, name: &str) -> Option<usize> {
200        self.rel_name_to_pos.get(name).copied()
201    }
202
203    /// Positions of rels incident to the given node position.
204    pub fn incident_rel_positions(&self, node_pos: usize) -> Vec<usize> {
205        let Some(node) = self.query_nodes.get(node_pos) else {
206            return Vec::new();
207        };
208        self.query_rels
209            .iter()
210            .enumerate()
211            .filter(|(_, r)| r.touches(&node.name))
212            .map(|(i, _)| i)
213            .collect()
214    }
215
216    /// Node positions connected to the given rel position.
217    pub fn rel_endpoint_positions(&self, rel_pos: usize) -> Option<(usize, usize)> {
218        let rel = self.query_rels.get(rel_pos)?;
219        let src = self.node_name_to_pos.get(&rel.src_node_name).copied()?;
220        let dst = self.node_name_to_pos.get(&rel.dst_node_name).copied()?;
221        Some((src, dst))
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    fn triangle_graph() -> QueryGraph {
230        let mut qg = QueryGraph::new();
231        qg.add_node(QueryNode::new("a", "a"));
232        qg.add_node(QueryNode::new("b", "b"));
233        qg.add_node(QueryNode::new("c", "c"));
234        qg.add_rel(QueryRel::new("e1", "e1", "a", "b"));
235        qg.add_rel(QueryRel::new("e2", "e2", "a", "c"));
236        qg.add_rel(QueryRel::new("e3", "e3", "b", "c"));
237        qg
238    }
239
240    #[test]
241    fn node_and_rel_positions() {
242        let qg = triangle_graph();
243        assert_eq!(qg.num_nodes(), 3);
244        assert_eq!(qg.num_rels(), 3);
245        assert_eq!(qg.node_pos("a"), Some(0));
246        assert_eq!(qg.rel_pos("e3"), Some(2));
247        assert_eq!(qg.num_variables(), 6);
248    }
249
250    #[test]
251    fn incident_rels() {
252        let qg = triangle_graph();
253        let incident = qg.incident_rel_positions(0);
254        assert_eq!(incident.len(), 2);
255        assert!(incident.contains(&0));
256        assert!(incident.contains(&1));
257    }
258
259    #[test]
260    fn rel_endpoints() {
261        let qg = triangle_graph();
262        assert_eq!(qg.rel_endpoint_positions(0), Some((0, 1)));
263        assert_eq!(qg.rel_endpoint_positions(2), Some((1, 2)));
264        assert_eq!(qg.rel_endpoint_positions(99), None);
265    }
266
267    #[test]
268    fn direction_roundtrip() {
269        assert_eq!(
270            "in".parse::<ExtendDirection>().unwrap(),
271            ExtendDirection::In
272        );
273        assert_eq!(
274            "both".parse::<ExtendDirection>().unwrap(),
275            ExtendDirection::Both
276        );
277        assert_eq!(ExtendDirection::Out.as_str(), "OUT");
278    }
279}
linkrs_query/planning/join_order/query_graph.rs

linkrs_query/planning/join_order/
query_graph.rs
