diskann 0.54.0

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

//! [`FlatIndex`] — the index wrapper for a [`DataProvider`](crate::provider::DataProvider)
//! over which we do flat search.
use std::num::NonZeroUsize;

use diskann_utils::future::SendFuture;

use crate::{
    ANNResult,
    error::{ErrorExt, IntoANNResult},
    flat::{DistancesUnordered, SearchStrategy},
    graph::SearchOutputBuffer,
    neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType},
    provider::DataProvider,
};

/// Statistics collected during a flat search.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SearchStats {
    /// The total number of distance computations performed during the scan.
    pub cmps: u32,

    /// The total number of results written to the output buffer.
    pub result_count: u32,
}

/// A thin wrapper around a [`DataProvider`] used for flat search.
#[derive(Debug)]
pub struct FlatIndex<P: DataProvider> {
    /// The backing provider.
    provider: P,
}

impl<P: DataProvider> FlatIndex<P> {
    /// Construct a new [`FlatIndex`] around `provider`.
    pub fn new(provider: P) -> Self {
        Self { provider }
    }

    /// Borrow the underlying provider.
    pub fn provider(&self) -> &P {
        &self.provider
    }

    /// Brute-force k-nearest-neighbor flat search.
    ///
    /// Streams every element produced by the strategy's visitor through the query
    /// computer, keeps the best `k` candidates in a [`NeighborPriorityQueue`], and
    /// writes the `(id, distance)` survivors into `output` in best-first order.
    pub fn knn_search<S, T, OB>(
        &self,
        k: NonZeroUsize,
        strategy: &S,
        context: &P::Context,
        query: T,
        output: &mut OB,
    ) -> impl SendFuture<ANNResult<SearchStats>>
    where
        S: SearchStrategy<P, T>,
        S::Id: NeighborPriorityQueueIdType,
        T: Send + Sync,
        OB: SearchOutputBuffer<S::Id> + Send + ?Sized,
    {
        async move {
            let mut visitor = strategy
                .create_visitor(&self.provider, context)
                .into_ann_result()?;

            let computer = strategy.build_query_computer(query).into_ann_result()?;

            let k = k.get();
            let mut queue = NeighborPriorityQueue::new(k);
            let mut cmps: u32 = 0;

            visitor
                .distances_unordered(&computer, |id, dist| {
                    cmps += 1;
                    queue.insert(Neighbor::new(id, dist));
                })
                .await
                .escalate("flat scan must complete to produce correct k-NN results")?;

            let result_count =
                output.extend(queue.iter().take(k).map(|n| (n.id, n.distance))) as u32;

            Ok(SearchStats { cmps, result_count })
        }
    }
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use crate::flat::{
        FlatIndex,
        test::{
            harness::KnnOracleRun,
            provider::{self as flat_provider},
        },
    };
    use crate::graph::test::synthetic::Grid;

    fn fixture(grid: Grid, size: usize) -> (FlatIndex<flat_provider::Provider>, usize) {
        let provider = flat_provider::Provider::grid(grid, size).unwrap();
        let len = provider.len();
        (FlatIndex::new(provider), len)
    }

    /// `knn_search` returns a `Send` future, and a shared `&FlatIndex` can serve
    /// many concurrent searches on a multi-threaded runtime, each producing the
    /// correct top-k independently.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn multithreaded_knn_search() {
        use std::sync::Arc;

        let (index, len) = fixture(Grid::Two, 4);
        let index = Arc::new(index);

        // Mix of corner, axis-aligned, and off-grid queries; k spans 1..=len.
        let cases: &[(&[f32], usize)] = &[
            (&[-1.0, -1.0], 1),
            (&[1.0, 1.0], len),
            (&[-1.0, 1.0], len / 2),
            (&[1.0, -1.0], len - 1),
            (&[0.0, 0.0], 3),
            (&[3.0, 3.0], len),
            (&[-2.0, 0.5], 2),
            (&[0.5, -0.5], len),
        ];

        let mut set = tokio::task::JoinSet::new();
        for (query, k) in cases {
            let index = Arc::clone(&index);
            let query: Vec<f32> = query.to_vec();
            let k = *k;
            set.spawn(async move {
                let outcome = KnnOracleRun::run(
                    &index,
                    &flat_provider::Strategy::new(index.provider().dim()),
                    &query,
                    k,
                )
                .await
                .expect("knn_search failed");
                (query, k, outcome)
            });
        }

        while let Some(joined) = set.join_next().await {
            let (query, k, outcome) = joined.expect("task panicked");
            assert_eq!(
                outcome.top_k, outcome.ground_truth,
                "query = {query:?}, k = {k}: top-k must match brute force",
            );
            assert_eq!(outcome.stats.cmps as usize, len);
            assert_eq!(outcome.stats.result_count as usize, k.min(len));
        }
    }

    ////////////
    // Errors //
    ////////////

    /// A transient error from the visitor's scan must escalate up through `knn_search`.
    #[test]
    fn transient_scan_error() {
        // The flat scan touches every id, so any transient id is guaranteed to be hit.
        for transient_ids in [&[0u32][..], &[3][..], &[1, 2, 5][..]] {
            let strategy =
                flat_provider::Strategy::with_transient(2, transient_ids.iter().copied());
            let (index, _) = fixture(Grid::Two, 3);
            let err = KnnOracleRun::run_sync(&index, &strategy, &[1.0, 0.0], 4)
                .expect_err("transient error during full scan must escalate");

            let msg = format!("{err}");
            assert!(
                transient_ids
                    .iter()
                    .any(|id| msg.contains(&format!("id {id}"))),
                "transients = {transient_ids:?}: expected error to name one of the \
                 transient ids, got: {msg}",
            );
        }
    }

    /// Run `knn_search` via the harness, assert it fails, and check the error
    /// message contains `expected_msg`.
    fn assert_search_error(strategy: &flat_provider::Strategy, query: &[f32], expected_msg: &str) {
        let (index, _) = fixture(Grid::Two, 3);
        let err = KnnOracleRun::run_sync(&index, strategy, query, 4)
            .expect_err("expected knn_search to fail");

        let msg = format!("{err}");
        assert!(
            msg.contains(expected_msg),
            "expected error containing {expected_msg:?}, got: {msg}",
        );
    }

    #[test]
    fn strategy_constructor_errors() {
        // Strategy/provider expect dim=2, query has dim=3.
        assert_search_error(
            &flat_provider::Strategy::new(2),
            &[0.0, 0.0, 0.0],
            "dimension mismatch",
        );

        // Strategy expects dim=5, provider has dim=2.
        assert_search_error(
            &flat_provider::Strategy::new(5),
            &[0.0, 0.0],
            "dimension mismatch",
        );
    }
}