1use crate::data::{AnyDataArray, FieldData};
2use crate::data::traits::DataObject;
3
4#[derive(Debug, Clone, Default)]
22pub struct Table {
23 columns: Vec<AnyDataArray>,
24 field_data: FieldData,
25}
26
27impl Table {
28 pub fn new() -> Self {
29 Self::default()
30 }
31
32 pub fn add_column(&mut self, column: AnyDataArray) {
34 if !self.columns.is_empty() {
35 assert_eq!(
36 column.num_tuples(),
37 self.num_rows(),
38 "column '{}' has {} rows, expected {}",
39 column.name(),
40 column.num_tuples(),
41 self.num_rows()
42 );
43 }
44 self.columns.push(column);
45 }
46
47 pub fn num_rows(&self) -> usize {
49 self.columns.first().map(|c| c.num_tuples()).unwrap_or(0)
50 }
51
52 pub fn num_columns(&self) -> usize {
54 self.columns.len()
55 }
56
57 pub fn column(&self, idx: usize) -> Option<&AnyDataArray> {
59 self.columns.get(idx)
60 }
61
62 pub fn column_by_name(&self, name: &str) -> Option<&AnyDataArray> {
64 self.columns.iter().find(|c| c.name() == name)
65 }
66
67 pub fn column_names(&self) -> Vec<&str> {
69 self.columns.iter().map(|c| c.name()).collect()
70 }
71
72 pub fn columns(&self) -> &[AnyDataArray] {
74 &self.columns
75 }
76
77 pub fn value_f64(&self, row: usize, column_name: &str) -> Option<f64> {
79 let col = self.column_by_name(column_name)?;
80 if row >= col.num_tuples() { return None; }
81 let mut buf = [0.0f64];
82 col.tuple_as_f64(row, &mut buf);
83 Some(buf[0])
84 }
85
86 pub fn filter_rows(&self, column_name: &str, predicate: impl Fn(f64) -> bool) -> Vec<usize> {
88 let Some(col) = self.column_by_name(column_name) else {
89 return Vec::new();
90 };
91 let mut result = Vec::new();
92 let mut buf = [0.0f64];
93 for i in 0..col.num_tuples() {
94 col.tuple_as_f64(i, &mut buf);
95 if predicate(buf[0]) {
96 result.push(i);
97 }
98 }
99 result
100 }
101
102 pub fn sort_by_column(&self, column_name: &str) -> Vec<usize> {
104 let Some(col) = self.column_by_name(column_name) else {
105 return Vec::new();
106 };
107 let mut indices: Vec<usize> = (0..col.num_tuples()).collect();
108 let mut buf = [0.0f64];
109 let mut values: Vec<f64> = Vec::with_capacity(col.num_tuples());
110 for i in 0..col.num_tuples() {
111 col.tuple_as_f64(i, &mut buf);
112 values.push(buf[0]);
113 }
114 indices.sort_by(|&a, &b| values[a].partial_cmp(&values[b]).unwrap_or(std::cmp::Ordering::Equal));
115 indices
116 }
117
118 pub fn select_rows(&self, indices: &[usize]) -> Table {
120 let mut result = Table::new();
121 for col in &self.columns {
122 let nc = col.num_components();
123 let mut data = Vec::with_capacity(indices.len() * nc);
124 let mut buf = vec![0.0f64; nc];
125 for &idx in indices {
126 col.tuple_as_f64(idx, &mut buf);
127 data.extend_from_slice(&buf);
128 }
129 result.columns.push(AnyDataArray::F64(
130 crate::data::DataArray::from_vec(col.name(), data, nc),
131 ));
132 }
133 result
134 }
135
136 pub fn remove_column(&mut self, name: &str) -> Option<AnyDataArray> {
138 if let Some(idx) = self.columns.iter().position(|c| c.name() == name) {
139 Some(self.columns.remove(idx))
140 } else {
141 None
142 }
143 }
144
145 pub fn with_column(mut self, column: AnyDataArray) -> Self {
147 self.add_column(column);
148 self
149 }
150
151 pub fn to_csv<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
165 let names: Vec<&str> = self.columns.iter().map(|c| c.name()).collect();
167 writeln!(w, "{}", names.join(","))?;
168
169 for row in 0..self.num_rows() {
171 let mut vals = Vec::with_capacity(self.columns.len());
172 for col in &self.columns {
173 let nc = col.num_components();
174 let mut buf = vec![0.0f64; nc];
175 col.tuple_as_f64(row, &mut buf);
176 if nc == 1 {
177 vals.push(format!("{}", buf[0]));
178 } else {
179 let components: Vec<String> = buf.iter().map(|v| format!("{v}")).collect();
180 vals.push(components.join(";"));
181 }
182 }
183 writeln!(w, "{}", vals.join(","))?;
184 }
185 Ok(())
186 }
187
188 pub fn from_csv<R: std::io::BufRead>(r: R) -> Result<Self, String> {
191 let mut lines = r.lines();
192
193 let header = lines.next()
195 .ok_or("empty CSV")?
196 .map_err(|e| e.to_string())?;
197 let col_names: Vec<&str> = header.trim().split(',').collect();
198 let ncols = col_names.len();
199 let mut columns: Vec<Vec<f64>> = vec![Vec::new(); ncols];
200
201 for line_result in lines {
203 let line = line_result.map_err(|e| e.to_string())?;
204 let trimmed = line.trim();
205 if trimmed.is_empty() { continue; }
206 let values: Vec<&str> = trimmed.split(',').collect();
207 for (i, val) in values.iter().enumerate().take(ncols) {
208 let v: f64 = val.trim().parse().unwrap_or(f64::NAN);
209 columns[i].push(v);
210 }
211 }
212
213 let mut table = Table::new();
214 for (i, name) in col_names.iter().enumerate() {
215 let arr = crate::data::DataArray::from_vec(name.trim(), columns[i].clone(), 1);
216 table.add_column(AnyDataArray::F64(arr));
217 }
218 Ok(table)
219 }
220
221 pub fn write_csv_file(&self, path: &std::path::Path) -> std::io::Result<()> {
223 let file = std::fs::File::create(path)?;
224 let mut w = std::io::BufWriter::new(file);
225 self.to_csv(&mut w)
226 }
227
228 pub fn read_csv_file(path: &std::path::Path) -> Result<Self, String> {
230 let file = std::fs::File::open(path).map_err(|e| e.to_string())?;
231 let reader = std::io::BufReader::new(file);
232 Self::from_csv(reader)
233 }
234}
235
236impl DataObject for Table {
237 fn field_data(&self) -> &FieldData {
238 &self.field_data
239 }
240
241 fn field_data_mut(&mut self) -> &mut FieldData {
242 &mut self.field_data
243 }
244}
245
246impl std::fmt::Display for Table {
247 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
248 write!(f, "Table: {} rows, {} columns [{}]",
249 self.num_rows(), self.num_columns(),
250 self.column_names().join(", "))
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use crate::data::DataArray;
258
259 #[test]
260 fn basic_table() {
261 let mut table = Table::new();
262 table.add_column(AnyDataArray::F64(DataArray::from_vec("x", vec![1.0, 2.0, 3.0], 1)));
263 table.add_column(AnyDataArray::F64(DataArray::from_vec("y", vec![4.0, 5.0, 6.0], 1)));
264
265 assert_eq!(table.num_rows(), 3);
266 assert_eq!(table.num_columns(), 2);
267 assert_eq!(table.column_names(), vec!["x", "y"]);
268 }
269
270 #[test]
271 fn column_lookup() {
272 let mut table = Table::new();
273 table.add_column(AnyDataArray::I32(DataArray::from_vec("ids", vec![10, 20, 30], 1)));
274 table.add_column(AnyDataArray::F64(DataArray::from_vec("values", vec![1.1, 2.2, 3.3], 1)));
275
276 let col = table.column_by_name("values").unwrap();
277 assert_eq!(col.num_tuples(), 3);
278 let mut buf = [0.0f64];
279 col.tuple_as_f64(1, &mut buf);
280 assert!((buf[0] - 2.2).abs() < 1e-10);
281 }
282
283 #[test]
284 fn value_f64_access() {
285 let table = Table::new()
286 .with_column(AnyDataArray::F64(DataArray::from_vec("x", vec![10.0, 20.0, 30.0], 1)));
287 assert_eq!(table.value_f64(1, "x"), Some(20.0));
288 assert_eq!(table.value_f64(5, "x"), None);
289 assert_eq!(table.value_f64(0, "missing"), None);
290 }
291
292 #[test]
293 fn filter_rows() {
294 let table = Table::new()
295 .with_column(AnyDataArray::F64(DataArray::from_vec("val", vec![1.0, 5.0, 2.0, 8.0, 3.0], 1)));
296 let big = table.filter_rows("val", |v| v > 3.0);
297 assert_eq!(big, vec![1, 3]);
298 }
299
300 #[test]
301 fn sort_by_column() {
302 let table = Table::new()
303 .with_column(AnyDataArray::F64(DataArray::from_vec("val", vec![3.0, 1.0, 4.0, 1.0, 5.0], 1)));
304 let sorted = table.sort_by_column("val");
305 assert_eq!(sorted[0], 1); assert_eq!(sorted[1], 3); }
308
309 #[test]
310 fn select_rows() {
311 let mut table = Table::new();
312 table.add_column(AnyDataArray::F64(DataArray::from_vec("x", vec![10.0, 20.0, 30.0, 40.0], 1)));
313 table.add_column(AnyDataArray::F64(DataArray::from_vec("y", vec![1.0, 2.0, 3.0, 4.0], 1)));
314 let sub = table.select_rows(&[0, 2]);
315 assert_eq!(sub.num_rows(), 2);
316 assert_eq!(sub.value_f64(0, "x"), Some(10.0));
317 assert_eq!(sub.value_f64(1, "x"), Some(30.0));
318 }
319
320 #[test]
321 fn remove_column() {
322 let mut table = Table::new();
323 table.add_column(AnyDataArray::F64(DataArray::from_vec("a", vec![1.0, 2.0], 1)));
324 table.add_column(AnyDataArray::F64(DataArray::from_vec("b", vec![3.0, 4.0], 1)));
325 assert_eq!(table.num_columns(), 2);
326 table.remove_column("a");
327 assert_eq!(table.num_columns(), 1);
328 assert!(table.column_by_name("b").is_some());
329 }
330
331 #[test]
332 fn multicomponent_columns() {
333 let mut table = Table::new();
334 table.add_column(AnyDataArray::F64(DataArray::from_vec(
335 "position",
336 vec![0.0, 0.0, 0.0, 1.0, 2.0, 3.0],
337 3,
338 )));
339 assert_eq!(table.num_rows(), 2);
340 assert_eq!(table.column(0).unwrap().num_components(), 3);
341 }
342
343 #[test]
344 fn csv_roundtrip() {
345 let table = Table::new()
346 .with_column(AnyDataArray::F64(DataArray::from_vec("x", vec![1.0, 2.0, 3.0], 1)))
347 .with_column(AnyDataArray::F64(DataArray::from_vec("y", vec![4.0, 5.0, 6.0], 1)));
348
349 let mut buf = Vec::new();
350 table.to_csv(&mut buf).unwrap();
351 let csv = String::from_utf8(buf.clone()).unwrap();
352 assert!(csv.starts_with("x,y\n"));
353
354 let loaded = Table::from_csv(std::io::BufReader::new(&buf[..])).unwrap();
355 assert_eq!(loaded.num_rows(), 3);
356 assert_eq!(loaded.num_columns(), 2);
357 assert_eq!(loaded.value_f64(0, "x"), Some(1.0));
358 assert_eq!(loaded.value_f64(2, "y"), Some(6.0));
359 }
360}