webdataset_core/
npy.rs

1//! Reading and writing NumPy's `.npy` and `.npz` container formats.
2//!
3//! Only the dense numeric subset that WebDataset actually stores is supported:
4//! C-ordered arrays of the dtypes in [`DType`]. Fortran-ordered arrays,
5//! structured dtypes, and object arrays are rejected with a clear error.
6
7use crate::error::{Error, Result};
8use crate::prelude::*;
9use crate::tensor::{DType, Tensor};
10
11const MAGIC: &[u8; 6] = b"\x93NUMPY";
12const ALIGNMENT: usize = 64;
13
14/// Decode a `.npy` byte string into a [`Tensor`].
15pub fn from_npy(data: &[u8]) -> Result<Tensor> {
16    if data.len() < 10 || &data[..6] != MAGIC {
17        return Err(Error::format("not a .npy file (bad magic)"));
18    }
19    let major = data[6];
20    let (header_len, header_start) = match major {
21        1 => (u16::from_le_bytes([data[8], data[9]]) as usize, 10),
22        2 | 3 => {
23            if data.len() < 12 {
24                return Err(Error::format("truncated .npy header"));
25            }
26            (u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize, 12)
27        }
28        other => return Err(Error::unsupported(format!(".npy version {other}"))),
29    };
30    let header_end = header_start + header_len;
31    if data.len() < header_end {
32        return Err(Error::format("truncated .npy header"));
33    }
34    let header = core::str::from_utf8(&data[header_start..header_end])?;
35    let (dtype, big_endian, fortran, shape) = parse_header(header)?;
36
37    let expected = shape.iter().product::<usize>() * dtype.size();
38    let body = &data[header_end..];
39    if body.len() < expected {
40        return Err(Error::format(format!(".npy body is {} bytes but shape {shape:?} needs {expected}", body.len())));
41    }
42    if fortran {
43        return Err(Error::unsupported("Fortran-ordered .npy arrays"));
44    }
45
46    let mut body = body[..expected].to_vec();
47    if big_endian != cfg!(target_endian = "big") && dtype.size() > 1 {
48        for chunk in body.chunks_exact_mut(dtype.size()) {
49            chunk.reverse();
50        }
51    }
52    Tensor::new(dtype, shape, body)
53}
54
55/// Encode a [`Tensor`] as a version 1.0 `.npy` byte string.
56pub fn to_npy(tensor: &Tensor) -> Vec<u8> {
57    let mut out = Vec::with_capacity(tensor.data().len() + ALIGNMENT + 16);
58    let (header, body) = npy_parts(tensor);
59    out.extend_from_slice(&header);
60    out.extend_from_slice(body);
61    out
62}
63
64/// Write a [`Tensor`] to `out` in version 1.0 `.npy` format.
65#[cfg(feature = "std")]
66pub fn write_npy(out: &mut impl std::io::Write, tensor: &Tensor) -> Result<()> {
67    let (header, body) = npy_parts(tensor);
68    out.write_all(&header)?;
69    out.write_all(body)?;
70    Ok(())
71}
72
73/// Read a `.npy` array from a stream.
74#[cfg(feature = "std")]
75pub fn read_npy(stream: &mut impl std::io::Read) -> Result<Tensor> {
76    let mut buf = Vec::new();
77    stream.read_to_end(&mut buf)?;
78    from_npy(&buf)
79}
80
81/// Build the header block and borrow the element data of a `.npy` encoding.
82fn npy_parts(tensor: &Tensor) -> (Vec<u8>, &[u8]) {
83    let order = if cfg!(target_endian = "big") { '>' } else { '<' };
84    let shape = match tensor.shape() {
85        [] => "()".to_string(),
86        [one] => format!("({one},)"),
87        many => format!("({})", many.iter().map(|d| d.to_string()).collect::<Vec<_>>().join(", ")),
88    };
89    let mut header =
90        format!("{{'descr': '{order}{}', 'fortran_order': False, 'shape': {shape}, }}", tensor.dtype().short_name());
91
92    // The header is padded so that the start of the data is 64-byte aligned.
93    let unpadded = MAGIC.len() + 2 + 2 + header.len() + 1;
94    let padding = (ALIGNMENT - unpadded % ALIGNMENT) % ALIGNMENT;
95    header.push_str(&" ".repeat(padding));
96    header.push('\n');
97
98    let mut block = Vec::with_capacity(unpadded + padding);
99    block.extend_from_slice(MAGIC);
100    block.extend_from_slice(&[1, 0]);
101    block.extend_from_slice(&(header.len() as u16).to_le_bytes());
102    block.extend_from_slice(header.as_bytes());
103    (block, tensor.data())
104}
105
106/// Parse the Python-literal header dict of a `.npy` file.
107fn parse_header(header: &str) -> Result<(DType, bool, bool, Vec<usize>)> {
108    let descr = extract(header, "descr").ok_or_else(|| Error::format("no descr in .npy header"))?;
109    let mut chars = descr.chars();
110    let (big_endian, name) = match chars.next() {
111        Some('<') | Some('|') | Some('=') => (false, chars.as_str()),
112        Some('>') => (true, chars.as_str()),
113        _ => (false, descr.as_str()),
114    };
115    let dtype = DType::parse(name)?;
116
117    let fortran = header.contains("'fortran_order': True");
118
119    let shape_text = header
120        .split_once("'shape':")
121        .and_then(|(_, rest)| rest.split_once('(').and_then(|(_, r)| r.split_once(')')).map(|(inner, _)| inner))
122        .ok_or_else(|| Error::format("no shape in .npy header"))?;
123    let shape = shape_text
124        .split(',')
125        .map(str::trim)
126        .filter(|s| !s.is_empty())
127        .map(|s| s.parse::<usize>().map_err(|_| Error::format(format!("bad shape entry {s:?}"))))
128        .collect::<Result<Vec<_>>>()?;
129
130    Ok((dtype, big_endian, fortran, shape))
131}
132
133/// Pull `'<key>': '<value>'` out of a Python-literal dict.
134fn extract(header: &str, key: &str) -> Option<String> {
135    let needle = format!("'{key}':");
136    let rest = header.split_once(&needle)?.1;
137    let rest = rest.trim_start();
138    let quote = rest.chars().next()?;
139    if quote != '\'' && quote != '"' {
140        return None;
141    }
142    let body = &rest[quote.len_utf8()..];
143    let end = body.find(quote)?;
144    Some(body[..end].to_string())
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn round_trips_a_matrix() {
153        let t = Tensor::from_f32_shaped(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]);
154        let bytes = to_npy(&t);
155        assert_eq!(&bytes[..6], MAGIC);
156        // The data must start on a 64 byte boundary.
157        let header_len = u16::from_le_bytes([bytes[8], bytes[9]]) as usize;
158        assert_eq!((10 + header_len) % ALIGNMENT, 0);
159
160        let back = from_npy(&bytes).unwrap();
161        assert_eq!(back, t);
162    }
163
164    #[test]
165    fn round_trips_a_vector_and_a_scalar() {
166        for shape in [vec![4], vec![1], vec![2, 2]] {
167            let n: usize = shape.iter().product();
168            let values: Vec<f32> = (0..n).map(|i| i as f32).collect();
169            let t = Tensor::from_f32_shaped(&values, shape);
170            assert_eq!(from_npy(&to_npy(&t)).unwrap(), t);
171        }
172    }
173
174    #[test]
175    fn rejects_junk() {
176        assert!(from_npy(b"not an npy file at all").is_err());
177        assert!(from_npy(b"").is_err());
178    }
179
180    #[test]
181    fn parses_a_real_numpy_header() {
182        let header = "{'descr': '<i8', 'fortran_order': False, 'shape': (3, 4), }";
183        let (dtype, big, fortran, shape) = parse_header(header).unwrap();
184        assert_eq!(dtype, DType::I64);
185        assert!(!big);
186        assert!(!fortran);
187        assert_eq!(shape, vec![3, 4]);
188    }
189
190    #[test]
191    fn rejects_fortran_order() {
192        let bytes = to_npy(&Tensor::from_f32(&[1.0, 2.0]));
193        let header_len = u16::from_le_bytes([bytes[8], bytes[9]]) as usize;
194        let header = core::str::from_utf8(&bytes[10..10 + header_len]).unwrap().replace("False", "True ");
195
196        let mut fortran = bytes[..10].to_vec();
197        fortran.extend_from_slice(header.as_bytes());
198        fortran.extend_from_slice(&bytes[10 + header_len..]);
199
200        assert_eq!(fortran.len(), bytes.len(), "only the flag should change");
201        let err = from_npy(&fortran).unwrap_err();
202        assert!(matches!(err, Error::Unsupported(_)), "{err}");
203    }
204}
webdataset_core/npy.rs

webdataset_core/
npy.rs
