oxirs_vec/
tree_indices_unified.rs

1//! Unified [`TreeIndex`] dispatcher over all tree-based index implementations.
2//!
3//! [`TreeIndex`] selects one concrete tree (ball / KD / VP / cover / random
4//! projection) based on [`TreeIndexConfig::tree_type`] and exposes a single
5//! [`VectorIndex`] interface over it.
6
7use crate::tree_indices_balltree::BallTree;
8use crate::tree_indices_covertree::CoverTree;
9use crate::tree_indices_kdtree::KdTree;
10use crate::tree_indices_rptree::RandomProjectionTree;
11use crate::tree_indices_types::{TreeIndexConfig, TreeType};
12use crate::tree_indices_vptree::VpTree;
13use crate::{Vector, VectorIndex};
14use anyhow::Result;
15
16/// Unified tree index interface
17pub struct TreeIndex {
18    tree_type: TreeType,
19    ball_tree: Option<BallTree>,
20    kd_tree: Option<KdTree>,
21    vp_tree: Option<VpTree>,
22    cover_tree: Option<CoverTree>,
23    rp_tree: Option<RandomProjectionTree>,
24}
25
26impl TreeIndex {
27    pub fn new(config: TreeIndexConfig) -> Self {
28        let tree_type = config.tree_type;
29
30        let (ball_tree, kd_tree, vp_tree, cover_tree, rp_tree) = match tree_type {
31            TreeType::BallTree => (Some(BallTree::new(config)), None, None, None, None),
32            TreeType::KdTree => (None, Some(KdTree::new(config)), None, None, None),
33            TreeType::VpTree => (None, None, Some(VpTree::new(config)), None, None),
34            TreeType::CoverTree => (None, None, None, Some(CoverTree::new(config)), None),
35            TreeType::RandomProjectionTree => (
36                None,
37                None,
38                None,
39                None,
40                Some(RandomProjectionTree::new(config)),
41            ),
42        };
43
44        Self {
45            tree_type,
46            ball_tree,
47            kd_tree,
48            vp_tree,
49            cover_tree,
50            rp_tree,
51        }
52    }
53
54    pub fn build(&mut self) -> Result<()> {
55        match self.tree_type {
56            TreeType::BallTree => self
57                .ball_tree
58                .as_mut()
59                .expect("ball_tree should be initialized for BallTree type")
60                .build(),
61            TreeType::KdTree => self
62                .kd_tree
63                .as_mut()
64                .expect("kd_tree should be initialized for KdTree type")
65                .build(),
66            TreeType::VpTree => self
67                .vp_tree
68                .as_mut()
69                .expect("vp_tree should be initialized for VpTree type")
70                .build(),
71            TreeType::CoverTree => self
72                .cover_tree
73                .as_mut()
74                .expect("cover_tree should be initialized for CoverTree type")
75                .build(),
76            TreeType::RandomProjectionTree => self
77                .rp_tree
78                .as_mut()
79                .expect("rp_tree should be initialized for RandomProjectionTree type")
80                .build(),
81        }
82    }
83
84    fn search_internal(&self, query: &[f32], k: usize) -> Vec<(usize, f32)> {
85        match self.tree_type {
86            TreeType::BallTree => self
87                .ball_tree
88                .as_ref()
89                .expect("ball_tree should be initialized for BallTree type")
90                .search(query, k),
91            TreeType::KdTree => self
92                .kd_tree
93                .as_ref()
94                .expect("kd_tree should be initialized for KdTree type")
95                .search(query, k),
96            TreeType::VpTree => self
97                .vp_tree
98                .as_ref()
99                .expect("vp_tree should be initialized for VpTree type")
100                .search(query, k),
101            TreeType::CoverTree => self
102                .cover_tree
103                .as_ref()
104                .expect("cover_tree should be initialized for CoverTree type")
105                .search(query, k),
106            TreeType::RandomProjectionTree => self
107                .rp_tree
108                .as_ref()
109                .expect("rp_tree should be initialized for RandomProjectionTree type")
110                .search(query, k),
111        }
112    }
113}
114
115impl VectorIndex for TreeIndex {
116    fn insert(&mut self, uri: String, vector: Vector) -> Result<()> {
117        let data = match self.tree_type {
118            TreeType::BallTree => {
119                &mut self
120                    .ball_tree
121                    .as_mut()
122                    .expect("ball_tree should be initialized for BallTree type")
123                    .data
124            }
125            TreeType::KdTree => {
126                &mut self
127                    .kd_tree
128                    .as_mut()
129                    .expect("kd_tree should be initialized for KdTree type")
130                    .data
131            }
132            TreeType::VpTree => {
133                &mut self
134                    .vp_tree
135                    .as_mut()
136                    .expect("vp_tree should be initialized for VpTree type")
137                    .data
138            }
139            TreeType::CoverTree => {
140                &mut self
141                    .cover_tree
142                    .as_mut()
143                    .expect("cover_tree should be initialized for CoverTree type")
144                    .data
145            }
146            TreeType::RandomProjectionTree => {
147                &mut self
148                    .rp_tree
149                    .as_mut()
150                    .expect("rp_tree should be initialized for RandomProjectionTree type")
151                    .data
152            }
153        };
154
155        data.push((uri, vector));
156        Ok(())
157    }
158
159    fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>> {
160        let query_f32 = query.as_f32();
161        let results = self.search_internal(&query_f32, k);
162
163        let data = match self.tree_type {
164            TreeType::BallTree => {
165                &self
166                    .ball_tree
167                    .as_ref()
168                    .expect("ball_tree should be initialized for BallTree type")
169                    .data
170            }
171            TreeType::KdTree => {
172                &self
173                    .kd_tree
174                    .as_ref()
175                    .expect("kd_tree should be initialized for KdTree type")
176                    .data
177            }
178            TreeType::VpTree => {
179                &self
180                    .vp_tree
181                    .as_ref()
182                    .expect("vp_tree should be initialized for VpTree type")
183                    .data
184            }
185            TreeType::CoverTree => {
186                &self
187                    .cover_tree
188                    .as_ref()
189                    .expect("cover_tree should be initialized for CoverTree type")
190                    .data
191            }
192            TreeType::RandomProjectionTree => {
193                &self
194                    .rp_tree
195                    .as_ref()
196                    .expect("rp_tree should be initialized for RandomProjectionTree type")
197                    .data
198            }
199        };
200
201        Ok(results
202            .into_iter()
203            .map(|(idx, dist)| (data[idx].0.clone(), dist))
204            .collect())
205    }
206
207    fn search_threshold(&self, query: &Vector, threshold: f32) -> Result<Vec<(String, f32)>> {
208        let query_f32 = query.as_f32();
209        let all_results = self.search_internal(&query_f32, 1000); // Search more broadly
210
211        let data = match self.tree_type {
212            TreeType::BallTree => {
213                &self
214                    .ball_tree
215                    .as_ref()
216                    .expect("ball_tree should be initialized for BallTree type")
217                    .data
218            }
219            TreeType::KdTree => {
220                &self
221                    .kd_tree
222                    .as_ref()
223                    .expect("kd_tree should be initialized for KdTree type")
224                    .data
225            }
226            TreeType::VpTree => {
227                &self
228                    .vp_tree
229                    .as_ref()
230                    .expect("vp_tree should be initialized for VpTree type")
231                    .data
232            }
233            TreeType::CoverTree => {
234                &self
235                    .cover_tree
236                    .as_ref()
237                    .expect("cover_tree should be initialized for CoverTree type")
238                    .data
239            }
240            TreeType::RandomProjectionTree => {
241                &self
242                    .rp_tree
243                    .as_ref()
244                    .expect("rp_tree should be initialized for RandomProjectionTree type")
245                    .data
246            }
247        };
248
249        Ok(all_results
250            .into_iter()
251            .filter(|(_, dist)| *dist <= threshold)
252            .map(|(idx, dist)| (data[idx].0.clone(), dist))
253            .collect())
254    }
255
256    fn get_vector(&self, uri: &str) -> Option<&Vector> {
257        let data = match self.tree_type {
258            TreeType::BallTree => {
259                &self
260                    .ball_tree
261                    .as_ref()
262                    .expect("ball_tree should be initialized for BallTree type")
263                    .data
264            }
265            TreeType::KdTree => {
266                &self
267                    .kd_tree
268                    .as_ref()
269                    .expect("kd_tree should be initialized for KdTree type")
270                    .data
271            }
272            TreeType::VpTree => {
273                &self
274                    .vp_tree
275                    .as_ref()
276                    .expect("vp_tree should be initialized for VpTree type")
277                    .data
278            }
279            TreeType::CoverTree => {
280                &self
281                    .cover_tree
282                    .as_ref()
283                    .expect("cover_tree should be initialized for CoverTree type")
284                    .data
285            }
286            TreeType::RandomProjectionTree => {
287                &self
288                    .rp_tree
289                    .as_ref()
290                    .expect("rp_tree should be initialized for RandomProjectionTree type")
291                    .data
292            }
293        };
294
295        data.iter().find(|(u, _)| u == uri).map(|(_, v)| v)
296    }
297}
298
299// Add rand to dependencies for VP-Tree and Random Projection Tree
300// Note: Replaced with scirs2_core::random
301
302// Placeholder for async task spawning - integrate with oxirs-core::parallel
303async fn spawn_task<F, T>(f: F) -> T
304where
305    F: FnOnce() -> T + Send + 'static,
306    T: Send + 'static,
307{
308    // In practice, this would use oxirs-core::parallel's task spawning
309    f()
310}
oxirs_vec/tree_indices_unified.rs

oxirs_vec/
tree_indices_unified.rs
