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    utils::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 crate::utils::aligned_file_reader::{
202        traits::AlignedReaderFactory, AlignedFileReaderFactory,
203    };
204    use diskann_utils::test_data_root;
205
206    use super::*;
207    use crate::data_model::{GraphLayoutVersion, GraphMetadata};
208
209    fn test_index_path() -> String {
210        test_data_root()
211            .join("disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index")
212            .to_string_lossy()
213            .to_string()
214    }
215
216    fn test_initialize_disk_sector_graph(
217        num_nodes_per_sector: u64,
218        num_sectors_per_node: usize,
219        sector_reader: <AlignedFileReaderFactory as AlignedReaderFactory>::AlignedReaderType,
220    ) -> DiskSectorGraph<<AlignedFileReaderFactory as AlignedReaderFactory>::AlignedReaderType>
221    {
222        DiskSectorGraph {
223            sectors_data: Poly::broadcast(0u8, 512, AlignedAllocator::A512).unwrap(),
224            sector_reader,
225            cur_sector_idx: 0,
226            num_nodes_per_sector,
227            node_len: 32,
228            max_n_batch_sector_read: 4,
229            num_sectors_per_node,
230            block_size: 64,
231        }
232    }
233
234    #[test]
235    fn test_new_disk_sector_graph_multi_node_per_sector() {
236        let metadata = GraphMetadata::new(1000, 32, 500, 32, 2, 20, 50, 1024, 256);
237        let header = GraphHeader::new(metadata, 64, GraphLayoutVersion::new(1, 0));
238        let reader = AlignedFileReaderFactory::new(test_index_path())
239            .build()
240            .unwrap();
241        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
242        assert_eq!(graph.sectors_data.len(), 128);
243        assert_eq!(graph.num_sectors_per_node, 1);
244        assert_eq!(graph.num_nodes_per_sector, 2);
245    }
246
247    #[test]
248    fn test_new_disk_sector_graph_multi_sector_per_node() {
249        let metadata = GraphMetadata::new(1000, 32, 500, 128, 0, 20, 50, 1024, 256);
250        let header = GraphHeader::new(metadata, 64, GraphLayoutVersion::new(1, 0));
251        let reader = AlignedFileReaderFactory::new(test_index_path())
252            .build()
253            .unwrap();
254        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
255        assert_eq!(graph.sectors_data.len(), 256);
256        assert_eq!(graph.num_sectors_per_node, 2);
257        assert_eq!(graph.num_nodes_per_sector, 0);
258    }
259
260    #[test]
261    fn test_new_disk_sector_graph_old_version_data() {
262        let metadata = GraphMetadata::new(1000, 32, 500, 128, 0, 20, 50, 1024, 256);
263        let header = GraphHeader::new(metadata, 9999, GraphLayoutVersion::new(0, 0));
264        let reader = AlignedFileReaderFactory::new(test_index_path())
265            .build()
266            .unwrap();
267        let graph = DiskSectorGraph::new(reader, &header, 2).unwrap();
268        assert_eq!(graph.block_size, DEFAULT_DISK_SECTOR_LEN);
269    }
270
271    #[test]
272    fn get_sector_buf_test() {
273        let reader = AlignedFileReaderFactory::new(test_index_path())
274            .build()
275            .unwrap();
276        let graph = test_initialize_disk_sector_graph(2, 1, reader);
277        let sector_buf = graph.get_sector_buf(0);
278        assert_eq!(sector_buf.len(), 64);
279    }
280
281    #[test]
282    fn get_node_offset_test_multi_node_per_sector() {
283        let reader = AlignedFileReaderFactory::new(test_index_path())
284            .build()
285            .unwrap();
286        let graph = test_initialize_disk_sector_graph(4, 1, reader);
287
288        assert_eq!(graph.get_node_offset(0), 0);
289        assert_eq!(graph.get_node_offset(1), 32);
290        assert_eq!(graph.get_node_offset(2), 64);
291        assert_eq!(graph.get_node_offset(3), 96);
292        assert_eq!(graph.get_node_offset(4), 0);
293        assert_eq!(graph.get_node_offset(5), 32);
294        assert_eq!(graph.get_node_offset(6), 64);
295        assert_eq!(graph.get_node_offset(7), 96);
296    }
297
298    #[test]
299    fn get_node_offset_test_multi_sector_per_node() {
300        let reader = AlignedFileReaderFactory::new(test_index_path())
301            .build()
302            .unwrap();
303        let graph = test_initialize_disk_sector_graph(0, 2, reader);
304
305        assert_eq!(graph.get_node_offset(0), 0);
306        assert_eq!(graph.get_node_offset(1), 0);
307        assert_eq!(graph.get_node_offset(2), 0);
308        assert_eq!(graph.get_node_offset(3), 0);
309        assert_eq!(graph.get_node_offset(4), 0);
310        assert_eq!(graph.get_node_offset(5), 0);
311    }
312
313    #[test]
314    fn node_sector_index_test_multi_node_per_sector() {
315        let reader = AlignedFileReaderFactory::new(test_index_path())
316            .build()
317            .unwrap();
318        let graph = test_initialize_disk_sector_graph(4, 1, reader);
319
320        assert_eq!(graph.node_sector_index(0), 1);
321        assert_eq!(graph.node_sector_index(3), 1);
322        assert_eq!(graph.node_sector_index(4), 2);
323        assert_eq!(graph.node_sector_index(5), 2);
324        assert_eq!(graph.node_sector_index(7), 2);
325        assert_eq!(graph.node_sector_index(8), 3);
326        assert_eq!(graph.node_sector_index(1023), 256);
327        assert_eq!(graph.node_sector_index(1024), 257);
328        assert_eq!(graph.node_sector_index(2047), 512);
329        assert_eq!(graph.node_sector_index(2048), 513);
330    }
331
332    #[test]
333    fn node_sector_index_test_multi_sector_per_node() {
334        let reader = AlignedFileReaderFactory::new(test_index_path())
335            .build()
336            .unwrap();
337        let graph = test_initialize_disk_sector_graph(0, 2, reader);
338
339        assert_eq!(graph.node_sector_index(0), 1);
340        assert_eq!(graph.node_sector_index(3), 7);
341        assert_eq!(graph.node_sector_index(4), 9);
342        assert_eq!(graph.node_sector_index(5), 11);
343        assert_eq!(graph.node_sector_index(7), 15);
344        assert_eq!(graph.node_sector_index(8), 17);
345        assert_eq!(graph.node_sector_index(1023), 2047);
346        assert_eq!(graph.node_sector_index(1024), 2049);
347        assert_eq!(graph.node_sector_index(2047), 4095);
348        assert_eq!(graph.node_sector_index(2048), 4097);
349    }
350
351    #[test]
352    fn test_read_graph_max_sectors() {
353        let reader = AlignedFileReaderFactory::new(test_index_path())
354            .build()
355            .unwrap();
356        let mut disk_sector_graph = test_initialize_disk_sector_graph(0, 2, reader);
357
358        // Try to read more sectors than the maximum allowed
359        let sectors_to_fetch = vec![1, 2, 3, 4, 5, 6];
360        let result = disk_sector_graph.read_graph(&sectors_to_fetch);
361
362        // Check that an error is returned
363        // Trying to read too many sectors. number of sectors to read: {}, max number of sectors can read: {}",
364        assert!(result.is_err());
365    }
366
367    #[test]
368    fn test_disk_sector_graph_deref() {
369        let reader = AlignedFileReaderFactory::new(test_index_path())
370            .build()
371            .unwrap();
372        let graph = test_initialize_disk_sector_graph(1, 1, reader);
373        let data = &graph;
374        assert_eq!(data.len(), 512);
375    }
376}
diskann_disk/search/provider/disk_sector_graph.rs

diskann_disk/search/provider/
disk_sector_graph.rs
