reifydb_core/value/column/
columns.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4use std::{
5	collections::HashMap,
6	ops::{Deref, Index, IndexMut},
7};
8
9use reifydb_type::{Fragment, RowNumber, Value};
10
11use crate::{
12	interface::resolved::{ResolvedRingBuffer, ResolvedTable, ResolvedView},
13	util::CowVec,
14	value::{
15		column::{Column, ColumnData, headers::ColumnHeaders},
16		container::UndefinedContainer,
17	},
18};
19
20#[derive(Debug, Clone)]
21pub struct Columns<'a> {
22	pub row_numbers: CowVec<RowNumber>,
23	pub columns: CowVec<Column<'a>>,
24}
25
26impl<'a> Deref for Columns<'a> {
27	type Target = [Column<'a>];
28
29	fn deref(&self) -> &Self::Target {
30		self.columns.deref()
31	}
32}
33
34impl<'a> Index<usize> for Columns<'a> {
35	type Output = Column<'a>;
36
37	fn index(&self, index: usize) -> &Self::Output {
38		self.columns.index(index)
39	}
40}
41
42impl<'a> IndexMut<usize> for Columns<'a> {
43	fn index_mut(&mut self, index: usize) -> &mut Self::Output {
44		&mut self.columns.make_mut()[index]
45	}
46}
47
48impl<'a> Columns<'a> {
49	pub fn new(columns: Vec<Column<'a>>) -> Self {
50		let n = columns.first().map_or(0, |c| c.data().len());
51		assert!(columns.iter().all(|c| c.data().len() == n));
52
53		Self {
54			row_numbers: CowVec::new(Vec::new()),
55			columns: CowVec::new(columns),
56		}
57	}
58
59	pub fn with_row_numbers(columns: Vec<Column<'a>>, row_numbers: Vec<RowNumber>) -> Self {
60		let n = columns.first().map_or(0, |c| c.data().len());
61		assert!(columns.iter().all(|c| c.data().len() == n));
62		assert_eq!(row_numbers.len(), n, "row_numbers length must match column data length");
63
64		Self {
65			row_numbers: CowVec::new(row_numbers),
66			columns: CowVec::new(columns),
67		}
68	}
69
70	pub fn single_row<'b>(rows: impl IntoIterator<Item = (&'b str, Value)>) -> Columns<'a> {
71		let mut columns = Vec::new();
72		let mut index = HashMap::new();
73
74		for (idx, (name, value)) in rows.into_iter().enumerate() {
75			let data = match value {
76				Value::Undefined => ColumnData::undefined(1),
77				Value::Boolean(v) => ColumnData::bool([v]),
78				Value::Float4(v) => ColumnData::float4([v.into()]),
79				Value::Float8(v) => ColumnData::float8([v.into()]),
80				Value::Int1(v) => ColumnData::int1([v]),
81				Value::Int2(v) => ColumnData::int2([v]),
82				Value::Int4(v) => ColumnData::int4([v]),
83				Value::Int8(v) => ColumnData::int8([v]),
84				Value::Int16(v) => ColumnData::int16([v]),
85				Value::Utf8(v) => ColumnData::utf8([v.clone()]),
86				Value::Uint1(v) => ColumnData::uint1([v]),
87				Value::Uint2(v) => ColumnData::uint2([v]),
88				Value::Uint4(v) => ColumnData::uint4([v]),
89				Value::Uint8(v) => ColumnData::uint8([v]),
90				Value::Uint16(v) => ColumnData::uint16([v]),
91				Value::Date(v) => ColumnData::date([v.clone()]),
92				Value::DateTime(v) => ColumnData::datetime([v.clone()]),
93				Value::Time(v) => ColumnData::time([v.clone()]),
94				Value::Duration(v) => ColumnData::duration([v.clone()]),
95				Value::RowNumber(v) => ColumnData::row_number([v]),
96				Value::IdentityId(v) => ColumnData::identity_id([v]),
97				Value::Uuid4(v) => ColumnData::uuid4([v]),
98				Value::Uuid7(v) => ColumnData::uuid7([v]),
99				Value::Blob(v) => ColumnData::blob([v.clone()]),
100				Value::Int(v) => ColumnData::int(vec![v]),
101				Value::Uint(v) => ColumnData::uint(vec![v]),
102				Value::Decimal(v) => ColumnData::decimal(vec![v]),
103				Value::Any(v) => ColumnData::any(vec![v]),
104			};
105
106			let column = Column {
107				name: Fragment::owned_internal(name.to_string()),
108				data,
109			};
110			index.insert(name, idx);
111			columns.push(column);
112		}
113
114		Self {
115			row_numbers: CowVec::new(Vec::new()),
116			columns: CowVec::new(columns),
117		}
118	}
119
120	pub fn apply_headers(&mut self, headers: &ColumnHeaders<'a>) {
121		// Apply the column names from headers to this Columns instance
122		for (i, name) in headers.columns.iter().enumerate() {
123			if i < self.len() {
124				let column = &mut self[i];
125				let data = std::mem::replace(column.data_mut(), ColumnData::undefined(0));
126
127				*column = Column {
128					name: name.clone(),
129					data,
130				};
131			}
132		}
133	}
134}
135
136impl<'a> Columns<'a> {
137	pub fn shape(&self) -> (usize, usize) {
138		let row_count = if !self.row_numbers.is_empty() {
139			self.row_numbers.len()
140		} else {
141			self.get(0).map(|c| c.data().len()).unwrap_or(0)
142		};
143		(row_count, self.len())
144	}
145
146	pub fn into_iter(self) -> impl Iterator<Item = Column<'a>> {
147		self.columns.into_iter()
148	}
149
150	pub fn is_empty(&self) -> bool {
151		self.shape().0 == 0
152	}
153
154	pub fn row(&self, i: usize) -> Vec<Value> {
155		self.iter().map(|c| c.data().get_value(i)).collect()
156	}
157
158	pub fn column(&self, name: &str) -> Option<&Column> {
159		self.iter().find(|col| col.name().text() == name)
160	}
161
162	pub fn row_count(&self) -> usize {
163		if !self.row_numbers.is_empty() {
164			self.row_numbers.len()
165		} else {
166			self.first().map_or(0, |col| col.data().len())
167		}
168	}
169
170	pub fn get_row(&self, index: usize) -> Vec<Value> {
171		self.iter().map(|col| col.data().get_value(index)).collect()
172	}
173}
174
175impl<'a> Column<'a> {
176	pub fn extend(&mut self, other: Column<'a>) -> crate::Result<()> {
177		self.data_mut().extend(other.data().clone())
178	}
179}
180
181impl<'a> Columns<'a> {
182	pub fn from_rows(names: &[&str], result_rows: &[Vec<Value>]) -> Self {
183		let column_count = names.len();
184
185		let mut columns: Vec<Column> = names
186			.iter()
187			.map(|name| Column {
188				name: Fragment::owned_internal(name.to_string()),
189				data: ColumnData::Undefined(UndefinedContainer::new(0)),
190			})
191			.collect();
192
193		for row in result_rows {
194			assert_eq!(row.len(), column_count, "row length does not match column count");
195			for (i, value) in row.iter().enumerate() {
196				columns[i].data_mut().push_value(value.clone());
197			}
198		}
199
200		Columns::new(columns)
201	}
202}
203
204impl<'a> Columns<'a> {
205	pub fn empty() -> Self {
206		Self {
207			row_numbers: CowVec::new(vec![]),
208			columns: CowVec::new(vec![]),
209		}
210	}
211
212	pub fn from_table(table: &ResolvedTable<'a>) -> Self {
213		let _source = table.clone();
214
215		let columns: Vec<Column> = table
216			.columns()
217			.iter()
218			.map(|col| {
219				let column_ident = Fragment::owned_internal(&col.name);
220				Column {
221					name: column_ident,
222					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
223				}
224			})
225			.collect();
226
227		Self {
228			row_numbers: CowVec::new(Vec::new()),
229			columns: CowVec::new(columns),
230		}
231	}
232
233	pub fn from_ring_buffer(ring_buffer: &ResolvedRingBuffer<'a>) -> Self {
234		let _source = ring_buffer.clone();
235
236		let columns: Vec<Column> = ring_buffer
237			.columns()
238			.iter()
239			.map(|col| {
240				let column_ident = Fragment::owned_internal(&col.name);
241				Column {
242					name: column_ident,
243					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
244				}
245			})
246			.collect();
247
248		Self {
249			row_numbers: CowVec::new(Vec::new()),
250			columns: CowVec::new(columns),
251		}
252	}
253
254	pub fn from_view(view: &ResolvedView<'a>) -> Self {
255		let _source = view.clone();
256
257		let columns: Vec<Column> = view
258			.columns()
259			.iter()
260			.map(|col| {
261				let column_ident = Fragment::owned_internal(&col.name);
262				Column {
263					name: column_ident,
264					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
265				}
266			})
267			.collect();
268
269		Self {
270			row_numbers: CowVec::new(Vec::new()),
271			columns: CowVec::new(columns),
272		}
273	}
274}
275
276#[cfg(test)]
277mod tests {
278	use reifydb_type::{Date, DateTime, Duration, Time};
279
280	use super::*;
281
282	#[test]
283	fn test_single_row_temporal_types() {
284		let date = Date::from_ymd(2025, 1, 15).unwrap();
285		let datetime = DateTime::from_timestamp(1642694400).unwrap();
286		let time = Time::from_hms(14, 30, 45).unwrap();
287		let duration = Duration::from_days(30);
288
289		let columns = Columns::single_row([
290			("date_col", Value::Date(date.clone())),
291			("datetime_col", Value::DateTime(datetime.clone())),
292			("time_col", Value::Time(time.clone())),
293			("interval_col", Value::Duration(duration.clone())),
294		]);
295
296		assert_eq!(columns.len(), 4);
297		assert_eq!(columns.shape(), (1, 4));
298
299		// Check that the values are correctly stored
300		assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
301		assert_eq!(columns.column("datetime_col").unwrap().data().get_value(0), Value::DateTime(datetime));
302		assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
303		assert_eq!(columns.column("interval_col").unwrap().data().get_value(0), Value::Duration(duration));
304	}
305
306	#[test]
307	fn test_single_row_mixed_types() {
308		let date = Date::from_ymd(2025, 7, 15).unwrap();
309		let time = Time::from_hms(9, 15, 30).unwrap();
310
311		let columns = Columns::single_row([
312			("bool_col", Value::Boolean(true)),
313			("int_col", Value::Int4(42)),
314			("str_col", Value::Utf8("hello".to_string())),
315			("date_col", Value::Date(date.clone())),
316			("time_col", Value::Time(time.clone())),
317			("undefined_col", Value::Undefined),
318		]);
319
320		assert_eq!(columns.len(), 6);
321		assert_eq!(columns.shape(), (1, 6));
322
323		// Check all values are correctly stored
324		assert_eq!(columns.column("bool_col").unwrap().data().get_value(0), Value::Boolean(true));
325		assert_eq!(columns.column("int_col").unwrap().data().get_value(0), Value::Int4(42));
326		assert_eq!(columns.column("str_col").unwrap().data().get_value(0), Value::Utf8("hello".to_string()));
327		assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
328		assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
329		assert_eq!(columns.column("undefined_col").unwrap().data().get_value(0), Value::Undefined);
330	}
331
332	#[test]
333	fn test_single_row_normal_column_names_work() {
334		let columns = Columns::single_row([("normal_column", Value::Int4(42))]);
335		assert_eq!(columns.len(), 1);
336		assert_eq!(columns.column("normal_column").unwrap().data().get_value(0), Value::Int4(42));
337	}
338}