diskann_disk/search/provider/aligned_file_reader/reader/
windows.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5use std::{ptr, thread, time::Duration};
6
7use diskann::{ANNError, ANNResult};
8
9use crate::search::provider::aligned_file_reader::{
10    platform::{
11        get_queued_completion_status, read_file_to_slice, FileHandle, IOCompletionPort, IOContext,
12        DWORD, OVERLAPPED, ULONG_PTR,
13    },
14    traits::AlignedFileReader,
15    AlignedRead, A512,
16};
17
18pub const MAX_IO_CONCURRENCY: usize = 128;
19pub const IO_COMPLETION_TIMEOUT: DWORD = u32::MAX; // Infinite timeout.
20pub const ASYNC_IO_COMPLETION_CHECK_INTERVAL: Duration = Duration::from_micros(5);
21
22/// AlignedFileReader for Windows.  When you modify this class run the benchmarks to make sure
23/// we don't regress on runtime.
24///
25/// # Run this before making your code change
26/// cargo bench --bench bench_main -p diskann -- --save-baseline prior_to_change
27///
28/// # Run this after making your code change to generate comparison metrics
29/// cargo bench --bench bench_main -p diskann -- --baseline prior_to_change
30pub struct WindowsAlignedFileReader {
31    io_context: IOContext,
32}
33
34impl WindowsAlignedFileReader {
35    pub fn new(fname: &str) -> ANNResult<Self> {
36        let mut io_context = IOContext::new();
37        tracing::debug!("Creating file handle for {}", fname);
38        match unsafe { FileHandle::new(fname) } {
39            Ok(file_handle) => io_context.file_handle = file_handle,
40            Err(err) => {
41                return Err(ANNError::log_io_error(err));
42            }
43        }
44
45        // Create a io completion port for the file handle, later it will be used to get the completion status.
46        match IOCompletionPort::new(&io_context.file_handle, None, 0, 0) {
47            Ok(io_completion_port) => io_context.io_completion_port = io_completion_port,
48            Err(err) => {
49                return Err(ANNError::log_io_error(err));
50            }
51        }
52
53        Ok(WindowsAlignedFileReader { io_context })
54    }
55}
56
57impl AlignedFileReader for WindowsAlignedFileReader {
58    /// Overlapped/`FILE_FLAG_NO_BUFFERING` I/O requires the buffer pointer to
59    /// be aligned to the device sector size in memory (512 bytes on typical
60    /// Windows volumes).
61    type Alignment = A512;
62
63    // Read the data from the file by sending concurrent io requests in batches.
64    fn read(&mut self, read_requests: &mut [AlignedRead<u8, A512>]) -> ANNResult<()> {
65        let n_requests = read_requests.len();
66        let n_batches = n_requests.div_ceil(MAX_IO_CONCURRENCY);
67        let ctx = &self.io_context;
68        let mut overlapped_in_out =
69            vec![unsafe { std::mem::zeroed::<OVERLAPPED>() }; MAX_IO_CONCURRENCY];
70
71        for batch_idx in 0..n_batches {
72            let batch_start = MAX_IO_CONCURRENCY * batch_idx;
73            let batch_size = std::cmp::min(n_requests - batch_start, MAX_IO_CONCURRENCY);
74
75            for j in 0..batch_size {
76                let req = &mut read_requests[batch_start + j];
77                let offset = req.offset();
78                let os = &mut overlapped_in_out[j];
79
80                match unsafe {
81                    read_file_to_slice(&ctx.file_handle, req.aligned_buf_mut(), os, offset)
82                } {
83                    Ok(_) => {}
84                    Err(error) => {
85                        return Err(ANNError::log_io_error(error));
86                    }
87                }
88            }
89
90            let mut n_read: DWORD = 0;
91            let mut n_complete: u64 = 0;
92            let mut completion_key: ULONG_PTR = 0;
93            let mut lp_os: *mut OVERLAPPED = ptr::null_mut();
94            while n_complete < batch_size as u64 {
95                match unsafe {
96                    get_queued_completion_status(
97                        &ctx.io_completion_port,
98                        &mut n_read,
99                        &mut completion_key,
100                        &mut lp_os,
101                        IO_COMPLETION_TIMEOUT,
102                    )
103                } {
104                    // An IO request completed.
105                    Ok(true) => n_complete += 1,
106                    // No IO request completed, continue to wait.
107                    Ok(false) => {
108                        thread::sleep(ASYNC_IO_COMPLETION_CHECK_INTERVAL);
109                    }
110                    // An error ocurred.
111                    Err(error) => return Err(ANNError::log_io_error(error)),
112                }
113            }
114        }
115
116        Ok(())
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use std::{
123        fs::File,
124        io::{BufReader, Read, Seek, SeekFrom},
125    };
126
127    use bincode::deserialize_from;
128    use diskann_utils::test_data_root;
129    use serde::{Deserialize, Serialize};
130
131    use super::*;
132    use crate::search::provider::aligned_file_reader::{AlignedRead, A512};
133    use diskann_quantization::alloc::{AlignedAllocator, Poly};
134
135    fn test_index_path() -> String {
136        test_data_root()
137            .join("disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index")
138            .to_string_lossy()
139            .to_string()
140    }
141
142    fn truth_node_data_path() -> String {
143        test_data_root()
144            .join("disk_index_misc/disk_index_node_data_aligned_reader_truth.bin")
145            .to_string_lossy()
146            .to_string()
147    }
148
149    const DEFAULT_DISK_SECTOR_LEN: usize = 4096;
150
151    #[derive(Debug, Serialize, Deserialize)]
152    struct NodeData {
153        num_neighbors: u32,
154        coordinates: Vec<f32>,
155        neighbors: Vec<u32>,
156    }
157
158    impl PartialEq for NodeData {
159        fn eq(&self, other: &Self) -> bool {
160            self.num_neighbors == other.num_neighbors
161                && self.coordinates == other.coordinates
162                && self.neighbors == other.neighbors
163        }
164    }
165
166    #[test]
167    fn test_new_aligned_file_reader() {
168        // Replace "test_file_path" with actual file path
169        let result = WindowsAlignedFileReader::new(&test_index_path());
170        assert!(result.is_ok());
171    }
172
173    #[test]
174    fn test_read() {
175        let mut reader = WindowsAlignedFileReader::new(&test_index_path()).unwrap();
176
177        let read_length = 512; // adjust according to your logic
178        let num_read = 10;
179        let mut aligned_mem =
180            Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
181
182        // create and add AlignedReads to the vector
183        let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();
184
185        let mut aligned_reads: Vec<AlignedRead<'_, u8, A512>> = mem_slices
186            .iter_mut()
187            .enumerate()
188            .map(|(i, slice)| {
189                let offset = (i * read_length) as u64;
190                AlignedRead::new(offset, slice).unwrap()
191            })
192            .collect();
193
194        let result = reader.read(&mut aligned_reads);
195        assert!(result.is_ok());
196
197        // Assert that the actual data is correct.
198        let mut file = File::open(test_index_path()).unwrap();
199        for current_read in aligned_reads {
200            let mut expected = vec![0; current_read.aligned_buf().len()];
201            file.seek(SeekFrom::Start(current_read.offset())).unwrap();
202            file.read_exact(&mut expected).unwrap();
203
204            assert_eq!(
205                expected,
206                current_read.aligned_buf(),
207                "aligned_buf did not contain the expected data"
208            );
209        }
210    }
211
212    #[test]
213    fn test_read_disk_index_by_sector() {
214        let mut reader = WindowsAlignedFileReader::new(&test_index_path()).unwrap();
215
216        let read_length = DEFAULT_DISK_SECTOR_LEN; // adjust according to your logic
217        let num_sector = 10;
218        let mut aligned_mem =
219            Poly::broadcast(0u8, read_length * num_sector, AlignedAllocator::A512).unwrap();
220
221        // Each slice will be used as the buffer for a read request of a sector.
222        let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();
223
224        let mut aligned_reads: Vec<AlignedRead<'_, u8, A512>> = mem_slices
225            .iter_mut()
226            .enumerate()
227            .map(|(sector_id, slice)| {
228                let offset = (sector_id * read_length) as u64;
229                AlignedRead::new(offset, slice).unwrap()
230            })
231            .collect();
232
233        let result = reader.read(&mut aligned_reads);
234        assert!(result.is_ok());
235
236        aligned_reads.iter().for_each(|read| {
237            assert_eq!(read.aligned_buf().len(), DEFAULT_DISK_SECTOR_LEN);
238        });
239
240        let disk_layout_meta = reconstruct_disk_meta(aligned_reads[0].aligned_buf_mut());
241        assert!(disk_layout_meta.len() > 9);
242
243        let dims = disk_layout_meta[1];
244        let num_pts = disk_layout_meta[0];
245        let node_len = disk_layout_meta[3];
246        let max_num_nodes_per_sector = disk_layout_meta[4];
247
248        assert!(node_len * max_num_nodes_per_sector < DEFAULT_DISK_SECTOR_LEN as u64);
249
250        let num_nbrs_start = (dims as usize) * std::mem::size_of::<f32>();
251        let nbrs_buf_start = num_nbrs_start + std::mem::size_of::<u32>();
252
253        let mut node_data_array = Vec::with_capacity(max_num_nodes_per_sector as usize * 9);
254
255        // Only validate the first 9 sectors with graph nodes.
256        (1..9).for_each(|sector_id| {
257            let sector_data = &mem_slices[sector_id];
258            for node_data in sector_data.chunks_exact(node_len as usize) {
259                // Extract coordinates data from the start of the node_data
260                let coordinates_end = (dims as usize) * std::mem::size_of::<f32>();
261                let coordinates = node_data[0..coordinates_end]
262                    .chunks_exact(std::mem::size_of::<f32>())
263                    .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
264                    .collect();
265
266                // Extract number of neighbors from the node_data
267                let neighbors_num = u32::from_le_bytes(
268                    node_data[num_nbrs_start..nbrs_buf_start]
269                        .try_into()
270                        .unwrap(),
271                );
272
273                let nbors_buf_end =
274                    nbrs_buf_start + (neighbors_num as usize) * std::mem::size_of::<u32>();
275
276                // Extract neighbors from the node data.
277                let mut neighbors = Vec::new();
278                for nbors_data in node_data[nbrs_buf_start..nbors_buf_end]
279                    .chunks_exact(std::mem::size_of::<u32>())
280                {
281                    let nbors_id = u32::from_le_bytes(nbors_data.try_into().unwrap());
282                    assert!(nbors_id < num_pts as u32);
283                    neighbors.push(nbors_id);
284                }
285
286                // Create NodeData struct and push it to the node_data_array
287                node_data_array.push(NodeData {
288                    num_neighbors: neighbors_num,
289                    coordinates,
290                    neighbors,
291                });
292            }
293        });
294
295        // Compare that each node read from the disk index are expected.
296        let node_data_truth_file = File::open(truth_node_data_path()).unwrap();
297        let reader = BufReader::new(node_data_truth_file);
298
299        let node_data_vec: Vec<NodeData> = deserialize_from(reader).unwrap();
300        for (node_from_node_data_file, node_from_disk_index) in
301            node_data_vec.iter().zip(node_data_array.iter())
302        {
303            // Verify that the NodeData from the file is equal to the NodeData in node_data_array
304            assert_eq!(node_from_node_data_file, node_from_disk_index);
305        }
306    }
307
308    #[test]
309    fn test_read_fail_invalid_file() {
310        let reader = WindowsAlignedFileReader::new("/invalid_path");
311        assert!(reader.is_err());
312    }
313
314    #[test]
315    #[allow(clippy::read_zero_byte_vec)]
316    fn test_read_no_requests() {
317        let mut reader = WindowsAlignedFileReader::new(&test_index_path()).unwrap();
318
319        let mut read_requests = Vec::<AlignedRead<u8, A512>>::new();
320        let result = reader.read(&mut read_requests);
321        assert!(result.is_ok());
322    }
323
324    fn reconstruct_disk_meta(buffer: &[u8]) -> Vec<u64> {
325        let size_of_u64 = std::mem::size_of::<u64>();
326
327        let num_values = buffer.len() / size_of_u64;
328        let mut disk_layout_meta = Vec::with_capacity(num_values);
329        let meta_data = &buffer[8..];
330
331        for chunk in meta_data.chunks_exact(size_of_u64) {
332            let value = u64::from_le_bytes(chunk.try_into().unwrap());
333            disk_layout_meta.push(value);
334        }
335
336        disk_layout_meta
337    }
338}
diskann_disk/search/provider/aligned_file_reader/reader/windows.rs

diskann_disk/search/provider/aligned_file_reader/reader/
windows.rs
