Skip to main content

linkrs_query/executor/traversal/
graph_reader.rs

1use crate::storage::QueryStorage;
2use linkrs_core::error::QueryError;
3use linkrs_core::types::storage_ids::VertexId;
4use linkrs_core::{Edge, EdgeDirection, Vertex};
5
6/// Identity of one edge inside a single path: source, destination, edge
7/// type and ranking. Parallel edges between the same pair stay distinct.
8pub type EdgeKey = (VertexId, VertexId, String, i64);
9
10/// Path-local identity of `edge`, the key trail (edge-unique) enumeration
11/// rejects when an edge repeats within one path.
12pub 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
21/// Resolve the endpoint tag pair declared by an edge type schema.
22///
23/// Returns `None` when the edge type is unknown or its schema declares no
24/// endpoint labels. The schema is the single source of truth for endpoint
25/// labels; there is no label guessing and no empty-tag fabrication.
26pub 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
38/// Resolve the tag used to materialize the neighbor on one side of `edge`.
39///
40/// A plan-provided tag wins when present. Otherwise the tag is derived from
41/// the edge type schema by the neighbor's endpoint side. Returns `None` when
42/// no tag can be determined; callers skip that hop instead of fabricating a
43/// vertex.
44pub 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
62/// Resolve the seed label for a bare-id traversal seed.
63///
64/// A plan-provided label wins. Otherwise every listed edge type must agree
65/// on the seed-side endpoint label (the source label for `Out`, the
66/// destination label for `In`, a single shared label for `Both`): the edge
67/// schemas then determine the seed domain instead of guessing it. Errors on
68/// unknown edge types, untyped schemas, or ambiguous domains.
69pub 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    /// Fetch the neighbor vertex on one side of `edge`, deriving its tag
134    /// from the edge type schema when the plan carries none.
135    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}