diskann_tools/utils/
gen_associated_data_from_range.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::io::Write;
7
8use diskann_providers::storage::StorageWriteProvider;
9use diskann_utils::io::Metadata;
10
11use super::{CMDResult, CMDToolError};
12
13pub fn gen_associated_data_from_range<S: StorageWriteProvider>(
14    storage_provider: &S,
15    associated_data_path: &str,
16    start: u32,
17    end: u32,
18) -> CMDResult<()> {
19    if end < start {
20        return Err(CMDToolError {
21            details: format!(
22                "invalid range: end ({end}) must be greater than or equal to start ({start})"
23            ),
24        });
25    }
26
27    // Calculate the number of integers and the number of integers in associated data
28    // Use checked arithmetic to avoid overflow when end == u32::MAX and start == 0
29    let num_ints = (end - start).checked_add(1).ok_or_else(|| CMDToolError {
30        details: format!("range [{start}, {end}] is too large: count overflows u32"),
31    })?;
32
33    let mut file = storage_provider.create_for_write(associated_data_path)?;
34
35    let int_length: u32 = 1;
36
37    // Write the number of integers and the length of each integer as little endian
38    Metadata::new(num_ints, int_length)?.write(&mut file)?;
39
40    // Write the integers from the range as little endian
41    for i in start..=end {
42        file.write_all(&i.to_le_bytes())?;
43    }
44
45    Ok(())
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use byteorder::{LittleEndian, ReadBytesExt};
52    use diskann_providers::storage::{StorageReadProvider, VirtualStorageProvider};
53
54    #[test]
55    fn test_gen_associated_data_from_range() {
56        let storage_provider = VirtualStorageProvider::new_memory();
57        let path = "/test_gen_associated_data_from_range.bin";
58
59        // Generate data from range 0 to 9
60        gen_associated_data_from_range(&storage_provider, path, 0, 9).unwrap();
61
62        // Read back and verify
63        let mut file = storage_provider.open_reader(path).unwrap();
64
65        // Read metadata
66        let num_ints = file.read_u32::<LittleEndian>().unwrap();
67        let int_length = file.read_u32::<LittleEndian>().unwrap();
68
69        assert_eq!(num_ints, 10);
70        assert_eq!(int_length, 1);
71
72        // Read integers
73        for expected in 0u32..=9 {
74            let actual = file.read_u32::<LittleEndian>().unwrap();
75            assert_eq!(actual, expected);
76        }
77    }
78
79    #[test]
80    fn test_gen_associated_data_from_range_single_value() {
81        let storage_provider = VirtualStorageProvider::new_memory();
82        let path = "/test_gen_associated_data_single.bin";
83
84        // Generate data for a single value
85        gen_associated_data_from_range(&storage_provider, path, 42, 42).unwrap();
86
87        let mut file = storage_provider.open_reader(path).unwrap();
88
89        let num_ints = file.read_u32::<LittleEndian>().unwrap();
90        let int_length = file.read_u32::<LittleEndian>().unwrap();
91
92        assert_eq!(num_ints, 1);
93        assert_eq!(int_length, 1);
94
95        let value = file.read_u32::<LittleEndian>().unwrap();
96        assert_eq!(value, 42);
97    }
98
99    #[test]
100    fn test_gen_associated_data_from_range_large() {
101        let storage_provider = VirtualStorageProvider::new_memory();
102        let path = "/test_gen_associated_data_large.bin";
103
104        // Generate data for range 100 to 199
105        gen_associated_data_from_range(&storage_provider, path, 100, 199).unwrap();
106
107        let mut file = storage_provider.open_reader(path).unwrap();
108
109        let num_ints = file.read_u32::<LittleEndian>().unwrap();
110        let int_length = file.read_u32::<LittleEndian>().unwrap();
111
112        assert_eq!(num_ints, 100);
113        assert_eq!(int_length, 1);
114
115        for expected in 100u32..=199 {
116            let actual = file.read_u32::<LittleEndian>().unwrap();
117            assert_eq!(actual, expected);
118        }
119    }
120
121    #[test]
122    fn test_gen_associated_data_from_range_end_less_than_start() {
123        let storage_provider = VirtualStorageProvider::new_memory();
124        let path = "/test.bin";
125
126        let result = gen_associated_data_from_range(&storage_provider, path, 10, 5);
127
128        let err = result.unwrap_err();
129        assert_eq!(
130            err.details,
131            "invalid range: end (5) must be greater than or equal to start (10)"
132        );
133        assert!(
134            !storage_provider.exists(path),
135            "File should not be created when range is invalid"
136        );
137    }
138
139    #[test]
140    fn test_gen_associated_data_from_range_max_overflow() {
141        let storage_provider = VirtualStorageProvider::new_memory();
142        let path = "/test.bin";
143
144        let result = gen_associated_data_from_range(&storage_provider, path, 0, u32::MAX);
145
146        let err = result.unwrap_err();
147        assert_eq!(
148            err.details,
149            format!("range [0, {}] is too large: count overflows u32", u32::MAX)
150        );
151        assert!(
152            !storage_provider.exists(path),
153            "File should not be created when range is too large"
154        );
155    }
156}
diskann_tools/utils/gen_associated_data_from_range.rs

diskann_tools/utils/
gen_associated_data_from_range.rs
