veclite_index/

hnsw.rs

use rand::Rng;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, HashSet};

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct OrderedFloat(pub f32);

impl Eq for OrderedFloat {}

impl PartialOrd for OrderedFloat {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for OrderedFloat {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.partial_cmp(&other.0).unwrap_or(Ordering::Equal)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HnswConfig {
    pub m: usize,
    pub m_max: usize,
    pub m_max0: usize,
    pub ef_construction: usize,
    pub ef_search: usize,
    pub ml: f32,
}

impl Default for HnswConfig {
    fn default() -> Self {
        Self {
            m: 16,
            m_max: 16,
            m_max0: 32,
            ef_construction: 200,
            ef_search: 64,
            ml: 1.0 / 16.0f32.ln(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HnswIndex {
    pub config: HnswConfig,
    pub entry_point: Option<usize>,
    pub max_layer: usize,
    // Node ID -> [Layer 0 connections, Layer 1 connections, ...]
    pub graph: HashMap<usize, Vec<Vec<usize>>>,
}

impl HnswIndex {
    pub fn new(config: HnswConfig) -> Self {
        Self {
            config,
            entry_point: None,
            max_layer: 0,
            graph: HashMap::new(),
        }
    }

    fn random_layer(&self) -> usize {
        let mut rng = rand::thread_rng();
        let unif: f32 = rng.gen_range(0.0..1.0);
        (-unif.ln() * self.config.ml).floor() as usize
    }

    pub fn insert<'a>(
        &mut self,
        node_id: usize,
        vector: &[f32],
        get_vector: &impl Fn(usize) -> &'a [f32],
        distance_fn: &impl Fn(&[f32], &[f32]) -> f32,
    ) {
        let l = self.random_layer();

        // Initialize layers for the new node
        self.graph.insert(node_id, vec![vec![]; l + 1]);

        let mut ep = if let Some(ep) = self.entry_point {
            ep
        } else {
            self.entry_point = Some(node_id);
            self.max_layer = l;
            return;
        };

        let max_layer = self.max_layer;

        let mut curr_node = ep;

        // Phase 1: Traverse from top layer to l+1
        for lc in (l + 1..=max_layer).rev() {
            let mut curr_dist = distance_fn(vector, get_vector(curr_node));
            let mut changed = true;
            while changed {
                changed = false;
                if let Some(neighbors) = self.graph.get(&curr_node).and_then(|g| g.get(lc)) {
                    for &neighbor in neighbors {
                        let dist = distance_fn(vector, get_vector(neighbor));
                        if dist < curr_dist {
                            curr_dist = dist;
                            curr_node = neighbor;
                            changed = true;
                        }
                    }
                }
            }
        }

        // Phase 2: Insert and connect at layers l down to 0
        ep = curr_node;
        for lc in (0..=l.min(max_layer)).rev() {
            let mut w = self.search_layer(
                vector,
                ep,
                self.config.ef_construction,
                lc,
                get_vector,
                distance_fn,
            );
            let neighbors = self.select_neighbors(&mut w, self.config.m);

            for &neighbor in &neighbors {
                self.graph.get_mut(&node_id).unwrap()[lc].push(neighbor);
                self.graph.get_mut(&neighbor).unwrap()[lc].push(node_id);

                let m_max = if lc == 0 {
                    self.config.m_max0
                } else {
                    self.config.m_max
                };
                let neighbor_conns = &mut self.graph.get_mut(&neighbor).unwrap()[lc];
                if neighbor_conns.len() > m_max {
                    // Shrink neighbor connections
                    let mut e_conn = BinaryHeap::new();
                    for &n2 in neighbor_conns.iter() {
                        let d = distance_fn(get_vector(neighbor), get_vector(n2));
                        // max heap, we want to keep closest, so push with dist.
                        // To keep smallest distances, we need to sort and take top, or use a min heap.
                        // wait, standard binary heap is max heap. If we want smallest distance, we can put standard dist to find the largest to drop, or just sort.
                        e_conn.push((OrderedFloat(-d), n2));
                    }
                    let mut new_conns = Vec::new();
                    while let Some((_, n)) = e_conn.pop() {
                        new_conns.push(n);
                        if new_conns.len() == m_max {
                            break;
                        }
                    }
                    *neighbor_conns = new_conns;
                }
            }
            ep = w
                .iter()
                .min_by_key(|(d, _)| OrderedFloat(*d))
                .map(|(_, id)| *id)
                .unwrap_or(ep);
        }

        if l > max_layer {
            self.max_layer = l;
            self.entry_point = Some(node_id);
        }
    }

    pub fn search<'a>(
        &self,
        query: &[f32],
        k: usize,
        ef_search: usize,
        get_vector: &impl Fn(usize) -> &'a [f32],
        distance_fn: &impl Fn(&[f32], &[f32]) -> f32,
    ) -> Vec<(usize, f32)> {
        let mut ep = if let Some(ep) = self.entry_point {
            ep
        } else {
            return Vec::new();
        };

        let max_layer = self.max_layer;

        for lc in (1..=max_layer).rev() {
            let mut curr_dist = distance_fn(query, get_vector(ep));
            let mut changed = true;
            while changed {
                changed = false;
                if let Some(neighbors) = self.graph.get(&ep).and_then(|g| g.get(lc)) {
                    for &neighbor in neighbors {
                        let dist = distance_fn(query, get_vector(neighbor));
                        if dist < curr_dist {
                            curr_dist = dist;
                            ep = neighbor;
                            changed = true;
                        }
                    }
                }
            }
        }

        let w = self.search_layer(query, ep, ef_search.max(k), 0, get_vector, distance_fn);

        let mut res = w.into_iter().collect::<Vec<_>>();
        res.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        res.truncate(k);
        res.into_iter().map(|(d, id)| (id, d)).collect()
    }

    fn search_layer<'a>(
        &self,
        query: &[f32],
        ep: usize,
        ef: usize,
        lc: usize,
        get_vector: &impl Fn(usize) -> &'a [f32],
        distance_fn: &impl Fn(&[f32], &[f32]) -> f32,
    ) -> Vec<(f32, usize)> {
        let mut v = HashSet::new();
        let mut c = BinaryHeap::new(); // Candidates to explore: min-heap (ordered by dist asc -> store -dist)
        let mut w = BinaryHeap::new(); // Nearest neighbors found: max-heap (ordered by dist desc -> store dist)

        let d = distance_fn(query, get_vector(ep));
        v.insert(ep);
        c.push((OrderedFloat(-d), ep));
        w.push((OrderedFloat(d), ep));

        while let Some((OrderedFloat(neg_c_dist), c_id)) = c.pop() {
            let c_dist = -neg_c_dist;
            let f_dist = w.peek().unwrap().0 .0;
            if c_dist > f_dist {
                break;
            }

            if let Some(neighbors) = self.graph.get(&c_id).and_then(|g| g.get(lc)) {
                for &e in neighbors {
                    if !v.contains(&e) {
                        v.insert(e);
                        let f_dist = w.peek().unwrap().0 .0;
                        let e_dist = distance_fn(query, get_vector(e));

                        if e_dist < f_dist || w.len() < ef {
                            c.push((OrderedFloat(-e_dist), e));
                            w.push((OrderedFloat(e_dist), e));
                            if w.len() > ef {
                                w.pop();
                            }
                        }
                    }
                }
            }
        }

        w.into_iter().map(|(OrderedFloat(d), id)| (d, id)).collect()
    }

    fn select_neighbors(&self, candidates: &mut [(f32, usize)], m: usize) -> Vec<usize> {
        candidates.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        candidates.iter().take(m).map(|(_, id)| *id).collect()
    }
}

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

    #[test]
    fn test_hnsw() {
        let config = HnswConfig::default();
        let mut index = HnswIndex::new(config);

        let vectors = vec![
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![1.0, 1.0],
            vec![-1.0, 0.0],
        ];

        let get_vector = |id: usize| vectors[id].as_slice();
        let distance_fn = |a: &[f32], b: &[f32]| {
            a.iter()
                .zip(b.iter())
                .map(|(x, y)| (x - y).powi(2))
                .sum::<f32>()
                .sqrt()
        };

        for (id, vec) in vectors.iter().enumerate() {
            index.insert(id, vec, &get_vector, &distance_fn);
        }

        let res = index.search(&[1.0, 0.1], 2, 10, &get_vector, &distance_fn);
        assert_eq!(res.len(), 2);
        assert_eq!(res[0].0, 0); // closest is [1.0, 0.0]
    }
}