dtrees_rs/reader/
data_reader.rs

1use super::{DataFormat, DataReaderError};
2use crate::bitsets::{BitCollection, Bitset, BitsetInit};
3use crate::cover::Cover;
4use std::collections::HashSet;
5use std::fs::File;
6use std::io::{BufRead, BufReader};
7use std::path::Path;
8
9/// Reads a text dataset of binary features into a [`Cover`].
10///
11/// The default format is one instance per line, whitespace separated, the
12/// label in column 0, `#` starting a comment, and no header. Feature values
13/// must be 0 or 1, and labels non-negative integers.
14pub struct DataReader {
15    format: DataFormat,
16    has_headers: bool,
17    comment_char: Option<char>,
18    label_column: Option<usize>,
19}
20
21impl Default for DataReader {
22    fn default() -> Self {
23        Self {
24            format: DataFormat::Space,
25            has_headers: false,
26            comment_char: Some('#'),
27            label_column: Some(0),
28        }
29    }
30}
31
32impl DataReader {
33    /// A reader with the default format.
34    pub fn new() -> Self {
35        Self::default()
36    }
37
38    /// Sets the column delimiter.
39    pub fn with_format(mut self, format: DataFormat) -> Self {
40        self.format = format;
41        self
42    }
43
44    /// Whether the first data line is a header to skip.
45    pub fn with_headers(mut self, has_headers: bool) -> Self {
46        self.has_headers = has_headers;
47        self
48    }
49
50    /// Lines starting with this character are ignored.
51    pub fn with_comment_char(mut self, comment_char: Option<char>) -> Self {
52        self.comment_char = comment_char;
53        self
54    }
55
56    /// Which column holds the label; `None` for unlabelled data.
57    pub fn with_label_column(mut self, label_column: Option<usize>) -> Self {
58        self.label_column = label_column;
59        self
60    }
61
62    /// Picks the delimiter from the file extension.
63    pub fn auto_detect_format(mut self, path: &Path) -> Self {
64        self.format = DataFormat::from_extension(path);
65        self
66    }
67
68    /// Reads a dataset file.
69    pub fn read_file(&self, path: &Path) -> Result<Cover, DataReaderError> {
70        let file = File::open(path)?;
71        let reader = BufReader::new(file);
72
73        let mut num_cols = 0;
74
75        let delimiter = self.format.delimiter();
76        let mut row_idx = 0;
77
78        let mut attributes: Vec<Bitset> = vec![];
79        let mut target = vec![];
80
81        for (i, line_result) in reader.lines().enumerate() {
82            let line = line_result?;
83
84            if line.trim().is_empty() {
85                continue;
86            }
87
88            if let Some(comment) = self.comment_char {
89                if line.trim().starts_with(comment) {
90                    continue;
91                }
92            }
93
94            if i == 0 && self.has_headers {
95                continue;
96            }
97
98            let tokens: Vec<&str> = line.split(delimiter).map(|s| s.trim()).collect();
99
100            if i <= 1 && num_cols == 0 {
101                let actual_cols = if self.label_column.is_some() {
102                    tokens.len() - 1
103                } else {
104                    tokens.len()
105                };
106
107                num_cols = num_cols.max(actual_cols);
108                attributes = vec![Bitset::new(BitsetInit::Empty(64)); num_cols]
109            }
110
111            for (col_idx, &token) in tokens.iter().enumerate() {
112                if Some(col_idx) == self.label_column {
113                    match token.parse::<usize>() {
114                        Ok(val) => target.push(val),
115                        Err(_) => {
116                            return Err(DataReaderError::Parse(format!(
117                                "Parse error at line {}, column {}: {}",
118                                i + 1,
119                                col_idx + 1,
120                                token
121                            )));
122                        }
123                    }
124                    continue;
125                }
126
127                let effective_col = if col_idx > self.label_column.unwrap_or(usize::MAX) {
128                    col_idx - 1
129                } else {
130                    col_idx
131                };
132
133                let capacity = attributes[effective_col].capacity();
134
135                if row_idx >= capacity {
136                    attributes[effective_col].resize(capacity * 2);
137                }
138
139                match token.parse::<usize>() {
140                    Ok(0) => {}
141
142                    Ok(1) => {
143                        attributes[effective_col].set(row_idx);
144                    }
145
146                    Ok(_) => {
147                        return Err(DataReaderError::Format(format!(
148                            "Non-binary value at line {}, column {}",
149                            i + 1,
150                            col_idx + 1
151                        )))
152                    }
153
154                    Err(_) => {
155                        return Err(DataReaderError::Parse(format!(
156                            "Parse error at line {}, column {}: {}",
157                            i + 1,
158                            col_idx + 1,
159                            token
160                        )));
161                    }
162                }
163            }
164
165            row_idx += 1;
166        }
167
168        for bitset in attributes.iter_mut() {
169            bitset.resize(row_idx);
170        }
171
172        let unique_targets = target.iter().copied().collect::<HashSet<usize>>().len();
173
174        let mut targets = vec![Bitset::new(BitsetInit::Empty(row_idx)); unique_targets];
175
176        for (tid, &t) in target.iter().enumerate() {
177            targets[t].set(tid);
178        }
179        Ok(Cover::new(attributes, targets, row_idx))
180    }
181}
182
183#[cfg(test)]
184mod data_reader_test {
185    use crate::reader::data_reader::DataReader;
186    use std::path::Path;
187
188    #[test]
189    fn load_small() {
190        let reader = DataReader::default();
191        let path = Path::new("test_data/anneal.txt");
192        let cover_result = reader.read_file(path);
193        let cover = cover_result.expect("the test data is readable");
194
195        assert_eq!(cover.num_labels, 2);
196        assert_eq!(cover.num_attributes, 93);
197        assert_eq!(cover.count(), 812);
198    }
199}
dtrees_rs/reader/data_reader.rs

dtrees_rs/reader/
data_reader.rs
