reifydb_core/value/column/
columns.rs1use 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 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 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 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}