cfpyo3_core/df/
io.rs

1use super::{ColumnsDtype, DataFrame, IndexDtype, COLUMNS_NBYTES, INDEX_NBYTES};
2use crate::{
3    as_data_slice_or_err,
4    toolkit::{
5        array::AFloat,
6        convert::{from_bytes, to_bytes, to_nbytes},
7    },
8};
9use anyhow::Result;
10use std::io::{Read, Write};
11
12mod buffer;
13mod bytes;
14mod fs;
15
16impl<'a, T: AFloat> DataFrame<'a, T> {
17    pub fn read(reader: &mut impl Read) -> Result<Self> {
18        let mut nbytes_buffer = [0u8; 8];
19        reader.read_exact(&mut nbytes_buffer)?;
20        let index_nbytes = i64::from_le_bytes(nbytes_buffer) as usize;
21        reader.read_exact(&mut nbytes_buffer)?;
22        let columns_nbytes = i64::from_le_bytes(nbytes_buffer) as usize;
23        let index_shape = index_nbytes / INDEX_NBYTES;
24        let columns_shape = columns_nbytes / COLUMNS_NBYTES;
25        let values_nbytes = to_nbytes::<T>(index_shape * columns_shape);
26        let mut index_buffer = vec![0u8; index_nbytes];
27        let mut columns_buffer = vec![0u8; columns_nbytes];
28        let mut values_buffer = vec![0u8; values_nbytes];
29        reader.read_exact(&mut index_buffer)?;
30        reader.read_exact(&mut columns_buffer)?;
31        reader.read_exact(&mut values_buffer)?;
32
33        let (index, columns, values) = unsafe {
34            (
35                from_bytes(index_buffer),
36                from_bytes(columns_buffer),
37                from_bytes(values_buffer),
38            )
39        };
40        DataFrame::from_vec(index, columns, values)
41    }
42
43    pub fn write(&self, writer: &mut impl Write) -> Result<()> {
44        let index = self.index();
45        let columns = self.columns();
46        let index_nbytes = to_nbytes::<IndexDtype>(index.len()) as i64;
47        let columns_nbytes = to_nbytes::<ColumnsDtype>(columns.len()) as i64;
48        writer.write_all(&index_nbytes.to_le_bytes())?;
49        writer.write_all(&columns_nbytes.to_le_bytes())?;
50        unsafe {
51            writer.write_all(to_bytes(as_data_slice_or_err!(index)))?;
52            writer.write_all(to_bytes(as_data_slice_or_err!(columns)))?;
53            writer.write_all(to_bytes(as_data_slice_or_err!(self.values())))?;
54        }
55        Ok(())
56    }
57}
58
59#[cfg(test)]
60pub(super) mod tests {
61    use super::*;
62    use crate::df::io::bytes::tests::get_test_df;
63    use std::fs::File;
64    use tempfile::tempdir;
65
66    #[test]
67    fn test_io() {
68        let df = get_test_df();
69        let dir = tempdir().unwrap();
70        let file_path = dir.path().join("test.cfdf");
71        let mut file = File::create(&file_path).unwrap();
72        df.write(&mut file).unwrap();
73        let mut file = File::open(&file_path).unwrap();
74        let loaded = DataFrame::<f32>::read(&mut file).unwrap();
75        assert_eq!(df.index(), loaded.index());
76        assert_eq!(df.columns(), loaded.columns());
77        assert_eq!(df.values(), loaded.values());
78        drop(file);
79        dir.close().unwrap();
80    }
81}