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

diskann_disk/utils/aligned_file_reader/
windows_aligned_file_reader.rs
