graphmind/vector/
manager.rs

1//! Manager for multiple vector indices
2//!
3//! Handles indexing for different node labels and property keys.
4
5use crate::graph::NodeId;
6use crate::vector::index::{DistanceMetric, VectorIndex, VectorResult};
7use std::collections::HashMap;
8use std::sync::{Arc, RwLock};
9
10/// Key for identifying a vector index: (Label, PropertyKey)
11#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
12pub struct IndexKey {
13    pub label: String,
14    pub property_key: String,
15}
16
17/// Manager for all vector indices in the system
18#[derive(Debug)]
19pub struct VectorIndexManager {
20    indices: RwLock<HashMap<IndexKey, Arc<RwLock<VectorIndex>>>>,
21}
22
23impl VectorIndexManager {
24    /// Create a new manager
25    pub fn new() -> Self {
26        Self {
27            indices: RwLock::new(HashMap::new()),
28        }
29    }
30
31    /// Create a new index
32    pub fn create_index(
33        &self,
34        label: &str,
35        property_key: &str,
36        dimensions: usize,
37        metric: DistanceMetric,
38    ) -> VectorResult<()> {
39        let key = IndexKey {
40            label: label.to_string(),
41            property_key: property_key.to_string(),
42        };
43
44        let index = VectorIndex::new(dimensions, metric);
45        let mut indices = self.indices.write().unwrap();
46        indices.insert(key, Arc::new(RwLock::new(index)));
47
48        Ok(())
49    }
50
51    /// Get an index
52    pub fn get_index(&self, label: &str, property_key: &str) -> Option<Arc<RwLock<VectorIndex>>> {
53        let key = IndexKey {
54            label: label.to_string(),
55            property_key: property_key.to_string(),
56        };
57
58        let indices = self.indices.read().unwrap();
59        indices.get(&key).cloned()
60    }
61
62    /// Add a vector to an index
63    pub fn add_vector(
64        &self,
65        label: &str,
66        property_key: &str,
67        node_id: NodeId,
68        vector: &Vec<f32>,
69    ) -> VectorResult<()> {
70        if let Some(index_lock) = self.get_index(label, property_key) {
71            let mut index = index_lock.write().unwrap();
72            index.add(node_id, vector)?;
73        }
74        Ok(())
75    }
76
77    /// Search an index
78    pub fn search(
79        &self,
80        label: &str,
81        property_key: &str,
82        query: &[f32],
83        k: usize,
84    ) -> VectorResult<Vec<(NodeId, f32)>> {
85        if let Some(index_lock) = self.get_index(label, property_key) {
86            let index = index_lock.read().unwrap();
87            return index.search(query, k);
88        }
89        Ok(Vec::new())
90    }
91
92    /// List all indices
93    pub fn list_indices(&self) -> Vec<IndexKey> {
94        let indices = self.indices.read().unwrap();
95        indices.keys().cloned().collect()
96    }
97
98    /// Save all indices to a directory
99    pub fn dump_all(&self, path: &std::path::Path) -> VectorResult<()> {
100        if !path.exists() {
101            std::fs::create_dir_all(path)?;
102        }
103
104        let indices = self.indices.read().unwrap();
105        let mut metadata = Vec::new();
106
107        for (key, index_lock) in indices.iter() {
108            let index = index_lock.read().unwrap();
109            let index_filename = format!("{}_{}.hnsw", key.label, key.property_key);
110            let index_path = path.join(&index_filename);
111            index.dump(&index_path)?;
112
113            metadata.push(serde_json::json!({
114                "label": key.label,
115                "property_key": key.property_key,
116                "dimensions": index.dimensions(),
117                "metric": index.metric(),
118                "filename": index_filename,
119            }));
120        }
121
122        let metadata_path = path.join("metadata.json");
123        let metadata_file = std::fs::File::create(metadata_path)?;
124        serde_json::to_writer_pretty(metadata_file, &metadata)
125            .map_err(|e| crate::vector::VectorError::IndexError(e.to_string()))?;
126
127        Ok(())
128    }
129
130    /// Load all indices from a directory
131    pub fn load_all(&self, path: &std::path::Path) -> VectorResult<()> {
132        if !path.exists() {
133            return Ok(());
134        }
135
136        let metadata_path = path.join("metadata.json");
137        if !metadata_path.exists() {
138            return Ok(());
139        }
140
141        let metadata_file = std::fs::File::open(metadata_path)?;
142        let metadata: Vec<serde_json::Value> = serde_json::from_reader(metadata_file)
143            .map_err(|e| crate::vector::VectorError::IndexError(e.to_string()))?;
144
145        let mut indices = self.indices.write().unwrap();
146        for item in metadata {
147            let label = item["label"].as_str().unwrap();
148            let property_key = item["property_key"].as_str().unwrap();
149            let dimensions = item["dimensions"].as_u64().unwrap() as usize;
150            let metric: DistanceMetric = serde_json::from_value(item["metric"].clone())
151                .map_err(|e| crate::vector::VectorError::IndexError(e.to_string()))?;
152            let filename = item["filename"].as_str().unwrap();
153
154            let index_path = path.join(filename);
155            let index = VectorIndex::load(&index_path, dimensions, metric)?;
156
157            let key = IndexKey {
158                label: label.to_string(),
159                property_key: property_key.to_string(),
160            };
161            indices.insert(key, Arc::new(RwLock::new(index)));
162        }
163
164        Ok(())
165    }
166}
167
168impl Default for VectorIndexManager {
169    fn default() -> Self {
170        Self::new()
171    }
172}
graphmind/vector/manager.rs

graphmind/vector/
manager.rs
