reifydb_value/value/frame/
extract.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{
5	error,
6	fmt::{self, Display, Formatter},
7};
8
9use super::frame::Frame;
10use crate::value::{
11	column_view::ColumnView,
12	try_from::{FromValueError, TryFromValue},
13};
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum FrameError {
17	ColumnNotFound {
18		name: String,
19	},
20
21	RowOutOfBounds {
22		row: usize,
23		len: usize,
24	},
25
26	ValueError {
27		column: String,
28		row: usize,
29		error: FromValueError,
30	},
31
32	InvalidColumn {
33		name: String,
34		message: String,
35	},
36}
37
38impl Display for FrameError {
39	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
40		match self {
41			FrameError::ColumnNotFound {
42				name,
43			} => {
44				write!(f, "column not found: {}", name)
45			}
46			FrameError::RowOutOfBounds {
47				row,
48				len,
49			} => {
50				write!(f, "row {} out of bounds (frame has {} rows)", row, len)
51			}
52			FrameError::ValueError {
53				column,
54				row,
55				error,
56			} => {
57				write!(f, "error extracting column '{}' row {}: {}", column, row, error)
58			}
59			FrameError::InvalidColumn {
60				name,
61				message,
62			} => {
63				write!(f, "column '{}' does not match its field: {}", name, message)
64			}
65		}
66	}
67}
68
69impl error::Error for FrameError {}
70
71impl Frame {
72	pub fn column(&self, name: &str) -> Result<Option<ColumnView<'_>>, FrameError> {
73		let schema = self.batch.schema_ref();
74		let Some(index) = schema.fields().iter().position(|field| field.name() == name) else {
75			return Ok(None);
76		};
77		ColumnView::try_from((self.batch.column(index), schema.field(index))).map(Some).map_err(|error| {
78			FrameError::InvalidColumn {
79				name: name.to_string(),
80				message: error.to_string(),
81			}
82		})
83	}
84
85	pub fn try_column(&self, name: &str) -> Result<ColumnView<'_>, FrameError> {
86		self.column(name)?.ok_or_else(|| FrameError::ColumnNotFound {
87			name: name.to_string(),
88		})
89	}
90
91	pub fn row_count(&self) -> usize {
92		self.batch.num_rows()
93	}
94
95	pub fn get<T: TryFromValue>(&self, column: &str, row: usize) -> Result<Option<T>, FrameError> {
96		let col = self.try_column(column)?;
97		let len = col.len();
98
99		if row >= len {
100			return Err(FrameError::RowOutOfBounds {
101				row,
102				len,
103			});
104		}
105
106		if !col.is_defined(row) {
107			return Ok(None);
108		}
109
110		let value = col.get_value(row);
111		T::try_from_value(&value).map(Some).map_err(|e| FrameError::ValueError {
112			column: column.to_string(),
113			row,
114			error: e,
115		})
116	}
117}
118
119#[cfg(test)]
120pub mod tests {
121	use std::sync::Arc;
122
123	use arrow_array::{ArrayRef, Int32Array, Int64Array, LargeStringArray, RecordBatch, UInt64Array};
124	use arrow_schema::{FieldRef, Schema};
125
126	use super::*;
127	use crate::value::{
128		system_columns::{SystemColumn, with_system_column},
129		value_type::{
130			ValueType,
131			field::{FieldType, to_field},
132		},
133	};
134
135	fn column(name: &str, value_type: ValueType, array: ArrayRef) -> (FieldRef, ArrayRef) {
136		let field_type = FieldType {
137			value_type: Some(value_type),
138			..FieldType::default()
139		};
140		(Arc::new(to_field(name, &field_type)), array)
141	}
142
143	fn make_test_frame() -> Frame {
144		let (fields, arrays): (Vec<FieldRef>, Vec<ArrayRef>) = vec![
145			column("id", ValueType::Int8, Arc::new(Int64Array::from(vec![1i64, 2, 3]))),
146			column(
147				"name",
148				ValueType::Utf8,
149				Arc::new(LargeStringArray::from(vec![
150					"Alice".to_string(),
151					"Bob".to_string(),
152					String::new(),
153				])),
154			),
155			column("score", ValueType::Int4, Arc::new(Int32Array::from(vec![100i32, 85, 92]))),
156		]
157		.into_iter()
158		.unzip();
159		let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap();
160		let batch = with_system_column(
161			batch,
162			SystemColumn::RowNumbers,
163			Arc::new(UInt64Array::from(vec![1u64, 2, 3])),
164		)
165		.unwrap();
166		Frame::from(batch)
167	}
168
169	#[test]
170	fn test_column_by_name() {
171		let frame = make_test_frame();
172		assert!(frame.column("id").unwrap().is_some());
173		assert!(frame.column("name").unwrap().is_some());
174		assert!(frame.column("nonexistent").unwrap().is_none());
175	}
176
177	#[test]
178	fn test_row_count() {
179		let frame = make_test_frame();
180		assert_eq!(frame.row_count(), 3);
181
182		let empty = Frame::from(RecordBatch::new_empty(Arc::new(Schema::empty())));
183		assert_eq!(empty.row_count(), 0);
184	}
185
186	#[test]
187	fn test_get_value() {
188		let frame = make_test_frame();
189
190		let id: Option<i64> = frame.get("id", 0).unwrap();
191		assert_eq!(id, Some(1i64));
192
193		let name: Option<String> = frame.get("name", 0).unwrap();
194		assert_eq!(name, Some("Alice".to_string()));
195
196		// The column has no bitvec, so an empty string must read as present, not as a none.
197		let name_at_2: Option<String> = frame.get("name", 2).unwrap();
198		assert_eq!(name_at_2, Some(String::new()));
199	}
200
201	#[test]
202	fn test_errors() {
203		let frame = make_test_frame();
204
205		// Each failure keeps its own variant, so a caller can tell a typo from a bad index from
206		// a type mismatch.
207		let err = frame.get::<i64>("nonexistent", 0).unwrap_err();
208		assert!(matches!(err, FrameError::ColumnNotFound { .. }));
209
210		let err = frame.get::<i64>("id", 100).unwrap_err();
211		assert!(matches!(err, FrameError::RowOutOfBounds { .. }));
212
213		let err = frame.get::<i32>("id", 0).unwrap_err();
214		assert!(matches!(err, FrameError::ValueError { .. }));
215	}
216}
reifydb_value/value/frame/extract.rs

reifydb_value/value/frame/
extract.rs
