burn_dataset/dataset/
in_memory.rs

1use std::{
2    fs::File,
3    io::{BufRead, BufReader},
4    path::Path,
5};
6
7use serde::de::DeserializeOwned;
8
9use crate::{Dataset, DatasetError};
10
11/// Dataset where all items are stored in ram.
12pub struct InMemDataset<I> {
13    items: Vec<I>,
14}
15
16impl<I> InMemDataset<I> {
17    /// Creates a new in memory dataset from the given items.
18    pub fn new(items: Vec<I>) -> Self {
19        InMemDataset { items }
20    }
21}
22
23impl<I> Dataset<I> for InMemDataset<I>
24where
25    I: Clone + Send + Sync,
26{
27    fn get(&self, index: usize) -> Result<I, DatasetError> {
28        match self.items.get(index) {
29            Some(item) => Ok(item.clone()),
30            None => panic!(
31                "Index out of bounds for InMemDataset: {index} >= {}",
32                self.items.len()
33            ),
34        }
35    }
36    fn len(&self) -> usize {
37        self.items.len()
38    }
39}
40
41impl<I> InMemDataset<I>
42where
43    I: Clone + DeserializeOwned,
44{
45    /// Create from a dataset. All items are loaded in memory.
46    pub fn from_dataset<E>(dataset: &impl Dataset<I, E>) -> Self
47    where
48        E: std::error::Error + Send + Sync + 'static,
49    {
50        let items: Vec<I> = dataset.iter().map(Result::unwrap).collect();
51        Self::new(items)
52    }
53
54    /// Create from a json rows file (one json per line).
55    ///
56    /// [Supported field types](https://docs.rs/serde_json/latest/serde_json/value/enum.Value.html)
57    pub fn from_json_rows<P: AsRef<Path>>(path: P) -> Result<Self, std::io::Error> {
58        let file = File::open(path)?;
59        let reader = BufReader::new(file);
60        let mut items = Vec::new();
61
62        for line in reader.lines() {
63            let item = serde_json::from_str(line.unwrap().as_str()).unwrap();
64            items.push(item);
65        }
66
67        let dataset = Self::new(items);
68
69        Ok(dataset)
70    }
71
72    /// Create from a csv file.
73    ///
74    /// The provided `csv::ReaderBuilder` can be configured to fit your csv format.
75    ///
76    /// The supported field types are: String, integer, float, and bool.
77    ///
78    /// See:
79    /// - [Reading with Serde](https://docs.rs/csv/latest/csv/tutorial/index.html#reading-with-serde)
80    /// - [Delimiters, quotes and variable length records](https://docs.rs/csv/latest/csv/tutorial/index.html#delimiters-quotes-and-variable-length-records)
81    pub fn from_csv<P: AsRef<Path>>(
82        path: P,
83        builder: &csv::ReaderBuilder,
84    ) -> Result<Self, std::io::Error> {
85        let mut rdr = builder.from_path(path)?;
86
87        let mut items = Vec::new();
88
89        for result in rdr.deserialize() {
90            let item: I = result?;
91            items.push(item);
92        }
93
94        let dataset = Self::new(items);
95
96        Ok(dataset)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102
103    use super::*;
104    #[cfg(feature = "sqlite")]
105    use crate::SqliteDataset;
106    use crate::test_data;
107
108    #[cfg(feature = "sqlite")]
109    use rstest::{fixture, rstest};
110    use serde::{Deserialize, Serialize};
111
112    #[cfg(feature = "sqlite")]
113    const DB_FILE: &str = "tests/data/sqlite-dataset.db";
114    const JSON_FILE: &str = "tests/data/dataset.json";
115    const CSV_FILE: &str = "tests/data/dataset.csv";
116    const CSV_FMT_FILE: &str = "tests/data/dataset-fmt.csv";
117
118    #[cfg(feature = "sqlite")]
119    type SqlDs = SqliteDataset<Sample>;
120
121    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
122    pub struct Sample {
123        column_str: String,
124        column_bytes: Vec<u8>,
125        column_int: i64,
126        column_bool: bool,
127        column_float: f64,
128    }
129
130    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
131    pub struct SampleCsv {
132        column_str: String,
133        column_int: i64,
134        column_bool: bool,
135        column_float: f64,
136    }
137
138    #[cfg(feature = "sqlite")]
139    #[fixture]
140    fn train_dataset() -> SqlDs {
141        SqliteDataset::from_db_file(DB_FILE, "train").unwrap()
142    }
143
144    #[cfg(feature = "sqlite")]
145    #[rstest]
146    pub fn from_dataset(train_dataset: SqlDs) {
147        let dataset = InMemDataset::from_dataset(&train_dataset);
148
149        let record_index: usize = 0;
150
151        assert_eq!(dataset.get(record_index).unwrap().column_str, "HI1");
152    }
153
154    #[cfg(feature = "sqlite")]
155    #[rstest]
156    #[should_panic(expected = "Index out of bounds")]
157    pub fn from_dataset_out_of_bounds(train_dataset: SqlDs) {
158        let non_existing_record_index: usize = 10;
159        train_dataset.get(non_existing_record_index).unwrap();
160    }
161
162    #[test]
163    pub fn from_json_rows() {
164        let dataset = InMemDataset::<Sample>::from_json_rows(JSON_FILE).unwrap();
165
166        let record_index: usize = 1;
167
168        assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2");
169        assert!(!dataset.get(record_index).unwrap().column_bool);
170    }
171
172    #[test]
173    #[should_panic(expected = "Index out of bounds")]
174    pub fn from_json_rows_out_of_bounds() {
175        let dataset = InMemDataset::<Sample>::from_json_rows(JSON_FILE).unwrap();
176        dataset.get(10).unwrap();
177    }
178
179    #[test]
180    pub fn from_csv_rows() {
181        let rdr = csv::ReaderBuilder::new();
182        let dataset = InMemDataset::<SampleCsv>::from_csv(CSV_FILE, &rdr).unwrap();
183
184        let record_index: usize = 1;
185
186        assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2");
187        assert_eq!(dataset.get(record_index).unwrap().column_int, 1);
188        assert!(!dataset.get(record_index).unwrap().column_bool);
189        assert_eq!(dataset.get(record_index).unwrap().column_float, 1.0);
190    }
191
192    #[test]
193    pub fn from_csv_rows_fmt() {
194        let mut rdr = csv::ReaderBuilder::new();
195        let rdr = rdr.delimiter(b' ').has_headers(false);
196        let dataset = InMemDataset::<SampleCsv>::from_csv(CSV_FMT_FILE, rdr).unwrap();
197
198        let record_index: usize = 1;
199
200        assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2");
201        assert_eq!(dataset.get(record_index).unwrap().column_int, 1);
202        assert!(!dataset.get(record_index).unwrap().column_bool);
203        assert_eq!(dataset.get(record_index).unwrap().column_float, 1.0);
204    }
205
206    #[test]
207    pub fn given_in_memory_dataset_when_iterate_should_iterate_though_all_items() {
208        let items_original = test_data::string_items();
209        let dataset = InMemDataset::new(items_original.clone());
210
211        let items: Vec<String> = dataset.iter().map(Result::unwrap).collect();
212
213        assert_eq!(items_original, items);
214    }
215}
burn_dataset/dataset/in_memory.rs

burn_dataset/dataset/
in_memory.rs
