1use ndarray::Array1;
2use serde::{Deserialize, Serialize};
3use std::cmp::Ordering;
4use vectradb_components::{VectorDocument, VectraDBError};
5
6pub mod hnsw;
8
9pub mod lsh;
11
12pub mod pq;
14
15pub use hnsw::HNSWIndex;
17pub use lsh::LSHIndex;
18pub use pq::PQIndex;
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct SearchConfig {
23 pub algorithm: SearchAlgorithm,
24 pub max_connections: usize,
25 pub search_ef: usize,
26 pub construction_ef: usize,
27 pub m: usize, pub ef_construction: usize, pub num_hashes: usize, pub num_buckets: usize, pub dimension: Option<usize>, pub num_subspaces: Option<usize>, pub codes_per_subspace: Option<usize>, }
35
36#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
37pub enum SearchAlgorithm {
38 HNSW,
39 LSH,
40 PQ,
41 Linear,
42}
43
44impl Default for SearchConfig {
45 fn default() -> Self {
46 Self {
47 algorithm: SearchAlgorithm::HNSW,
48 max_connections: 16,
49 search_ef: 50,
50 construction_ef: 200,
51 m: 16,
52 ef_construction: 200,
53 num_hashes: 10,
54 num_buckets: 1000,
55 dimension: Some(384),
56 num_subspaces: Some(8),
57 codes_per_subspace: Some(256),
58 }
59 }
60}
61
62#[derive(Debug, Clone)]
64pub struct SearchResult {
65 pub id: String,
66 pub distance: f32,
67 pub similarity: f32,
68}
69
70impl Ord for SearchResult {
71 fn cmp(&self, other: &Self) -> Ordering {
72 other.distance.total_cmp(&self.distance)
74 }
75}
76
77impl PartialOrd for SearchResult {
78 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
79 Some(self.cmp(other))
80 }
81}
82
83impl PartialEq for SearchResult {
84 fn eq(&self, other: &Self) -> bool {
85 self.distance == other.distance && self.id == other.id
86 }
87}
88
89impl Eq for SearchResult {}
90
91pub trait AdvancedSearch {
93 fn search(&self, query: &Array1<f32>, k: usize) -> Result<Vec<SearchResult>, VectraDBError>;
94 fn insert(&mut self, document: VectorDocument) -> Result<(), VectraDBError>;
95 fn remove(&mut self, id: &str) -> Result<(), VectraDBError>;
96 fn update(&mut self, id: &str, document: VectorDocument) -> Result<(), VectraDBError>;
97 fn build_index(&mut self, documents: Vec<VectorDocument>) -> Result<(), VectraDBError>;
98 fn get_stats(&self) -> SearchStats;
99}
100
101#[derive(Debug, Clone)]
103pub struct SearchStats {
104 pub total_vectors: usize,
105 pub index_size_bytes: usize,
106 pub average_search_time_ms: f64,
107 pub construction_time_ms: f64,
108}
109
110impl Default for SearchStats {
111 fn default() -> Self {
112 Self {
113 total_vectors: 0,
114 index_size_bytes: 0,
115 average_search_time_ms: 0.0,
116 construction_time_ms: 0.0,
117 }
118 }
119}