diskann-disk 0.56.0

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

use std::{marker::PhantomData, pin::Pin, sync::Arc};

use diskann::{
    graph::{
        glue::{InsertStrategy, PruneStrategy},
        Config, DiskANNIndex,
    },
    provider::DefaultContext,
    utils::VectorRepr,
    ANNError, ANNResult,
};
use diskann_providers::storage::{DynWriteProvider, WriteProviderWrapper};
use diskann_providers::{
    index::diskann_async,
    model::graph::provider::async_::{
        common::{
            FullPrecision, NoDeletes, NoStore, Quantized as DefaultQuantized, SetElementHelper,
            VectorStore,
        },
        inmem::{
            spherical, DefaultProvider, DefaultProviderParameters, FullPrecisionProvider,
            SetStartPoints,
        },
    },
    storage::{DiskGraphOnly, SaveWith},
};
use diskann_quantization::spherical::iface;
use diskann_utils::future::{AsyncFriendly, SendFuture};

use super::quantizer::BuildQuantizer;

/// Builder facade for in memory index construction and persistence.
///
/// Thread safety:
/// Implementors must be `Send` and `Sync`. Methods can be called from many tasks.
pub(super) trait InmemIndexBuilder<T: Sized>: Send + Sync {
    /// Return the total capacity of the provider, **excluding** start points.
    fn capacity(&self) -> usize;

    /// Return the total capacity of the provider, **including** start points.
    fn total_points(&self) -> usize;

    /// Set a single start point to search.
    ///
    /// The slice must match the underlying vector type, else `WrongDataType` is returned.
    fn set_start_point(&self, start_point: &[T]) -> ANNResult<()>;

    /// Insert a vector with a `id`.
    ///
    /// The slice must match the underlying vector type.
    fn insert_vector<'a>(
        &'a self,
        id: u32,
        vector: &'a [T],
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>>;

    /// Prune the built graph over `[range.start, range.end)`.
    fn final_prune(
        &self,
        range: core::ops::Range<u32>,
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + '_>>;

    /// Persist only the graph file set.
    fn save_graph<'a>(
        &'a self,
        storage_provider: &'a dyn DynWriteProvider,
        start_point_and_path: &'a (u32, DiskGraphOnly),
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>>;

    /// Return the number of vector reads for full_precision and quantized stores respectively.
    #[cfg(debug_assertions)]
    fn counts_for_get_vector(&self) -> (usize, usize);

    /// Count the number of nodes in the graph reachable from the given `start_points`.
    ///
    /// This function has a large memory footprint for large graphs and should not be called
    /// frequently. This is mainly for analysis and sanity tests.
    #[cfg(debug_assertions)]
    fn count_reachable_nodes(&self) -> Pin<Box<dyn SendFuture<ANNResult<usize>> + '_>>;
}

//////////////////////////////////
// FullPrecision Implementation //
//////////////////////////////////

impl<T> InmemIndexBuilder<T> for DiskANNIndex<FullPrecisionProvider<T>>
where
    T: VectorRepr,
{
    fn capacity(&self) -> usize {
        self.provider().capacity()
    }

    fn total_points(&self) -> usize {
        self.provider().total_points()
    }

    fn set_start_point(&self, start_point: &[T]) -> ANNResult<()> {
        self.provider()
            .set_start_points(std::iter::once(start_point))
    }

    fn insert_vector<'a>(
        &'a self,
        id: u32,
        vector: &'a [T],
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>> {
        Box::pin(async move {
            self.insert(&FullPrecision, &DefaultContext, &id, vector)
                .await
        })
    }

    fn final_prune(
        &self,
        range: core::ops::Range<u32>,
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + '_>> {
        Box::pin(async move {
            self.prune_range(&FullPrecision, &DefaultContext, range)
                .await
        })
    }

    fn save_graph<'a>(
        &'a self,
        storage_provider: &'a dyn DynWriteProvider,
        start_point_and_path: &'a (u32, DiskGraphOnly),
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>> {
        Box::pin(async move {
            let wrapper = WriteProviderWrapper::new(storage_provider);
            self.save_with(&wrapper, start_point_and_path).await
        })
    }

    #[cfg(debug_assertions)]
    fn counts_for_get_vector(&self) -> (usize, usize) {
        self.provider().counts_for_get_vector()
    }

    #[cfg(debug_assertions)]
    fn count_reachable_nodes(&self) -> Pin<Box<dyn SendFuture<ANNResult<usize>> + '_>> {
        Box::pin(async move {
            let provider = self.provider();
            let start_points = provider.starting_points()?;
            let mut neighbor_accessor = provider.neighbors();
            self.count_reachable_nodes(&start_points, &mut neighbor_accessor)
                .await
        })
    }
}

//////////////////////////
// Quant Implementation //
//////////////////////////

pub(super) struct QuantInMemBuilder<T, Q, S>
where
    Q: AsyncFriendly,
{
    index: DiskANNIndex<DefaultProvider<NoStore, Q>>,
    strategy: S,
    _vector_data_type: PhantomData<T>,
}

impl<T, Q, S> QuantInMemBuilder<T, Q, S>
where
    Q: AsyncFriendly,
{
    pub fn new(index: DiskANNIndex<DefaultProvider<NoStore, Q>>, strategy: S) -> Self {
        Self {
            index,
            strategy,
            _vector_data_type: PhantomData,
        }
    }

    fn index(&self) -> &DiskANNIndex<DefaultProvider<NoStore, Q>> {
        &self.index
    }
}

impl<T, Q, S> InmemIndexBuilder<T> for QuantInMemBuilder<T, Q, S>
where
    T: VectorRepr,
    Q: AsyncFriendly + VectorStore + SetElementHelper<T>,
    S: Send
        + Sync
        + for<'a> InsertStrategy<'a, DefaultProvider<NoStore, Q>, &'a [T]>
        + PruneStrategy<DefaultProvider<NoStore, Q>>,
    DefaultProvider<NoStore, Q>: SaveWith<(u32, u32, DiskGraphOnly), Error = ANNError>,
{
    fn capacity(&self) -> usize {
        self.index().provider().capacity()
    }

    fn total_points(&self) -> usize {
        self.index().provider().total_points()
    }

    fn set_start_point(&self, start_point: &[T]) -> ANNResult<()> {
        self.index()
            .provider()
            .set_start_points(std::iter::once(start_point))
    }

    fn insert_vector<'a>(
        &'a self,
        id: u32,
        vector: &'a [T],
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>> {
        Box::pin(async move {
            self.index()
                .insert(&self.strategy, &DefaultContext, &id, vector)
                .await
        })
    }

    fn final_prune(
        &self,
        range: core::ops::Range<u32>,
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + '_>> {
        Box::pin(async move {
            self.index()
                .prune_range(&self.strategy, &DefaultContext, range)
                .await
        })
    }

    fn save_graph<'a>(
        &'a self,
        storage_provider: &'a dyn DynWriteProvider,
        start_point_and_path: &'a (u32, DiskGraphOnly),
    ) -> Pin<Box<dyn SendFuture<ANNResult<()>> + 'a>> {
        Box::pin(async move {
            let wrapper = WriteProviderWrapper::new(storage_provider);
            self.index().save_with(&wrapper, start_point_and_path).await
        })
    }

    #[cfg(debug_assertions)]
    fn counts_for_get_vector(&self) -> (usize, usize) {
        self.index().provider().counts_for_get_vector()
    }

    #[cfg(debug_assertions)]
    fn count_reachable_nodes(&self) -> Pin<Box<dyn SendFuture<ANNResult<usize>> + '_>> {
        Box::pin(async move {
            let provider = self.index().provider();
            let start_points = provider.starting_points()?;
            let mut neighbor_accessor = provider.neighbors();
            self.index()
                .count_reachable_nodes(&start_points, &mut neighbor_accessor)
                .await
        })
    }
}

/// Create a new in-memory index builder for vectors of type `T`.
///
/// Chooses the builder implementation based on the given `BuildQuantizer`.
/// - `NoQuant` uses a plain index with no quantization.
/// - `Scalar1Bit`, `Spherical1Bit`, and `PQ` create quantized only indexes backed by
///   `QuantInMemBuilder`.
///
/// # Parameters
/// * `config` – Index configuration.
/// * `build_quantizer` – Quantization strategy to apply.
///
/// # Returns
/// An `Arc<dyn InmemIndexBuilder>` wrapped in `ANNResult`.
///
/// # Errors
/// Returns an error if the underlying index creation fails.
pub(super) fn new_inmem_index_builder<T>(
    config: Config,
    params: DefaultProviderParameters,
    build_quantizer: &BuildQuantizer,
) -> ANNResult<Arc<dyn InmemIndexBuilder<T>>>
where
    T: VectorRepr,
{
    match &build_quantizer {
        BuildQuantizer::NoQuant(_) => diskann_async::new_index::<T, _>(config, params, NoDeletes)
            .map(|index| index as Arc<dyn InmemIndexBuilder<T>>),
        BuildQuantizer::Scalar1Bit(q) => {
            let index = diskann_async::new_quant_only_index(config, params, q.clone(), NoDeletes)?;
            Ok(Arc::new(QuantInMemBuilder::new(index, DefaultQuantized)))
        }
        BuildQuantizer::Spherical1Bit(q) => {
            let quantizer = q.try_clone().map_err(spherical::AllocatorError::from)?;
            let plan = iface::Impl::<1>::new(quantizer).map_err(spherical::AllocatorError::from)?;
            let index = diskann_async::new_quant_only_index(config, params, plan, NoDeletes)?;
            Ok(Arc::new(QuantInMemBuilder::new(
                index,
                spherical::Quantized::build(),
            )))
        }
        BuildQuantizer::PQ(table) => {
            let index =
                diskann_async::new_quant_only_index(config, params, table.clone(), NoDeletes)?;
            Ok(Arc::new(QuantInMemBuilder::new(index, DefaultQuantized)))
        }
    }
}