Skip to main content

orbital_data/
dataset.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::projection::ProjectionError;
6use crate::{DataRecord, DataSchema, DataValue};
7
8/// The shared shape. Tables render it; charts build series by field key.
9#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
10pub struct Dataset {
11    pub schema: DataSchema,
12    pub records: Vec<DataRecord>,
13}
14
15impl Dataset {
16    pub fn new(schema: DataSchema, records: Vec<DataRecord>) -> Self {
17        Self { schema, records }
18    }
19
20    pub fn from_records(schema: DataSchema, records: Vec<DataRecord>) -> Self {
21        Self::new(schema, records)
22    }
23
24    /// Build a dataset from text field definitions and `(id, cells)` rows.
25    pub fn from_text_rows(
26        fields: impl IntoIterator<Item = (String, String)>,
27        rows: impl IntoIterator<Item = (String, HashMap<String, String>)>,
28    ) -> Self {
29        let schema = DataSchema::from_text_fields(fields);
30        let records = rows
31            .into_iter()
32            .map(|(id, cells)| DataRecord::from_text_map(id, cells))
33            .collect();
34        Self::new(schema, records)
35    }
36
37    pub fn field_keys(&self) -> impl Iterator<Item = &str> {
38        self.schema.fields.iter().map(|f| f.key.as_str())
39    }
40
41    /// Whether `field` is declared in the dataset schema.
42    pub fn has_field(&self, field: &str) -> bool {
43        self.schema.fields.iter().any(|f| f.key == field)
44    }
45
46    /// Extract a numeric column for chart value axes.
47    ///
48    /// Accepts [`DataValue::Number`]. [`DataValue::Null`] maps to `f64::NAN` as a gap.
49    pub fn column_as_numbers(&self, field: &str) -> Result<Vec<f64>, ProjectionError> {
50        if self.records.is_empty() {
51            return Err(ProjectionError::EmptyDataset);
52        }
53        if !self.has_field(field) {
54            return Err(ProjectionError::UnknownField {
55                field: field.to_string(),
56            });
57        }
58
59        self.records
60            .iter()
61            .map(|record| match record.get(field) {
62                None | Some(DataValue::Null) => Ok(f64::NAN),
63                Some(DataValue::Number(n)) => Ok(*n),
64                Some(other) => Err(ProjectionError::TypeMismatch {
65                    field: field.to_string(),
66                    expected: "number",
67                    got: other.data_type(),
68                }),
69            })
70            .collect()
71    }
72
73    /// Extract a categorical column for band/point axes.
74    ///
75    /// Accepts text, category, date, number, and bool variants (stringified).
76    pub fn column_as_categories(&self, field: &str) -> Result<Vec<String>, ProjectionError> {
77        if self.records.is_empty() {
78            return Err(ProjectionError::EmptyDataset);
79        }
80        if !self.has_field(field) {
81            return Err(ProjectionError::UnknownField {
82                field: field.to_string(),
83            });
84        }
85
86        self.records
87            .iter()
88            .map(|record| match record.get(field) {
89                None | Some(DataValue::Null) => Ok(String::new()),
90                Some(value) => Ok(value.display_string()),
91            })
92            .collect()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use std::collections::HashMap;
100
101    fn sample_dataset() -> Dataset {
102        let schema = DataSchema::from_text_fields([
103            ("quarter".into(), "Quarter".into()),
104            ("revenue".into(), "Revenue".into()),
105            ("cost".into(), "Cost".into()),
106        ]);
107        let records = vec![
108            DataRecord::new(
109                "1",
110                HashMap::from([
111                    ("quarter".into(), DataValue::Category("Q1".into())),
112                    ("revenue".into(), DataValue::Number(100.0)),
113                    ("cost".into(), DataValue::Number(60.0)),
114                ]),
115            ),
116            DataRecord::new(
117                "2",
118                HashMap::from([
119                    ("quarter".into(), DataValue::Category("Q2".into())),
120                    ("revenue".into(), DataValue::Number(120.0)),
121                    ("cost".into(), DataValue::Number(70.0)),
122                ]),
123            ),
124        ];
125        Dataset::from_records(schema, records)
126    }
127
128    #[test]
129    fn column_as_numbers_happy_path() {
130        let dataset = sample_dataset();
131        let nums = dataset.column_as_numbers("revenue").unwrap();
132        assert_eq!(nums, vec![100.0, 120.0]);
133    }
134
135    #[test]
136    fn column_as_numbers_null_becomes_nan() {
137        let mut dataset = sample_dataset();
138        dataset.records[0]
139            .values
140            .insert("revenue".into(), DataValue::Null);
141        let nums = dataset.column_as_numbers("revenue").unwrap();
142        assert!(nums[0].is_nan());
143        assert_eq!(nums[1], 120.0);
144    }
145
146    #[test]
147    fn column_as_numbers_type_mismatch() {
148        let dataset = sample_dataset();
149        let err = dataset.column_as_numbers("quarter").unwrap_err();
150        assert!(matches!(
151            err,
152            ProjectionError::TypeMismatch {
153                field,
154                expected: "number",
155                ..
156            } if field == "quarter"
157        ));
158    }
159
160    #[test]
161    fn column_as_numbers_unknown_field() {
162        let dataset = sample_dataset();
163        let err = dataset.column_as_numbers("missing").unwrap_err();
164        assert!(matches!(err, ProjectionError::UnknownField { .. }));
165    }
166
167    #[test]
168    fn column_as_numbers_empty_dataset() {
169        let dataset = Dataset::default();
170        let err = dataset.column_as_numbers("revenue").unwrap_err();
171        assert!(matches!(err, ProjectionError::EmptyDataset));
172    }
173
174    #[test]
175    fn column_as_categories_happy_path() {
176        let dataset = sample_dataset();
177        let cats = dataset.column_as_categories("quarter").unwrap();
178        assert_eq!(cats, vec!["Q1", "Q2"]);
179    }
180
181    #[test]
182    fn has_field() {
183        let dataset = sample_dataset();
184        assert!(dataset.has_field("quarter"));
185        assert!(!dataset.has_field("missing"));
186    }
187}