diskann_disk/search/provider/
disk_sector_graph.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5#![warn(missing_docs)]
6
7//! Sector graph
8use std::ops::Deref;
9
10use diskann::{ANNError, ANNResult};
11use diskann_quantization::alloc::{AlignedAllocator, Poly};
12
13use crate::{
14    data_model::GraphHeader,
15    search::provider::aligned_file_reader::{traits::AlignedFileReader, AlignedRead, Alignment},
16};
17
18const DEFAULT_DISK_SECTOR_LEN: usize = 4096;
19
20/// Sector graph read from disk index
21pub struct DiskSectorGraph<AlignedReaderType: AlignedFileReader> {
22    /// Ensure `sector_reader` is dropped before `sectors_data` by placing it before `sectors_data`.
23    /// Graph storage to read sectors
24    sector_reader: AlignedReaderType,
25    /// Sector bytes from disk
26    /// One sector has num_nodes_per_sector nodes
27    /// Each node's layout: {full precision vector:[T; DIM]}{num_nbrs: u32}{neighbors: [u32; num_nbrs]}
28    /// The fp vector is not aligned
29    ///
30    /// index info for multi-node sectors
31    /// node `i` is in sector: [i / num_nodes_per_sector]
32    /// offset in sector: [(i % num_nodes_per_sector) * node_len]
33    ///
34    /// index info for multi-sector nodes
35    /// node `i` is in sector: [i * max_node_len.div_ceil(block_size)]
36    /// offset in sector: [0]
37    sectors_data: Poly<[u8], AlignedAllocator>,
38    /// Current sector index into which the next read reads data
39    cur_sector_idx: u64,
40
41    /// 0 for multi-sector nodes, >0 for multi-node sectors
42    num_nodes_per_sector: u64,
43
44    node_len: u64,
45
46    max_n_batch_sector_read: usize,
47
48    num_sectors_per_node: usize,
49
50    block_size: usize,
51}
52
53impl<AlignedReaderType: AlignedFileReader> DiskSectorGraph<AlignedReaderType> {
54    /// Create SectorGraph instance
55    pub fn new(
56        sector_reader: AlignedReaderType,
57        header: &GraphHeader,
58        max_n_batch_sector_read: usize,
59    ) -> ANNResult<Self> {
60        let mut block_size = header.block_size() as usize;
61        let version = header.layout_version();
62        if (version.major_version() == 0 && version.minor_version() == 0) || block_size == 0 {
63            block_size = DEFAULT_DISK_SECTOR_LEN;
64        }
65
66        let num_nodes_per_sector = header.metadata().num_nodes_per_block;
67        let node_len = header.metadata().node_len;
68        let num_sectors_per_node = if num_nodes_per_sector > 0 {
69            1
70        } else {
71            (node_len as usize).div_ceil(block_size)
72        };
73
74        Ok(Self {
75            sector_reader,
76            sectors_data: Poly::broadcast(
77                0u8,
78                max_n_batch_sector_read * num_sectors_per_node * block_size,
79                AlignedAllocator::new(AlignedReaderType::Alignment::VALUE),
80            )
81            .map_err(ANNError::log_index_error)?,
82            cur_sector_idx: 0,
83            num_nodes_per_sector,
84            node_len,
85            max_n_batch_sector_read,
86            num_sectors_per_node,
87            block_size,
88        })
89    }
90
91    /// Reconfigure SectorGraph if the max number of sectors to read is larger than the current one
92    pub fn reconfigure(&mut self, max_n_batch_sector_read: usize) -> ANNResult<()> {
93        if max_n_batch_sector_read > self.max_n_batch_sector_read {
94            self.max_n_batch_sector_read = max_n_batch_sector_read;
95            self.sectors_data = Poly::broadcast(
96                0u8,
97                max_n_batch_sector_read * self.num_sectors_per_node * self.block_size,
98                AlignedAllocator::new(AlignedReaderType::Alignment::VALUE),
99            )
100            .map_err(ANNError::log_index_error)?;
101        }
102        Ok(())
103    }
104
105    /// Reset SectorGraph
106    pub fn reset(&mut self) {
107        self.cur_sector_idx = 0;
108    }
109
110    /// Read sectors into sectors_data
111    /// They are in the same order as sectors_to_fetch
112    pub fn read_graph(&mut self, sectors_to_fetch: &[u64]) -> ANNResult<()> {
113        let cur_sector_idx_usize: usize = self.cur_sector_idx.try_into()?;
114        if sectors_to_fetch.len() > self.max_n_batch_sector_read - cur_sector_idx_usize {
115            return Err(ANNError::log_index_error(format_args!(
116                "Trying to read too many sectors. number of sectors to read: {}, max number of sectors can read: {}",
117                sectors_to_fetch.len(),
118                self.max_n_batch_sector_read - cur_sector_idx_usize,
119            )));
120        }
121
122        let len_per_node = self.num_sectors_per_node * self.block_size;
123        if len_per_node == 0 {
124            return Err(ANNError::log_index_error(format_args!(
125                "len_per_node is 0 (num_sectors_per_node={}, block_size={})",
126                self.num_sectors_per_node, self.block_size,
127            )));
128        }
129        let range = cur_sector_idx_usize * len_per_node
130            ..(cur_sector_idx_usize + sectors_to_fetch.len()) * len_per_node;
131        debug_assert!(
132            range.len() % len_per_node == 0,
133            "range length {} is not divisible by {}",
134            range.len(),
135            len_per_node
136        );
137        let mut sector_slices: Vec<&mut [u8]> =
138            self.sectors_data[range].chunks_mut(len_per_node).collect();
139        let mut read_requests: Vec<AlignedRead<'_, u8, AlignedReaderType::Alignment>> =
140            Vec::with_capacity(sector_slices.len());
141        for (local_sector_idx, slice) in sector_slices.iter_mut().enumerate() {
142            let sector_id = sectors_to_fetch[local_sector_idx];
143            read_requests.push(AlignedRead::new(sector_id * self.block_size as u64, slice)?);
144        }
145
146        self.sector_reader.read(&mut read_requests)?;
147        self.cur_sector_idx += sectors_to_fetch.len() as u64;
148
149        Ok(())
150    }
151
152    #[inline]
153    /// Get node data by local index.
154    pub fn node_disk_buf(&self, node_index_local: usize, vertex_id: u32) -> &[u8] {
155        // get sector_buf where this node is located
156        let sector_buf = self.get_sector_buf(node_index_local);
157        let node_offset = self.get_node_offset(vertex_id);
158        &sector_buf[node_offset..node_offset + self.node_len as usize]
159    }
160
161    /// Get sector data by local index
162    #[inline]
163    fn get_sector_buf(&self, local_sector_idx: usize) -> &[u8] {
164        let len_per_node = self.num_sectors_per_node * self.block_size;
165        &self.sectors_data[local_sector_idx * len_per_node..(local_sector_idx + 1) * len_per_node]
166    }
167
168    /// Get offset of node in sectors_data
169    #[inline]
170    fn get_node_offset(&self, vertex_id: u32) -> usize {
171        if self.num_nodes_per_sector == 0 {
172            // multi-sector node
173            0
174        } else {
175            // multi node in a sector
176            (vertex_id as u64 % self.num_nodes_per_sector * self.node_len) as usize
177        }
178    }
179
180    #[inline]
181    /// Gets the index for the sector that contains the node with the given vertex_id
182    pub fn node_sector_index(&self, vertex_id: u32) -> u64 {
183        1 + if self.num_nodes_per_sector > 0 {
184            vertex_id as u64 / self.num_nodes_per_sector
185        } else {
186            vertex_id as u64 * self.num_sectors_per_node as u64
187        }
188    }
189}
190
191impl<AlignedReaderType: AlignedFileReader> Deref for DiskSectorGraph<AlignedReaderType> {
192    type Target = [u8];
193
194    fn deref(&self) -> &Self::Target {
195        &self.sectors_data
196    }
197}
198
199#[cfg(test)]
200mod disk_sector_graph_test {
201    use diskann_utils::test_data_root;
202
203    use super::*;
204    use crate::{
205        data_model::{GraphLayoutVersion, GraphMetadata},
206        search::provider::aligned_file_reader::{
207            traits::AlignedReaderFactory, AlignedFileReaderFactory,
208        },
209    };
210
211    fn test_index_path() -> String {
212        test_data_root()
213            .join("disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index")
214            .to_string_lossy()
215            .to_string()
216    }
217
218    fn test_initialize_disk_sector_graph(
219        num_nodes_per_sector: u64,
220        num_sectors_per_node: usize,
221        sector_reader: <AlignedFileReaderFactory as AlignedReaderFactory>::AlignedReaderType,
222    ) -> DiskSectorGraph<<AlignedFileReaderFactory as AlignedReaderFactory>::AlignedReaderType>
223    {
224        DiskSectorGraph {
225            sectors_data: Poly::broadcast(0u8, 512, AlignedAllocator::A512).unwrap(),
226            sector_reader,
227            cur_sector_idx: 0,
228            num_nodes_per_sector,
229            node_len: 32,
230            max_n_batch_sector_read: 4,
231            num_sectors_per_node,
232            block_size: 64,
233        }
234    }
235
236    #[test]
237    fn test_new_disk_sector_graph_multi_node_per_sector() {
238        let metadata = GraphMetadata::new(1000, 32, 500, 32, 2, 20, 50, 1024, 256);
239        let header = GraphHeader::new(metadata, 64, GraphLayoutVersion::new(1, 0));
240        let reader = AlignedFileReaderFactory::new(test_index_path())
241            .build()
242            .unwrap();
243        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
244        assert_eq!(graph.sectors_data.len(), 128);
245        assert_eq!(graph.num_sectors_per_node, 1);
246        assert_eq!(graph.num_nodes_per_sector, 2);
247    }
248
249    #[test]
250    fn test_new_disk_sector_graph_multi_sector_per_node() {
251        let metadata = GraphMetadata::new(1000, 32, 500, 128, 0, 20, 50, 1024, 256);
252        let header = GraphHeader::new(metadata, 64, GraphLayoutVersion::new(1, 0));
253        let reader = AlignedFileReaderFactory::new(test_index_path())
254            .build()
255            .unwrap();
256        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
257        assert_eq!(graph.sectors_data.len(), 256);
258        assert_eq!(graph.num_sectors_per_node, 2);
259        assert_eq!(graph.num_nodes_per_sector, 0);
260    }
261
262    #[test]
263    fn test_new_disk_sector_graph_old_version_data() {
264        let metadata = GraphMetadata::new(1000, 32, 500, 128, 0, 20, 50, 1024, 256);
265        let header = GraphHeader::new(metadata, 9999, GraphLayoutVersion::new(0, 0));
266        let reader = AlignedFileReaderFactory::new(test_index_path())
267            .build()
268            .unwrap();
269        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
270        assert_eq!(graph.block_size, DEFAULT_DISK_SECTOR_LEN);
271    }
272
273    #[test]
274    fn get_sector_buf_test() {
275        let reader = AlignedFileReaderFactory::new(test_index_path())
276            .build()
277            .unwrap();
278        let graph = test_initialize_disk_sector_graph(2, 1, reader);
279        let sector_buf = graph.get_sector_buf(0);
280        assert_eq!(sector_buf.len(), 64);
281    }
282
283    #[test]
284    fn get_node_offset_test_multi_node_per_sector() {
285        let reader = AlignedFileReaderFactory::new(test_index_path())
286            .build()
287            .unwrap();
288        let graph = test_initialize_disk_sector_graph(4, 1, reader);
289
290        assert_eq!(graph.get_node_offset(0), 0);
291        assert_eq!(graph.get_node_offset(1), 32);
292        assert_eq!(graph.get_node_offset(2), 64);
293        assert_eq!(graph.get_node_offset(3), 96);
294        assert_eq!(graph.get_node_offset(4), 0);
295        assert_eq!(graph.get_node_offset(5), 32);
296        assert_eq!(graph.get_node_offset(6), 64);
297        assert_eq!(graph.get_node_offset(7), 96);
298    }
299
300    #[test]
301    fn get_node_offset_test_multi_sector_per_node() {
302        let reader = AlignedFileReaderFactory::new(test_index_path())
303            .build()
304            .unwrap();
305        let graph = test_initialize_disk_sector_graph(0, 2, reader);
306
307        assert_eq!(graph.get_node_offset(0), 0);
308        assert_eq!(graph.get_node_offset(1), 0);
309        assert_eq!(graph.get_node_offset(2), 0);
310        assert_eq!(graph.get_node_offset(3), 0);
311        assert_eq!(graph.get_node_offset(4), 0);
312        assert_eq!(graph.get_node_offset(5), 0);
313    }
314
315    #[test]
316    fn node_sector_index_test_multi_node_per_sector() {
317        let reader = AlignedFileReaderFactory::new(test_index_path())
318            .build()
319            .unwrap();
320        let graph = test_initialize_disk_sector_graph(4, 1, reader);
321
322        assert_eq!(graph.node_sector_index(0), 1);
323        assert_eq!(graph.node_sector_index(3), 1);
324        assert_eq!(graph.node_sector_index(4), 2);
325        assert_eq!(graph.node_sector_index(5), 2);
326        assert_eq!(graph.node_sector_index(7), 2);
327        assert_eq!(graph.node_sector_index(8), 3);
328        assert_eq!(graph.node_sector_index(1023), 256);
329        assert_eq!(graph.node_sector_index(1024), 257);
330        assert_eq!(graph.node_sector_index(2047), 512);
331        assert_eq!(graph.node_sector_index(2048), 513);
332    }
333
334    #[test]
335    fn node_sector_index_test_multi_sector_per_node() {
336        let reader = AlignedFileReaderFactory::new(test_index_path())
337            .build()
338            .unwrap();
339        let graph = test_initialize_disk_sector_graph(0, 2, reader);
340
341        assert_eq!(graph.node_sector_index(0), 1);
342        assert_eq!(graph.node_sector_index(3), 7);
343        assert_eq!(graph.node_sector_index(4), 9);
344        assert_eq!(graph.node_sector_index(5), 11);
345        assert_eq!(graph.node_sector_index(7), 15);
346        assert_eq!(graph.node_sector_index(8), 17);
347        assert_eq!(graph.node_sector_index(1023), 2047);
348        assert_eq!(graph.node_sector_index(1024), 2049);
349        assert_eq!(graph.node_sector_index(2047), 4095);
350        assert_eq!(graph.node_sector_index(2048), 4097);
351    }
352
353    #[test]
354    fn test_read_graph_max_sectors() {
355        let reader = AlignedFileReaderFactory::new(test_index_path())
356            .build()
357            .unwrap();
358        let mut disk_sector_graph = test_initialize_disk_sector_graph(0, 2, reader);
359
360        // Try to read more sectors than the maximum allowed
361        let sectors_to_fetch = vec![1, 2, 3, 4, 5, 6];
362        let result = disk_sector_graph.read_graph(&sectors_to_fetch);
363
364        // Check that an error is returned
365        // Trying to read too many sectors. number of sectors to read: {}, max number of sectors can read: {}",
366        assert!(result.is_err());
367    }
368
369    #[test]
370    fn test_disk_sector_graph_deref() {
371        let reader = AlignedFileReaderFactory::new(test_index_path())
372            .build()
373            .unwrap();
374        let graph = test_initialize_disk_sector_graph(1, 1, reader);
375        let data = &graph;
376        assert_eq!(data.len(), 512);
377    }
378
379    #[test]
380    fn test_reconfigure_grows_buffer() {
381        let reader = AlignedFileReaderFactory::new(test_index_path())
382            .build()
383            .unwrap();
384        let mut graph = test_initialize_disk_sector_graph(2, 1, reader);
385        assert_eq!(graph.max_n_batch_sector_read, 4);
386
387        // Reconfigure to larger batch — buffer must grow beyond initial 512 bytes
388        graph.reconfigure(16).unwrap();
389        assert_eq!(graph.max_n_batch_sector_read, 16);
390        assert_eq!(graph.sectors_data.len(), 16 * 64);
391    }
392
393    #[test]
394    fn test_reconfigure_noop_for_smaller_size() {
395        let reader = AlignedFileReaderFactory::new(test_index_path())
396            .build()
397            .unwrap();
398        let mut graph = test_initialize_disk_sector_graph(2, 1, reader);
399        let original_len = graph.sectors_data.len();
400
401        // Reconfigure with same or smaller size should be a no-op
402        graph.reconfigure(4).unwrap();
403        assert_eq!(graph.max_n_batch_sector_read, 4);
404        assert_eq!(graph.sectors_data.len(), original_len);
405
406        graph.reconfigure(2).unwrap();
407        assert_eq!(graph.max_n_batch_sector_read, 4);
408        assert_eq!(graph.sectors_data.len(), original_len);
409    }
410
411    #[test]
412    fn test_new_disk_sector_graph_zero_block_size_defaults() {
413        let metadata = GraphMetadata::new(1000, 32, 500, 32, 2, 20, 50, 1024, 256);
414        // block_size = 0 should fall back to DEFAULT_DISK_SECTOR_LEN regardless of version
415        let header = GraphHeader::new(metadata, 0, GraphLayoutVersion::new(1, 0));
416        let reader = AlignedFileReaderFactory::new(test_index_path())
417            .build()
418            .unwrap();
419        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
420        assert_eq!(graph.block_size, DEFAULT_DISK_SECTOR_LEN);
421    }
422}
diskann_disk/search/provider/disk_sector_graph.rs

diskann_disk/search/provider/
disk_sector_graph.rs
