nodedb_vector/hnsw/

search.rs

// SPDX-License-Identifier: Apache-2.0

//! HNSW search algorithm (Malkov & Yashunin, Algorithm 2).
//!
//! Beam search through the multi-layer graph with optional Roaring bitmap
//! pre-filtering for efficient filtered k-NN queries.

use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// Issue a non-faulting prefetch hint for the memory region starting at `ptr`
/// into all cache levels (T0). No-op on architectures without a prefetch
/// intrinsic.
#[inline(always)]
fn prefetch_t0(ptr: *const u8) {
    #[cfg(target_arch = "x86_64")]
    {
        // SAFETY: _mm_prefetch is a pure hint — it never faults even on an
        // invalid address, and it does not dereference the pointer.
        unsafe {
            core::arch::x86_64::_mm_prefetch::<{ core::arch::x86_64::_MM_HINT_T0 }>(
                ptr as *const i8,
            );
        }
    }
    // aarch64 hardware prefetchers handle sequential + pointer-chasing patterns
    // well without explicit hints. wasm32 and all other targets: intentional no-op.
    #[cfg(not(target_arch = "x86_64"))]
    let _ = ptr;
}

use roaring::RoaringBitmap;

use crate::hnsw::graph::{Candidate, HnswIndex, SearchResult};

impl HnswIndex {
    /// K-NN search: find the `k` closest vectors to `query`.
    ///
    /// `ef` controls the search beam width (higher = better recall, slower).
    /// Must be >= k. Typical values: ef = 2*k to 10*k.
    pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim, "query dimension mismatch");
        if self.is_empty() {
            return Vec::new();
        }

        /// Maximum beam width to prevent runaway search cost.
        const MAX_EF: usize = 8192;
        let ef = ef.max(k).min(MAX_EF);
        let Some(ep) = self.entry_point else {
            return Vec::new();
        };

        // Phase 1: Greedy descent from top layer to layer 1.
        let mut current_ep = ep;
        for layer in (1..=self.max_layer).rev() {
            let results = search_layer(self, query, current_ep, 1, layer, None, 0);
            if let Some(nearest) = results.first() {
                current_ep = nearest.id;
            }
        }

        // Phase 2: Beam search at layer 0.
        let results = search_layer(self, query, current_ep, ef, 0, None, 0);

        results
            .into_iter()
            .take(k)
            .map(|c| SearchResult {
                id: c.id,
                distance: c.dist,
            })
            .collect()
    }

    /// Filtered K-NN search with Roaring bitmap pre-filtering.
    pub fn search_filtered(
        &self,
        query: &[f32],
        k: usize,
        ef: usize,
        filter: &RoaringBitmap,
    ) -> Vec<SearchResult> {
        self.search_filtered_offset(query, k, ef, filter, 0)
    }

    /// Filtered K-NN search where the bitmap is keyed in a shifted ID space.
    ///
    /// `id_offset` is added to local node IDs before testing `filter.contains`.
    /// Used by multi-segment collections where the bitmap holds GLOBAL ids
    /// and each segment's HNSW nodes are numbered starting at `base_id`.
    pub fn search_filtered_offset(
        &self,
        query: &[f32],
        k: usize,
        ef: usize,
        filter: &RoaringBitmap,
        id_offset: u32,
    ) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim, "query dimension mismatch");
        if self.is_empty() {
            return Vec::new();
        }

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

        let mut current_ep = ep;
        for layer in (1..=self.max_layer).rev() {
            let results = search_layer(self, query, current_ep, 1, layer, None, 0);
            if let Some(nearest) = results.first() {
                current_ep = nearest.id;
            }
        }

        let results = search_layer(self, query, current_ep, ef, 0, Some(filter), id_offset);

        results
            .into_iter()
            .take(k)
            .map(|c| SearchResult {
                id: c.id,
                distance: c.dist,
            })
            .collect()
    }

    /// Deserialize a Roaring bitmap from bytes and perform filtered search.
    pub fn search_with_bitmap_bytes(
        &self,
        query: &[f32],
        k: usize,
        ef: usize,
        bitmap_bytes: &[u8],
    ) -> Vec<SearchResult> {
        self.search_with_bitmap_bytes_offset(query, k, ef, bitmap_bytes, 0)
    }

    /// Deserialize a Roaring bitmap and search with an ID offset applied
    /// before testing membership. See `search_filtered_offset` for rationale.
    pub fn search_with_bitmap_bytes_offset(
        &self,
        query: &[f32],
        k: usize,
        ef: usize,
        bitmap_bytes: &[u8],
        id_offset: u32,
    ) -> Vec<SearchResult> {
        match RoaringBitmap::deserialize_from(bitmap_bytes) {
            Ok(bitmap) => self.search_filtered_offset(query, k, ef, &bitmap, id_offset),
            Err(_) => self.search(query, k, ef),
        }
    }
}

/// Unified HNSW beam search on a single layer with optional pre-filter.
///
/// When `filter` is `None`, all non-deleted nodes enter the result set.
/// When `filter` is `Some`, only nodes present in the bitmap enter results,
/// but all nodes are still traversed for graph connectivity.
///
/// Scratch buffers are drawn from `index.arena` (a `RefCell`-guarded
/// `BeamSearchArena`).  The arena is reset at entry and its `Vec`/`HashSet`
/// capacity grows to the high-water mark over successive calls, giving
/// amortised zero-allocation steady state.
pub(crate) fn search_layer(
    index: &HnswIndex,
    query: &[f32],
    entry_point: u32,
    ef: usize,
    layer: usize,
    filter: Option<&RoaringBitmap>,
    id_offset: u32,
) -> Vec<Candidate> {
    let mut arena = index.arena.borrow_mut();

    // Reset scratch buffers — retains Vec/HashSet capacity from prior calls.
    arena.reset();

    arena.visited.insert(entry_point);

    let ep_dist = index.dist_to_node(query, entry_point);
    let ep_candidate = Candidate {
        dist: ep_dist,
        id: entry_point,
    };

    // Strategy (c): take the arena Vecs, build BinaryHeaps on them, do the
    // search, then move the Vecs back.  The Vecs retain their allocated
    // capacity across the take/into_vec round-trip, so steady-state searches
    // incur zero heap allocation.
    let mut cand_vec = std::mem::take(&mut arena.candidates);
    cand_vec.push(Reverse(ep_candidate));
    let mut candidates = BinaryHeap::from(cand_vec);

    let mut res_vec = std::mem::take(&mut arena.results);

    let passes_filter = |id: u32| -> bool {
        if index.nodes[id as usize].deleted {
            return false;
        }
        match filter {
            Some(f) => f.contains(id + id_offset),
            None => true,
        }
    };

    if passes_filter(entry_point) {
        res_vec.push(ep_candidate);
    }
    let mut results = BinaryHeap::from(res_vec);

    while let Some(Reverse(current)) = candidates.pop() {
        if let Some(worst) = results.peek()
            && current.dist > worst.dist
            && results.len() >= ef
        {
            break;
        }

        // Prefetch the vector of the next candidate before touching this
        // iteration's neighbor list, so it lands in cache by the time the
        // inner loop calls dist_to_node on it.
        if let Some(Reverse(next)) = candidates.peek()
            && let Some(node) = index.nodes.get(next.id as usize)
            && let Some(v) = node.vector.first()
        {
            prefetch_t0(v as *const f32 as *const u8);
        }

        let neighbors = index.neighbors_at(current.id, layer);
        if neighbors.is_empty() && layer >= index.node_num_layers(current.id) {
            continue;
        }

        for &neighbor_id in neighbors {
            if !arena.visited.insert(neighbor_id) {
                continue;
            }

            let dist = index.dist_to_node(query, neighbor_id);
            let neighbor = Candidate {
                dist,
                id: neighbor_id,
            };

            let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist);
            let should_explore = dist < worst_dist || results.len() < ef;

            if should_explore {
                candidates.push(Reverse(neighbor));
            }

            if passes_filter(neighbor_id) {
                results.push(neighbor);
                if results.len() > ef {
                    results.pop();
                }
            }
        }
    }

    // Restore the candidates Vec to the arena before releasing the borrow.
    arena.candidates = candidates.into_vec();

    let mut result_vec = results.into_vec();
    // arena.results stays empty — reset() will clear it on the next call.
    drop(arena);

    result_vec.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist));
    result_vec
}

#[cfg(test)]
mod tests {
    use crate::distance::DistanceMetric;
    use crate::hnsw::{HnswIndex, HnswParams};
    use roaring::RoaringBitmap;

    fn build_index(n: usize, dim: usize) -> HnswIndex {
        let mut idx = HnswIndex::with_seed(
            dim,
            HnswParams {
                m: 16,
                m0: 32,
                ef_construction: 100,
                metric: DistanceMetric::L2,
            },
            42,
        );
        for i in 0..n {
            let v: Vec<f32> = (0..dim).map(|d| (i * dim + d) as f32).collect();
            idx.insert(v).unwrap();
        }
        idx
    }

    #[test]
    fn search_empty_index() {
        let idx = HnswIndex::new(3, HnswParams::default());
        let results = idx.search(&[1.0, 2.0, 3.0], 5, 50);
        assert!(results.is_empty());
    }

    #[test]
    fn search_single_element() {
        let mut idx = HnswIndex::with_seed(
            2,
            HnswParams {
                m: 4,
                m0: 8,
                ef_construction: 16,
                metric: DistanceMetric::L2,
            },
            1,
        );
        idx.insert(vec![1.0, 0.0]).unwrap();
        let results = idx.search(&[1.0, 0.0], 1, 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, 0);
        assert!(results[0].distance < 1e-6);
    }

    #[test]
    fn search_finds_exact_match() {
        let idx = build_index(50, 3);
        let query = idx.get_vector(25).unwrap().to_vec();
        let results = idx.search(&query, 1, 50);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, 25);
        assert!(results[0].distance < 1e-6);
    }

    #[test]
    fn search_returns_sorted_by_distance() {
        let idx = build_index(100, 4);
        let query = vec![50.0, 50.0, 50.0, 50.0];
        let results = idx.search(&query, 10, 64);
        assert_eq!(results.len(), 10);
        for w in results.windows(2) {
            assert!(w[0].distance <= w[1].distance);
        }
    }

    #[test]
    fn search_k_larger_than_index() {
        let idx = build_index(5, 2);
        let results = idx.search(&[0.0, 0.0], 20, 50);
        assert_eq!(results.len(), 5);
    }

    #[test]
    fn search_recall_at_10() {
        let idx = build_index(500, 3);
        let query = vec![100.0, 100.0, 100.0];
        let results = idx.search(&query, 10, 128);

        let mut truth: Vec<(u32, f32)> = (0..500)
            .map(|i| {
                let v = idx.get_vector(i).unwrap();
                let d = crate::distance::l2_squared(&query, v);
                (i, d)
            })
            .collect();
        truth.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        let truth_top10: std::collections::HashSet<u32> = truth[..10].iter().map(|t| t.0).collect();

        let found: std::collections::HashSet<u32> = results.iter().map(|r| r.id).collect();
        let recall = found.intersection(&truth_top10).count() as f64 / 10.0;
        assert!(recall >= 0.8, "recall@10 = {recall:.2}, expected >= 0.80");
    }

    #[test]
    fn search_excludes_tombstoned() {
        let mut idx = build_index(20, 3);
        idx.delete(0);
        let results = idx.search(&[0.0, 0.0, 0.0], 5, 32);
        for r in &results {
            assert_ne!(r.id, 0, "tombstoned node appeared in results");
        }
    }

    #[test]
    fn search_filtered_respects_bitmap() {
        let idx = build_index(50, 3);
        let mut filter = RoaringBitmap::new();
        for i in (0..50u32).step_by(2) {
            filter.insert(i);
        }
        let results = idx.search_filtered(&[0.0, 0.0, 0.0], 5, 64, &filter);
        assert_eq!(results.len(), 5);
        for r in &results {
            assert!(r.id % 2 == 0, "got odd id {}", r.id);
        }
    }

    #[test]
    fn search_filtered_empty_returns_empty() {
        let idx = build_index(20, 3);
        let filter = RoaringBitmap::new();
        let results = idx.search_filtered(&[0.0, 0.0, 0.0], 5, 64, &filter);
        assert!(results.is_empty());
    }

    #[test]
    fn bitmap_bytes_roundtrip() {
        let idx = build_index(50, 3);
        let mut filter = RoaringBitmap::new();
        for i in 0..25u32 {
            filter.insert(i);
        }
        let mut bytes = Vec::new();
        filter.serialize_into(&mut bytes).unwrap();
        let results = idx.search_with_bitmap_bytes(&[0.0, 0.0, 0.0], 5, 32, &bytes);
        for r in &results {
            assert!(r.id < 25, "got filtered-out node {}", r.id);
        }
    }
}