nodedb_vector/

flat.rs

// SPDX-License-Identifier: Apache-2.0

//! Flat (brute-force) vector index for small collections.
//!
//! Simple linear scan over all stored vectors. No graph overhead, exact
//! results. Automatically used when a collection has fewer than
//! `DEFAULT_FLAT_INDEX_THRESHOLD` vectors (default 10K). Also serves as the
//! search method for growing segments before HNSW construction.
//!
//! Complexity: O(N × D) per query where N = vectors, D = dimensions.

use roaring::RoaringBitmap;

use crate::distance::{DistanceMetric, distance};
use crate::hnsw::SearchResult;

/// Default threshold below which collections use flat index instead of HNSW.
pub const DEFAULT_FLAT_INDEX_THRESHOLD: usize = 10_000;

/// Flat vector index: append-only buffer with brute-force search.
pub struct FlatIndex {
    dim: usize,
    metric: DistanceMetric,
    /// Vectors stored contiguously for cache-friendly sequential scan.
    data: Vec<f32>,
    /// Tombstone bitmap: `deleted[i]` = true means vector i is soft-deleted.
    deleted: Vec<bool>,
    /// Number of live (non-deleted) vectors.
    live_count: usize,
}

impl FlatIndex {
    /// Create a new empty flat index.
    pub fn new(dim: usize, metric: DistanceMetric) -> Self {
        Self {
            dim,
            metric,
            data: Vec::new(),
            deleted: Vec::new(),
            live_count: 0,
        }
    }

    /// Insert a vector. Returns the assigned vector ID.
    pub fn insert(&mut self, vector: Vec<f32>) -> u32 {
        assert_eq!(
            vector.len(),
            self.dim,
            "dimension mismatch: expected {}, got {}",
            self.dim,
            vector.len()
        );
        let id = self.len() as u32;
        self.data.extend_from_slice(&vector);
        self.deleted.push(false);
        self.live_count += 1;
        id
    }

    /// Soft-delete a vector by ID.
    pub fn delete(&mut self, id: u32) -> bool {
        let idx = id as usize;
        if idx < self.deleted.len() && !self.deleted[idx] {
            self.deleted[idx] = true;
            self.live_count -= 1;
            true
        } else {
            false
        }
    }

    /// Brute-force k-NN search with an explicit distance metric override.
    /// Overrides the `self.metric` configured at collection creation time.
    pub fn search_with_metric(
        &self,
        query: &[f32],
        top_k: usize,
        metric: DistanceMetric,
    ) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim);
        let n = self.len();
        if n == 0 || top_k == 0 {
            return Vec::new();
        }

        let mut candidates: Vec<SearchResult> = Vec::with_capacity(n.min(top_k * 2));
        for i in 0..n {
            if self.deleted[i] {
                continue;
            }
            let start = i * self.dim;
            let vec_slice = &self.data[start..start + self.dim];
            let dist = distance(query, vec_slice, metric);
            candidates.push(SearchResult {
                id: i as u32,
                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)
        });
        candidates
    }

    /// Brute-force k-NN search. Exact results — no approximation.
    pub fn search(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim);
        let n = self.len();
        if n == 0 || top_k == 0 {
            return Vec::new();
        }

        let mut candidates: Vec<SearchResult> = Vec::with_capacity(n.min(top_k * 2));
        for i in 0..n {
            if self.deleted[i] {
                continue;
            }
            let start = i * self.dim;
            let vec_slice = &self.data[start..start + self.dim];
            let dist = distance(query, vec_slice, self.metric);
            candidates.push(SearchResult {
                id: i as u32,
                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)
        });
        candidates
    }

    /// Search with a pre-filter bitmap (byte-array format).
    pub fn search_filtered(&self, query: &[f32], top_k: usize, bitmap: &[u8]) -> Vec<SearchResult> {
        self.search_filtered_offset(query, top_k, bitmap, 0)
    }

    /// Filtered search with an explicit metric override.
    pub fn search_filtered_offset_with_metric(
        &self,
        query: &[f32],
        top_k: usize,
        bitmap: &[u8],
        id_offset: u32,
        metric: DistanceMetric,
    ) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim);
        let n = self.len();
        if n == 0 || top_k == 0 {
            return Vec::new();
        }

        let parsed = RoaringBitmap::deserialize_from(bitmap).ok();

        let mut candidates: Vec<SearchResult> = Vec::with_capacity(top_k * 2);
        for i in 0..n {
            if self.deleted[i] {
                continue;
            }
            if let Some(ref bm) = parsed {
                let global = (i as u32).saturating_add(id_offset);
                if !bm.contains(global) {
                    continue;
                }
            }
            let start = i * self.dim;
            let vec_slice = &self.data[start..start + self.dim];
            let dist = distance(query, vec_slice, metric);
            candidates.push(SearchResult {
                id: i as u32,
                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)
        });
        candidates
    }

    /// Search with a pre-filter bitmap applying a global id offset.
    ///
    /// `bitmap` is a serialized `RoaringBitmap` (matching the HNSW filter
    /// format). Bit `i + id_offset` tests local id `i`. Used by multi-segment
    /// collections where the bitmap holds GLOBAL vector ids. If the bytes
    /// fail to deserialize, the search degrades to unfiltered.
    pub fn search_filtered_offset(
        &self,
        query: &[f32],
        top_k: usize,
        bitmap: &[u8],
        id_offset: u32,
    ) -> Vec<SearchResult> {
        assert_eq!(query.len(), self.dim);
        let n = self.len();
        if n == 0 || top_k == 0 {
            return Vec::new();
        }

        let parsed = RoaringBitmap::deserialize_from(bitmap).ok();

        let mut candidates: Vec<SearchResult> = Vec::with_capacity(top_k * 2);
        for i in 0..n {
            if self.deleted[i] {
                continue;
            }
            if let Some(ref bm) = parsed {
                let global = (i as u32).saturating_add(id_offset);
                if !bm.contains(global) {
                    continue;
                }
            }
            let start = i * self.dim;
            let vec_slice = &self.data[start..start + self.dim];
            let dist = distance(query, vec_slice, self.metric);
            candidates.push(SearchResult {
                id: i as u32,
                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)
        });
        candidates
    }

    pub fn len(&self) -> usize {
        self.deleted.len()
    }

    pub fn live_count(&self) -> usize {
        self.live_count
    }

    pub fn is_empty(&self) -> bool {
        self.live_count == 0
    }

    pub fn get_vector(&self, id: u32) -> Option<&[f32]> {
        let idx = id as usize;
        if idx < self.deleted.len() && !self.deleted[idx] {
            let start = idx * self.dim;
            Some(&self.data[start..start + self.dim])
        } else {
            None
        }
    }

    /// Raw access bypassing tombstone filter — used by snapshot/restore.
    pub fn get_vector_raw(&self, id: u32) -> Option<&[f32]> {
        let idx = id as usize;
        if idx < self.deleted.len() {
            let start = idx * self.dim;
            Some(&self.data[start..start + self.dim])
        } else {
            None
        }
    }

    /// Whether the given local id has been tombstoned.
    pub fn is_deleted(&self, id: u32) -> bool {
        let idx = id as usize;
        idx < self.deleted.len() && self.deleted[idx]
    }

    /// Insert a vector that is already tombstoned (for checkpoint restore).
    pub fn insert_tombstoned(&mut self, vector: Vec<f32>) -> u32 {
        assert_eq!(
            vector.len(),
            self.dim,
            "dimension mismatch: expected {}, got {}",
            self.dim,
            vector.len()
        );
        let id = self.len() as u32;
        self.data.extend_from_slice(&vector);
        self.deleted.push(true);
        // No live_count increment — it's dead on arrival.
        id
    }

    pub fn dim(&self) -> usize {
        self.dim
    }

    pub fn metric(&self) -> DistanceMetric {
        self.metric
    }

    pub fn tombstone_count(&self) -> usize {
        self.len().saturating_sub(self.live_count)
    }
}

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

    #[test]
    fn insert_and_search() {
        let mut idx = FlatIndex::new(3, DistanceMetric::L2);
        for i in 0..100u32 {
            idx.insert(vec![i as f32, 0.0, 0.0]);
        }
        assert_eq!(idx.len(), 100);
        assert_eq!(idx.live_count(), 100);

        let results = idx.search(&[50.0, 0.0, 0.0], 3);
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].id, 50);
        assert!(results[0].distance < 0.01);
    }

    #[test]
    fn delete_excludes_from_search() {
        let mut idx = FlatIndex::new(2, DistanceMetric::L2);
        idx.insert(vec![0.0, 0.0]);
        idx.insert(vec![1.0, 0.0]);
        idx.insert(vec![2.0, 0.0]);

        assert!(idx.delete(1));
        assert_eq!(idx.live_count(), 2);

        let results = idx.search(&[1.0, 0.0], 3);
        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r.id != 1));
    }

    #[test]
    fn exact_results() {
        let mut idx = FlatIndex::new(2, DistanceMetric::Cosine);
        idx.insert(vec![1.0, 0.0]);
        idx.insert(vec![0.0, 1.0]);
        idx.insert(vec![1.0, 1.0]);

        let results = idx.search(&[1.0, 0.0], 1);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, 0);
    }

    #[test]
    fn empty_search() {
        let idx = FlatIndex::new(3, DistanceMetric::L2);
        let results = idx.search(&[1.0, 0.0, 0.0], 5);
        assert!(results.is_empty());
    }

    #[test]
    fn filtered_search() {
        let mut idx = FlatIndex::new(2, DistanceMetric::L2);
        for i in 0..8u32 {
            idx.insert(vec![i as f32, 0.0]);
        }
        let bitmap = vec![0b11001100u8];
        let results = idx.search_filtered(&[3.0, 0.0], 2, &bitmap);
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, 3);
        assert_eq!(results[1].id, 2);
    }
}