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 {
22	pub row_numbers: CowVec<RowNumber>,
23	pub columns: CowVec<Column>,
24}
25
26impl Deref for Columns {
27	type Target = [Column];
28
29	fn deref(&self) -> &Self::Target {
30		self.columns.deref()
31	}
32}
33
34impl Index<usize> for Columns {
35	type Output = Column;
36
37	fn index(&self, index: usize) -> &Self::Output {
38		self.columns.index(index)
39	}
40}
41
42impl IndexMut<usize> for Columns {
43	fn index_mut(&mut self, index: usize) -> &mut Self::Output {
44		&mut self.columns.make_mut()[index]
45	}
46}
47
48impl Columns {
49	pub fn new(columns: Vec<Column>) -> 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>, 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 {
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::IdentityId(v) => ColumnData::identity_id([v]),
96				Value::Uuid4(v) => ColumnData::uuid4([v]),
97				Value::Uuid7(v) => ColumnData::uuid7([v]),
98				Value::Blob(v) => ColumnData::blob([v.clone()]),
99				Value::Int(v) => ColumnData::int(vec![v]),
100				Value::Uint(v) => ColumnData::uint(vec![v]),
101				Value::Decimal(v) => ColumnData::decimal(vec![v]),
102				Value::Any(v) => ColumnData::any(vec![v]),
103			};
104
105			let column = Column {
106				name: Fragment::internal(name.to_string()),
107				data,
108			};
109			index.insert(name, idx);
110			columns.push(column);
111		}
112
113		Self {
114			row_numbers: CowVec::new(Vec::new()),
115			columns: CowVec::new(columns),
116		}
117	}
118
119	pub fn apply_headers(&mut self, headers: &ColumnHeaders) {
120		// Apply the column names from headers to this Columns instance
121		for (i, name) in headers.columns.iter().enumerate() {
122			if i < self.len() {
123				let column = &mut self[i];
124				let data = std::mem::replace(column.data_mut(), ColumnData::undefined(0));
125
126				*column = Column {
127					name: name.clone(),
128					data,
129				};
130			}
131		}
132	}
133}
134
135impl Columns {
136	pub fn shape(&self) -> (usize, usize) {
137		let row_count = if !self.row_numbers.is_empty() {
138			self.row_numbers.len()
139		} else {
140			self.get(0).map(|c| c.data().len()).unwrap_or(0)
141		};
142		(row_count, self.len())
143	}
144
145	pub fn into_iter(self) -> impl Iterator<Item = Column> {
146		self.columns.into_iter()
147	}
148
149	pub fn is_empty(&self) -> bool {
150		self.shape().0 == 0
151	}
152
153	pub fn row(&self, i: usize) -> Vec<Value> {
154		self.iter().map(|c| c.data().get_value(i)).collect()
155	}
156
157	pub fn column(&self, name: &str) -> Option<&Column> {
158		self.iter().find(|col| col.name().text() == name)
159	}
160
161	pub fn row_count(&self) -> usize {
162		if !self.row_numbers.is_empty() {
163			self.row_numbers.len()
164		} else {
165			self.first().map_or(0, |col| col.data().len())
166		}
167	}
168
169	pub fn get_row(&self, index: usize) -> Vec<Value> {
170		self.iter().map(|col| col.data().get_value(index)).collect()
171	}
172}
173
174impl Column {
175	pub fn extend(&mut self, other: Column) -> crate::Result<()> {
176		self.data_mut().extend(other.data().clone())
177	}
178}
179
180impl Columns {
181	pub fn from_rows(names: &[&str], result_rows: &[Vec<Value>]) -> Self {
182		let column_count = names.len();
183
184		let mut columns: Vec<Column> = names
185			.iter()
186			.map(|name| Column {
187				name: Fragment::internal(name.to_string()),
188				data: ColumnData::Undefined(UndefinedContainer::new(0)),
189			})
190			.collect();
191
192		for row in result_rows {
193			assert_eq!(row.len(), column_count, "row length does not match column count");
194			for (i, value) in row.iter().enumerate() {
195				columns[i].data_mut().push_value(value.clone());
196			}
197		}
198
199		Columns::new(columns)
200	}
201}
202
203impl Columns {
204	pub fn empty() -> Self {
205		Self {
206			row_numbers: CowVec::new(vec![]),
207			columns: CowVec::new(vec![]),
208		}
209	}
210
211	pub fn from_table(table: &ResolvedTable) -> Self {
212		let _source = table.clone();
213
214		let columns: Vec<Column> = table
215			.columns()
216			.iter()
217			.map(|col| {
218				let column_ident = Fragment::internal(&col.name);
219				Column {
220					name: column_ident,
221					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
222				}
223			})
224			.collect();
225
226		Self {
227			row_numbers: CowVec::new(Vec::new()),
228			columns: CowVec::new(columns),
229		}
230	}
231
232	pub fn from_ringbuffer(ringbuffer: &ResolvedRingBuffer) -> Self {
233		let _source = ringbuffer.clone();
234
235		let columns: Vec<Column> = ringbuffer
236			.columns()
237			.iter()
238			.map(|col| {
239				let column_ident = Fragment::internal(&col.name);
240				Column {
241					name: column_ident,
242					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
243				}
244			})
245			.collect();
246
247		Self {
248			row_numbers: CowVec::new(Vec::new()),
249			columns: CowVec::new(columns),
250		}
251	}
252
253	pub fn from_view(view: &ResolvedView) -> Self {
254		let _source = view.clone();
255
256		let columns: Vec<Column> = view
257			.columns()
258			.iter()
259			.map(|col| {
260				let column_ident = Fragment::internal(&col.name);
261				Column {
262					name: column_ident,
263					data: ColumnData::with_capacity(col.constraint.get_type(), 0),
264				}
265			})
266			.collect();
267
268		Self {
269			row_numbers: CowVec::new(Vec::new()),
270			columns: CowVec::new(columns),
271		}
272	}
273}
274
275#[cfg(test)]
276mod tests {
277	use reifydb_type::{Date, DateTime, Duration, Time};
278
279	use super::*;
280
281	#[test]
282	fn test_single_row_temporal_types() {
283		let date = Date::from_ymd(2025, 1, 15).unwrap();
284		let datetime = DateTime::from_timestamp(1642694400).unwrap();
285		let time = Time::from_hms(14, 30, 45).unwrap();
286		let duration = Duration::from_days(30);
287
288		let columns = Columns::single_row([
289			("date_col", Value::Date(date.clone())),
290			("datetime_col", Value::DateTime(datetime.clone())),
291			("time_col", Value::Time(time.clone())),
292			("interval_col", Value::Duration(duration.clone())),
293		]);
294
295		assert_eq!(columns.len(), 4);
296		assert_eq!(columns.shape(), (1, 4));
297
298		// Check that the values are correctly stored
299		assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
300		assert_eq!(columns.column("datetime_col").unwrap().data().get_value(0), Value::DateTime(datetime));
301		assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
302		assert_eq!(columns.column("interval_col").unwrap().data().get_value(0), Value::Duration(duration));
303	}
304
305	#[test]
306	fn test_single_row_mixed_types() {
307		let date = Date::from_ymd(2025, 7, 15).unwrap();
308		let time = Time::from_hms(9, 15, 30).unwrap();
309
310		let columns = Columns::single_row([
311			("bool_col", Value::Boolean(true)),
312			("int_col", Value::Int4(42)),
313			("str_col", Value::Utf8("hello".to_string())),
314			("date_col", Value::Date(date.clone())),
315			("time_col", Value::Time(time.clone())),
316			("undefined_col", Value::Undefined),
317		]);
318
319		assert_eq!(columns.len(), 6);
320		assert_eq!(columns.shape(), (1, 6));
321
322		// Check all values are correctly stored
323		assert_eq!(columns.column("bool_col").unwrap().data().get_value(0), Value::Boolean(true));
324		assert_eq!(columns.column("int_col").unwrap().data().get_value(0), Value::Int4(42));
325		assert_eq!(columns.column("str_col").unwrap().data().get_value(0), Value::Utf8("hello".to_string()));
326		assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
327		assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
328		assert_eq!(columns.column("undefined_col").unwrap().data().get_value(0), Value::Undefined);
329	}
330
331	#[test]
332	fn test_single_row_normal_column_names_work() {
333		let columns = Columns::single_row([("normal_column", Value::Int4(42))]);
334		assert_eq!(columns.len(), 1);
335		assert_eq!(columns.column("normal_column").unwrap().data().get_value(0), Value::Int4(42));
336	}
337}