nuvec_io/
idx.rs

1//! The IDX format (MNIST and Fashion-MNIST): a big-endian header (two zero bytes, a type
2//! code, the number of dimensions, then one `u32` per dimension) followed by row-major
3//! big-endian data.
4
5use std::path::Path;
6
7use nuvec_core::{Array, Error, Result};
8
9fn invalid(msg: impl Into<String>) -> Error {
10    Error::InvalidArgument(format!("idx: {}", msg.into()))
11}
12
13/// Parses an IDX file's bytes. Element types map to `uint8`, `int8`, `int16`, `int32`,
14/// `float32` and `float64`.
15pub fn parse_idx(bytes: &[u8]) -> Result<Array> {
16    let header = bytes
17        .get(..4)
18        .ok_or_else(|| invalid("file shorter than its header"))?;
19    if header[..2] != [0, 0] {
20        return Err(invalid("bad magic number"));
21    }
22    let ndim = usize::from(header[3]);
23    let dims_end = 4 + 4 * ndim;
24    let shape: Vec<usize> = bytes
25        .get(4..dims_end)
26        .ok_or_else(|| invalid("truncated dimensions"))?
27        .as_chunks::<4>()
28        .0
29        .iter()
30        .map(|&d| u32::from_be_bytes(d) as usize)
31        .collect();
32    let count = shape
33        .iter()
34        .try_fold(1usize, |n, &d| n.checked_mul(d))
35        .ok_or_else(|| invalid("dimensions overflow"))?;
36    let data = &bytes[dims_end..];
37    macro_rules! decode {
38        ($t:ty) => {{
39            let size = std::mem::size_of::<$t>();
40            if data.len() != count * size {
41                return Err(invalid(format!(
42                    "expected {} data bytes for shape {shape:?}, got {}",
43                    count * size,
44                    data.len()
45                )));
46            }
47            let values: Vec<$t> = data
48                .chunks_exact(size)
49                .map(|c| <$t>::from_be_bytes(c.try_into().expect("element bytes")))
50                .collect();
51            Array::from_slice(&values, &shape)
52        }};
53    }
54    match header[2] {
55        0x08 => decode!(u8),
56        0x09 => decode!(i8),
57        0x0B => decode!(i16),
58        0x0C => decode!(i32),
59        0x0D => decode!(f32),
60        0x0E => decode!(f64),
61        code => Err(invalid(format!("unknown type code {code:#04x}"))),
62    }
63}
64
65/// Reads an (uncompressed) IDX file.
66pub fn read_idx(path: impl AsRef<Path>) -> Result<Array> {
67    let path = path.as_ref();
68    let bytes = std::fs::read(path).map_err(|e| Error::Io(format!("{}: {e}", path.display())))?;
69    parse_idx(&bytes)
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    fn file(code: u8, shape: &[u32], data: &[u8]) -> Vec<u8> {
77        let mut b = vec![0, 0, code, shape.len() as u8];
78        for d in shape {
79            b.extend(d.to_be_bytes());
80        }
81        b.extend(data);
82        b
83    }
84
85    #[test]
86    fn images_and_labels() {
87        let images = parse_idx(&file(
88            0x08,
89            &[2, 2, 3],
90            &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 255],
91        ))
92        .unwrap();
93        assert_eq!(images.shape(), &[2, 2, 3]);
94        assert_eq!(images.to_vec::<u8>().unwrap()[11], 255);
95        let labels = parse_idx(&file(0x08, &[3], &[7, 2, 1])).unwrap();
96        assert_eq!(labels.to_vec::<u8>().unwrap(), [7, 2, 1]);
97    }
98
99    #[test]
100    fn big_endian_wider_types() {
101        let mut data = Vec::new();
102        data.extend((-2i32).to_be_bytes());
103        data.extend(300i32.to_be_bytes());
104        assert_eq!(
105            parse_idx(&file(0x0C, &[2], &data))
106                .unwrap()
107                .to_vec::<i32>()
108                .unwrap(),
109            [-2, 300]
110        );
111        let data: Vec<u8> = [1.5f32, -0.25]
112            .iter()
113            .flat_map(|v| v.to_be_bytes())
114            .collect();
115        assert_eq!(
116            parse_idx(&file(0x0D, &[2], &data))
117                .unwrap()
118                .to_vec::<f32>()
119                .unwrap(),
120            [1.5, -0.25]
121        );
122    }
123
124    #[test]
125    fn malformed_files_are_errors() {
126        assert!(parse_idx(&[0, 0]).is_err());
127        assert!(parse_idx(&[1, 0, 8, 1, 0, 0, 0, 1, 5]).is_err());
128        assert!(parse_idx(&file(0x08, &[3], &[1, 2])).is_err());
129        assert!(parse_idx(&file(0x42, &[1], &[1])).is_err());
130        assert!(matches!(
131            read_idx("/nonexistent/train-images-idx3-ubyte"),
132            Err(Error::Io(_))
133        ));
134    }
135}
nuvec_io/idx.rs

nuvec_io/
idx.rs
