ggplot_rs/data/
dataframe.rs

1use indexmap::IndexMap;
2
3use super::{GroupKey, Value};
4
5/// Internal columnar DataFrame for data storage and manipulation.
6#[derive(Clone, Debug)]
7pub struct DataFrame {
8    columns: IndexMap<String, Vec<Value>>,
9    nrows: usize,
10    /// Problems found while assembling the frame (e.g. mismatched column
11    /// lengths). Reported by [`validate`](Self::validate), and turned into a
12    /// `GGError::ValidationError` when a plot using this frame is built.
13    issues: Vec<String>,
14}
15
16impl DataFrame {
17    /// Create an empty DataFrame.
18    pub fn new() -> Self {
19        DataFrame {
20            columns: IndexMap::new(),
21            nrows: 0,
22            issues: Vec::new(),
23        }
24    }
25
26    /// Get a column by name.
27    pub fn column(&self, name: &str) -> Option<&[Value]> {
28        self.columns.get(name).map(|v| v.as_slice())
29    }
30
31    /// Get number of rows.
32    pub fn nrows(&self) -> usize {
33        self.nrows
34    }
35
36    /// Get number of columns.
37    pub fn ncols(&self) -> usize {
38        self.columns.len()
39    }
40
41    /// Get column names.
42    pub fn column_names(&self) -> Vec<&str> {
43        self.columns.keys().map(|s| s.as_str()).collect()
44    }
45
46    /// Check if a column exists.
47    pub fn has_column(&self, name: &str) -> bool {
48        self.columns.contains_key(name)
49    }
50
51    /// Add (or replace) a column.
52    ///
53    /// Never panics. If the length doesn't match the existing rows, the shorter
54    /// side is padded with [`Value::Na`] and the mismatch is recorded as an
55    /// [issue](Self::issues): building or rendering a plot from this frame then
56    /// fails with `GGError::ValidationError`. Use
57    /// [`try_add_column`](Self::try_add_column) to reject the column instead.
58    pub fn add_column(&mut self, name: String, mut values: Vec<Value>) {
59        if self.columns.is_empty() {
60            self.nrows = values.len();
61        } else if values.len() != self.nrows {
62            self.issues
63                .push(Self::mismatch_message(&name, values.len(), self.nrows));
64            if values.len() < self.nrows {
65                values.resize(self.nrows, Value::Na);
66            } else {
67                let n = values.len();
68                for col in self.columns.values_mut() {
69                    col.resize(n, Value::Na);
70                }
71                self.nrows = n;
72            }
73        }
74        self.columns.insert(name, values);
75    }
76
77    /// Add (or replace) a column, rejecting a length mismatch with a
78    /// `GGError::ValidationError` (the frame is left unchanged).
79    pub fn try_add_column(
80        &mut self,
81        name: String,
82        values: Vec<Value>,
83    ) -> Result<(), crate::plot::GGError> {
84        if !self.columns.is_empty() && values.len() != self.nrows {
85            return Err(crate::plot::GGError::ValidationError(
86                Self::mismatch_message(&name, values.len(), self.nrows),
87            ));
88        }
89        self.add_column(name, values);
90        Ok(())
91    }
92
93    fn mismatch_message(name: &str, len: usize, nrows: usize) -> String {
94        format!("column '{name}' has {len} values but the data has {nrows} rows")
95    }
96
97    /// Problems recorded while assembling this frame (empty when well-formed).
98    pub fn issues(&self) -> &[String] {
99        &self.issues
100    }
101
102    /// `Ok` when the frame is well-formed, else a `GGError::ValidationError`
103    /// describing every recorded [issue](Self::issues).
104    pub fn validate(&self) -> Result<(), crate::plot::GGError> {
105        if self.issues.is_empty() {
106            Ok(())
107        } else {
108            Err(crate::plot::GGError::ValidationError(
109                self.issues.join("; "),
110            ))
111        }
112    }
113
114    /// Get a mutable reference to a column.
115    pub fn column_mut(&mut self, name: &str) -> Option<&mut Vec<Value>> {
116        self.columns.get_mut(name)
117    }
118
119    /// Remove a column, returning its values if it existed. Row count
120    /// (`nrows`) is unaffected (B3-5's D3 "mapped-only columns" drop).
121    pub fn remove_column(&mut self, name: &str) -> Option<Vec<Value>> {
122        self.columns.shift_remove(name)
123    }
124
125    /// Group by one or more key columns. Returns a Vec of DataFrames, one per group.
126    pub fn group_by(&self, keys: &[&str]) -> Vec<DataFrame> {
127        if self.nrows == 0 {
128            return vec![];
129        }
130
131        // Build group keys for each row
132        // Borrowing keys (no per-row String clones for text columns); a missing
133        // value is its own group, distinct from the literal string "NA".
134        let key_cols: Vec<Option<&Vec<Value>>> =
135            keys.iter().map(|k| self.columns.get(*k)).collect();
136        let mut group_map: IndexMap<Vec<GroupKey<'_>>, Vec<usize>> = IndexMap::new();
137
138        for i in 0..self.nrows {
139            let key: Vec<GroupKey<'_>> = key_cols
140                .iter()
141                .map(|col| col.map_or(GroupKey::Na, |c| c[i].group_key()))
142                .collect();
143            group_map.entry(key).or_default().push(i);
144        }
145
146        group_map
147            .into_values()
148            .map(|indices| {
149                let mut df = DataFrame::new();
150                for (name, col) in &self.columns {
151                    let values: Vec<Value> = indices.iter().map(|&i| col[i].clone()).collect();
152                    df.add_column(name.clone(), values);
153                }
154                df
155            })
156            .collect()
157    }
158
159    /// Vertically stack another DataFrame onto this one.
160    pub fn vstack(&mut self, other: &DataFrame) {
161        if other.nrows == 0 {
162            return;
163        }
164        if self.columns.is_empty() {
165            *self = other.clone();
166            return;
167        }
168
169        // Add columns from other that we have
170        for (name, col) in &self.columns {
171            if let Some(other_col) = other.columns.get(name) {
172                // Will extend below
173                let _ = (col, other_col);
174            }
175        }
176
177        // Also add columns from other that we don't have (fill with NA)
178        for name in other.columns.keys() {
179            if !self.columns.contains_key(name) {
180                self.columns
181                    .insert(name.clone(), vec![Value::Na; self.nrows]);
182            }
183        }
184
185        let old_nrows = self.nrows;
186        self.nrows += other.nrows;
187
188        for (name, col) in &mut self.columns {
189            if let Some(other_col) = other.columns.get(name) {
190                col.extend(other_col.iter().cloned());
191            } else {
192                col.extend(std::iter::repeat_with(|| Value::Na).take(other.nrows));
193            }
194            debug_assert_eq!(col.len(), old_nrows + other.nrows);
195        }
196    }
197
198    /// Select a subset of columns.
199    pub fn select(&self, columns: &[&str]) -> DataFrame {
200        let mut df = DataFrame::new();
201        for &col_name in columns {
202            if let Some(col) = self.columns.get(col_name) {
203                df.add_column(col_name.to_string(), col.clone());
204            }
205        }
206        df
207    }
208
209    /// Get a single row as a map.
210    pub fn row(&self, idx: usize) -> IndexMap<String, Value> {
211        assert!(
212            idx < self.nrows,
213            "Row index {idx} out of bounds ({} rows)",
214            self.nrows
215        );
216        let mut map = IndexMap::new();
217        for (name, col) in &self.columns {
218            map.insert(name.clone(), col[idx].clone());
219        }
220        map
221    }
222
223    /// Sort by a column (ascending). Returns a new DataFrame.
224    pub fn sort_by(&self, column: &str) -> DataFrame {
225        let col = match self.columns.get(column) {
226            Some(c) => c,
227            None => return self.clone(),
228        };
229
230        let mut indices: Vec<usize> = (0..self.nrows).collect();
231        indices.sort_by(|&a, &b| {
232            let va = col[a].as_f64().unwrap_or(f64::NAN);
233            let vb = col[b].as_f64().unwrap_or(f64::NAN);
234            va.total_cmp(&vb)
235        });
236
237        let mut df = DataFrame::new();
238        for (name, c) in &self.columns {
239            let values: Vec<Value> = indices.iter().map(|&i| c[i].clone()).collect();
240            df.add_column(name.clone(), values);
241        }
242        df
243    }
244
245    /// Create from rows (list of maps).
246    pub fn from_rows(rows: Vec<IndexMap<String, Value>>) -> Self {
247        if rows.is_empty() {
248            return DataFrame::new();
249        }
250
251        // Collect all column names from all rows
252        let mut col_names: IndexMap<String, ()> = IndexMap::new();
253        for row in &rows {
254            for key in row.keys() {
255                col_names.entry(key.clone()).or_default();
256            }
257        }
258
259        let mut df = DataFrame::new();
260        for name in col_names.keys() {
261            let values: Vec<Value> = rows
262                .iter()
263                .map(|row| row.get(name).cloned().unwrap_or(Value::Na))
264                .collect();
265            df.add_column(name.clone(), values);
266        }
267        df
268    }
269
270    /// Get all unique values in a column.
271    pub fn unique_values(&self, column: &str) -> Vec<Value> {
272        let col = match self.columns.get(column) {
273            Some(c) => c,
274            None => return vec![],
275        };
276        let mut seen = std::collections::HashSet::new();
277        let mut result = Vec::new();
278        for v in col {
279            if seen.insert(v.group_key()) {
280                result.push(v.clone());
281            }
282        }
283        result
284    }
285}
286
287impl DataFrame {
288    /// Load a DataFrame from a CSV file.
289    /// First row is treated as column headers.
290    /// Values are parsed as Float if possible, otherwise kept as strings.
291    /// The literal string "NA" is parsed as Value::Na.
292    pub fn from_csv(path: &str) -> Result<Self, std::io::Error> {
293        let content = std::fs::read_to_string(path)?;
294        let mut lines = content.lines();
295
296        let header = match lines.next() {
297            Some(h) => h,
298            None => return Ok(DataFrame::new()),
299        };
300
301        let col_names: Vec<&str> = header.split(',').map(|s| s.trim()).collect();
302        let mut columns: Vec<Vec<Value>> = vec![Vec::new(); col_names.len()];
303
304        for line in lines {
305            let line = line.trim();
306            if line.is_empty() {
307                continue;
308            }
309            let fields: Vec<&str> = line.split(',').collect();
310            for (i, field) in fields.iter().enumerate() {
311                if i >= col_names.len() {
312                    continue;
313                }
314                let field = field.trim();
315                let val = if field == "NA" || field == "na" {
316                    Value::Na
317                } else if let Ok(f) = field.parse::<f64>() {
318                    Value::Float(f)
319                } else {
320                    Value::Str(field.to_string())
321                };
322                columns[i].push(val);
323            }
324            // Pad missing columns with NA
325            for col in columns.iter_mut().skip(fields.len()) {
326                col.push(Value::Na);
327            }
328        }
329
330        let mut df = DataFrame::new();
331        for (i, name) in col_names.iter().enumerate() {
332            if !columns[i].is_empty() {
333                df.add_column(name.to_string(), std::mem::take(&mut columns[i]));
334            }
335        }
336
337        Ok(df)
338    }
339}
340
341impl Default for DataFrame {
342    fn default() -> Self {
343        Self::new()
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350
351    #[test]
352    fn mismatched_column_lengths_pad_and_record_an_issue() {
353        let mut df = DataFrame::new();
354        df.add_column("x".into(), vec![Value::Float(1.0), Value::Float(2.0)]);
355        df.add_column("y".into(), vec![Value::Float(3.0)]);
356        assert_eq!(df.nrows(), 2);
357        assert_eq!(df.column("y").unwrap()[1], Value::Na);
358        df.add_column("z".into(), vec![Value::Float(0.0); 4]);
359        assert_eq!(df.nrows(), 4);
360        assert!(df
361            .column_names()
362            .iter()
363            .all(|c| df.column(c).unwrap().len() == 4));
364        assert_eq!(df.issues().len(), 2);
365        let err = df.validate().unwrap_err().to_string();
366        assert!(err.contains("'y'") && err.contains("'z'"), "{err}");
367
368        let mut ok = DataFrame::new();
369        ok.add_column("x".into(), vec![Value::Float(1.0)]);
370        assert!(ok.validate().is_ok());
371        assert!(ok
372            .try_add_column("y".into(), vec![Value::Na, Value::Na])
373            .is_err());
374        assert!(!ok.has_column("y"), "rejected column is not added");
375        assert!(ok.try_add_column("y".into(), vec![Value::Na]).is_ok());
376        assert!(ok.validate().is_ok());
377    }
378
379    #[test]
380    fn group_by_keeps_na_apart_from_literal_na() {
381        let mut df = DataFrame::new();
382        df.add_column(
383            "g".into(),
384            vec![
385                Value::Na,
386                Value::Str("NA".into()),
387                Value::Na,
388                Value::Str("NA".into()),
389            ],
390        );
391        df.add_column("v".into(), (0..4).map(Value::Integer).collect());
392        let groups = df.group_by(&["g"]);
393        assert_eq!(groups.len(), 2);
394        assert!(groups[0].column("g").unwrap().iter().all(Value::is_na));
395        assert!(groups[1]
396            .column("g")
397            .unwrap()
398            .iter()
399            .all(|v| v.as_str() == Some("NA")));
400        assert_eq!(df.unique_values("g").len(), 2);
401    }
402
403    #[test]
404    fn test_add_column_and_access() {
405        let mut df = DataFrame::new();
406        df.add_column("x".into(), vec![Value::Float(1.0), Value::Float(2.0)]);
407        df.add_column("y".into(), vec![Value::Float(3.0), Value::Float(4.0)]);
408
409        assert_eq!(df.nrows(), 2);
410        assert_eq!(df.ncols(), 2);
411        assert!(df.has_column("x"));
412        assert!(!df.has_column("z"));
413    }
414
415    #[test]
416    fn test_group_by() {
417        let mut df = DataFrame::new();
418        df.add_column(
419            "cat".into(),
420            vec![
421                Value::Str("a".into()),
422                Value::Str("b".into()),
423                Value::Str("a".into()),
424            ],
425        );
426        df.add_column(
427            "val".into(),
428            vec![Value::Float(1.0), Value::Float(2.0), Value::Float(3.0)],
429        );
430
431        let groups = df.group_by(&["cat"]);
432        assert_eq!(groups.len(), 2);
433        assert_eq!(groups[0].nrows(), 2); // "a" group
434        assert_eq!(groups[1].nrows(), 1); // "b" group
435    }
436
437    #[test]
438    fn test_vstack() {
439        let mut df1 = DataFrame::new();
440        df1.add_column("x".into(), vec![Value::Float(1.0)]);
441
442        let mut df2 = DataFrame::new();
443        df2.add_column("x".into(), vec![Value::Float(2.0)]);
444
445        df1.vstack(&df2);
446        assert_eq!(df1.nrows(), 2);
447    }
448
449    #[test]
450    fn test_sort_by() {
451        let mut df = DataFrame::new();
452        df.add_column(
453            "x".into(),
454            vec![Value::Float(3.0), Value::Float(1.0), Value::Float(2.0)],
455        );
456        let sorted = df.sort_by("x");
457        let col = sorted.column("x").unwrap();
458        assert_eq!(col[0].as_f64(), Some(1.0));
459        assert_eq!(col[1].as_f64(), Some(2.0));
460        assert_eq!(col[2].as_f64(), Some(3.0));
461    }
462}
ggplot_rs/data/dataframe.rs

ggplot_rs/data/
dataframe.rs
