nodedb_vector/codec_index/

search.rs

// SPDX-License-Identifier: Apache-2.0

//! Search algorithm for `HnswCodecIndex<C>`.
//!
//! Two-stage HNSW search:
//! - Phase 1 (layers max..1): greedy descent using `fast_symmetric_distance`.
//! - Phase 2 (layer 0): ef-wide beam search using `fast_symmetric_distance`
//!   for navigation, then a final rerank pass with `exact_asymmetric_distance`.

use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashSet};

use nodedb_codec::vector_quant::codec::VectorCodec;

use super::graph::HnswCodecIndex;

/// A single result from a codec-index search.
#[derive(Debug, Clone)]
pub struct CodecSearchResult {
    /// Caller-supplied id from `HnswCodecIndex::insert`.
    pub id: u32,
    /// `exact_asymmetric_distance` between the query and this vector.
    pub distance: f32,
}

/// Internal candidate during beam search.
#[derive(Clone, Copy, PartialEq)]
struct Cand {
    dist: f32,
    /// Dense index into `HnswCodecIndex::nodes`.
    idx: u32,
}

impl Eq for Cand {}

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

impl Ord for Cand {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.dist
            .partial_cmp(&other.dist)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(self.idx.cmp(&other.idx))
    }
}

impl<C: VectorCodec> HnswCodecIndex<C> {
    /// K-NN search returning up to `k` results.
    ///
    /// `ef_search` controls the beam width at layer 0 (must be >= k).
    ///
    /// The returned results are sorted ascending by `exact_asymmetric_distance`.
    pub fn search(&self, query: &[f32], k: usize, ef_search: usize) -> Vec<CodecSearchResult> {
        if self.is_empty() {
            return Vec::new();
        }

        let Some(ep) = self.entry_point else {
            return Vec::new();
        };

        let ef = ef_search.max(k);

        // Precompute the query forms used by the two phases.
        let q_encoded = self.codec.encode(query);
        let q_prepared = self.codec.prepare_query(query);

        // Phase 1: greedy descent through layers max_layer..1.
        let mut cur_ep = ep;
        for layer in (1..=self.max_layer).rev() {
            cur_ep = self.greedy_nearest_search(&q_encoded, cur_ep, layer);
        }

        // Phase 2: ef-wide beam search at layer 0.
        let candidates = self.search_layer_0(&q_encoded, cur_ep, ef);

        // Rerank top ef_search candidates with exact asymmetric distance.
        let mut reranked: Vec<(f32, u32)> = candidates
            .into_iter()
            .take(ef)
            .map(|c| {
                let asym = self
                    .codec
                    .exact_asymmetric_distance(&q_prepared, &self.nodes[c.idx as usize].quantized);
                (asym, self.nodes[c.idx as usize].id)
            })
            .collect();

        reranked.sort_unstable_by(|a, b| a.0.total_cmp(&b.0));
        reranked.truncate(k);

        reranked
            .into_iter()
            .map(|(distance, id)| CodecSearchResult { id, distance })
            .collect()
    }

    /// Greedy single-nearest descent at `layer` using the pre-encoded query.
    fn greedy_nearest_search(&self, q_enc: &C::Quantized, ep_idx: u32, layer: usize) -> u32 {
        let mut best_idx = ep_idx;
        let mut best_dist = self
            .codec
            .fast_symmetric_distance(q_enc, &self.nodes[ep_idx as usize].quantized);

        loop {
            let mut improved = false;
            for &nb in self.neighbors_at(best_idx, layer) {
                if self.nodes[nb as usize].deleted {
                    continue;
                }
                let d = self
                    .codec
                    .fast_symmetric_distance(q_enc, &self.nodes[nb as usize].quantized);
                if d < best_dist {
                    best_dist = d;
                    best_idx = nb;
                    improved = true;
                }
            }
            if !improved {
                break;
            }
        }

        best_idx
    }

    /// Beam search at layer 0 using `fast_symmetric_distance`.
    ///
    /// Returns candidates sorted ascending by symmetric distance (used only
    /// for routing; final ranking is done by `exact_asymmetric_distance` in
    /// the caller).
    fn search_layer_0(&self, q_enc: &C::Quantized, ep_idx: u32, ef: usize) -> Vec<Cand> {
        let mut visited: HashSet<u32> = HashSet::new();
        visited.insert(ep_idx);

        let ep_dist = self
            .codec
            .fast_symmetric_distance(q_enc, &self.nodes[ep_idx as usize].quantized);
        let ep_cand = Cand {
            dist: ep_dist,
            idx: ep_idx,
        };

        let mut candidates: BinaryHeap<Reverse<Cand>> = BinaryHeap::new();
        candidates.push(Reverse(ep_cand));

        let mut results: BinaryHeap<Cand> = BinaryHeap::new();
        if !self.nodes[ep_idx as usize].deleted {
            results.push(ep_cand);
        }

        while let Some(Reverse(cur)) = candidates.pop() {
            let worst = results.peek().map_or(f32::INFINITY, |w| w.dist);
            if cur.dist > worst && results.len() >= ef {
                break;
            }

            for &nb in self.neighbors_at(cur.idx, 0) {
                if !visited.insert(nb) {
                    continue;
                }
                let d = self
                    .codec
                    .fast_symmetric_distance(q_enc, &self.nodes[nb as usize].quantized);
                let worst_now = results.peek().map_or(f32::INFINITY, |w| w.dist);
                if d < worst_now || results.len() < ef {
                    candidates.push(Reverse(Cand { dist: d, idx: nb }));
                }
                if !self.nodes[nb as usize].deleted {
                    results.push(Cand { dist: d, idx: nb });
                    if results.len() > ef {
                        results.pop();
                    }
                }
            }
        }

        let mut out: Vec<Cand> = results.into_vec();
        out.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist));
        out
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use nodedb_codec::vector_quant::{bbq::BbqCodec, rabitq::RaBitQCodec};

    use crate::{codec_index::HnswCodecIndex, distance::l2_squared, quantize::Sq8Codec};

    // ── helpers ───────────────────────────────────────────────────────────────

    fn xorshift(state: &mut u64) -> u64 {
        *state ^= *state << 13;
        *state ^= *state >> 7;
        *state ^= *state << 17;
        *state
    }

    fn rand_vec(state: &mut u64, dim: usize) -> Vec<f32> {
        (0..dim)
            .map(|_| (xorshift(state) as f32 / u64::MAX as f32) * 2.0 - 1.0)
            .collect()
    }

    /// Brute-force top-k by L2-squared.
    fn ground_truth(vecs: &[Vec<f32>], query: &[f32], k: usize) -> Vec<u32> {
        let mut scored: Vec<(f32, u32)> = vecs
            .iter()
            .enumerate()
            .map(|(i, v)| (l2_squared(query, v), i as u32))
            .collect();
        scored.sort_unstable_by(|a, b| a.0.total_cmp(&b.0));
        scored.into_iter().take(k).map(|(_, id)| id).collect()
    }

    // ── Sq8 round-trip ────────────────────────────────────────────────────────

    #[test]
    fn sq8_top1_exact_match() {
        let dim = 8;
        let n = 50usize;
        let mut state = 0xDEAD_BEEF_u64;
        let vecs: Vec<Vec<f32>> = (0..n).map(|_| rand_vec(&mut state, dim)).collect();
        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
        let codec = Sq8Codec::calibrate(&refs, dim);
        let mut idx: HnswCodecIndex<Sq8Codec> = HnswCodecIndex::new(dim, 8, 100, codec, 7);
        for (i, v) in vecs.iter().enumerate() {
            idx.insert(i as u32, v);
        }
        // Query with vector 17 — top-1 should return id 17.
        let query = vecs[17].clone();
        let results = idx.search(&query, 1, 50);
        assert_eq!(results.len(), 1, "expected 1 result");
        assert_eq!(
            results[0].id, 17,
            "top-1 should be the queried vector itself"
        );
        assert!(
            results[0].distance < 0.1,
            "distance to self should be near 0, got {}",
            results[0].distance
        );
    }

    // ── RaBitQ recall ─────────────────────────────────────────────────────────

    #[test]
    fn rabitq_recall_at_least_60_percent() {
        let dim = 16;
        let n = 100usize;
        let k = 5usize;
        let mut state = 0xCAFE_BABE_u64;
        let vecs: Vec<Vec<f32>> = (0..n).map(|_| rand_vec(&mut state, dim)).collect();
        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
        let codec = RaBitQCodec::calibrate(&refs, dim, 0xABCD_1234);
        let mut idx: HnswCodecIndex<RaBitQCodec> = HnswCodecIndex::new(dim, 8, 150, codec, 99);
        for (i, v) in vecs.iter().enumerate() {
            idx.insert(i as u32, v);
        }

        let n_queries = 10usize;
        let mut total_hits = 0usize;
        let mut total = 0usize;
        for _qi in 0..n_queries {
            let query = rand_vec(&mut state, dim);
            let truth: std::collections::HashSet<u32> =
                ground_truth(&vecs, &query, k).into_iter().collect();
            let results = idx.search(&query, k, k * 4);
            let found: std::collections::HashSet<u32> = results.iter().map(|r| r.id).collect();
            total_hits += found.intersection(&truth).count();
            total += k;
        }
        let recall = total_hits as f64 / total as f64;
        // 1-bit codecs at low dim (D=16) are inherently approximate; the
        // O(1/√D) bound only bites at higher dimensions. This is a sanity
        // test: at D=16 / n=100 / k=5, RaBitQ typically lands ≥ 30%.
        assert!(
            recall >= 0.30,
            "RaBitQ recall@{k} = {recall:.2}, expected >= 0.30"
        );
    }

    // ── BBQ recall ────────────────────────────────────────────────────────────

    #[test]
    fn bbq_recall_at_least_60_percent() {
        let dim = 16;
        let n = 100usize;
        let k = 5usize;
        let mut state = 0x1234_5678_u64;
        let vecs: Vec<Vec<f32>> = (0..n).map(|_| rand_vec(&mut state, dim)).collect();
        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
        let codec = BbqCodec::calibrate(&refs, dim, 3);
        let mut idx: HnswCodecIndex<BbqCodec> = HnswCodecIndex::new(dim, 8, 150, codec, 42);
        for (i, v) in vecs.iter().enumerate() {
            idx.insert(i as u32, v);
        }

        let n_queries = 10usize;
        let mut total_hits = 0usize;
        let mut total = 0usize;
        for _qi in 0..n_queries {
            let query = rand_vec(&mut state, dim);
            let truth: std::collections::HashSet<u32> =
                ground_truth(&vecs, &query, k).into_iter().collect();
            let results = idx.search(&query, k, k * 4);
            let found: std::collections::HashSet<u32> = results.iter().map(|r| r.id).collect();
            total_hits += found.intersection(&truth).count();
            total += k;
        }
        let recall = total_hits as f64 / total as f64;
        // BBQ at D=16 is approximate (corrective factors help vs raw binary
        // but the 1-bit code itself is information-bounded). Sanity test
        // threshold; production uses BBQ + oversample ×3 rerank pass.
        assert!(
            recall >= 0.40,
            "BBQ recall@{k} = {recall:.2}, expected >= 0.40"
        );
    }

    // ── Edge cases ────────────────────────────────────────────────────────────

    #[test]
    fn empty_index_returns_empty() {
        let codec = {
            let vecs: Vec<Vec<f32>> = (0..5).map(|i| vec![i as f32; 4]).collect();
            let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
            Sq8Codec::calibrate(&refs, 4)
        };
        let idx: HnswCodecIndex<Sq8Codec> = HnswCodecIndex::new(4, 8, 50, codec, 1);
        let results = idx.search(&[0.0, 0.0, 0.0, 0.0], 5, 20);
        assert!(results.is_empty(), "empty index must return no results");
    }

    #[test]
    fn single_vector_index_always_returns_it() {
        let dim = 4;
        let vecs = [vec![1.0f32, 2.0, 3.0, 4.0]];
        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
        let codec = Sq8Codec::calibrate(&refs, dim);
        let mut idx: HnswCodecIndex<Sq8Codec> = HnswCodecIndex::new(dim, 8, 50, codec, 5);
        idx.insert(0, &vecs[0]);
        // Query with a completely different vector.
        let results = idx.search(&[10.0, 20.0, 30.0, 40.0], 1, 10);
        assert_eq!(results.len(), 1, "single-node index must return 1 result");
        assert_eq!(results[0].id, 0, "the only node must be returned");
    }
}