reifydb_value/value/frame/
row.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{collections::HashMap, rc::Rc};
5
6use super::{extract::FrameError, frame::Frame};
7use crate::{
8	error::Error,
9	value::{column_view::ColumnView, try_from::TryFromValue},
10};
11
12#[derive(Debug)]
13struct ColumnIndex<'a> {
14	by_name: HashMap<String, usize>,
15	views: Vec<Result<ColumnView<'a>, Error>>,
16}
17
18impl<'a> ColumnIndex<'a> {
19	fn new(frame: &'a Frame) -> Self {
20		let schema = frame.batch.schema_ref();
21		let mut by_name = HashMap::with_capacity(schema.fields().len());
22		let mut views = Vec::with_capacity(schema.fields().len());
23		for (idx, (field, array)) in schema.fields().iter().zip(frame.batch.columns()).enumerate() {
24			by_name.insert(field.name().clone(), idx);
25			views.push(ColumnView::try_from((array, field.as_ref())));
26		}
27		Self {
28			by_name,
29			views,
30		}
31	}
32
33	fn get(&self, name: &str) -> Option<&Result<ColumnView<'a>, Error>> {
34		self.by_name.get(name).map(|&idx| &self.views[idx])
35	}
36}
37
38#[derive(Debug)]
39pub struct FrameRow<'a> {
40	index: Rc<ColumnIndex<'a>>,
41	row_idx: usize,
42}
43
44impl<'a> FrameRow<'a> {
45	pub fn get<T: TryFromValue>(&self, column: &str) -> Result<Option<T>, FrameError> {
46		let col = match self.index.get(column) {
47			Some(Ok(view)) => view,
48			Some(Err(error)) => {
49				return Err(FrameError::InvalidColumn {
50					name: column.to_string(),
51					message: error.to_string(),
52				});
53			}
54			None => {
55				return Err(FrameError::ColumnNotFound {
56					name: column.to_string(),
57				});
58			}
59		};
60
61		if !col.is_defined(self.row_idx) {
62			return Ok(None);
63		}
64
65		let value = col.get_value(self.row_idx);
66		T::try_from_value(&value).map(Some).map_err(|e| FrameError::ValueError {
67			column: column.to_string(),
68			row: self.row_idx,
69			error: e,
70		})
71	}
72}
73
74pub struct FrameRows<'a> {
75	index: Rc<ColumnIndex<'a>>,
76	current: usize,
77	len: usize,
78}
79
80impl<'a> FrameRows<'a> {
81	pub(super) fn new(frame: &'a Frame) -> Self {
82		Self {
83			index: Rc::new(ColumnIndex::new(frame)),
84			current: 0,
85			len: frame.batch.num_rows(),
86		}
87	}
88}
89
90impl<'a> Iterator for FrameRows<'a> {
91	type Item = FrameRow<'a>;
92
93	fn next(&mut self) -> Option<Self::Item> {
94		if self.current >= self.len {
95			return None;
96		}
97
98		let row = FrameRow {
99			index: Rc::clone(&self.index),
100			row_idx: self.current,
101		};
102
103		self.current += 1;
104		Some(row)
105	}
106
107	fn size_hint(&self) -> (usize, Option<usize>) {
108		let remaining = self.len.saturating_sub(self.current);
109		(remaining, Some(remaining))
110	}
111}
112
113impl ExactSizeIterator for FrameRows<'_> {}
114
115impl<'a> DoubleEndedIterator for FrameRows<'a> {
116	fn next_back(&mut self) -> Option<Self::Item> {
117		if self.current >= self.len {
118			return None;
119		}
120
121		self.len -= 1;
122
123		Some(FrameRow {
124			index: Rc::clone(&self.index),
125			row_idx: self.len,
126		})
127	}
128}
129
130impl Frame {
131	pub fn rows(&self) -> FrameRows<'_> {
132		FrameRows::new(self)
133	}
134}
135
136#[cfg(test)]
137pub mod tests {
138	use std::sync::Arc;
139
140	use arrow_array::{ArrayRef, Int64Array, LargeStringArray, RecordBatch, UInt64Array};
141	use arrow_schema::{FieldRef, Schema};
142
143	use super::*;
144	use crate::value::{
145		system_columns::{SystemColumn, with_system_column},
146		value_type::{
147			ValueType,
148			field::{FieldType, to_field},
149		},
150	};
151
152	fn column(name: &str, value_type: ValueType, array: ArrayRef) -> (FieldRef, ArrayRef) {
153		let field_type = FieldType {
154			value_type: Some(value_type),
155			..FieldType::default()
156		};
157		(Arc::new(to_field(name, &field_type)), array)
158	}
159
160	fn make_test_frame() -> Frame {
161		let (fields, arrays): (Vec<FieldRef>, Vec<ArrayRef>) = vec![
162			column("id", ValueType::Int8, Arc::new(Int64Array::from(vec![1i64, 2, 3]))),
163			column(
164				"name",
165				ValueType::Utf8,
166				Arc::new(LargeStringArray::from(vec![
167					"Alice".to_string(),
168					"Bob".to_string(),
169					String::new(),
170				])),
171			),
172		]
173		.into_iter()
174		.unzip();
175		let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap();
176		let batch = with_system_column(
177			batch,
178			SystemColumn::RowNumbers,
179			Arc::new(UInt64Array::from(vec![100u64, 200, 300])),
180		)
181		.unwrap();
182		Frame::from(batch)
183	}
184
185	#[test]
186	fn test_rows_iterator() {
187		let frame = make_test_frame();
188		let rows: Vec<_> = frame.rows().collect();
189
190		assert_eq!(rows.len(), 3);
191		assert_eq!(rows[0].get::<i64>("id").unwrap(), Some(1i64));
192		assert_eq!(rows[1].get::<i64>("id").unwrap(), Some(2i64));
193		assert_eq!(rows[2].get::<i64>("id").unwrap(), Some(3i64));
194	}
195
196	#[test]
197	fn test_row_get() {
198		let frame = make_test_frame();
199		let mut rows = frame.rows();
200
201		let row0 = rows.next().unwrap();
202		assert_eq!(row0.get::<i64>("id").unwrap(), Some(1i64));
203		assert_eq!(row0.get::<String>("name").unwrap(), Some("Alice".to_string()));
204
205		let row2 = rows.nth(1).unwrap(); // Skip to index 2
206		assert_eq!(row2.get::<i64>("id").unwrap(), Some(3i64));
207		assert_eq!(row2.get::<String>("name").unwrap(), Some(String::new())); // All values are defined
208	}
209
210	#[test]
211	fn test_exact_size_iterator() {
212		let frame = make_test_frame();
213		let rows = frame.rows();
214
215		assert_eq!(rows.len(), 3);
216	}
217
218	#[test]
219	fn test_double_ended_iterator() {
220		let frame = make_test_frame();
221		let mut rows = frame.rows();
222
223		let last = rows.next_back().unwrap();
224		assert_eq!(last.get::<i64>("id").unwrap(), Some(3i64));
225
226		let first = rows.next().unwrap();
227		assert_eq!(first.get::<i64>("id").unwrap(), Some(1i64));
228	}
229}
reifydb_value/value/frame/row.rs

reifydb_value/value/frame/
row.rs
