1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::projection::ProjectionError;
6use crate::{DataRecord, DataSchema, DataValue};
7
8#[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 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 pub fn has_field(&self, field: &str) -> bool {
43 self.schema.fields.iter().any(|f| f.key == field)
44 }
45
46 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 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}