diskann_disk/storage/
disk_index_reader.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5use std::{marker::PhantomData, sync::Arc};
6
7use diskann::ANNResult;
8use diskann_providers::storage::StorageReadProvider;
9use diskann_providers::{storage::PQStorage, utils::load_metadata_from_file};
10
11use crate::search::pq::PQData;
12use tracing::info;
13
14/// This struct is used by the DiskIndexSearcher to read the index data from storage. Noted that the index data here is different from index graph,
15/// It includes the PQ data, pivot table, and the warmup query data.
16/// The Storage acts as a provider to read the data from storage system.
17/// The storage provider should be provided as a generic type and be specified by the caller when it initializes the DiskIndexSearcher.
18pub struct DiskIndexReader<VectorType> {
19    phantom: PhantomData<VectorType>,
20
21    pq_data: Arc<PQData>,
22
23    num_points: usize,
24}
25
26impl<VectorType> DiskIndexReader<VectorType> {
27    /// Create DiskIndexReader instance
28    pub fn new<Storage: StorageReadProvider>(
29        pq_pivot_path: String,
30        pq_compressed_data_path: String,
31        storage_provider: &Storage,
32    ) -> ANNResult<Self> {
33        let pq_storage = PQStorage::new(&pq_pivot_path, &pq_compressed_data_path, None);
34        let pq_pivot_table = pq_storage.load_pq_pivots_bin::<Storage>(
35            &pq_pivot_path,
36            0, // Use 0 to infer num_pq_chunks from the file
37            storage_provider,
38        )?;
39
40        // Auto-detect number of points from compressed PQ file metadata
41        let metadata = load_metadata_from_file(storage_provider, &pq_compressed_data_path)?;
42
43        let pq_compressed_data = PQStorage::load_pq_compressed_vectors_bin::<Storage>(
44            &pq_compressed_data_path,
45            metadata.npoints(),
46            pq_pivot_table.get_num_chunks(),
47            storage_provider,
48        )?;
49        info!(
50            "Loaded PQ centroids and in-memory compressed vectors. #points:{} #pq_chunks: {}",
51            metadata.npoints(),
52            pq_pivot_table.get_num_chunks()
53        );
54
55        Ok(DiskIndexReader {
56            phantom: PhantomData,
57            pq_data: Arc::<PQData>::new(PQData::new(pq_pivot_table, pq_compressed_data)?),
58            num_points: metadata.npoints(),
59        })
60    }
61
62    pub fn get_pq_data(&self) -> Arc<PQData> {
63        Arc::clone(&self.pq_data)
64    }
65
66    pub fn get_num_points(&self) -> usize {
67        self.num_points
68    }
69}
70
71#[cfg(test)]
72mod disk_index_storage_test {
73    use diskann::ANNErrorKind;
74    use diskann_providers::storage::VirtualStorageProvider;
75    use diskann_utils::test_data_root;
76    use vfs::OverlayFS;
77
78    use super::*;
79
80    #[test]
81    fn load_pivot_test() {
82        let pivot_file_prefix: &str = "/sift/siftsmall_learn";
83        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
84        let storage = DiskIndexReader::<f32>::new::<VirtualStorageProvider<OverlayFS>>(
85            pivot_file_prefix.to_string() + "_pq_pivots.bin",
86            pivot_file_prefix.to_string() + "_pq_compressed.bin",
87            &storage_provider,
88        )
89        .unwrap();
90
91        // Creating the backend storage is sufficient to verify the constraints on the
92        // PQ schema as both `FixedChunkPQTable` and the possible alternatives (such as
93        // `quantization::TransposedTable`) check for the well-formedness of the schema.
94        let _: Arc<PQData> = storage.get_pq_data();
95    }
96
97    #[test]
98    fn load_pivot_file_not_exist_test() {
99        let pivot_file_prefix: &str = "/sift/siftsmall_learn_file_not_exist";
100        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
101        let err = match DiskIndexReader::<f32>::new::<VirtualStorageProvider<OverlayFS>>(
102            pivot_file_prefix.to_string() + "_pq_pivots.bin",
103            pivot_file_prefix.to_string() + "_pq_compressed.bin",
104            &storage_provider,
105        ) {
106            Ok(_) => panic!("this function should not have succeeded"),
107            Err(err) => err,
108        };
109        assert_eq!(err.kind(), ANNErrorKind::PQError);
110        assert!(err.to_string().contains("PQ k-means pivot file not found"));
111    }
112
113    #[test]
114    fn test_get_num_points() {
115        let pivot_file_prefix: &str = "/sift/siftsmall_learn";
116        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
117        let storage = DiskIndexReader::<f32>::new::<VirtualStorageProvider<OverlayFS>>(
118            pivot_file_prefix.to_string() + "_pq_pivots.bin",
119            pivot_file_prefix.to_string() + "_pq_compressed.bin",
120            &storage_provider,
121        )
122        .unwrap();
123
124        let num_points = storage.get_num_points();
125        assert_eq!(num_points, 25000);
126    }
127}
diskann_disk/storage/disk_index_reader.rs

diskann_disk/storage/
disk_index_reader.rs
