linkrs_query/executor/traversal/
graph_reader.rs1use crate::storage::QueryStorage;
2use linkrs_core::error::QueryError;
3use linkrs_core::types::storage_ids::VertexId;
4use linkrs_core::{Edge, EdgeDirection, Vertex};
5
6pub type EdgeKey = (VertexId, VertexId, String, i64);
9
10pub fn edge_key(edge: &Edge) -> EdgeKey {
13 (
14 *edge.src(),
15 *edge.dst(),
16 edge.edge_type().to_string(),
17 edge.ranking(),
18 )
19}
20
21pub fn resolve_edge_endpoint_tags(
27 storage: &dyn QueryStorage,
28 space_name: &str,
29 edge_type: &str,
30) -> Option<(String, String)> {
31 let info = storage.get_edge_type(space_name, edge_type).ok()??;
32 if info.src_tag_name.is_empty() || info.dst_tag_name.is_empty() {
33 return None;
34 }
35 Some((info.src_tag_name, info.dst_tag_name))
36}
37
38pub fn resolve_neighbor_tag(
45 storage: &dyn QueryStorage,
46 space_name: &str,
47 edge: &Edge,
48 neighbor_id: &VertexId,
49 override_tag: &str,
50) -> Option<String> {
51 if !override_tag.is_empty() {
52 return Some(override_tag.to_string());
53 }
54 let (src_tag, dst_tag) = resolve_edge_endpoint_tags(storage, space_name, edge.edge_type())?;
55 if neighbor_id == edge.dst() {
56 Some(dst_tag)
57 } else {
58 Some(src_tag)
59 }
60}
61
62pub fn resolve_seed_tag(
70 storage: &dyn QueryStorage,
71 space_name: &str,
72 edge_types: &[String],
73 direction: EdgeDirection,
74 override_tag: &str,
75) -> Result<String, QueryError> {
76 if !override_tag.is_empty() {
77 return Ok(override_tag.to_string());
78 }
79 if edge_types.is_empty() {
80 return Err(QueryError::execution(
81 "Traversal from a bare id requires an edge type to determine the seed label"
82 .to_string(),
83 ));
84 }
85 let mut seed_tags: Vec<String> = Vec::new();
86 for edge_type in edge_types {
87 let info = storage
88 .get_edge_type(space_name, edge_type)
89 .map_err(|e| QueryError::execution(format!("Storage error: {e}")))?
90 .ok_or_else(|| QueryError::execution(format!("Edge type '{edge_type}' not found")))?;
91 let side_tags: Vec<&str> = match direction {
92 EdgeDirection::Out => vec![info.src_tag_name.as_str()],
93 EdgeDirection::In => vec![info.dst_tag_name.as_str()],
94 EdgeDirection::Both => {
95 if info.src_tag_name != info.dst_tag_name {
96 return Err(QueryError::execution(format!(
97 "Bidirectional traversal over '{edge_type}' cannot determine the seed label: endpoint labels differ"
98 )));
99 }
100 vec![info.src_tag_name.as_str()]
101 }
102 };
103 for tag in side_tags {
104 if !tag.is_empty() && !seed_tags.iter().any(|t| t == tag) {
105 seed_tags.push(tag.to_string());
106 }
107 }
108 }
109 match seed_tags.as_slice() {
110 [tag] => Ok(tag.clone()),
111 _ => Err(QueryError::execution(
112 "Traversal from a bare id requires exactly one seed label".to_string(),
113 )),
114 }
115}
116
117pub struct TraversalGraphReader<'a> {
118 storage: &'a dyn QueryStorage,
119}
120
121impl<'a> TraversalGraphReader<'a> {
122 pub fn new(storage: &'a dyn QueryStorage) -> Self {
123 Self { storage }
124 }
125
126 pub fn get_vertex(&self, space_name: &str, tag: &str, vertex_id: &VertexId) -> Option<Vertex> {
127 if tag.is_empty() {
128 return None;
129 }
130 self.storage.get_vertex(space_name, tag, vertex_id).ok()?
131 }
132
133 pub fn get_neighbor_vertex(
136 &self,
137 space_name: &str,
138 edge: &Edge,
139 neighbor_id: &VertexId,
140 override_tag: &str,
141 ) -> Option<Vertex> {
142 let tag = resolve_neighbor_tag(self.storage, space_name, edge, neighbor_id, override_tag)?;
143 self.get_vertex(space_name, &tag, neighbor_id)
144 }
145
146 pub fn get_edges(
147 &self,
148 space_name: &str,
149 vertex_id: &VertexId,
150 direction: EdgeDirection,
151 edge_types: &[String],
152 ) -> Vec<Edge> {
153 self.storage
154 .get_node_edges(space_name, vertex_id, direction, edge_types)
155 .unwrap_or_default()
156 }
157
158 pub fn filter_edges<'b>(&self, edges: &'b [Edge], edge_types: &[String]) -> Vec<&'b Edge> {
159 if edge_types.is_empty() {
160 edges.iter().collect()
161 } else {
162 edges
163 .iter()
164 .filter(|e| edge_types.contains(&e.edge_type))
165 .collect()
166 }
167 }
168
169 pub fn get_neighbor_id(
170 &self,
171 edge: &Edge,
172 current_id: &VertexId,
173 direction: EdgeDirection,
174 ) -> VertexId {
175 match direction {
176 EdgeDirection::Out => *edge.dst(),
177 EdgeDirection::In => *edge.src(),
178 EdgeDirection::Both => {
179 if edge.src() == current_id {
180 *edge.dst()
181 } else {
182 *edge.src()
183 }
184 }
185 }
186 }
187
188 pub fn read_neighbors(
189 &self,
190 space_name: &str,
191 tag: &str,
192 vertex_id: &VertexId,
193 direction: EdgeDirection,
194 edge_types: &[String],
195 ) -> Vec<(Vertex, Edge)> {
196 let edges = self.get_edges(space_name, vertex_id, direction, edge_types);
197 if edges.is_empty() {
198 return Vec::new();
199 }
200 let mut neighbor_ids = Vec::with_capacity(edges.len());
201 for edge in &edges {
202 neighbor_ids.push(self.get_neighbor_id(edge, vertex_id, direction));
203 }
204 let vertices = self
205 .storage
206 .get_vertices_batch(space_name, tag, &neighbor_ids)
207 .unwrap_or_default();
208 let mut result = Vec::with_capacity(edges.len());
209 for (edge, vertex) in edges.into_iter().zip(vertices) {
210 if let Some(vertex) = vertex {
211 result.push((vertex, edge));
212 }
213 }
214 result
215 }
216}