1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161

use core::mem;
use std::collections::HashSet;
use std::sync::{Arc, RwLock};

use hnsw::{Hnsw, Searcher};
use rand_pcg::Pcg64;
use space::{MetricPoint, Neighbor};

const MAX_REMOVED_BEFORE_REBUILD: usize = 100;

#[derive(Clone)]
pub struct Euclidean(Vec<f32>);

impl MetricPoint for Euclidean {
    fn distance(&self, rhs: &Self) -> u32 {
        space::f32_metric(
            self.0
                .iter()
                .zip(rhs.0.iter())
                .map(|(&a, &b)| (a - b).powi(2))
                .sum::<f32>()
                .sqrt(),
        )
    }
}

#[derive(Debug)]
pub struct AnnNeighbor {
    pub id: String,
    pub index: usize,
    pub distance: u32,
}

#[derive(Clone)]
pub struct VectorIndex {
    pub searcher: Arc<RwLock<Searcher>>,
    pub hnsw: Arc<RwLock<Hnsw<Euclidean, Pcg64, 12, 24>>>,
    pub vectors: Arc<RwLock<Vec<(String, Vec<f32>)>>>,
    pub removed: Arc<RwLock<HashSet<String>>>,
}

impl VectorIndex {
    pub fn new() -> Self {
        VectorIndex {
            searcher: Arc::new(RwLock::new(Searcher::default())),
            hnsw: Arc::new(RwLock::new(Hnsw::new())),
            vectors: Arc::new(RwLock::new(Vec::new())),
            removed: Arc::new(Default::default()),
        }
    }

    pub fn insert(&self, v: Vec<f32>, id: String) {
        let mut hnsw = self.hnsw.write().unwrap();
        let mut searcher = self.searcher.write().unwrap();
        let mut vectors = self.vectors.write().unwrap();
        vectors.push((id, v.clone()));
        hnsw.insert(Euclidean(v), &mut searcher);
    }

    pub fn search(&self, v: &[f32]) -> Vec<AnnNeighbor> {
        let mut neighbors = [Neighbor::invalid(); 8];
        let mut searcher = self.searcher.write().unwrap();
        let hnsw = self.hnsw.write().unwrap();
        let removed = self.removed.read().unwrap();
        let vectors = self.vectors.read().unwrap();
        hnsw.nearest(&Euclidean(v.to_vec()), 8, &mut searcher, &mut neighbors);
        println!("Neighbors {:?}", neighbors);
        neighbors
            .iter()
            .filter(|e| {
                let id = e.index;
                let id_val = vectors[id].0.clone();
                !removed.contains(&id_val)
            })
            .cloned()
            .map(|n| AnnNeighbor {
                id: vectors[n.index].0.clone(),
                index: n.index,
                distance: n.distance,
            })
            .collect()
    }

    pub fn remove(&self, id: String) {
        let mut removed = self.removed.write().unwrap();
        removed.insert(id);
        if removed.len() > MAX_REMOVED_BEFORE_REBUILD {
            self.rebuild();
        }
    }

    pub fn rebuild(&self) {
        let new_index = VectorIndex::new();
        let vectors = self.vectors.read().unwrap();
        let removed = self.removed.read().unwrap();
        for (id, v) in vectors.iter() {
            if !removed.contains(id) {
                new_index.insert(v.clone(), id.clone());
            }
        }

        drop(vectors);
        drop(removed);

        // replace all data structures inplace
        let mut searcher_old = self.searcher.write().unwrap();
        let mut hnsw_old = self.hnsw.write().unwrap();
        let mut vectors_old = self.vectors.write().unwrap();
        let mut removed_old = self.removed.write().unwrap();

        let searcher_new = new_index.searcher.read().unwrap();
        mem::replace(&mut *searcher_old, (*searcher_new).clone());

        let hnsw_new = new_index.hnsw.read().unwrap();
        mem::replace(&mut *hnsw_old, (*hnsw_new).clone());

        let vectors_new = new_index.vectors.read().unwrap();
        mem::replace(&mut *vectors_old, (*vectors_new).clone());

        removed_old.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_vector_index() {
        let index = VectorIndex::new();
        let features: [&[f32]; 9] = [
            &[0.0, 0.0, 0.0, 0.0],
            &[0.0, 0.0, 0.0, 1.0],
            &[0.0, 0.0, 1.0, 1.0],
            &[0.0, 1.0, 1.0, 0.0],
            &[1.0, 1.0, 0.0, 0.0],
            &[0.0, 0.0, 1.0, 1.0],
            &[0.0, 1.0, 1.0, 0.0],
            &[1.0, 1.0, 0.0, 0.0],
            &[1.0, 0.0, 0.0, 1.0],
        ];

        let mut i = 0;
        for &feature in &features {
            let v = feature.to_vec();
            index.insert(v, i.to_string());
            i += 1;
        }

        let neighbors = index.search(&features[0].to_vec().clone());
        assert_eq!(neighbors[0].id, "0".to_string());

        index.remove("0".to_string());
        let neighbors = index.search(&features[0].to_vec().clone());
        assert_eq!(neighbors[0].id, "1".to_string());

        index.rebuild();
        let neighbors = index.search(&features[0].to_vec().clone());
        assert_eq!(neighbors[0].id, "1".to_string());
    }
}