Method search

nodedb-vector/src/ivf.rs:130–189

Search: find top-k nearest neighbors.

(&self, query: &[f32], top_k: usize)

    /// Search: find top-k nearest neighbors.
    pub fn search(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim);
        if self.centroids.is_empty() || self.count == 0 {
            return Vec::new();
        }

        let pq = match &self.pq {
            Some(p) => p,
            None => return Vec::new(),
        };

        let nprobe = self.params.nprobe.min(self.centroids.len());
        let mut centroid_dists: Vec<(usize, f32)> = self
            .centroids
            .iter()
            .enumerate()
            .map(|(i, c)| (i, distance(query, c, self.params.metric)))
            .collect();
        centroid_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

        let mut candidates: Vec<SearchResult> = Vec::new();

        for &(cell_idx, _) in centroid_dists.iter().take(nprobe) {
            let residual_query: Vec<f32> = query
                .iter()
                .zip(&self.centroids[cell_idx])
                .map(|(q, c)| q - c)
                .collect();
            let table = match pq.build_distance_table(&residual_query) {
                Ok(t) => t,
                Err(e) => {
                    tracing::warn!(error = %e, "IVF PQ build_distance_table budget exhausted; skipping cell");
                    continue;
                }
            };

            for (id, code) in &self.cells[cell_idx] {
                let dist = pq.asymmetric_distance(&table, code);
                candidates.push(SearchResult {
                    id: *id,
                    distance: dist,
                });
            }
        }

        if candidates.len() > top_k {
            candidates.select_nth_unstable_by(top_k, |a, b| {
                a.distance
                    .partial_cmp(&b.distance)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            candidates.truncate(top_k);
        }
        candidates.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });