From ba7a26348eac899449a0160fb904f2522606efa8 Mon Sep 17 00:00:00 2001 From: Sanhaoji2 Date: Tue, 10 Feb 2026 16:48:42 +0800 Subject: [PATCH 1/4] Add API to load pre-built index to wrapped_async --- Cargo.lock | 1 + diskann-providers/Cargo.toml | 1 + diskann-providers/src/index/wrapped_async.rs | 51 +++++++++++++++++++- 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 426fe8795..62d62c80a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -612,6 +612,7 @@ dependencies = [ "diskann-utils", "diskann-vector", "diskann-wide", + "futures-executor", "futures-util", "half", "hashbrown 0.16.1", diff --git a/diskann-providers/Cargo.toml b/diskann-providers/Cargo.toml index d23fcc038..14d2c0d70 100644 --- a/diskann-providers/Cargo.toml +++ b/diskann-providers/Cargo.toml @@ -40,6 +40,7 @@ bf-tree = { workspace = true, optional = true } anyhow.workspace = true prost = "0.14.1" futures-util.workspace = true +futures-executor = "0.3" polonius-the-crab = { version = "0.5.0", optional = true } serde_json = { workspace = true, optional = true } vfs = { workspace = true, optional = true } diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index d3d37416d..7f7ffb2db 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -4,9 +4,10 @@ */ use std::{num::NonZeroUsize, sync::Arc}; +use futures_executor::block_on; use diskann::{ - ANNResult, + ANNError, ANNResult, graph::{ self, ConsolidateKind, InplaceDeleteMethod, SearchParams, glue::{ @@ -20,6 +21,10 @@ use diskann::{ utils::{ONE, async_tools::VectorIdBoxSlice}, }; + +use crate::{model::IndexConfiguration, + storage::{StorageReadProvider, LoadWith, AsyncQuantLoadContext,AsyncIndexMetadata, file_storage_provider::FileStorageProvider}}; + pub struct DiskANNIndex { /// The underlying async DiskANNIndex. pub inner: Arc>, @@ -317,3 +322,47 @@ where self.handle.block_on(self.inner.get_degree_stats(accessor)) } } + +pub trait SyncLoadWith: Sized { + + fn load_with(path: &str, index_config: IndexConfiguration) -> Self; +} + +// Load static memory index from pre-built files synchronously +impl SyncLoadWith for DiskANNIndex +where + DP: DataProvider + LoadWith, +{ + fn load_with(path: &str, index_config: IndexConfiguration) -> Self { + let storage = FileStorageProvider; + let data_provider = create_data_provider(&storage, path, &index_config); + + match data_provider { + Ok(dp) => DiskANNIndex::::new_with_current_thread_runtime(index_config.config, dp), + Err(e) => panic!("Failed to create data provider: {}", e), + } + } +} + +pub fn create_data_provider<'a, P, DP>( + provider: &P, + path: &'a str, + index_config: &'a IndexConfiguration, +) -> ANNResult +where + P: StorageReadProvider, + DP: DataProvider + LoadWith, +{ + let pq_context = 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: false, // only support in-memory index loading here + prefetch_cache_line_level: index_config.prefetch_cache_line_level, + }; + + let data_provider = DP::load_with(provider, &pq_context); + + block_on(data_provider) +} From 498622a27537b2f093a8796184b7e13b3a8c1956 Mon Sep 17 00:00:00 2001 From: Sanhaoji2 Date: Tue, 10 Feb 2026 19:40:26 +0800 Subject: [PATCH 2/4] Fix compile issue --- diskann-providers/src/index/wrapped_async.rs | 31 ++++++++++++-------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index 7f7ffb2db..168ee0e31 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -3,8 +3,8 @@ * Licensed under the MIT license. */ -use std::{num::NonZeroUsize, sync::Arc}; use futures_executor::block_on; +use std::{num::NonZeroUsize, sync::Arc}; use diskann::{ ANNError, ANNResult, @@ -21,9 +21,13 @@ use diskann::{ utils::{ONE, async_tools::VectorIdBoxSlice}, }; - -use crate::{model::IndexConfiguration, - storage::{StorageReadProvider, LoadWith, AsyncQuantLoadContext,AsyncIndexMetadata, file_storage_provider::FileStorageProvider}}; +use crate::{ + model::IndexConfiguration, + storage::{ + AsyncIndexMetadata, AsyncQuantLoadContext, LoadWith, StorageReadProvider, + file_storage_provider::FileStorageProvider, + }, +}; pub struct DiskANNIndex { /// The underlying async DiskANNIndex. @@ -323,23 +327,26 @@ where } } -pub trait SyncLoadWith: Sized { - - fn load_with(path: &str, index_config: IndexConfiguration) -> Self; +pub trait SyncLoadWith: Sized { + fn load_with(path: &str, index_config: IndexConfiguration) -> ANNResult; } // Load static memory index from pre-built files synchronously -impl SyncLoadWith for DiskANNIndex +impl SyncLoadWith for DiskANNIndex where DP: DataProvider + LoadWith, { - fn load_with(path: &str, index_config: IndexConfiguration) -> Self { + fn load_with(path: &str, index_config: IndexConfiguration) -> ANNResult> { let storage = FileStorageProvider; let data_provider = create_data_provider(&storage, path, &index_config); match data_provider { - Ok(dp) => DiskANNIndex::::new_with_current_thread_runtime(index_config.config, dp), - Err(e) => panic!("Failed to create data provider: {}", e), + Ok(dp) => { + let index = + DiskANNIndex::::new_with_current_thread_runtime(index_config.config, dp); + Ok(index) + } + Err(e) => return Err(e.into()), } } } @@ -363,6 +370,6 @@ where }; let data_provider = DP::load_with(provider, &pq_context); - + block_on(data_provider) } From 124e0f98bfb99a4252316bd3940db83846f0b1ec Mon Sep 17 00:00:00 2001 From: Sanhaoji2 Date: Tue, 10 Feb 2026 19:43:59 +0800 Subject: [PATCH 3/4] Fix compile issue --- diskann-providers/src/index/wrapped_async.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index 168ee0e31..cd660db87 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -346,7 +346,7 @@ where DiskANNIndex::::new_with_current_thread_runtime(index_config.config, dp); Ok(index) } - Err(e) => return Err(e.into()), + Err(e) => return Err(e), } } } From e3001a4dafc5fdc2a0339824e8c5f811be645c3a Mon Sep 17 00:00:00 2001 From: Sanhaoji2 Date: Tue, 10 Feb 2026 19:47:19 +0800 Subject: [PATCH 4/4] fix compile issue --- diskann-providers/src/index/wrapped_async.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index cd660db87..69c4495b8 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -346,7 +346,7 @@ where DiskANNIndex::::new_with_current_thread_runtime(index_config.config, dp); Ok(index) } - Err(e) => return Err(e), + Err(e) => Err(e), } } }