legume_numeric/matrix/
tensor_io.rs

1use crate::matrix::common_io::{read_lines_of_types, write_lines, Delimiter};
2use crate::matrix::parquet::*;
3use crate::matrix::traits::*;
4
5use candle_core::{Device, Tensor};
6
7impl IoOps for Tensor {
8    type Scalar = f32;
9    type Mat = Self;
10
11    fn read_data(
12        file_path: &str,
13        delim: impl Into<Delimiter>,
14        skip: Option<usize>,
15        row_name_index: Option<usize>,
16        column_indices: Option<&[usize]>,
17        column_names: Option<&[Box<str>]>,
18    ) -> anyhow::Result<MatWithNames<Self::Mat>> {
19        let (rows, cols, data) = Self::read_data_vec_with_indices_names(
20            file_path,
21            delim,
22            skip,
23            row_name_index,
24            column_indices,
25            column_names,
26        )?;
27
28        let nrows = rows.len();
29        let ncols = cols.len();
30        Ok(MatWithNames {
31            rows,
32            cols,
33            mat: Tensor::from_vec(data, (nrows, ncols), &Device::Cpu)?,
34        })
35    }
36
37    fn read_file_delim(
38        tsv_file: &str,
39        delim: impl Into<Delimiter>,
40        skip: Option<usize>,
41    ) -> anyhow::Result<Self::Mat> {
42        let hdr_line = match skip {
43            Some(skip) => skip as i64,
44            None => -1, // no skipping
45        };
46
47        let data = read_lines_of_types::<f32>(tsv_file, delim, hdr_line)?.lines;
48
49        if data.is_empty() {
50            return Err(anyhow::anyhow!("No data in file"));
51        }
52
53        let ncols = data[0].len();
54        let nrows = data.len();
55        let data = data.into_iter().flatten().collect::<Vec<_>>();
56
57        Ok(Tensor::from_vec(data, (nrows, ncols), &Device::Cpu)?)
58    }
59
60    fn write_file_delim(&self, file: &str, delim: &str) -> anyhow::Result<()> {
61        let dims = self.dims();
62
63        if dims.len() != 2 {
64            return Err(anyhow::anyhow!("Expected 2 dimensions, got {}", dims.len()));
65        }
66
67        let lines: Vec<Box<str>> = (0..dims[0])
68            .map(|i| {
69                let row = self.narrow(0, i, 1).expect("failed to narrow in");
70                let flatten_row = row.flatten_to(1).expect("flatten");
71                let row_vec = flatten_row.to_vec1::<f32>().expect("to_vec1");
72                row_vec
73                    .iter()
74                    .map(|&x| format!("{}", x))
75                    .collect::<Vec<_>>()
76                    .join(delim)
77                    .into_boxed_str()
78            })
79            .collect();
80
81        write_lines(&lines, file)?;
82
83        Ok(())
84    }
85
86    fn to_parquet_with_names(
87        &self,
88        file_path: &str,
89        row_names: (Option<&[Box<str>]>, Option<&str>),
90        column_names: Option<&[Box<str>]>,
91    ) -> anyhow::Result<()> {
92        let dims = self.dims();
93
94        if dims.len() != 2 {
95            return Err(anyhow::anyhow!("expected 2 dimensions, got {}", dims.len()));
96        }
97
98        let (nrows, ncols) = (dims[0], dims[1]);
99
100        let (row_names_slice, row_column_name) = row_names;
101
102        let writer = ParquetWriter::new(
103            file_path,
104            (nrows, ncols),
105            (row_names_slice, column_names),
106            None,
107            row_column_name,
108        )?;
109        let row_names = writer.row_names_vec();
110
111        if row_names.len() != nrows {
112            return Err(anyhow::anyhow!("row names don't match"));
113        }
114
115        let mut writer = writer.get_writer()?;
116        let mut row_group_writer = writer.next_row_group()?;
117        parquet_add_bytearray(&mut row_group_writer, row_names)?;
118
119        let tensor = self.to_dtype(candle_core::DType::F32)?;
120
121        for j in 0..ncols {
122            let data_j = tensor.narrow(1, j, 1)?.flatten_all()?.to_vec1::<f32>()?;
123            parquet_add_numeric_column(&mut row_group_writer, &data_j)?;
124        }
125        row_group_writer.close()?;
126        writer.close()?;
127        Ok(())
128    }
129
130    fn from_parquet_with_indices_names(
131        file_path: &str,
132        row_name_index: Option<usize>,
133        column_indices: Option<&[usize]>,
134        column_names: Option<&[Box<str>]>,
135    ) -> anyhow::Result<MatWithNames<Self>> {
136        let parquet = ParquetReader::new(file_path, row_name_index, column_indices, column_names)?;
137
138        let nrows = parquet.row_names.len();
139        let ncols = parquet.column_names.len();
140
141        let data: Vec<f32> = parquet
142            .row_major_data
143            .into_iter()
144            .map(|x| x as f32)
145            .collect();
146
147        Ok(MatWithNames {
148            rows: parquet.row_names,
149            cols: parquet.column_names,
150            mat: Tensor::from_vec(data, (nrows, ncols), &Device::Cpu)?,
151        })
152    }
153}
legume_numeric/matrix/tensor_io.rs

legume_numeric/matrix/
tensor_io.rs
