Return indexed vectors after disk search

diskann-disk/src/search/provider/disk_provider.rs

@@ -850,6 +850,18 @@ pub struct SearchResultItem<AssociatedData> {
     pub distance: f32,
 }
 
+pub struct SearchResultWithIndexedVectors<AssociatedData, VectorData> {
+    pub results: Vec<SearchResultItemWithIndexedVector<AssociatedData, VectorData>>,
+    pub stats: SearchResultStats,
+}
+
+pub struct SearchResultItemWithIndexedVector<AssociatedData, VectorData> {
+    pub vertex_id: u32,
+    pub data: AssociatedData,
+    pub distance: f32,
+    pub indexed_vector: Box<[VectorData]>,
+}
+
 impl<Data, ProviderFactory> DiskIndexSearcher<Data, ProviderFactory>
 where
     Data: GraphDataType<VectorIdType = u32>,
@@ -1095,6 +1107,67 @@ where
         Ok(search_result)
     }
 
+    /// Perform a disk search and return valid hits with their indexed vectors.
+    pub fn search_with_indexed_vectors(
+        &self,
+        query: &[Data::VectorDataType],
+        return_list_size: u32,
+        search_list_size: u32,
+        beam_width: Option<usize>,
+        mode: SearchMode<'_>,
+    ) -> ANNResult<SearchResultWithIndexedVectors<Data::AssociatedDataType, Data::VectorDataType>>
+    {
+        let SearchResult { mut results, stats } =
+            self.search(query, return_list_size, search_list_size, beam_width, mode)?;
+        results.truncate(stats.result_count as usize);
+        if results.is_empty() {
+            return Ok(SearchResultWithIndexedVectors {
+                results: Vec::new(),
+                stats,
+            });
+        }
+
+        let provider = self.index.provider();
+        let mut scratch = PoolOption::try_pooled(
+            &self.scratch_pool,
+            &DiskSearchScratchArgs {
+                graph_degree: provider.graph_header.max_degree::<Data::VectorDataType>()?,
+                pq_dim: provider.pq_data.get_dim(),
+                num_pq_chunks: provider.pq_data.get_num_chunks(),
+                num_pq_centers: provider.pq_data.get_num_centers(),
+                vertex_factory: &self.vertex_provider_factory,
+                graph_header: &provider.graph_header,
+            },
+        )?;
+
+        const BATCH_SIZE: usize = 128;
+        let mut indexed_vectors = Vec::with_capacity(results.len());
+        for batch in results.chunks(BATCH_SIZE) {
+            let ids = batch
+                .iter()
+                .map(|result| result.vertex_id)
+                .collect::<Vec<_>>();
+            ensure_vertex_loaded(&mut scratch.vertex_provider, &ids)?;
+            for id in ids {
+                indexed_vectors.push(Box::from(scratch.vertex_provider.get_vector(&id)?));
+            }
+        }
+
+        let results = results
+            .into_iter()
+            .zip(indexed_vectors)
+            .map(
+                |(result, indexed_vector)| SearchResultItemWithIndexedVector {
+                    indexed_vector,
+                    vertex_id: result.vertex_id,
+                    data: result.data,
+                    distance: result.distance,
+                },
+            )
+            .collect();
+        Ok(SearchResultWithIndexedVectors { results, stats })
+    }
+
     /// Perform a raw search on the disk index.
     /// This is a lower-level API that allows more control over the search parameters and output buffers.
     #[allow(clippy::too_many_arguments)]
@@ -1535,6 +1608,76 @@ mod disk_provider_tests {
         result.into_inner().into_vec()
     }
 
+    #[rstest]
+    #[case(CachingStrategy::None)]
+    #[case(CachingStrategy::StaticCacheWithBfsNodes(32))]
+    fn test_search_with_indexed_vectors(#[case] caching_strategy: CachingStrategy) {
+        let storage_provider = Arc::new(VirtualStorageProvider::new_overlay(test_data_root()));
+        let searcher = create_disk_index_searcher::<GraphDataF32VectorUnitData>(
+            CreateDiskIndexSearcherParams {
+                max_thread_num: 1,
+                pq_pivot_file_path: TEST_PQ_PIVOT_128DIM,
+                pq_compressed_file_path: TEST_PQ_COMPRESSED_128DIM,
+                index_path: TEST_INDEX_128DIM,
+                index_path_prefix: TEST_INDEX_PREFIX_128DIM,
+                caching_strategy,
+                ..Default::default()
+            },
+            &storage_provider,
+        );
+        let query = vec![0.1f32; 128];
+        let legacy = searcher
+            .search(&query, 10, 20, None, SearchMode::graph())
+            .unwrap();
+        let indexed = searcher
+            .search_with_indexed_vectors(&query, 10, 20, None, SearchMode::graph())
+            .unwrap();
+
+        assert_eq!(indexed.results.len(), indexed.stats.result_count as usize);
+        assert!(indexed.results.iter().zip(&legacy.results).all(|(a, b)| (
+            a.vertex_id,
+            a.distance
+        ) == (
+            b.vertex_id,
+            b.distance
+        )));
+        let source =
+            read_bin::<f32>(&mut storage_provider.open_reader(TEST_DATA_FILE).unwrap()).unwrap();
+        assert!(indexed
+            .results
+            .iter()
+            .all(|result| result.indexed_vector.as_ref() == source.row(result.vertex_id as usize)));
+
+        let batched = searcher
+            .search_with_indexed_vectors(&query, 129, 129, None, SearchMode::graph())
+            .unwrap();
+        assert_eq!(batched.results.len(), 129);
+        assert!(batched
+            .results
+            .iter()
+            .all(|result| result.indexed_vector.as_ref() == source.row(result.vertex_id as usize)));
+
+        let partial = searcher
+            .search_with_indexed_vectors(
+                &query,
+                10,
+                10,
+                None,
+                SearchMode::flat_filtered(|id| matches!(*id, 72 | 87 | 170)),
+            )
+            .unwrap();
+        assert_eq!(partial.results.len(), 3);
+        assert_eq!(partial.stats.result_count, 3);
+
+        let empty = searcher
+            .search_with_indexed_vectors(&query, 10, 10, None, SearchMode::flat_filtered(|_| false))
+            .unwrap();
+        assert!(empty.results.is_empty());
+        assert!(searcher
+            .search_with_indexed_vectors(&query, 2, 1, None, SearchMode::graph())
+            .is_err());
+    }
+
     struct TestDiskSearchParams<'a, StorageType> {
         storage_provider: &'a StorageType,
         index_search_engine: &'a DiskIndexSearcher<