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;
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    let mut file = storage_provider.create_for_write(associated_data_path)?;
20
21    // Calculate the number of integers and the number of integers in associated data
22    let num_ints = end - start + 1;
23    let int_length: u32 = 1;
24
25    // Write the number of integers and the length of each integer as little endian
26    Metadata::new(num_ints, int_length)?.write(&mut file)?;
27
28    // Write the integers from the range as little endian
29    for i in start..=end {
30        file.write_all(&i.to_le_bytes())?;
31    }
32
33    Ok(())
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use byteorder::{LittleEndian, ReadBytesExt};
40    use diskann_providers::storage::{StorageReadProvider, VirtualStorageProvider};
41
42    #[test]
43    fn test_gen_associated_data_from_range() {
44        let storage_provider = VirtualStorageProvider::new_memory();
45        let path = "/test_gen_associated_data_from_range.bin";
46
47        // Generate data from range 0 to 9
48        gen_associated_data_from_range(&storage_provider, path, 0, 9).unwrap();
49
50        // Read back and verify
51        let mut file = storage_provider.open_reader(path).unwrap();
52
53        // Read metadata
54        let num_ints = file.read_u32::<LittleEndian>().unwrap();
55        let int_length = file.read_u32::<LittleEndian>().unwrap();
56
57        assert_eq!(num_ints, 10);
58        assert_eq!(int_length, 1);
59
60        // Read integers
61        for expected in 0u32..=9 {
62            let actual = file.read_u32::<LittleEndian>().unwrap();
63            assert_eq!(actual, expected);
64        }
65    }
66
67    #[test]
68    fn test_gen_associated_data_from_range_single_value() {
69        let storage_provider = VirtualStorageProvider::new_memory();
70        let path = "/test_gen_associated_data_single.bin";
71
72        // Generate data for a single value
73        gen_associated_data_from_range(&storage_provider, path, 42, 42).unwrap();
74
75        let mut file = storage_provider.open_reader(path).unwrap();
76
77        let num_ints = file.read_u32::<LittleEndian>().unwrap();
78        let int_length = file.read_u32::<LittleEndian>().unwrap();
79
80        assert_eq!(num_ints, 1);
81        assert_eq!(int_length, 1);
82
83        let value = file.read_u32::<LittleEndian>().unwrap();
84        assert_eq!(value, 42);
85    }
86
87    #[test]
88    fn test_gen_associated_data_from_range_large() {
89        let storage_provider = VirtualStorageProvider::new_memory();
90        let path = "/test_gen_associated_data_large.bin";
91
92        // Generate data for range 100 to 199
93        gen_associated_data_from_range(&storage_provider, path, 100, 199).unwrap();
94
95        let mut file = storage_provider.open_reader(path).unwrap();
96
97        let num_ints = file.read_u32::<LittleEndian>().unwrap();
98        let int_length = file.read_u32::<LittleEndian>().unwrap();
99
100        assert_eq!(num_ints, 100);
101        assert_eq!(int_length, 1);
102
103        for expected in 100u32..=199 {
104            let actual = file.read_u32::<LittleEndian>().unwrap();
105            assert_eq!(actual, expected);
106        }
107    }
108}
diskann_tools/utils/gen_associated_data_from_range.rs

diskann_tools/utils/
gen_associated_data_from_range.rs
