diff --git a/Cargo.lock b/Cargo.lock index e80330d7d..3d7a91009 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -513,6 +513,7 @@ dependencies = [ "diskann", "diskann-benchmark-runner", "diskann-disk", + "diskann-inmem", "diskann-label-filter", "diskann-providers", "diskann-quantization", @@ -578,6 +579,7 @@ dependencies = [ "cfg-if", "criterion", "diskann", + "diskann-inmem", "diskann-linalg", "diskann-platform", "diskann-providers", @@ -602,6 +604,45 @@ dependencies = [ "vfs", ] +[[package]] +name = "diskann-inmem" +version = "0.41.0" +dependencies = [ + "anyhow", + "approx", + "arc-swap", + "bincode", + "bytemuck", + "byteorder", + "cfg-if", + "diskann", + "diskann-linalg", + "diskann-platform", + "diskann-providers", + "diskann-quantization", + "diskann-utils", + "diskann-vector", + "diskann-wide", + "futures-util", + "half", + "hashbrown 0.16.1", + "num-traits", + "once_cell", + "proptest", + "prost", + "rand 0.9.2", + "rand_distr", + "rayon", + "rstest", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.17", + "tokio", + "tracing", + "vfs", +] + [[package]] name = "diskann-label-filter" version = "0.41.0" @@ -663,6 +704,7 @@ dependencies = [ "cfg-if", "criterion", "diskann", + "diskann-inmem", "diskann-linalg", "diskann-platform", "diskann-quantization", diff --git a/Cargo.toml b/Cargo.toml index 8f4c20aeb..55a53e444 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ "diskann", # Providers "diskann-providers", + "diskann-inmem", "diskann-disk", "diskann-label-filter", # Infrastructure @@ -55,6 +56,7 @@ diskann-platform = { path = "diskann-platform", version = "0.41.0" } diskann = { path = "diskann", version = "0.41.0" } # Providers diskann-providers = { path = "diskann-providers", default-features = false, version = "0.41.0" } +diskann-inmem = { path = "diskann-inmem", version = "0.41.0" } diskann-disk = { path = "diskann-disk", version = "0.41.0" } diskann-label-filter = { path = "diskann-label-filter", version = "0.41.0" } # Infra diff --git a/diskann-benchmark/Cargo.toml b/diskann-benchmark/Cargo.toml index 4810ca401..35d6e8913 100644 --- a/diskann-benchmark/Cargo.toml +++ b/diskann-benchmark/Cargo.toml @@ -15,6 +15,7 @@ bytemuck.workspace = true roaring.workspace = true clap = { workspace = true, features = ["derive"] } diskann-providers = { workspace = true } +diskann-inmem = { workspace = true } diskann.workspace = true diskann-utils.workspace = true half = { workspace = true, features = ["rand_distr", "num-traits"] } diff --git a/diskann-benchmark/src/backend/index/benchmarks.rs b/diskann-benchmark/src/backend/index/benchmarks.rs index 8ff8b2f08..41c8a2de8 100644 --- a/diskann-benchmark/src/backend/index/benchmarks.rs +++ b/diskann-benchmark/src/backend/index/benchmarks.rs @@ -18,12 +18,12 @@ use diskann_benchmark_runner::{ Any, Checkpoint, }; use diskann_providers::{ - index::diskann_async, model::{ configuration::IndexConfiguration, - graph::provider::async_::{common, inmem::DefaultProvider}, + graph::provider::async_::common, }, }; +use diskann_inmem::{diskann_async, DefaultProvider}; use diskann_utils::{ future::AsyncFriendly, views::{Matrix, MatrixView}, diff --git a/diskann-benchmark/src/backend/index/build.rs b/diskann-benchmark/src/backend/index/build.rs index 94d342437..d23de62bb 100644 --- a/diskann-benchmark/src/backend/index/build.rs +++ b/diskann-benchmark/src/backend/index/build.rs @@ -23,10 +23,10 @@ use diskann_providers::{ self, model::{ configuration::IndexConfiguration, - graph::provider::async_::inmem::{DefaultProvider, SetStartPoints}, }, storage::{AsyncIndexMetadata, LoadWith, SaveWith}, }; +use diskann_inmem::{DefaultProvider, SetStartPoints}; use diskann_utils::{ future::AsyncFriendly, views::{Matrix, MatrixView}, diff --git a/diskann-benchmark/src/backend/index/product.rs b/diskann-benchmark/src/backend/index/product.rs index 9141a6367..3ba5a6fc5 100644 --- a/diskann-benchmark/src/backend/index/product.rs +++ b/diskann-benchmark/src/backend/index/product.rs @@ -37,9 +37,9 @@ mod imp { use std::{io::Write, sync::Arc}; use diskann_providers::{ - index::diskann_async::{self}, model::{graph::provider::async_::common, IndexConfiguration}, }; + use diskann_inmem::diskann_async; use diskann_utils::views::{Matrix, MatrixView}; use diskann_benchmark_runner::{ diff --git a/diskann-benchmark/src/backend/index/scalar.rs b/diskann-benchmark/src/backend/index/scalar.rs index bce32071d..eef0146a2 100644 --- a/diskann-benchmark/src/backend/index/scalar.rs +++ b/diskann-benchmark/src/backend/index/scalar.rs @@ -82,12 +82,12 @@ mod imp { Any, Checkpoint, Output, }; use diskann_providers::{ - index::diskann_async::{self}, model::{ configuration::IndexConfiguration, - graph::provider::async_::{common, inmem}, + graph::provider::async_::common, }, }; + use diskann_inmem::{self as inmem, diskann_async}; use diskann_utils::views::{Matrix, MatrixView}; use half::f16; diff --git a/diskann-benchmark/src/backend/index/spherical.rs b/diskann-benchmark/src/backend/index/spherical.rs index a656548d5..2ca825fed 100644 --- a/diskann-benchmark/src/backend/index/spherical.rs +++ b/diskann-benchmark/src/backend/index/spherical.rs @@ -64,9 +64,9 @@ mod imp { Any, Checkpoint, Output, }; use diskann_providers::{ - index::diskann_async::{self}, - model::graph::provider::async_::{common::NoDeletes, inmem}, + model::graph::provider::async_::common::NoDeletes, }; + use diskann_inmem::{self as inmem, diskann_async}; use diskann_quantization::alloc::GlobalAllocator; use diskann_utils::views::Matrix; use rand::SeedableRng; diff --git a/diskann-benchmark/src/backend/index/update.rs b/diskann-benchmark/src/backend/index/update.rs index cb90707a2..fa82d2779 100644 --- a/diskann-benchmark/src/backend/index/update.rs +++ b/diskann-benchmark/src/backend/index/update.rs @@ -15,7 +15,7 @@ use diskann::{ utils::async_tools, }; use diskann_benchmark_runner::utils::{percentiles, MicroSeconds}; -use diskann_providers::model::graph::provider::async_::inmem::DefaultProvider; +use diskann_inmem::DefaultProvider; use diskann_utils::views::Matrix; use serde::Serialize; diff --git a/diskann-benchmark/src/inputs/async_.rs b/diskann-benchmark/src/inputs/async_.rs index fccd5e28d..76ef9cafb 100644 --- a/diskann-benchmark/src/inputs/async_.rs +++ b/diskann-benchmark/src/inputs/async_.rs @@ -17,10 +17,10 @@ use diskann_benchmark_runner::{ use diskann_providers::{ model::{ configuration::IndexConfiguration, - graph::provider::async_::inmem::DefaultProviderParameters, }, utils::load_metadata_from_file, }; +use diskann_inmem::DefaultProviderParameters; use serde::{Deserialize, Serialize}; use crate::{ diff --git a/diskann-disk/Cargo.toml b/diskann-disk/Cargo.toml index 548042255..92a715ecc 100644 --- a/diskann-disk/Cargo.toml +++ b/diskann-disk/Cargo.toml @@ -16,6 +16,7 @@ targets = ["x86_64-unknown-linux-gnu", "aarch64-pc-windows-msvc", "i686-pc-windo diskann = { workspace = true } diskann-utils = { workspace = true } diskann-providers = { workspace = true } +diskann-inmem = { workspace = true } diskann-vector = { workspace = true } diskann-linalg = { workspace = true } diskann-quantization = { workspace = true, features = ["rayon"] } diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index 4c6788636..28f731322 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -19,7 +19,6 @@ use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ model::{ graph::{ - provider::async_::inmem::DefaultProviderParameters, traits::{AdHoc, GraphDataType}, }, IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, @@ -30,6 +29,7 @@ use diskann_providers::{ VectorDataIterator, MAX_MEDOID_SAMPLE_SIZE, }, }; +use diskann_inmem::DefaultProviderParameters; use tokio::task::JoinSet; use tracing::{debug, info}; diff --git a/diskann-inmem/Cargo.toml b/diskann-inmem/Cargo.toml new file mode 100644 index 000000000..b71397214 --- /dev/null +++ b/diskann-inmem/Cargo.toml @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +[package] +name = "diskann-inmem" +version.workspace = true +authors.workspace = true +description.workspace = true +documentation.workspace = true +license.workspace = true +edition = "2024" +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +bincode.workspace = true +bytemuck = { workspace = true, features = ["must_cast"]} +byteorder.workspace = true +half = { workspace = true, features = ["bytemuck", "num-traits"] } +hashbrown.workspace = true +num-traits.workspace = true +once_cell.workspace = true +rand.workspace = true +rayon.workspace = true +serde = { workspace = true, features = ["derive"] } +thiserror.workspace = true +diskann-platform = { workspace = true } +diskann-vector = { workspace = true } +rand_distr.workspace = true +cfg-if.workspace = true +diskann-wide = { workspace = true } +arc-swap.workspace = true +tracing.workspace = true +diskann-linalg = { workspace = true } +vfs = { workspace = true } +diskann = { workspace = true } +diskann-utils = { workspace = true } +diskann-quantization = { workspace = true, features = ["rayon"] } +tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } +anyhow.workspace = true +prost = "0.14.1" +futures-util.workspace = true +serde_json.workspace = true +diskann-providers = { workspace = true } + +[dev-dependencies] +approx.workspace = true +tempfile.workspace = true +rstest.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["full"] } +proptest.workspace = true + +[package.metadata.docs.rs] +default-target = "x86_64-pc-windows-msvc" +targets = [ + "x86_64-unknown-linux-gnu", + "aarch64-pc-windows-msvc", + "i686-pc-windows-msvc", + "x86_64-pc-windows-msvc", +] + +[features] +default = [] +experimental_diversity_search = ["diskann/experimental_diversity_search"] + +# Some 'cfg's in the source tree will be flagged by `cargo clippy -j 2 --workspace --no-deps --all-targets -- -D warnings` +[lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage)'] } + +[lints.clippy] +# Commenting out for now. It would give warnings. +# undocumented_unsafe_blocks = "warn" +uninlined_format_args = "allow" diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-inmem/src/diskann_async.rs similarity index 98% rename from diskann-providers/src/index/diskann_async.rs rename to diskann-inmem/src/diskann_async.rs index adfd6695a..8642923dd 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-inmem/src/diskann_async.rs @@ -11,16 +11,11 @@ use diskann::{ utils::VectorRepr, }; use diskann_utils::future::AsyncFriendly; +use diskann_providers::model::graph::provider::async_::common::{CreateDeleteProvider, CreateVectorStore, NoDeletes, NoStore}; -use crate::model::{ - self, - graph::provider::async_::{ - common::{CreateDeleteProvider, CreateVectorStore, NoDeletes, NoStore}, - inmem::{ - CreateFullPrecision, DefaultProvider, DefaultProviderParameters, DefaultQuant, - FullPrecisionProvider, - }, - }, +use crate::{ + CreateFullPrecision, DefaultProvider, DefaultProviderParameters, DefaultQuant, + FullPrecisionProvider, }; ///////////////////////// @@ -197,20 +192,20 @@ pub(crate) mod tests { use rstest::rstest; use super::*; - use crate::{ + use diskann_providers::{ model::graph::provider::{ async_::{ TableDeleteProviderAsync, common::{FullPrecision, Hybrid, NoDeletes, Quantized, TableBasedDeletes}, - inmem::{self, DefaultQuant, SetStartPoints}, }, layers::BetaFilter, }, test_utils::{ assert_range_results_exactly_match, assert_top_k_exactly_match, groundtruth, is_match, }, - utils::{self, VectorDataIterator, create_rnd_from_seed_in_tests, file_util}, + utils::{self as providers_utils, VectorDataIterator, create_rnd_from_seed_in_tests, file_util}, }; + use crate::{self as inmem, DefaultQuant, SetStartPoints}; // Callbacks for use with `simplified_builder`. fn no_modify(_: &mut diskann::graph::config::Builder) {} @@ -2416,7 +2411,7 @@ pub(crate) mod tests { let index = new_quant_index::( config, params, - inmem::WithBits::<$nbits>::new(quantizer), + crate::WithBits::<$nbits>::new(quantizer), NoDeletes, ) .unwrap(); @@ -2520,7 +2515,7 @@ pub(crate) mod tests { new_quant_only_index( config, params, - inmem::WithBits::<$nbits>::new(quantizer), + crate::WithBits::<$nbits>::new(quantizer), NoDeletes, ) .unwrap(), @@ -2636,7 +2631,7 @@ pub(crate) mod tests { let build_fn = async |index: Arc>, data: Arc>| { let ctx = &DefaultContext; - let strategy = inmem::spherical::Quantized::build(); + let strategy = crate::spherical::Quantized::build(); for (i, vector) in data.row_iter().enumerate() { index .insert(strategy, ctx, &(i as u32), vector) @@ -2690,7 +2685,7 @@ pub(crate) mod tests { // Quantized Search let mut output = search_output_buffer::IdDistance::new(&mut ids, &mut distances); - let strategy = inmem::spherical::Quantized::search( + let strategy = crate::spherical::Quantized::search( diskann_quantization::spherical::iface::QueryLayout::FourBitTransposed, ); @@ -2708,7 +2703,7 @@ pub(crate) mod tests { } // Ensure that the query computer used for insertion uses the `SameAsData` layout. - let strategy = inmem::spherical::Quantized::build(); + let strategy = crate::spherical::Quantized::build(); let accessor = strategy.search_accessor(index.provider(), ctx).unwrap(); let computer = accessor.build_query_computer(data.row(0)).unwrap(); assert_eq!( @@ -2759,7 +2754,7 @@ pub(crate) mod tests { let build_fn = async |index: Arc>, data: Arc>| { let ctx = &DefaultContext; - let strategy = inmem::spherical::Quantized::build(); + let strategy = crate::spherical::Quantized::build(); for (i, vector) in data.row_iter().enumerate() { index .insert(strategy, ctx, &(i as u32), vector) @@ -2798,7 +2793,7 @@ pub(crate) mod tests { for (q, query) in data.row_iter().enumerate() { // Quantized Search let mut output = search_output_buffer::IdDistance::new(&mut ids, &mut distances); - let strategy = inmem::spherical::Quantized::search( + let strategy = crate::spherical::Quantized::search( diskann_quantization::spherical::iface::QueryLayout::FourBitTransposed, ); @@ -2818,9 +2813,9 @@ pub(crate) mod tests { } // Ensure that the query computer used for insertion uses the `SameAsData` layout. - let strategy = inmem::spherical::Quantized::build(); - let accessor = , + let strategy = crate::spherical::Quantized::build(); + let accessor = , [f32], _, >>::search_accessor(&strategy, index.provider(), ctx) @@ -2995,7 +2990,7 @@ pub(crate) mod tests { .to_path_buf(); let storage = VirtualStorageProvider::new_overlay(workspace_root); - let mut iter = VectorDataIterator::<_, crate::model::graph::traits::AdHoc>::new( + let mut iter = VectorDataIterator::<_, diskann_providers::model::graph::traits::AdHoc>::new( file, None, &storage, ) .unwrap(); @@ -3589,7 +3584,7 @@ pub(crate) mod tests { // This is the two level index. let (index, data) = init_from_file( - inmem::test::Flaky::new(9), + crate::test::Flaky::new(9), parameters, SIFTSMALL, 8, @@ -3712,7 +3707,7 @@ pub(crate) mod tests { populate_data(&index.data_provider, ctx, &vectors).await; let r = index - .consolidate_vector(&inmem::test::SuperFlaky, ctx, 0) + .consolidate_vector(&crate::test::SuperFlaky, ctx, 0) .await .unwrap(); assert_eq!(r, ConsolidateKind::FailedVectorRetrieval); @@ -3758,13 +3753,13 @@ pub(crate) mod tests { let index_sat = create_retry_saturated_index(NonZeroU32::new(1).unwrap(), true) .await .unwrap(); - let mut accessor_sat = inmem::FullAccessor::new(index_sat.provider()); + let mut accessor_sat = crate::FullAccessor::new(index_sat.provider()); let res_sat = index_sat.get_degree_stats(&mut accessor_sat).await.unwrap(); let index_unsat = create_retry_saturated_index(NonZeroU32::new(1).unwrap(), false) .await .unwrap(); - let mut accessor_unsat = inmem::FullAccessor::new(index_unsat.provider()); + let mut accessor_unsat = crate::FullAccessor::new(index_unsat.provider()); let res_unsat = index_sat .get_degree_stats(&mut accessor_unsat) .await @@ -3780,13 +3775,13 @@ pub(crate) mod tests { let index_sat = create_retry_saturated_index(NonZeroU32::new(3).unwrap(), false) .await .unwrap(); - let mut accessor_sat = inmem::FullAccessor::new(index_sat.provider()); + let mut accessor_sat = crate::FullAccessor::new(index_sat.provider()); let res_sat = index_sat.get_degree_stats(&mut accessor_sat).await.unwrap(); let index_unsat = create_retry_saturated_index(NonZeroU32::new(1).unwrap(), false) .await .unwrap(); - let mut accessor_unsat = inmem::FullAccessor::new(index_unsat.provider()); + let mut accessor_unsat = crate::FullAccessor::new(index_unsat.provider()); let res_unsat = index_sat .get_degree_stats(&mut accessor_unsat) .await diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs b/diskann-inmem/src/full_precision.rs similarity index 99% rename from diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs rename to diskann-inmem/src/full_precision.rs index 4c4f393a8..8925c9da3 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs +++ b/diskann-inmem/src/full_precision.rs @@ -24,19 +24,20 @@ use diskann::{ use diskann_utils::future::AsyncFriendly; use diskann_vector::{DistanceFunction, distance::Metric}; -use crate::model::graph::{ +use diskann_providers::model::graph::{ provider::async_::{ FastMemoryVectorProviderAsync, SimpleNeighborProviderAsync, common::{ CreateVectorStore, FullPrecision, Internal, NoDeletes, NoStore, Panics, PrefetchCacheLineLevel, SetElementHelper, }, - inmem::DefaultProvider, postprocess::{AsDeletionCheck, DeletionCheck, RemoveDeletedIdsAndCopy}, }, traits::AdHoc, }; +use crate::DefaultProvider; + /// A type alias for the DefaultProvider with full-precision as the primary vector store. pub type FullPrecisionProvider = DefaultProvider, Q, D, Ctx>; diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs b/diskann-inmem/src/lib.rs similarity index 68% rename from diskann-providers/src/model/graph/provider/async_/inmem/mod.rs rename to diskann-inmem/src/lib.rs index 4f45b684c..76119318d 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/mod.rs +++ b/diskann-inmem/src/lib.rs @@ -11,7 +11,8 @@ pub use provider::{DefaultProvider, DefaultProviderParameters, SetStartPoints}; // Extensions mod scalar; -pub use scalar::{SQError, SQStore, WithBits}; +pub use scalar::{SQStore, WithBits}; +pub use diskann_providers::storage::SQError; #[cfg(not(test))] mod product; @@ -23,9 +24,15 @@ mod full_precision; pub use full_precision::{ CreateFullPrecision, FullAccessor, FullPrecisionProvider, FullPrecisionStore, }; -pub(super) use full_precision::{GetFullPrecision, Rerank}; +pub use full_precision::{GetFullPrecision, Rerank}; #[cfg(test)] pub mod product; #[cfg(test)] pub(crate) mod test; + +// Helper functions for creating inmem indexes +pub mod diskann_async; + +// Storage implementations for inmem providers +pub mod storage; diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/product.rs b/diskann-inmem/src/product.rs similarity index 99% rename from diskann-providers/src/model/graph/provider/async_/inmem/product.rs rename to diskann-inmem/src/product.rs index 2886d1d20..236364fe3 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/product.rs +++ b/diskann-inmem/src/product.rs @@ -20,7 +20,7 @@ use diskann::{ use diskann_utils::future::AsyncFriendly; use diskann_vector::distance::Metric; -use crate::model::{ +use diskann_providers::model::{ graph::{ provider::async_::{ FastMemoryQuantVectorProviderAsync, FastMemoryVectorProviderAsync, @@ -30,10 +30,6 @@ use crate::model::{ VectorStore, }, distances, - inmem::{ - DefaultProvider, FullPrecisionProvider, FullPrecisionStore, GetFullPrecision, - Rerank, - }, postprocess::{AsDeletionCheck, DeletionCheck, RemoveDeletedIdsAndCopy}, }, traits::AdHoc, @@ -41,6 +37,11 @@ use crate::model::{ pq::{self, FixedChunkPQTable}, }; +use crate::{ + DefaultProvider, FullPrecisionProvider, FullPrecisionStore, GetFullPrecision, + Rerank, +}; + /// The default quant provider. pub type DefaultQuant = FastMemoryQuantVectorProviderAsync; diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs b/diskann-inmem/src/provider.rs similarity index 99% rename from diskann-providers/src/model/graph/provider/async_/inmem/provider.rs rename to diskann-inmem/src/provider.rs index b4245355c..e3b492c82 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/provider.rs +++ b/diskann-inmem/src/provider.rs @@ -5,7 +5,7 @@ use std::{fmt::Debug, future::Future, num::NonZeroUsize}; -use crate::storage::{StorageReadProvider, StorageWriteProvider}; +use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; #[cfg(test)] use diskann::neighbor::Neighbor; use diskann::{ @@ -20,7 +20,7 @@ use diskann::{ use diskann_utils::future::AsyncFriendly; use diskann_vector::distance::Metric; -use crate::{ +use diskann_providers::{ model::graph::provider::async_::{ SimpleNeighborProviderAsync, StartPoints, TableDeleteProviderAsync, common::{ diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs b/diskann-inmem/src/scalar.rs similarity index 96% rename from diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs rename to diskann-inmem/src/scalar.rs index 149b54019..7e379cd2c 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs +++ b/diskann-inmem/src/scalar.rs @@ -5,7 +5,7 @@ use std::{future::Future, sync::Mutex}; -use crate::storage::{StorageReadProvider, StorageWriteProvider}; +use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider, SQError}; use diskann::{ ANNError, ANNResult, graph::glue::{ @@ -33,7 +33,7 @@ use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::M use thiserror::Error; use super::{DefaultProvider, GetFullPrecision, Rerank}; -use crate::{ +use diskann_providers::{ common::IgnoreLockPoison, model::graph::{ provider::async_::{ @@ -42,7 +42,6 @@ use crate::{ AlignedMemoryVectorStore, CreateVectorStore, NoStore, Quantized, SetElementHelper, TestCallCount, VectorStore, }, - inmem::{FullPrecisionProvider, FullPrecisionStore}, postprocess::{AsDeletionCheck, DeletionCheck, RemoveDeletedIdsAndCopy}, }, traits::AdHoc, @@ -50,6 +49,8 @@ use crate::{ storage::{self, AsyncIndexMetadata, AsyncQuantLoadContext, LoadWith, SaveWith}, }; +use crate::{FullPrecisionProvider, FullPrecisionStore}; + type CVRef<'a, const NBITS: usize> = CompensatedVectorRef<'a, NBITS>; /// A thin wrapper around [`ScalarQuantizer`] that encodes the number of bits desired for @@ -794,36 +795,6 @@ impl storage::bin::GetData for SQStore { } } -#[derive(Debug, Error)] -pub enum SQError { - #[error("Issue with canonical layout of data: {0:?}")] - CanonicalLayoutError(#[from] NotCanonical), - - #[error("Input contains NaN values.")] - InputContainsNaN(#[from] InputContainsNaN), - - #[error("Input full-precision conversion error : {0}")] - FullPrecisionConversionErr(String), - - #[error("Mean Norm is missing in the quantizer.")] - MeanNormMissing(#[from] MeanNormMissing), - - #[error("Unsupported distance metric: {0:?}")] - UnsupportedDistanceMetric(Metric), - - #[error("Error while loading quantizer proto struct from file: {0:?}")] - ProtoStorageError(#[from] crate::storage::protos::ProtoStorageError), - - #[error("Error while converting proto struct to Scalar Qunatizer: {0:?}")] - QuantizerDecodeError(#[from] crate::storage::protos::ProtoConversionError), -} - -impl From for ANNError { - #[cold] - fn from(err: SQError) -> Self { - ANNError::log_sq_error(err) - } -} #[cfg(test)] mod tests { diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs b/diskann-inmem/src/spherical.rs similarity index 99% rename from diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs rename to diskann-inmem/src/spherical.rs index b705b43d3..f5c28a3fd 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs +++ b/diskann-inmem/src/spherical.rs @@ -30,7 +30,7 @@ use diskann_vector::distance::Metric; use thiserror::Error; use super::{GetFullPrecision, Rerank}; -use crate::{ +use diskann_providers::{ common::IgnoreLockPoison, model::graph::{ provider::async_::{ @@ -40,7 +40,6 @@ use crate::{ TestCallCount, VectorStore, }, distances::UnwrapErr, - inmem::{DefaultProvider, FullPrecisionProvider, FullPrecisionStore}, postprocess::{AsDeletionCheck, DeletionCheck, RemoveDeletedIdsAndCopy}, }, traits::AdHoc, @@ -48,6 +47,8 @@ use crate::{ utils::{Bridge, BridgeErr}, }; +use crate::{DefaultProvider, FullPrecisionProvider, FullPrecisionStore}; + ///////////////////// // Error Promotion // ///////////////////// diff --git a/diskann-inmem/src/storage.rs b/diskann-inmem/src/storage.rs new file mode 100644 index 000000000..a2610a072 --- /dev/null +++ b/diskann-inmem/src/storage.rs @@ -0,0 +1,208 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Storage implementations for inmem providers. + +use std::num::NonZeroUsize; + +use diskann::{ + ANNError, ANNResult, graph::DiskANNIndex, provider::DataProvider, utils::VectorRepr, +}; +use diskann_utils::future::AsyncFriendly; +use diskann_providers::storage::{ + StorageReadProvider, StorageWriteProvider, AsyncIndexMetadata, AsyncQuantLoadContext, + DiskGraphOnly, LoadWith, SaveWith, +}; +use diskann_providers::model::{ + configuration::IndexConfiguration, + graph::provider::async_::{ + FastMemoryQuantVectorProviderAsync, TableDeleteProviderAsync, common, + }, +}; + +use crate::{DefaultProvider, FullPrecisionStore}; + +impl SaveWith for DiskANNIndex> +where + U: AsyncFriendly, + V: AsyncFriendly, + D: AsyncFriendly, + DefaultProvider: SaveWith<(u32, AsyncIndexMetadata), Error = ANNError>, +{ + type Ok = (); + type Error = ANNError; + + async fn save_with

(&self, provider: &P, ctx_prefix: &AsyncIndexMetadata) -> ANNResult<()> + where + P: StorageWriteProvider, + { + let start_id = get_and_validate_single_starting_point(&self.data_provider)?; + + self.data_provider + .save_with(provider, &(start_id, ctx_prefix.clone())) + .await?; + + Ok(()) + } +} + +// This implementation saves only graph and not the vector/quant data. +impl SaveWith<(u32, DiskGraphOnly)> for DiskANNIndex> +where + U: AsyncFriendly, + V: AsyncFriendly, + D: AsyncFriendly, + DefaultProvider: SaveWith<(u32, u32, DiskGraphOnly), Error = ANNError>, +{ + type Ok = (); + type Error = ANNError; + + async fn save_with

(&self, provider: &P, ctx_prefix: &(u32, DiskGraphOnly)) -> ANNResult<()> + where + P: StorageWriteProvider, + { + let start_id = get_and_validate_single_starting_point(&self.data_provider)?; + + self.data_provider + .save_with(provider, &(start_id, ctx_prefix.0, ctx_prefix.1.clone())) + .await?; + Ok(()) + } +} + +/// Creates a `AsyncQuantLoadContext` from an `IndexConfiguration` with the given path and disk index flag. +pub fn create_load_context( + path: &str, + index_config: &IndexConfiguration, + is_disk_index: bool, +) -> ANNResult { + Ok(AsyncQuantLoadContext { + metadata: AsyncIndexMetadata::new(path), + num_frozen_points: index_config.num_frozen_pts, + metric: index_config.dist_metric, + prefetch_lookahead: index_config.prefetch_lookahead.map(|x| x.get()), + is_disk_index, + prefetch_cache_line_level: index_config.prefetch_cache_line_level, + }) +} + +impl<'a, DP> LoadWith<(&'a str, IndexConfiguration)> for DiskANNIndex +where + DP: DataProvider + LoadWith, +{ + type Error = ANNError; + async fn load_with

( + provider: &P, + (path, index_config): &(&'a str, IndexConfiguration), + ) -> ANNResult + where + P: StorageReadProvider, + { + let pq_context = create_load_context(path, index_config, false)?; + + let data_provider = DP::load_with(provider, &pq_context).await?; + let num_threads = index_config.num_threads; + Ok(Self::new( + index_config.config.clone(), + data_provider, + NonZeroUsize::new(num_threads), + )) + } +} + +pub async fn load_pq_index( + provider: &P, + path: &str, + config: IndexConfiguration, +) -> ANNResult>> +where + P: StorageReadProvider, + T: VectorRepr, +{ + DiskANNIndex::load_with(provider, &(path, config)).await +} + +pub async fn load_pq_index_with_deletes( + provider: &P, + path: &str, + config: IndexConfiguration, +) -> ANNResult< + DiskANNIndex< + crate::DefaultProvider< + FullPrecisionStore, + FastMemoryQuantVectorProviderAsync, + TableDeleteProviderAsync, + >, + >, +> +where + P: StorageReadProvider, + T: VectorRepr, +{ + DiskANNIndex::load_with(provider, &(path, config)).await +} + +pub async fn load_fp_index( + provider: &P, + path: &str, + config: IndexConfiguration, +) -> ANNResult>> +where + P: StorageReadProvider, + T: VectorRepr, + Q: AsyncFriendly, + crate::FullPrecisionProvider: LoadWith, +{ + DiskANNIndex::load_with(provider, &(path, config)).await +} + +pub async fn load_index( + provider: &P, + path: &str, + config: IndexConfiguration, +) -> ANNResult>> +where + P: StorageReadProvider, + U: AsyncFriendly, + V: AsyncFriendly, + crate::DefaultProvider: LoadWith, +{ + DiskANNIndex::load_with(provider, &(path, config)).await +} + +pub async fn load_index_with_deletes( + provider: &P, + path: &str, + config: IndexConfiguration, +) -> ANNResult< + DiskANNIndex>, +> +where + P: StorageReadProvider, + T: VectorRepr, +{ + DiskANNIndex::load_with(provider, &(path, config)).await +} + +/// Extracts and validates single start point from DefaultProvider. +/// +/// # Errors +/// - Returns an error if the number of start points is not exactly 1 +fn get_and_validate_single_starting_point( + data_provider: &DefaultProvider, +) -> ANNResult { + use diskann::provider::DataProvider; + use diskann::ANNErrorKind; + + let start_points: Vec = data_provider.start_points(); + if start_points.len() != 1 { + return Err(ANNError::log_error(ANNErrorKind::InvalidStartPoint, format!( + "Index must have exactly 1 start point for saving. Found {} start points: {:?}", + start_points.len(), + start_points + ))); + } + Ok(start_points[0]) +} diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/test.rs b/diskann-inmem/src/test.rs similarity index 98% rename from diskann-providers/src/model/graph/provider/async_/inmem/test.rs rename to diskann-inmem/src/test.rs index 9c6f908b1..e1a457e91 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/test.rs +++ b/diskann-inmem/src/test.rs @@ -21,9 +21,10 @@ use diskann::{ }; use super::{DefaultProvider, DefaultQuant}; -use crate::model::graph::provider::async_::{ - SimpleNeighborProviderAsync, TableDeleteProviderAsync, inmem::FullPrecisionStore, +use diskann_providers::model::graph::provider::async_::{ + SimpleNeighborProviderAsync, TableDeleteProviderAsync, }; +use crate::FullPrecisionStore; /// A full-precision accessor that spuriously fails for non-start points with a controllable /// frequency. diff --git a/diskann-providers/Cargo.toml b/diskann-providers/Cargo.toml index c0deaeca6..65441596c 100644 --- a/diskann-providers/Cargo.toml +++ b/diskann-providers/Cargo.toml @@ -55,6 +55,7 @@ rstest.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["full"] } proptest.workspace = true +diskann-inmem = { workspace = true } [[bench]] name = "bench_main" diff --git a/diskann-providers/src/index/mod.rs b/diskann-providers/src/index/mod.rs index 5c46ec76d..34139cc54 100644 --- a/diskann-providers/src/index/mod.rs +++ b/diskann-providers/src/index/mod.rs @@ -3,6 +3,5 @@ * Licensed under the MIT license. */ -pub mod diskann_async; - +// diskann_async module has been moved to diskann-inmem crate pub mod wrapped_async; diff --git a/diskann-providers/src/model/graph/provider/async_/mod.rs b/diskann-providers/src/model/graph/provider/async_/mod.rs index 3d89359e2..182f8062d 100644 --- a/diskann-providers/src/model/graph/provider/async_/mod.rs +++ b/diskann-providers/src/model/graph/provider/async_/mod.rs @@ -29,8 +29,8 @@ pub use fast_memory_vector_provider::FastMemoryVectorProviderAsync; pub mod fast_memory_quant_vector_provider; pub use fast_memory_quant_vector_provider::FastMemoryQuantVectorProviderAsync; -// The default `inmem` data provider for the async index. -pub mod inmem; +// The default `inmem` data provider has been moved to the diskann-inmem crate. +// Import from diskann_inmem crate instead of this module path. // Bf-tree based data provider for the async index #[cfg(feature = "bf_tree")] diff --git a/diskann-providers/src/storage/index_storage.rs b/diskann-providers/src/storage/index_storage.rs index 674726e5f..81dd68e4a 100644 --- a/diskann-providers/src/storage/index_storage.rs +++ b/diskann-providers/src/storage/index_storage.rs @@ -3,511 +3,6 @@ * Licensed under the MIT license. */ -use std::num::NonZeroUsize; - -use super::{StorageReadProvider, StorageWriteProvider}; -use diskann::{ - ANNError, ANNResult, graph::DiskANNIndex, provider::DataProvider, utils::VectorRepr, -}; -use diskann_utils::future::AsyncFriendly; - -use super::{AsyncIndexMetadata, AsyncQuantLoadContext, DiskGraphOnly, LoadWith, SaveWith}; -use crate::model::{ - configuration::IndexConfiguration, - graph::provider::async_::{ - FastMemoryQuantVectorProviderAsync, TableDeleteProviderAsync, common, - inmem::{self, DefaultProvider, FullPrecisionStore}, - }, -}; - -impl SaveWith for DiskANNIndex> -where - U: AsyncFriendly, - V: AsyncFriendly, - D: AsyncFriendly, - DefaultProvider: SaveWith<(u32, AsyncIndexMetadata), Error = ANNError>, -{ - type Ok = (); - type Error = ANNError; - - async fn save_with

(&self, provider: &P, ctx_prefix: &AsyncIndexMetadata) -> ANNResult<()> - where - P: StorageWriteProvider, - { - let start_id = get_and_validate_single_starting_point(&self.data_provider)?; - - self.data_provider - .save_with(provider, &(start_id, ctx_prefix.clone())) - .await?; - - Ok(()) - } -} - -// This implementation saves only graph and not the vector/quant data. -impl SaveWith<(u32, DiskGraphOnly)> for DiskANNIndex> -where - U: AsyncFriendly, - V: AsyncFriendly, - D: AsyncFriendly, - DefaultProvider: SaveWith<(u32, u32, DiskGraphOnly), Error = ANNError>, -{ - type Ok = (); - type Error = ANNError; - - async fn save_with

(&self, provider: &P, ctx_prefix: &(u32, DiskGraphOnly)) -> ANNResult<()> - where - P: StorageWriteProvider, - { - let start_id = get_and_validate_single_starting_point(&self.data_provider)?; - - self.data_provider - .save_with(provider, &(start_id, ctx_prefix.0, ctx_prefix.1.clone())) - .await?; - Ok(()) - } -} - -/// Creates a `AsyncQuantLoadContext` from an `IndexConfiguration` with the given path and disk index flag. -pub fn create_load_context( - path: &str, - index_config: &IndexConfiguration, - is_disk_index: bool, -) -> ANNResult { - Ok(AsyncQuantLoadContext { - metadata: AsyncIndexMetadata::new(path), - num_frozen_points: index_config.num_frozen_pts, - metric: index_config.dist_metric, - prefetch_lookahead: index_config.prefetch_lookahead.map(|x| x.get()), - is_disk_index, - prefetch_cache_line_level: index_config.prefetch_cache_line_level, - }) -} - -impl<'a, DP> LoadWith<(&'a str, IndexConfiguration)> for DiskANNIndex -where - DP: DataProvider + LoadWith, -{ - type Error = ANNError; - async fn load_with

( - provider: &P, - (path, index_config): &(&'a str, IndexConfiguration), - ) -> ANNResult - where - P: StorageReadProvider, - { - let pq_context = create_load_context(path, index_config, false)?; - - let data_provider = DP::load_with(provider, &pq_context).await?; - let num_threads = index_config.num_threads; - Ok(Self::new( - index_config.config.clone(), - data_provider, - NonZeroUsize::new(num_threads), - )) - } -} - -pub async fn load_pq_index( - provider: &P, - path: &str, - config: IndexConfiguration, -) -> ANNResult>> -where - P: StorageReadProvider, - T: VectorRepr, -{ - DiskANNIndex::load_with(provider, &(path, config)).await -} - -pub async fn load_pq_index_with_deletes( - provider: &P, - path: &str, - config: IndexConfiguration, -) -> ANNResult< - DiskANNIndex< - inmem::DefaultProvider< - FullPrecisionStore, - FastMemoryQuantVectorProviderAsync, - TableDeleteProviderAsync, - >, - >, -> -where - P: StorageReadProvider, - T: VectorRepr, -{ - DiskANNIndex::load_with(provider, &(path, config)).await -} - -pub async fn load_fp_index( - provider: &P, - path: &str, - config: IndexConfiguration, -) -> ANNResult>> -where - P: StorageReadProvider, - T: VectorRepr, - Q: AsyncFriendly, - inmem::FullPrecisionProvider: LoadWith, -{ - DiskANNIndex::load_with(provider, &(path, config)).await -} - -pub async fn load_index( - provider: &P, - path: &str, - config: IndexConfiguration, -) -> ANNResult>> -where - P: StorageReadProvider, - U: AsyncFriendly, - V: AsyncFriendly, - inmem::DefaultProvider: LoadWith, -{ - DiskANNIndex::load_with(provider, &(path, config)).await -} - -pub async fn load_index_with_deletes( - provider: &P, - path: &str, - config: IndexConfiguration, -) -> ANNResult< - DiskANNIndex>, -> -where - P: StorageReadProvider, - T: VectorRepr, -{ - DiskANNIndex::load_with(provider, &(path, config)).await -} - -/// Retrieves starting points and enforces that there is exactly one starting point. -/// -/// This helper function: -/// 1. Retrieves the starting points from the data provider -/// 2. Validates there is exactly one starting point -/// 3. Returns the single start point if valid -/// -/// Returns an error if there are multiple starting points or no starting points. -fn get_and_validate_single_starting_point( - data_provider: &DefaultProvider, -) -> ANNResult { - let start_ids = data_provider.starting_points()?; - - let num_starting_points = start_ids.len(); - if num_starting_points > 1 { - return Err(ANNError::log_index_error(format_args!( - "ERROR: Save index does not support multiple starting points. Found {} starting points.", - num_starting_points - ))); - } - - start_ids - .first() - .cloned() - .ok_or_else(|| ANNError::log_index_error("ERROR: No starting points found")) -} -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use std::{num::NonZeroUsize, sync::Arc}; - - use crate::storage::VirtualStorageProvider; - use diskann::{ - graph::{AdjacencyList, config, glue::InsertStrategy}, - provider::{Accessor, SetElement}, - utils::{IntoUsize, ONE}, - }; - use diskann_utils::{ - Reborrow, - views::{Matrix, MatrixView}, - }; - use diskann_vector::distance::Metric; - use vfs::MemoryFS; - - use super::*; - use crate::{ - index::diskann_async::{self, MemoryIndex}, - model::graph::provider::async_::{ - SimpleNeighborProviderAsync, - common::{FullPrecision, NoDeletes, NoStore, TableBasedDeletes}, - inmem::{self}, - }, - utils::{create_rnd_from_seed_in_tests, file_util}, - }; - - async fn build_index( - index: &Arc>, - strategy: S, - data: MatrixView<'_, f32>, - ) where - DP: DataProvider + SetElement<[f32]>, - DP::Context: Default, - S: InsertStrategy + Clone, - { - let ctx = &DP::Context::default(); - for (i, v) in data.row_iter().enumerate() { - index - .insert(strategy.clone(), ctx, &(i as u32), v) - .await - .unwrap(); - } - } - - // Our test strategy here is to basically build one main index using quantization - // and to save that. - // - // We will the try reloading with the following flavors: - // 1. Without quant, with delete set. - // 2. Without quant, without delete set. - // 3. With quant, with delete set. - // 4. With quant, without delete set. - #[tokio::test] - async fn test_save_and_load() { - let save_path = "/index"; - let file_path = "/test_data/sift/siftsmall_learn_256pts.fbin"; - let train_data = { - let workspace_root = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .parent() - .unwrap() - .to_path_buf(); - let storage = VirtualStorageProvider::new_overlay(workspace_root); - let (train_data, npoints, dim) = file_util::load_bin(&storage, file_path, 0).unwrap(); - Matrix::::try_from(train_data.into(), npoints, dim).unwrap() - }; - - let pq_bytes = 8; - let pq_table = diskann_async::train_pq( - train_data.as_view(), - pq_bytes, - &mut create_rnd_from_seed_in_tests(0xe3c52ef001bc7ade), - 2, - ) - .unwrap(); - - let (config, parameters) = diskann_async::simplified_builder( - 20, - 32, - Metric::L2, - train_data.ncols(), - train_data.nrows(), - |_| {}, - ) - .unwrap(); - - let index = diskann_async::new_quant_index::( - config, - parameters, - pq_table, - TableBasedDeletes, - ) - .unwrap(); - - build_index(&index, FullPrecision, train_data.as_view()).await; - - // Check that all nodes are reachable. - { - let count = index - .count_reachable_nodes( - &index.provider().starting_points().unwrap(), - &mut index.provider().neighbors(), - ) - .await - .unwrap(); - assert_eq!(count, train_data.nrows() + 1); - } - - // Save the resulting index. - let provider = VirtualStorageProvider::new(MemoryFS::new()); - index - .save_with(&provider, &AsyncIndexMetadata::new(save_path.to_string())) - .await - .unwrap(); - - // Convert into the full index configuration. - let config = IndexConfiguration::new( - Metric::L2, - train_data.ncols(), - train_data.nrows(), - ONE, - 1, - config::Builder::new( - 30, - config::MaxDegree::default_slack(), - 20, - Metric::L2.into(), - ) - .build() - .unwrap(), - ); - - let id_iter = index.data_provider.iter(); - - // Without Quant, With Delete Set. - { - let reloaded = load_index_with_deletes::(&provider, save_path, config.clone()) - .await - .unwrap(); - - assert_eq!(id_iter, reloaded.data_provider.iter()); - check_accessor_equal( - inmem::FullAccessor::new(index.provider()), - inmem::FullAccessor::new(reloaded.provider()), - id_iter.clone(), - ) - .await; - - check_graphs_equal( - &index.provider().neighbor_provider, - &reloaded.provider().neighbor_provider, - id_iter.clone(), - ) - } - - // Without Quant, Without Delete Set. - { - let reloaded = load_fp_index::(&provider, save_path, config.clone()) - .await - .unwrap(); - - assert_eq!(id_iter, reloaded.data_provider.iter()); - check_accessor_equal( - inmem::FullAccessor::new(index.provider()), - inmem::FullAccessor::new(reloaded.provider()), - id_iter.clone(), - ) - .await; - - check_graphs_equal( - &index.provider().neighbor_provider, - &reloaded.provider().neighbor_provider, - id_iter.clone(), - ) - } - - // With Quant, With Delete Set. - { - let reloaded = - load_pq_index_with_deletes::(&provider, save_path, config.clone()) - .await - .unwrap(); - - assert_eq!(id_iter, reloaded.data_provider.iter()); - check_accessor_equal( - inmem::FullAccessor::new(index.provider()), - inmem::FullAccessor::new(reloaded.provider()), - index.data_provider.iter(), - ) - .await; - - check_accessor_equal( - inmem::product::QuantAccessor::new(index.provider()), - inmem::product::QuantAccessor::new(reloaded.provider()), - index.data_provider.iter(), - ) - .await; - - check_graphs_equal( - &index.provider().neighbor_provider, - &reloaded.provider().neighbor_provider, - id_iter.clone(), - ) - } - - // With Quant, Without Delete Set. - { - let reloaded = load_pq_index::(&provider, save_path, config.clone()) - .await - .unwrap(); - - assert_eq!(id_iter, reloaded.data_provider.iter()); - check_accessor_equal( - inmem::FullAccessor::new(index.provider()), - inmem::FullAccessor::new(reloaded.provider()), - index.data_provider.iter(), - ) - .await; - - check_accessor_equal( - inmem::product::QuantAccessor::new(index.provider()), - inmem::product::QuantAccessor::new(reloaded.provider()), - index.data_provider.iter(), - ) - .await; - - check_graphs_equal( - &index.provider().neighbor_provider, - &reloaded.provider().neighbor_provider, - id_iter.clone(), - ) - } - } - - async fn check_accessor_equal(mut left: A, mut right: B, itr: Itr) - where - A: for<'a> Accessor = &'a T>, - B: for<'a> Accessor = &'a T>, - T: PartialEq + std::fmt::Debug + ?Sized, - Itr: Iterator, - { - for i in itr { - assert_eq!( - left.get_element(i).await.unwrap().reborrow(), - right.get_element(i).await.unwrap().reborrow(), - "failed for index {}", - i - ); - } - } - - fn check_graphs_equal( - left: &SimpleNeighborProviderAsync, - right: &SimpleNeighborProviderAsync, - itr: Itr, - ) where - Itr: Iterator, - { - let mut lv = AdjacencyList::new(); - let mut rv = AdjacencyList::new(); - for i in itr { - left.get_neighbors_sync(i.into_usize(), &mut lv).unwrap(); - right.get_neighbors_sync(i.into_usize(), &mut rv).unwrap(); - assert_eq!(lv, rv, "failed for index {}", i); - } - } - - fn create_test_index(num_start_points: usize) -> MemoryIndex { - let (config, mut parameters) = - diskann_async::simplified_builder(20, 32, Metric::L2, 3, 5, |_| {}).unwrap(); - - parameters.frozen_points = NonZeroUsize::new(num_start_points).unwrap(); - diskann_async::new_index::(config, parameters, NoDeletes).unwrap() - } - - #[tokio::test] - async fn test_validate_single_starting_point() { - // Test case 1: Single start point should succeed - { - let index = create_test_index(1); - let result = get_and_validate_single_starting_point(&index.data_provider); - assert!(result.is_ok(), "Failed to validate single start point"); - } - - // Test case 2: Multiple start points should fail - { - let index = create_test_index(2); - let result = get_and_validate_single_starting_point(&index.data_provider); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("not support multiple starting points") - ); - } - } -} +// Storage implementations for DiskANN indexes. +// The inmem-specific implementations have been moved to diskann-inmem crate. +// Tests that use inmem types should also be moved to diskann-inmem. diff --git a/diskann-providers/src/storage/mod.rs b/diskann-providers/src/storage/mod.rs index 1233b11f6..38f1fd7cd 100644 --- a/diskann-providers/src/storage/mod.rs +++ b/diskann-providers/src/storage/mod.rs @@ -39,7 +39,10 @@ pub use path_utility::{ }; pub mod index_storage; -pub use index_storage::{ +// Storage implementations that were specific to inmem have been moved to diskann-inmem crate. +// Re-export these functions from diskann-inmem in dev dependencies for tests. +#[cfg(test)] +pub use diskann_inmem::storage::{ create_load_context, load_fp_index, load_index_with_deletes, load_pq_index, load_pq_index_with_deletes, }; diff --git a/diskann-providers/src/storage/sq_storage.rs b/diskann-providers/src/storage/sq_storage.rs index 1b7be50e9..14f591dd1 100644 --- a/diskann-providers/src/storage/sq_storage.rs +++ b/diskann-providers/src/storage/sq_storage.rs @@ -9,7 +9,6 @@ use super::{StorageReadProvider, StorageWriteProvider}; use diskann_quantization::scalar::ScalarQuantizer; use super::protos; -use crate::model::graph::provider::async_::inmem::SQError; /// The suffix for the compressed SQ vectors file. const COMPRESSED_DATA_FILE_NAME_SUFFIX: &str = "sq_compressed.bin"; @@ -140,3 +139,34 @@ mod tests { assert_eq!(quantizer.compare(&loaded_quantizer), Ok(())); } } + +#[derive(Debug, thiserror::Error)] +pub enum SQError { + #[error("Issue with canonical layout of data: {0:?}")] + CanonicalLayoutError(#[from] diskann_quantization::meta::NotCanonical), + + #[error("Input contains NaN values.")] + InputContainsNaN(#[from] diskann_quantization::scalar::InputContainsNaN), + + #[error("Input full-precision conversion error : {0}")] + FullPrecisionConversionErr(String), + + #[error("Mean Norm is missing in the quantizer.")] + MeanNormMissing(#[from] diskann_quantization::scalar::MeanNormMissing), + + #[error("Unsupported distance metric: {0:?}")] + UnsupportedDistanceMetric(diskann_vector::distance::Metric), + + #[error("Error while loading quantizer proto struct from file: {0:?}")] + ProtoStorageError(#[from] crate::storage::protos::ProtoStorageError), + + #[error("Error while converting proto struct to Scalar Quantizer: {0:?}")] + QuantizerDecodeError(#[from] crate::storage::protos::ProtoConversionError), +} + +impl From for diskann::ANNError { + #[cold] + fn from(err: SQError) -> Self { + diskann::ANNError::log_sq_error(err) + } +}