omendb_core/vector/hnsw/
graph_storage.rs

1//! Graph storage abstraction for HNSW index
2//!
3//! Provides a unified API for the in-memory neighbor list storage.
4//! Persistence is handled by serializing the entire `HNSWIndex` to .omen format.
5
6use super::storage::NeighborLists;
7use serde::{Deserialize, Serialize};
8
9/// Graph storage backend for HNSW index
10///
11/// Wraps `NeighborLists` for in-memory neighbor storage.
12/// Persistence is handled externally by `.omen` format serialization.
13#[derive(Debug, Serialize, Deserialize)]
14pub struct GraphStorage(NeighborLists);
15
16impl GraphStorage {
17    /// Create new storage with max levels
18    #[must_use]
19    pub fn new(max_levels: usize) -> Self {
20        Self(NeighborLists::new(max_levels))
21    }
22
23    /// Create storage with pre-allocated capacity
24    #[must_use]
25    pub fn with_capacity(num_nodes: usize, max_levels: usize, m: usize) -> Self {
26        Self(NeighborLists::with_capacity(num_nodes, max_levels, m))
27    }
28
29    /// Create from existing neighbor lists (used when loading from persistence)
30    #[must_use]
31    pub fn from_neighbor_lists(lists: NeighborLists) -> Self {
32        Self(lists)
33    }
34
35    /// Get neighbors for a node at a specific level
36    #[inline]
37    #[must_use]
38    pub fn get_neighbors(&self, node_id: u32, level: u8) -> Vec<u32> {
39        self.0.get_neighbors(node_id, level)
40    }
41
42    /// Execute closure with read access to neighbors (zero-copy)
43    #[inline]
44    pub fn with_neighbors<F, R>(&self, node_id: u32, level: u8, f: F) -> R
45    where
46        F: FnOnce(&[u32]) -> R,
47    {
48        self.0.with_neighbors(node_id, level, f)
49    }
50
51    /// Set neighbors for a node at a specific level
52    #[inline]
53    pub fn set_neighbors(&mut self, node_id: u32, level: u8, neighbors: Vec<u32>) {
54        self.0.set_neighbors(node_id, level, neighbors);
55    }
56
57    /// Add bidirectional link between two nodes
58    #[inline]
59    pub fn add_bidirectional_link(&mut self, node_a: u32, node_b: u32, level: u8) {
60        self.0.add_bidirectional_link(node_a, node_b, level);
61    }
62
63    /// Add bidirectional link (parallel version)
64    #[inline]
65    pub fn add_bidirectional_link_parallel(&self, node_a: u32, node_b: u32, level: u8) {
66        self.0
67            .add_bidirectional_link_parallel(node_a, node_b, level);
68    }
69
70    /// Remove unidirectional link (parallel version)
71    #[inline]
72    pub fn remove_link_parallel(&self, node_a: u32, node_b: u32, level: u8) {
73        self.0.remove_link_parallel(node_a, node_b, level);
74    }
75
76    /// Set neighbors (parallel version)
77    #[inline]
78    pub fn set_neighbors_parallel(&self, node_id: u32, level: u8, neighbors: Vec<u32>) {
79        self.0.set_neighbors_parallel(node_id, level, neighbors);
80    }
81
82    /// Get `M_max` (max neighbors per node)
83    #[must_use]
84    pub fn m_max(&self) -> usize {
85        self.0.m_max()
86    }
87
88    /// Get memory usage in bytes
89    #[must_use]
90    pub fn memory_usage(&self) -> usize {
91        self.0.memory_usage()
92    }
93
94    /// Prefetch neighbor list into CPU cache
95    ///
96    /// Hints to CPU that we'll need neighbor data soon. Only beneficial on
97    /// x86/ARM servers - disabled on Apple Silicon where DMP handles this.
98    #[inline]
99    pub fn prefetch(&self, node_id: u32, level: u8) {
100        self.0.prefetch(node_id, level);
101    }
102
103    /// Reorder graph using BFS for cache locality
104    pub fn reorder_bfs(&mut self, entry_point: u32, start_level: u8) -> Vec<u32> {
105        self.0.reorder_bfs(entry_point, start_level)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_graph_storage_new() {
115        let storage = GraphStorage::new(8);
116        assert_eq!(storage.m_max(), 32);
117    }
118
119    #[test]
120    fn test_graph_storage_get_set_neighbors() {
121        let mut storage = GraphStorage::new(8);
122
123        storage.set_neighbors(0, 0, vec![1, 2, 3]);
124        storage.set_neighbors(0, 1, vec![4, 5]);
125
126        assert_eq!(storage.get_neighbors(0, 0), vec![1, 2, 3]);
127        assert_eq!(storage.get_neighbors(0, 1), vec![4, 5]);
128        assert_eq!(storage.get_neighbors(99, 0), Vec::<u32>::new());
129    }
130
131    #[test]
132    fn test_graph_storage_add_bidirectional_link() {
133        let mut storage = GraphStorage::new(8);
134
135        storage.add_bidirectional_link(0, 1, 0);
136
137        let neighbors_0 = storage.get_neighbors(0, 0);
138        let neighbors_1 = storage.get_neighbors(1, 0);
139
140        assert!(neighbors_0.contains(&1));
141        assert!(neighbors_1.contains(&0));
142    }
143
144    #[test]
145    fn test_graph_storage_serialization() {
146        let mut storage = GraphStorage::new(8);
147        storage.set_neighbors(0, 0, vec![1, 2, 3]);
148
149        let serialized = bincode::serialize(&storage).unwrap();
150        let deserialized: GraphStorage = bincode::deserialize(&serialized).unwrap();
151
152        assert_eq!(deserialized.get_neighbors(0, 0), vec![1, 2, 3]);
153    }
154}
omendb_core/vector/hnsw/graph_storage.rs

omendb_core/vector/hnsw/
graph_storage.rs
