reifydb_type/value/frame/
extract.rs1use std::fmt::{Display, Formatter};
5
6use super::{Frame, FrameColumn};
7use crate::{FromValueError, TryFromValue, TryFromValueCoerce};
8
9#[derive(Debug, Clone, PartialEq)]
11pub enum FrameError {
12 ColumnNotFound {
14 name: String,
15 },
16 RowOutOfBounds {
18 row: usize,
19 len: usize,
20 },
21 ValueError {
23 column: String,
24 row: usize,
25 error: FromValueError,
26 },
27}
28
29impl Display for FrameError {
30 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
31 match self {
32 FrameError::ColumnNotFound {
33 name,
34 } => {
35 write!(f, "column not found: {}", name)
36 }
37 FrameError::RowOutOfBounds {
38 row,
39 len,
40 } => {
41 write!(f, "row {} out of bounds (frame has {} rows)", row, len)
42 }
43 FrameError::ValueError {
44 column,
45 row,
46 error,
47 } => {
48 write!(f, "error extracting column '{}' row {}: {}", column, row, error)
49 }
50 }
51 }
52}
53
54impl std::error::Error for FrameError {}
55
56impl Frame {
57 pub fn column(&self, name: &str) -> Option<&FrameColumn> {
59 self.columns.iter().find(|c| c.name == name)
60 }
61
62 pub fn try_column(&self, name: &str) -> Result<&FrameColumn, FrameError> {
64 self.column(name).ok_or_else(|| FrameError::ColumnNotFound {
65 name: name.to_string(),
66 })
67 }
68
69 pub fn row_count(&self) -> usize {
71 self.columns.first().map(|c| c.data.len()).unwrap_or(0)
72 }
73
74 pub fn get<T: TryFromValue>(&self, column: &str, row: usize) -> Result<Option<T>, FrameError> {
79 let col = self.try_column(column)?;
80 let len = col.data.len();
81
82 if row >= len {
83 return Err(FrameError::RowOutOfBounds {
84 row,
85 len,
86 });
87 }
88
89 if !col.data.is_defined(row) {
91 return Ok(None);
92 }
93
94 let value = col.data.get_value(row);
95 T::try_from_value(&value).map(Some).map_err(|e| FrameError::ValueError {
96 column: column.to_string(),
97 row,
98 error: e,
99 })
100 }
101
102 pub fn get_coerce<T: TryFromValueCoerce>(&self, column: &str, row: usize) -> Result<Option<T>, FrameError> {
106 let col = self.try_column(column)?;
107 let len = col.data.len();
108
109 if row >= len {
110 return Err(FrameError::RowOutOfBounds {
111 row,
112 len,
113 });
114 }
115
116 if !col.data.is_defined(row) {
118 return Ok(None);
119 }
120
121 let value = col.data.get_value(row);
122 T::try_from_value_coerce(&value).map(Some).map_err(|e| FrameError::ValueError {
123 column: column.to_string(),
124 row,
125 error: e,
126 })
127 }
128
129 pub fn column_values<T: TryFromValue>(&self, name: &str) -> Result<Vec<Option<T>>, FrameError> {
133 let col = self.try_column(name)?;
134 (0..col.data.len())
135 .map(|row| {
136 if !col.data.is_defined(row) {
137 Ok(None)
138 } else {
139 let value = col.data.get_value(row);
140 T::try_from_value(&value).map(Some).map_err(|e| FrameError::ValueError {
141 column: name.to_string(),
142 row,
143 error: e,
144 })
145 }
146 })
147 .collect()
148 }
149
150 pub fn column_values_coerce<T: TryFromValueCoerce>(&self, name: &str) -> Result<Vec<Option<T>>, FrameError> {
154 let col = self.try_column(name)?;
155 (0..col.data.len())
156 .map(|row| {
157 if !col.data.is_defined(row) {
158 Ok(None)
159 } else {
160 let value = col.data.get_value(row);
161 T::try_from_value_coerce(&value).map(Some).map_err(|e| FrameError::ValueError {
162 column: name.to_string(),
163 row,
164 error: e,
165 })
166 }
167 })
168 .collect()
169 }
170}
171
172#[cfg(test)]
173mod tests {
174 use super::*;
175 use crate::{
176 BitVec,
177 value::{
178 container::{NumberContainer, Utf8Container},
179 frame::FrameColumnData,
180 },
181 };
182
183 fn make_test_frame() -> Frame {
184 Frame::with_row_numbers(
185 vec![
186 FrameColumn {
187 namespace: None,
188 source: None,
189 name: "id".to_string(),
190 data: FrameColumnData::Int8(NumberContainer::from_vec(vec![1i64, 2, 3])),
191 },
192 FrameColumn {
193 namespace: None,
194 source: None,
195 name: "name".to_string(),
196 data: FrameColumnData::Utf8(Utf8Container::new(
197 vec!["Alice".to_string(), "Bob".to_string(), String::new()],
198 BitVec::from_slice(&[true, true, false]),
199 )),
200 },
201 FrameColumn {
202 namespace: None,
203 source: None,
204 name: "score".to_string(),
205 data: FrameColumnData::Int4(NumberContainer::from_vec(vec![100i32, 85, 92])),
206 },
207 ],
208 vec![1.into(), 2.into(), 3.into()],
209 )
210 }
211
212 #[test]
213 fn test_column_by_name() {
214 let frame = make_test_frame();
215 assert!(frame.column("id").is_some());
216 assert!(frame.column("name").is_some());
217 assert!(frame.column("nonexistent").is_none());
218 }
219
220 #[test]
221 fn test_row_count() {
222 let frame = make_test_frame();
223 assert_eq!(frame.row_count(), 3);
224
225 let empty = Frame::new(vec![]);
226 assert_eq!(empty.row_count(), 0);
227 }
228
229 #[test]
230 fn test_get_value() {
231 let frame = make_test_frame();
232
233 let id: Option<i64> = frame.get("id", 0).unwrap();
235 assert_eq!(id, Some(1i64));
236
237 let name: Option<String> = frame.get("name", 0).unwrap();
239 assert_eq!(name, Some("Alice".to_string()));
240
241 let name_undefined: Option<String> = frame.get("name", 2).unwrap();
243 assert_eq!(name_undefined, None);
244 }
245
246 #[test]
247 fn test_get_coerce() {
248 let frame = make_test_frame();
249
250 let score: Option<i64> = frame.get_coerce("score", 0).unwrap();
252 assert_eq!(score, Some(100i64));
253
254 let score_f64: Option<f64> = frame.get_coerce("score", 1).unwrap();
256 assert_eq!(score_f64, Some(85.0f64));
257 }
258
259 #[test]
260 fn test_column_values() {
261 let frame = make_test_frame();
262
263 let ids: Vec<Option<i64>> = frame.column_values("id").unwrap();
264 assert_eq!(ids, vec![Some(1), Some(2), Some(3)]);
265
266 let names: Vec<Option<String>> = frame.column_values("name").unwrap();
267 assert_eq!(names, vec![Some("Alice".to_string()), Some("Bob".to_string()), None]);
268 }
269
270 #[test]
271 fn test_column_values_coerce() {
272 let frame = make_test_frame();
273
274 let scores: Vec<Option<i64>> = frame.column_values_coerce("score").unwrap();
276 assert_eq!(scores, vec![Some(100), Some(85), Some(92)]);
277 }
278
279 #[test]
280 fn test_errors() {
281 let frame = make_test_frame();
282
283 let err = frame.get::<i64>("nonexistent", 0).unwrap_err();
285 assert!(matches!(err, FrameError::ColumnNotFound { .. }));
286
287 let err = frame.get::<i64>("id", 100).unwrap_err();
289 assert!(matches!(err, FrameError::RowOutOfBounds { .. }));
290
291 let err = frame.get::<i32>("id", 0).unwrap_err();
293 assert!(matches!(err, FrameError::ValueError { .. }));
294 }
295}