Method search

nodedb-vector/src/codec_index/search.rs:57–100

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`.

(&self, query: &[f32], k: usize, ef_search: usize)

    ///
    /// 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 {