Skip to main content

reifydb_value/value/frame/
extract.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2026 ReifyDB
3
4use 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		// Get strict-typed value
209		let id: Option<i64> = frame.get("id", 0).unwrap();
210		assert_eq!(id, Some(1i64));
211
212		// Get string value
213		let name: Option<String> = frame.get("name", 0).unwrap();
214		assert_eq!(name, Some("Alice".to_string()));
215
216		// All values are defined (no bitvec), empty string at index 2
217		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		// Int4 coerced to i64
226		let score: Option<i64> = frame.get_coerce("score", 0).unwrap();
227		assert_eq!(score, Some(100i64));
228
229		// Int4 coerced to f64
230		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		// Int4 coerced to Vec<Option<i64>>
250		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		// Column not found
259		let err = frame.get::<i64>("nonexistent", 0).unwrap_err();
260		assert!(matches!(err, FrameError::ColumnNotFound { .. }));
261
262		// Row out of bounds
263		let err = frame.get::<i64>("id", 100).unwrap_err();
264		assert!(matches!(err, FrameError::RowOutOfBounds { .. }));
265
266		// ValueType mismatch (strict)
267		let err = frame.get::<i32>("id", 0).unwrap_err();
268		assert!(matches!(err, FrameError::ValueError { .. }));
269	}
270}