nuvec_core/
file.rs

1//! Arrays backed by a file's bytes, for loaders (safetensors, npy, ...).
2
3use std::path::Path;
4use std::sync::Arc;
5
6use crate::array::Data;
7use crate::buffer::Buffer;
8use crate::{Array, DType, Error, Result, shape};
9
10/// A file's bytes in array storage: on a Mac with a GPU, the file is memory-mapped and
11/// wrapped as GPU-visible memory without a copy, so arrays made from it are views whose
12/// pages are read when first touched (by the CPU or the GPU). Elsewhere the file is read
13/// into memory. The mapping lives as long as any array made from it.
14///
15/// ```
16/// use nuvec_core::{DType, MappedFile};
17///
18/// # fn main() -> nuvec_core::Result<()> {
19/// # if cfg!(miri) { return Ok(()); } // no file system under Miri's isolation
20/// let path = std::env::temp_dir().join("nuvec-mapped-file-doc.bin");
21/// let bytes: Vec<u8> = [1.0f32, 2.0, 3.0, 4.0].iter().flat_map(|v| v.to_le_bytes()).collect();
22/// std::fs::write(&path, &bytes).unwrap();
23/// let file = MappedFile::open(&path)?;
24/// let a = file.array(4, DType::F32, &[3])?; // bytes 4..16, no copy
25/// assert_eq!(a.to_vec::<f32>()?, vec![2.0, 3.0, 4.0]);
26/// # Ok(())
27/// # }
28/// ```
29#[derive(Clone)]
30pub struct MappedFile {
31    buffer: Arc<Buffer>,
32}
33
34impl std::fmt::Debug for MappedFile {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("MappedFile")
37            .field("len", &self.len())
38            .finish()
39    }
40}
41
42impl MappedFile {
43    /// Maps (or reads) the file at `path`.
44    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
45        Ok(Self {
46            buffer: Arc::new(Buffer::map_file(path.as_ref())?),
47        })
48    }
49
50    /// The file's size in bytes.
51    pub fn len(&self) -> usize {
52        self.buffer.len()
53    }
54
55    /// Whether the file is empty.
56    pub fn is_empty(&self) -> bool {
57        self.len() == 0
58    }
59
60    /// The file's bytes.
61    pub fn bytes(&self) -> &[u8] {
62        self.buffer.bytes()
63    }
64
65    /// Bytes `offset..offset + len` as storage of their own: one copy into an aligned buffer,
66    /// for files whose arrays are not aligned to their element sizes (safetensors written
67    /// without header padding, as MLX writes them). Arrays made from the copy are views of
68    /// it, and the file's mapping is not kept.
69    pub fn section(&self, offset: usize, len: usize) -> Result<Self> {
70        if offset.checked_add(len).is_none_or(|end| end > self.len()) {
71            return Err(Error::Io(format!(
72                "{len} bytes at offset {offset} exceed the file's {} bytes",
73                self.len()
74            )));
75        }
76        let mut copy = Buffer::try_alloc(len)?;
77        copy.bytes_mut()
78            .copy_from_slice(&self.bytes()[offset..offset + len]);
79        Ok(Self {
80            buffer: Arc::new(copy),
81        })
82    }
83
84    /// A row-major array of `dtype` and `shape` stored at byte `offset`: a view of the file
85    /// if `offset` is a multiple of the element size, otherwise a copy.
86    pub fn array(&self, offset: usize, dtype: DType, shape: &[usize]) -> Result<Array> {
87        let nbytes = shape::checked_nbytes(shape, dtype.size())?;
88        if offset
89            .checked_add(nbytes)
90            .is_none_or(|end| end > self.len())
91        {
92            return Err(Error::Io(format!(
93                "{nbytes} bytes at offset {offset} exceed the file's {} bytes",
94                self.len()
95            )));
96        }
97        if offset.is_multiple_of(dtype.size()) {
98            let data = Data {
99                buffer: self.buffer.clone(),
100                offset: (offset / dtype.size()) as isize,
101            };
102            return Array::leaf(Ok(data), dtype, shape.into());
103        }
104        let mut copy = Buffer::try_alloc(nbytes)?;
105        copy.bytes_mut()
106            .copy_from_slice(&self.bytes()[offset..offset + nbytes]);
107        Array::from_buffer(copy, dtype, shape)
108    }
109}
nuvec_core/file.rs

nuvec_core/
file.rs
