reifydb_type/value/frame/
extract.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the MIT, see license.md file
3
4use std::fmt::{Display, Formatter};
5
6use super::{Frame, FrameColumn};
7use crate::{FromValueError, TryFromValue, TryFromValueCoerce};
8
9/// Error type for Frame extraction operations
10#[derive(Debug, Clone, PartialEq)]
11pub enum FrameError {
12	/// Column not found by name
13	ColumnNotFound {
14		name: String,
15	},
16	/// Row index out of bounds
17	RowOutOfBounds {
18		row: usize,
19		len: usize,
20	},
21	/// Value extraction error
22	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	/// Get a column by name
58	pub fn column(&self, name: &str) -> Option<&FrameColumn> {
59		self.columns.iter().find(|c| c.name == name)
60	}
61
62	/// Get a column by name, returning an error if not found
63	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	/// Get the number of rows in the frame
70	pub fn row_count(&self) -> usize {
71		self.columns.first().map(|c| c.data.len()).unwrap_or(0)
72	}
73
74	/// Extract a single value by column name and row index (strict type matching).
75	///
76	/// Returns `Ok(None)` for Undefined values.
77	/// Returns `Err` for missing columns, out of bounds rows, or type mismatches.
78	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		// Check if value is undefined
90		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	/// Extract a single value with widening coercion.
103	///
104	/// Returns `Ok(None)` for Undefined values.
105	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		// Check if value is undefined
117		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	/// Extract an entire column as `Vec<Option<T>>` (strict type matching).
130	///
131	/// Undefined values become `None`, type mismatches return an error.
132	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	/// Extract an entire column with widening coercion.
151	///
152	/// Undefined values become `None`, incompatible types return an error.
153	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		// Get strict-typed value
234		let id: Option<i64> = frame.get("id", 0).unwrap();
235		assert_eq!(id, Some(1i64));
236
237		// Get string value
238		let name: Option<String> = frame.get("name", 0).unwrap();
239		assert_eq!(name, Some("Alice".to_string()));
240
241		// Get undefined value
242		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		// Int4 coerced to i64
251		let score: Option<i64> = frame.get_coerce("score", 0).unwrap();
252		assert_eq!(score, Some(100i64));
253
254		// Int4 coerced to f64
255		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		// Int4 coerced to Vec<Option<i64>>
275		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		// Column not found
284		let err = frame.get::<i64>("nonexistent", 0).unwrap_err();
285		assert!(matches!(err, FrameError::ColumnNotFound { .. }));
286
287		// Row out of bounds
288		let err = frame.get::<i64>("id", 100).unwrap_err();
289		assert!(matches!(err, FrameError::RowOutOfBounds { .. }));
290
291		// Type mismatch (strict)
292		let err = frame.get::<i32>("id", 0).unwrap_err();
293		assert!(matches!(err, FrameError::ValueError { .. }));
294	}
295}