linkrs_query/planning/statements/paths/
shortest_path_planner.rs

1//! Shortest Path Planner
2//!
3//! Responsible for planning shortest path queries, supporting algorithms such as BFS (Breadth-First Search).
4
5use crate::optimizer::cost::CostModelConfig;
6use crate::planning::statements::seeks::seek_strategy_base::{
7    NodePattern, SeekStrategyContext, SeekStrategySelector, SeekStrategyType,
8};
9use crate::storage::StorageReader;
10use linkrs_core::types::graph_schema::EdgeDirection;
11use linkrs_core::types::VertexId;
12use linkrs_core::{Edge, StorageError, Value, Vertex};
13use std::collections::{HashMap, HashSet, VecDeque};
14use std::sync::Arc;
15
16pub type PlannerError = StorageError;
17
18#[derive(Debug)]
19pub struct ShortestPathPlanner;
20
21impl Default for ShortestPathPlanner {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl ShortestPathPlanner {
28    pub fn new() -> Self {
29        Self
30    }
31
32    pub fn plan_shortest_path(
33        &self,
34        start: &NodePattern,
35        end: &NodePattern,
36        edge_pattern: &EdgePattern,
37        space_id: u64,
38    ) -> Result<ShortestPathPlan, PlannerError> {
39        // Fail fast on explicit endpoint ids: an invalid start or end is a
40        // planning error, while a valid but unreachable pair still yields an
41        // empty path set from `find_shortest_path`.
42        if let Some(vid) = &start.vid {
43            VertexId::try_from(vid)?;
44        }
45        if let Some(vid) = &end.vid {
46            VertexId::try_from(vid)?;
47        }
48        let bfs_config = BfsConfig {
49            max_iterations: 10000,
50            max_path_length: 100,
51            direction: self.extract_direction(edge_pattern)?,
52            edge_types: edge_pattern.types.clone().unwrap_or_default(),
53        };
54
55        let start_context = SeekStrategyContext::new(space_id, start.clone(), vec![]);
56        let selector = SeekStrategySelector::new(CostModelConfig::default());
57        let start_strategy = selector.select_strategy(&start_context);
58
59        let start_finder = match start_strategy {
60            SeekStrategyType::VertexSeek => StartVidSource::VertexSeek(start.clone()),
61            SeekStrategyType::IndexSeek => StartVidSource::IndexScan(start.clone()),
62            SeekStrategyType::PropIndexSeek => StartVidSource::PropIndexScan(start.clone()),
63            SeekStrategyType::VariablePropIndexSeek => {
64                StartVidSource::VariablePropIndexScan(start.clone())
65            }
66            SeekStrategyType::EdgeSeek => StartVidSource::EdgeScan(start.clone()),
67            SeekStrategyType::ScanSeek => StartVidSource::FullScan(start.clone()),
68        };
69
70        let end_condition = EndCondition {
71            pattern: end.clone(),
72        };
73
74        Ok(ShortestPathPlan {
75            start: start_finder,
76            end: end_condition,
77            bfs_config,
78        })
79    }
80
81    fn extract_direction(&self, edge: &EdgePattern) -> Result<EdgeDirection, PlannerError> {
82        Ok(match edge.direction {
83            Some(ref dir) => *dir,
84            None => EdgeDirection::Both,
85        })
86    }
87}
88
89#[derive(Debug, Clone)]
90pub struct EdgePattern {
91    pub types: Option<Vec<String>>,
92    pub direction: Option<EdgeDirection>,
93    pub properties: Vec<(Arc<str>, Value)>,
94}
95
96#[derive(Debug)]
97pub enum StartVidSource {
98    VertexSeek(NodePattern),
99    IndexScan(NodePattern),
100    PropIndexScan(NodePattern),
101    VariablePropIndexScan(NodePattern),
102    EdgeScan(NodePattern),
103    FullScan(NodePattern),
104}
105
106#[derive(Debug, Clone)]
107pub struct EndCondition {
108    pub pattern: NodePattern,
109}
110
111#[derive(Debug)]
112pub struct BfsConfig {
113    pub max_iterations: usize,
114    pub max_path_length: usize,
115    pub direction: EdgeDirection,
116    pub edge_types: Vec<String>,
117}
118
119#[derive(Debug)]
120pub struct ShortestPathPlan {
121    pub start: StartVidSource,
122    pub end: EndCondition,
123    pub bfs_config: BfsConfig,
124}
125
126#[derive(Debug)]
127pub struct ShortestPathResult {
128    pub paths: Vec<ShortestPath>,
129    pub nodes_visited: usize,
130    pub edges_explored: usize,
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Hash)]
134pub struct ShortestPath {
135    pub vertices: Vec<Value>,
136    pub edges: Vec<Edge>,
137}
138
139impl ShortestPathPlanner {
140    pub fn find_shortest_path<S: StorageReader>(
141        &self,
142        storage: &S,
143        plan: &ShortestPathPlan,
144    ) -> Result<ShortestPathResult, PlannerError> {
145        let start_vids = self.resolve_start_vids(storage, &plan.start)?;
146        let end_pattern = &plan.end.pattern;
147
148        let mut all_paths = Vec::new();
149        let mut total_nodes_visited = 0;
150        let mut total_edges_explored = 0;
151
152        for start_vid in start_vids {
153            if let Some(end_vid) = self.resolve_end_vid(storage, end_pattern)? {
154                match self.bfs_search(storage, &start_vid, &end_vid, &plan.bfs_config) {
155                    Ok(Some(path)) => {
156                        total_nodes_visited += path.vertices.len();
157                        total_edges_explored += path.edges.len();
158                        all_paths.push(path);
159                    }
160                    Ok(None) => continue,
161                    Err(e) => return Err(e),
162                }
163            }
164        }
165
166        all_paths.sort_by_key(|a| a.vertices.len());
167
168        Ok(ShortestPathResult {
169            paths: all_paths,
170            nodes_visited: total_nodes_visited,
171            edges_explored: total_edges_explored,
172        })
173    }
174
175    fn resolve_start_vids<S: StorageReader>(
176        &self,
177        storage: &S,
178        start: &StartVidSource,
179    ) -> Result<Vec<Value>, PlannerError> {
180        match start {
181            StartVidSource::VertexSeek(pattern) => match &pattern.vid {
182                Some(vid) => {
183                    let vid_clone: Value = vid.clone();
184                    Ok(vec![vid_clone])
185                }
186                None => self.scan_matching_vertices(storage, pattern),
187            },
188            StartVidSource::IndexScan(pattern) => self.scan_matching_vertices(storage, pattern),
189            StartVidSource::PropIndexScan(pattern) => self.scan_matching_vertices(storage, pattern),
190            StartVidSource::VariablePropIndexScan(pattern) => {
191                self.scan_matching_vertices(storage, pattern)
192            }
193            StartVidSource::EdgeScan(pattern) => self.scan_matching_vertices(storage, pattern),
194            StartVidSource::FullScan(pattern) => self.scan_matching_vertices(storage, pattern),
195        }
196    }
197
198    fn resolve_end_vid<S: StorageReader>(
199        &self,
200        storage: &S,
201        pattern: &NodePattern,
202    ) -> Result<Option<Value>, PlannerError> {
203        match &pattern.vid {
204            Some(vid) => {
205                let vid_clone: Value = vid.clone();
206                Ok(Some(vid_clone))
207            }
208            None => {
209                let vertices = self.scan_matching_vertices(storage, pattern)?;
210                Ok(vertices.first().cloned())
211            }
212        }
213    }
214
215    fn scan_matching_vertices<S: StorageReader>(
216        &self,
217        storage: &S,
218        pattern: &NodePattern,
219    ) -> Result<Vec<Value>, PlannerError> {
220        let vertices = storage.scan_vertices("default")?;
221        let mut matching: Vec<Value> = Vec::new();
222
223        for vertex in vertices {
224            if self.vertex_matches_pattern(&vertex, pattern) {
225                matching.push(Value::from(*vertex.vid()));
226            }
227        }
228
229        Ok(matching)
230    }
231
232    fn vertex_matches_pattern(&self, vertex: &Vertex, pattern: &NodePattern) -> bool {
233        if !pattern.labels.is_empty() {
234            let distinct = pattern
235                .labels
236                .iter()
237                .collect::<std::collections::HashSet<_>>();
238            if distinct.len() > 1 {
239                return false;
240            }
241            let has_all_labels = pattern.labels.iter().all(|label| vertex.tag.name == *label);
242            if !has_all_labels {
243                return false;
244            }
245        }
246
247        for (prop_name, prop_value) in &pattern.properties {
248            let found = vertex
249                .properties()
250                .iter()
251                .any(|(name, value)| name == prop_name && value == prop_value);
252            if !found {
253                return false;
254            }
255        }
256
257        true
258    }
259
260    fn bfs_search<S: StorageReader>(
261        &self,
262        storage: &S,
263        start: &Value,
264        end: &Value,
265        config: &BfsConfig,
266    ) -> Result<Option<ShortestPath>, StorageError> {
267        let mut queue = VecDeque::new();
268        let mut visited: HashSet<Value> = HashSet::new();
269        let mut parent_map: HashMap<Value, (Value, Edge)> = HashMap::new();
270
271        queue.push_back(start.clone());
272        visited.insert(start.clone());
273
274        while let Some(current) = queue.pop_front() {
275            if current == *end {
276                return Ok(self.reconstruct_path(start, end, &parent_map));
277            }
278
279            if parent_map.len() >= config.max_iterations {
280                continue;
281            }
282
283            let current_vid = VertexId::try_from(&current)
284                .map_err(|e| StorageError::invalid_input(e.to_string()))?;
285            let edges = storage.get_node_edges(
286                "default",
287                &current_vid,
288                config.direction,
289                &config.edge_types,
290            )?;
291
292            for edge in edges {
293                let neighbor = if Value::from(edge.src) == current {
294                    Value::from(edge.dst)
295                } else {
296                    Value::from(edge.src)
297                };
298
299                if visited.insert(neighbor.clone()) {
300                    parent_map.insert(neighbor.clone(), (current.clone(), edge.clone()));
301                    queue.push_back(neighbor);
302                }
303            }
304        }
305
306        Ok(None)
307    }
308
309    fn reconstruct_path(
310        &self,
311        start: &Value,
312        end: &Value,
313        parent_map: &HashMap<Value, (Value, Edge)>,
314    ) -> Option<ShortestPath> {
315        let mut vertices = Vec::new();
316        let mut edges = Vec::new();
317        let mut current = end.clone();
318
319        let mut path = Vec::new();
320        path.push(current.clone());
321
322        while let Some((parent, edge)) = parent_map.get(&current) {
323            vertices.push(current.clone());
324            edges.push(edge.clone());
325            current = parent.clone();
326            path.push(current.clone());
327        }
328
329        vertices.push(start.clone());
330        vertices.reverse();
331        edges.reverse();
332
333        Some(ShortestPath { vertices, edges })
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340
341    #[test]
342    fn test_shortest_path_planner_new() {
343        let _planner = ShortestPathPlanner::new();
344        // The test has been successful; reaching this point indicates that the goal has been achieved.
345    }
346
347    #[test]
348    fn test_edge_pattern() {
349        let _pattern = EdgePattern {
350            types: Some(vec!["follows".to_string()]),
351            direction: Some(EdgeDirection::Out),
352            properties: vec![],
353        };
354        // The test has been successful; reaching this point indicates that the goal has been achieved.
355    }
356}
linkrs_query/planning/statements/paths/shortest_path_planner.rs

linkrs_query/planning/statements/paths/
shortest_path_planner.rs
