diskann_disk/storage/
disk_index_reader.rs

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

diskann_disk/storage/
disk_index_reader.rs
