Skip to main content

vtk_pure_rs/data/
table.rs

1use crate::data::{AnyDataArray, FieldData};
2use crate::data::traits::DataObject;
3
4/// Columnar data table (rows × named columns).
5///
6/// Analogous to VTK's `vtkTable`. Each column is an `AnyDataArray` where
7/// the number of tuples is the number of rows.
8///
9/// # Examples
10///
11/// ```
12/// use crate::data::{Table, AnyDataArray, DataArray};
13///
14/// let table = Table::new()
15///     .with_column(AnyDataArray::F64(DataArray::from_vec("x", vec![1.0, 2.0, 3.0], 1)))
16///     .with_column(AnyDataArray::F64(DataArray::from_vec("y", vec![4.0, 5.0, 6.0], 1)));
17/// assert_eq!(table.num_rows(), 3);
18/// assert_eq!(table.num_columns(), 2);
19/// assert_eq!(table.value_f64(1, "x"), Some(2.0));
20/// ```
21#[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    /// Add a column. All columns must have the same number of tuples (rows).
33    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    /// Number of rows.
48    pub fn num_rows(&self) -> usize {
49        self.columns.first().map(|c| c.num_tuples()).unwrap_or(0)
50    }
51
52    /// Number of columns.
53    pub fn num_columns(&self) -> usize {
54        self.columns.len()
55    }
56
57    /// Get column by index.
58    pub fn column(&self, idx: usize) -> Option<&AnyDataArray> {
59        self.columns.get(idx)
60    }
61
62    /// Get column by name.
63    pub fn column_by_name(&self, name: &str) -> Option<&AnyDataArray> {
64        self.columns.iter().find(|c| c.name() == name)
65    }
66
67    /// Get column names.
68    pub fn column_names(&self) -> Vec<&str> {
69        self.columns.iter().map(|c| c.name()).collect()
70    }
71
72    /// Iterate over columns.
73    pub fn columns(&self) -> &[AnyDataArray] {
74        &self.columns
75    }
76
77    /// Get a single scalar value at (row, column_name).
78    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    /// Get row indices where a scalar column satisfies a predicate.
87    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    /// Get row indices sorted by a column's values (ascending).
103    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    /// Extract a subset of rows by indices, returning a new Table.
119    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    /// Remove a column by name. Returns the removed column if found.
137    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    /// Builder: add a column.
146    pub fn with_column(mut self, column: AnyDataArray) -> Self {
147        self.add_column(column);
148        self
149    }
150
151    /// Write the table as CSV to a writer.
152    ///
153    /// ```
154    /// use crate::data::{Table, AnyDataArray, DataArray};
155    ///
156    /// let table = Table::new()
157    ///     .with_column(AnyDataArray::F64(DataArray::from_vec("x", vec![1.0, 2.0], 1)))
158    ///     .with_column(AnyDataArray::F64(DataArray::from_vec("y", vec![3.0, 4.0], 1)));
159    /// let mut buf = Vec::new();
160    /// table.to_csv(&mut buf).unwrap();
161    /// let csv = String::from_utf8(buf).unwrap();
162    /// assert!(csv.contains("x,y"));
163    /// ```
164    pub fn to_csv<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
165        // Header
166        let names: Vec<&str> = self.columns.iter().map(|c| c.name()).collect();
167        writeln!(w, "{}", names.join(","))?;
168
169        // Rows
170        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    /// Read a table from CSV. First line is header (column names).
189    /// All values are parsed as f64.
190    pub fn from_csv<R: std::io::BufRead>(r: R) -> Result<Self, String> {
191        let mut lines = r.lines();
192
193        // Header
194        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        // Data rows
202        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    /// Write CSV to a file path.
222    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    /// Read CSV from a file path.
229    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); // val=1.0
306        assert_eq!(sorted[1], 3); // val=1.0
307    }
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}