polars_io/
mmap.rs

1use std::fs::File;
2use std::io::{BufReader, Cursor, Read, Seek};
3
4use polars_buffer::Buffer;
5use polars_core::config::verbose;
6use polars_utils::mmap::MMapSemaphore;
7
8/// Trait used to get a hold to file handler or to the underlying bytes
9/// without performing a Read.
10pub trait MmapBytesReader: Read + Seek + Send + Sync {
11    fn to_file(&self) -> Option<&File> {
12        None
13    }
14
15    fn to_bytes(&self) -> Option<&[u8]> {
16        None
17    }
18}
19
20impl MmapBytesReader for File {
21    fn to_file(&self) -> Option<&File> {
22        Some(self)
23    }
24}
25
26impl MmapBytesReader for BufReader<File> {
27    fn to_file(&self) -> Option<&File> {
28        Some(self.get_ref())
29    }
30}
31
32impl MmapBytesReader for BufReader<&File> {
33    fn to_file(&self) -> Option<&File> {
34        Some(self.get_ref())
35    }
36}
37
38impl<T> MmapBytesReader for Cursor<T>
39where
40    T: AsRef<[u8]> + Send + Sync,
41{
42    fn to_bytes(&self) -> Option<&[u8]> {
43        Some(self.get_ref().as_ref())
44    }
45}
46
47impl<T: MmapBytesReader + ?Sized> MmapBytesReader for Box<T> {
48    fn to_file(&self) -> Option<&File> {
49        T::to_file(self)
50    }
51
52    fn to_bytes(&self) -> Option<&[u8]> {
53        T::to_bytes(self)
54    }
55}
56
57impl<T: MmapBytesReader> MmapBytesReader for &mut T {
58    fn to_file(&self) -> Option<&File> {
59        T::to_file(self)
60    }
61
62    fn to_bytes(&self) -> Option<&[u8]> {
63        T::to_bytes(self)
64    }
65}
66
67// Handle various forms of input bytes.
68// TODO: we should just remove this and use Buffer.
69pub enum ReaderBytes<'a> {
70    Borrowed(&'a [u8]),
71    Owned(Buffer<u8>),
72}
73
74impl std::ops::Deref for ReaderBytes<'_> {
75    type Target = [u8];
76    fn deref(&self) -> &[u8] {
77        match self {
78            Self::Borrowed(ref_bytes) => ref_bytes,
79            Self::Owned(vec) => vec,
80        }
81    }
82}
83
84impl<'a, T: 'a + MmapBytesReader> From<&'a mut T> for ReaderBytes<'a> {
85    fn from(m: &'a mut T) -> Self {
86        match m.to_bytes() {
87            // , but somehow bchk doesn't see that lifetime is 'a.
88            Some(s) => {
89                let s = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(s) };
90                ReaderBytes::Borrowed(s)
91            },
92            None => {
93                if let Some(f) = m.to_file() {
94                    let mmap = MMapSemaphore::new_from_file(f).unwrap();
95                    ReaderBytes::Owned(Buffer::from_owner(mmap))
96                } else {
97                    if verbose() {
98                        eprintln!("could not memory map file; read to buffer.")
99                    }
100                    let mut buf = vec![];
101                    m.read_to_end(&mut buf).expect("could not read");
102                    ReaderBytes::Owned(Buffer::from_vec(buf))
103                }
104            },
105        }
106    }
107}
polars_io/mmap.rs

polars_io/
mmap.rs
