diskann_disk/utils/aligned_file_reader/
aligned_read.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use diskann::{ANNError, ANNResult};
7
8pub const DISK_IO_ALIGNMENT: usize = 512;
9
10/// Aligned read struct for disk IO, it takes the ownership of the AlignedBoxedSlice and returns the AlignedBoxWithSlice data immutably.
11pub struct AlignedRead<'a, T> {
12    /// where to read from
13    /// offset needs to be aligned with DISK_IO_ALIGNMENT
14    offset: u64,
15
16    /// where to read into
17    /// aligned_buf and its len need to be aligned with DISK_IO_ALIGNMENT
18    aligned_buf: &'a mut [T],
19}
20
21impl<'a, T> AlignedRead<'a, T> {
22    pub fn new(offset: u64, aligned_buf: &'a mut [T]) -> ANNResult<Self> {
23        Self::assert_is_aligned(offset as usize)?;
24        Self::assert_is_aligned(std::mem::size_of_val(aligned_buf))?;
25
26        Ok(Self {
27            offset,
28            aligned_buf,
29        })
30    }
31
32    fn assert_is_aligned(val: usize) -> ANNResult<()> {
33        match val % DISK_IO_ALIGNMENT {
34            0 => Ok(()),
35            _ => Err(ANNError::log_disk_io_request_alignment_error(format!(
36                "The offset or length of AlignedRead request is not {} bytes aligned",
37                DISK_IO_ALIGNMENT
38            ))),
39        }
40    }
41
42    /// where to read from
43    /// offset needs to be aligned with DISK_IO_ALIGNMENT
44    pub fn offset(&self) -> u64 {
45        self.offset
46    }
47
48    pub fn aligned_buf(&self) -> &[T] {
49        self.aligned_buf
50    }
51
52    /// where to read into
53    /// aligned_buf and its len need to be aligned with DISK_IO_ALIGNMENT
54    pub fn aligned_buf_mut(&mut self) -> &mut [T] {
55        self.aligned_buf
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    #[test]
64    fn test_aligned_read_valid() {
65        let mut buffer = vec![0u8; 512];
66        let aligned_read = AlignedRead::new(0, &mut buffer);
67
68        assert!(aligned_read.is_ok());
69        let aligned_read = aligned_read.unwrap();
70        assert_eq!(aligned_read.offset(), 0);
71        assert_eq!(aligned_read.aligned_buf().len(), 512);
72    }
73
74    #[test]
75    fn test_aligned_read_valid_offset() {
76        let mut buffer = vec![0u8; 1024];
77        let aligned_read = AlignedRead::new(512, &mut buffer);
78
79        assert!(aligned_read.is_ok());
80        let aligned_read = aligned_read.unwrap();
81        assert_eq!(aligned_read.offset(), 512);
82    }
83
84    #[test]
85    fn test_aligned_read_invalid_offset() {
86        let mut buffer = vec![0u8; 512];
87        let aligned_read = AlignedRead::new(100, &mut buffer);
88
89        assert!(aligned_read.is_err());
90    }
91
92    #[test]
93    fn test_aligned_read_invalid_buffer_size() {
94        let mut buffer = vec![0u8; 100];
95        let aligned_read = AlignedRead::new(0, &mut buffer);
96
97        assert!(aligned_read.is_err());
98    }
99
100    #[test]
101    fn test_aligned_read_buffer_access() {
102        let mut buffer = vec![42u8; 512];
103        let mut aligned_read = AlignedRead::new(0, &mut buffer).unwrap();
104
105        // Test immutable access
106        assert_eq!(aligned_read.aligned_buf()[0], 42);
107
108        // Test mutable access
109        aligned_read.aligned_buf_mut()[0] = 100;
110        assert_eq!(aligned_read.aligned_buf()[0], 100);
111    }
112}
diskann_disk/utils/aligned_file_reader/aligned_read.rs

diskann_disk/utils/aligned_file_reader/
aligned_read.rs
