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