Skip to main content

vectradb_search/
lib.rs

1use ndarray::Array1;
2use serde::{Deserialize, Serialize};
3use std::cmp::Ordering;
4use vectradb_components::{VectorDocument, VectraDBError};
5
6/// HNSW (Hierarchical Navigable Small World) search implementation
7pub mod hnsw;
8
9/// LSH (Locality Sensitive Hashing) search implementation  
10pub mod lsh;
11
12/// Product Quantization search implementation
13pub mod pq;
14
15/// Re-export search algorithms
16pub use hnsw::HNSWIndex;
17pub use lsh::LSHIndex;
18pub use pq::PQIndex;
19
20/// Search configuration for different algorithms
21#[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,                          // For HNSW
28    pub ef_construction: usize,            // For HNSW
29    pub num_hashes: usize,                 // For LSH
30    pub num_buckets: usize,                // For LSH
31    pub dimension: Option<usize>,          // Vector dimension
32    pub num_subspaces: Option<usize>,      // For PQ
33    pub codes_per_subspace: Option<usize>, // For PQ
34}
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/// Search result with distance and metadata
63#[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        // Reverse for max-heap (min distance = max priority)
73        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
91/// Trait for advanced search algorithms
92pub 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/// Search algorithm statistics
102#[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}