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    /// The statement's command verb (`INSERT`, `UPDATE`, `DELETE`, `UPSERT`,
70    /// `MERGE`, `TRUNCATE`, ...). `None` when the statement produced no DML
71    /// outcome.
72    #[serde(default)]
73    pub command: Option<String>,
74}
75
76impl QueryResult {
77    /// Empty result set with no columns.
78    pub fn empty() -> Self {
79        Self {
80            columns: Vec::new(),
81            rows: Vec::new(),
82            rows_affected: 0,
83            command: None,
84        }
85    }
86
87    /// Number of rows in the result set.
88    pub fn row_count(&self) -> usize {
89        self.rows.len()
90    }
91
92    /// Get a single row as a `HashMap<column_name, value>`.
93    pub fn row_as_map(&self, index: usize) -> Option<HashMap<&str, &Value>> {
94        let row = self.rows.get(index)?;
95        Some(
96            self.columns
97                .iter()
98                .zip(row.iter())
99                .map(|(col, val)| (col.as_str(), val))
100                .collect(),
101        )
102    }
103}
104
105impl SubGraph {
106    /// Empty subgraph.
107    pub fn empty() -> Self {
108        Self {
109            nodes: Vec::new(),
110            edges: Vec::new(),
111        }
112    }
113
114    /// Number of nodes in the subgraph.
115    pub fn node_count(&self) -> usize {
116        self.nodes.len()
117    }
118
119    /// Number of edges in the subgraph.
120    pub fn edge_count(&self) -> usize {
121        self.edges.len()
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn search_result() {
131        let r = SearchResult {
132            id: "vec-1".into(),
133            node_id: Some(NodeId::try_new("node-1").expect("test fixture")),
134            distance: 0.123,
135            metadata: HashMap::new(),
136        };
137        assert_eq!(r.id, "vec-1");
138        assert!(r.distance < 1.0);
139    }
140
141    #[test]
142    fn subgraph_construction() {
143        let mut sg = SubGraph::empty();
144        sg.nodes.push(SubGraphNode {
145            id: NodeId::try_new("a").expect("test fixture"),
146            depth: 0,
147            properties: HashMap::new(),
148        });
149        sg.nodes.push(SubGraphNode {
150            id: NodeId::try_new("b").expect("test fixture"),
151            depth: 1,
152            properties: HashMap::new(),
153        });
154        sg.edges.push(SubGraphEdge {
155            id: EdgeId::try_first(
156                NodeId::try_new("a").expect("test fixture"),
157                NodeId::try_new("b").expect("test fixture"),
158                "KNOWS",
159            )
160            .expect("test fixture"),
161            from: NodeId::try_new("a").expect("test fixture"),
162            to: NodeId::try_new("b").expect("test fixture"),
163            label: "KNOWS".into(),
164            properties: HashMap::new(),
165        });
166        assert_eq!(sg.node_count(), 2);
167        assert_eq!(sg.edge_count(), 1);
168    }
169
170    #[test]
171    fn query_result_row_as_map() {
172        let qr = QueryResult {
173            columns: vec!["name".into(), "age".into()],
174            rows: vec![vec![Value::String("Alice".into()), Value::Integer(30)]],
175            rows_affected: 0,
176            command: None,
177        };
178        let row = qr.row_as_map(0).unwrap();
179        assert_eq!(row["name"].as_str(), Some("Alice"));
180        assert_eq!(row["age"].as_i64(), Some(30));
181    }
182
183    #[test]
184    fn query_result_empty() {
185        let qr = QueryResult::empty();
186        assert_eq!(qr.row_count(), 0);
187        assert!(qr.row_as_map(0).is_none());
188    }
189}