Skip to main content

nodedb_types/
result.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Query result types returned by the `NodeDb` trait methods.
4
5use std::collections::HashMap;
6
7use serde::{Deserialize, Serialize};
8
9use crate::id::{EdgeId, NodeId};
10use crate::value::Value;
11
12/// Result of a k-NN vector search.
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct SearchResult {
15    /// Document/vector ID.
16    pub id: String,
17    /// Corresponding graph node ID (if the vector is linked to a node).
18    pub node_id: Option<NodeId>,
19    /// Distance from the query vector under the configured metric.
20    pub distance: f32,
21    /// Optional metadata fields attached to this vector.
22    pub metadata: HashMap<String, Value>,
23}
24
25/// Result of a graph traversal (BFS/DFS).
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
27pub struct SubGraph {
28    /// Nodes discovered during traversal, keyed by node ID.
29    pub nodes: Vec<SubGraphNode>,
30    /// Edges traversed.
31    pub edges: Vec<SubGraphEdge>,
32}
33
34/// A node in a traversal result.
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
36pub struct SubGraphNode {
37    /// Node identifier.
38    pub id: NodeId,
39    /// Depth at which this node was discovered (0 = start node).
40    pub depth: u8,
41    /// Optional properties attached to this node.
42    pub properties: HashMap<String, Value>,
43}
44
45/// An edge in a traversal result.
46#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
47pub struct SubGraphEdge {
48    /// Edge identifier.
49    pub id: EdgeId,
50    /// Source node.
51    pub from: NodeId,
52    /// Target node.
53    pub to: NodeId,
54    /// Edge label/type (e.g., "KNOWS", "RELATES_TO").
55    pub label: String,
56    /// Optional properties attached to this edge.
57    pub properties: HashMap<String, Value>,
58}
59
60/// Result of a SQL query or multi-modal query.
61#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
62pub struct QueryResult {
63    /// Column names.
64    pub columns: Vec<String>,
65    /// Rows of values. Each row has one value per column.
66    pub rows: Vec<Vec<Value>>,
67    /// Number of rows affected (for INSERT/UPDATE/DELETE).
68    pub rows_affected: u64,
69}
70
71impl QueryResult {
72    /// Empty result set with no columns.
73    pub fn empty() -> Self {
74        Self {
75            columns: Vec::new(),
76            rows: Vec::new(),
77            rows_affected: 0,
78        }
79    }
80
81    /// Number of rows in the result set.
82    pub fn row_count(&self) -> usize {
83        self.rows.len()
84    }
85
86    /// Get a single row as a `HashMap<column_name, value>`.
87    pub fn row_as_map(&self, index: usize) -> Option<HashMap<&str, &Value>> {
88        let row = self.rows.get(index)?;
89        Some(
90            self.columns
91                .iter()
92                .zip(row.iter())
93                .map(|(col, val)| (col.as_str(), val))
94                .collect(),
95        )
96    }
97}
98
99impl SubGraph {
100    /// Empty subgraph.
101    pub fn empty() -> Self {
102        Self {
103            nodes: Vec::new(),
104            edges: Vec::new(),
105        }
106    }
107
108    /// Number of nodes in the subgraph.
109    pub fn node_count(&self) -> usize {
110        self.nodes.len()
111    }
112
113    /// Number of edges in the subgraph.
114    pub fn edge_count(&self) -> usize {
115        self.edges.len()
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn search_result() {
125        let r = SearchResult {
126            id: "vec-1".into(),
127            node_id: Some(NodeId::try_new("node-1").expect("test fixture")),
128            distance: 0.123,
129            metadata: HashMap::new(),
130        };
131        assert_eq!(r.id, "vec-1");
132        assert!(r.distance < 1.0);
133    }
134
135    #[test]
136    fn subgraph_construction() {
137        let mut sg = SubGraph::empty();
138        sg.nodes.push(SubGraphNode {
139            id: NodeId::try_new("a").expect("test fixture"),
140            depth: 0,
141            properties: HashMap::new(),
142        });
143        sg.nodes.push(SubGraphNode {
144            id: NodeId::try_new("b").expect("test fixture"),
145            depth: 1,
146            properties: HashMap::new(),
147        });
148        sg.edges.push(SubGraphEdge {
149            id: EdgeId::try_first(
150                NodeId::try_new("a").expect("test fixture"),
151                NodeId::try_new("b").expect("test fixture"),
152                "KNOWS",
153            )
154            .expect("test fixture"),
155            from: NodeId::try_new("a").expect("test fixture"),
156            to: NodeId::try_new("b").expect("test fixture"),
157            label: "KNOWS".into(),
158            properties: HashMap::new(),
159        });
160        assert_eq!(sg.node_count(), 2);
161        assert_eq!(sg.edge_count(), 1);
162    }
163
164    #[test]
165    fn query_result_row_as_map() {
166        let qr = QueryResult {
167            columns: vec!["name".into(), "age".into()],
168            rows: vec![vec![Value::String("Alice".into()), Value::Integer(30)]],
169            rows_affected: 0,
170        };
171        let row = qr.row_as_map(0).unwrap();
172        assert_eq!(row["name"].as_str(), Some("Alice"));
173        assert_eq!(row["age"].as_i64(), Some(30));
174    }
175
176    #[test]
177    fn query_result_empty() {
178        let qr = QueryResult::empty();
179        assert_eq!(qr.row_count(), 0);
180        assert!(qr.row_as_map(0).is_none());
181    }
182}