1use std::{
5 hash::Hash,
6 ops::{Index, IndexMut},
7};
8
9use indexmap::IndexMap;
10use reifydb_codec::row::{
11 bytes::EncodedBytes,
12 shape::{RowFamily, RowShape},
13};
14use reifydb_value::{
15 Result,
16 fragment::Fragment,
17 reifydb_assertions,
18 value::{
19 Value,
20 constraint::Constraint,
21 datetime::{CREATED_AT_COLUMN_NAME, DateTime, TIME_COLUMN_NAME, UPDATED_AT_COLUMN_NAME},
22 partition::Partition,
23 row_number::{ROW_NUMBER_COLUMN_NAME, RowNumber},
24 system_columns::{RowStamps, SystemColumns},
25 value_type::ValueType,
26 },
27};
28use serde::{Deserialize, Serialize};
29
30use crate::{
31 interface::catalog::column::Column as CatalogColumn,
32 return_internal_error,
33 row::Row,
34 value::column::{ColumnBuffer, ColumnWithName, data::Column, headers::ColumnHeaders},
35};
36
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct Columns {
39 pub system: SystemColumns,
40 pub columns: Vec<ColumnBuffer>,
41 pub names: Vec<Fragment>,
42}
43
44impl Columns {
45 #[inline]
46 pub fn row_numbers(&self) -> &[RowNumber] {
47 self.system.row_numbers()
48 }
49
50 #[inline]
51 pub fn partitions(&self) -> &[Partition] {
52 self.system.partitions()
53 }
54
55 #[inline]
56 pub fn created_at(&self) -> &[DateTime] {
57 self.system.created_at()
58 }
59
60 #[inline]
61 pub fn updated_at(&self) -> &[DateTime] {
62 self.system.updated_at()
63 }
64
65 #[inline]
66 pub fn time(&self) -> &[DateTime] {
67 self.system.time()
68 }
69
70 pub fn system_column(&self, name: &str) -> Option<ColumnBuffer> {
71 let name = name.strip_prefix('#').unwrap_or(name);
72
73 if name == ROW_NUMBER_COLUMN_NAME && !self.row_numbers().is_empty() {
74 let values: Vec<u64> = self.row_numbers().iter().map(|r| r.value()).collect();
75 return Some(ColumnBuffer::uint8(values));
76 }
77 if name == CREATED_AT_COLUMN_NAME && !self.created_at().is_empty() {
78 return Some(ColumnBuffer::datetime(self.created_at().to_vec()));
79 }
80 if name == UPDATED_AT_COLUMN_NAME && !self.updated_at().is_empty() {
81 return Some(ColumnBuffer::datetime(self.updated_at().to_vec()));
82 }
83 if name == TIME_COLUMN_NAME && !self.time().is_empty() {
84 return Some(ColumnBuffer::datetime(self.time().to_vec()));
85 }
86 None
87 }
88}
89
90#[derive(Debug, Clone, Copy)]
91pub struct ColumnRef<'a> {
92 name: &'a Fragment,
93 data: &'a ColumnBuffer,
94}
95
96impl Index<usize> for Columns {
97 type Output = ColumnBuffer;
98
99 fn index(&self, index: usize) -> &Self::Output {
100 &self.columns[index]
101 }
102}
103
104impl IndexMut<usize> for Columns {
105 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
106 &mut self.columns[index]
107 }
108}
109
110impl<'a> ColumnRef<'a> {
111 pub fn new(name: &'a Fragment, data: &'a ColumnBuffer) -> Self {
112 Self {
113 name,
114 data,
115 }
116 }
117
118 pub fn name(&self) -> &'a Fragment {
119 self.name
120 }
121
122 pub fn data(&self) -> &'a ColumnBuffer {
123 self.data
124 }
125
126 pub fn get_type(&self) -> ValueType {
127 self.data.get_type()
128 }
129
130 pub fn column(&self) -> Column {
131 Column::from_column_buffer(self.data.clone())
132 }
133
134 pub fn with_new_data(&self, data: ColumnBuffer) -> ColumnWithName {
135 ColumnWithName::new(self.name.clone(), data)
136 }
137}
138
139fn value_to_buffer(value: Value) -> ColumnBuffer {
140 match value {
141 Value::None {
142 inner,
143 } => ColumnBuffer::none_typed(inner, 1),
144 Value::Boolean(v) => ColumnBuffer::bool([v]),
145 Value::Float4(v) => ColumnBuffer::float4([v.into()]),
146 Value::Float8(v) => ColumnBuffer::float8([v.into()]),
147 Value::Int1(v) => ColumnBuffer::int1([v]),
148 Value::Int2(v) => ColumnBuffer::int2([v]),
149 Value::Int4(v) => ColumnBuffer::int4([v]),
150 Value::Int8(v) => ColumnBuffer::int8([v]),
151 Value::Int16(v) => ColumnBuffer::int16([v]),
152 Value::Utf8(v) => ColumnBuffer::utf8([v]),
153 Value::Uint1(v) => ColumnBuffer::uint1([v]),
154 Value::Uint2(v) => ColumnBuffer::uint2([v]),
155 Value::Uint4(v) => ColumnBuffer::uint4([v]),
156 Value::Uint8(v) => ColumnBuffer::uint8([v]),
157 Value::Uint16(v) => ColumnBuffer::uint16([v]),
158 Value::Date(v) => ColumnBuffer::date([v]),
159 Value::DateTime(v) => ColumnBuffer::datetime([v]),
160 Value::Time(v) => ColumnBuffer::time([v]),
161 Value::Duration(v) => ColumnBuffer::duration([v]),
162 Value::IdentityId(v) => ColumnBuffer::identity_id([v]),
163 Value::Uuid4(v) => ColumnBuffer::uuid4([v]),
164 Value::Uuid7(v) => ColumnBuffer::uuid7([v]),
165 Value::Blob(v) => ColumnBuffer::blob([v]),
166 Value::Int(v) => ColumnBuffer::int(vec![v]),
167 Value::Uint(v) => ColumnBuffer::uint(vec![v]),
168 Value::Decimal(v) => ColumnBuffer::decimal(vec![v]),
169 Value::DictionaryId(v) => ColumnBuffer::dictionary_id(vec![v]),
170 Value::Any(v) => ColumnBuffer::any(vec![*v]),
171 Value::Type(v) => ColumnBuffer::any(vec![Value::Type(v)]),
172 Value::List(v) => ColumnBuffer::any(vec![Value::List(v)]),
173 Value::Record(v) => ColumnBuffer::any(vec![Value::Record(v)]),
174 Value::Tuple(v) => ColumnBuffer::any(vec![Value::Tuple(v)]),
175 }
176}
177
178impl Columns {
179 pub fn scalar_value(&self) -> Value {
180 reifydb_assertions! {
181 assert_eq!(self.len(), 1, "scalar_value() requires exactly 1 column, got {}", self.len());
182 assert_eq!(
183 self.row_count(),
184 1,
185 "scalar_value() requires exactly 1 row, got {}",
186 self.row_count()
187 );
188 }
189 self.columns[0].get_value(0)
190 }
191
192 pub fn new(columns: Vec<ColumnWithName>) -> Self {
193 let n = columns.first().map_or(0, |c| c.data.len());
194 assert!(columns.iter().all(|c| c.data.len() == n));
195
196 let mut names = Vec::with_capacity(columns.len());
197 let mut buffers = Vec::with_capacity(columns.len());
198 for c in columns {
199 names.push(c.name);
200 buffers.push(c.data);
201 }
202
203 Self {
204 system: SystemColumns::empty(),
205 columns: buffers,
206 names,
207 }
208 }
209
210 pub fn with_system(columns: Vec<ColumnWithName>, system: SystemColumns) -> Self {
211 let n = columns.first().map_or(0, |c| c.data.len());
212 assert!(columns.iter().all(|c| c.data.len() == n));
213 system.assert_invariants(n, "Columns::with_system");
214
215 let mut names = Vec::with_capacity(columns.len());
216 let mut buffers = Vec::with_capacity(columns.len());
217 for c in columns {
218 names.push(c.name);
219 buffers.push(c.data);
220 }
221
222 Self {
223 system,
224 columns: buffers,
225 names,
226 }
227 }
228
229 pub fn single_row<'b>(rows: impl IntoIterator<Item = (&'b str, Value)>) -> Columns {
230 let mut names = Vec::new();
231 let mut buffers = Vec::new();
232 for (name, value) in rows {
233 names.push(Fragment::internal(name));
234 buffers.push(value_to_buffer(value));
235 }
236 Self {
237 system: SystemColumns::empty(),
238 columns: buffers,
239 names,
240 }
241 }
242
243 pub fn with_row_numbers(mut self, row_numbers: Vec<RowNumber>) -> Self {
244 let n = row_numbers.len();
245 let now = DateTime::default();
246 let keep = |existing: &[DateTime]| {
247 if existing.len() == n {
248 existing.to_vec()
249 } else {
250 vec![now; n]
251 }
252 };
253 self.system = SystemColumns::new(
254 row_numbers,
255 self.system.partitions().to_vec(),
256 keep(self.system.created_at()),
257 keep(self.system.updated_at()),
258 keep(self.system.time()),
259 );
260 self
261 }
262
263 pub fn from_catalog_columns(cols: &[CatalogColumn]) -> Self {
264 let mut names = Vec::with_capacity(cols.len());
265 let mut buffers = Vec::with_capacity(cols.len());
266 for col in cols {
267 names.push(Fragment::internal(&col.name));
268 buffers.push(ColumnBuffer::with_capacity(col.constraint.get_type(), 0));
269 }
270 Self {
271 system: SystemColumns::empty(),
272 columns: buffers,
273 names,
274 }
275 }
276
277 pub fn apply_headers(&mut self, headers: &ColumnHeaders) {
278 let n = self.len();
279 let names = &mut self.names;
280 for (i, name) in headers.columns.iter().enumerate() {
281 if i < n {
282 names[i] = name.clone();
283 }
284 }
285 }
286}
287
288impl Columns {
289 pub fn number(&self) -> RowNumber {
290 assert_eq!(self.row_count(), 1, "number() requires exactly 1 row, got {}", self.row_count());
291 if self.row_numbers().is_empty() {
292 RowNumber(0)
293 } else {
294 self.row_numbers()[0]
295 }
296 }
297
298 pub fn shape(&self) -> (usize, usize) {
299 let row_count = if !self.row_numbers().is_empty() {
300 self.row_numbers().len()
301 } else {
302 self.columns.first().map(|c| c.len()).unwrap_or(0)
303 };
304 (row_count, self.len())
305 }
306
307 pub fn heap_size(&self) -> usize {
308 let data: usize = self.columns.iter().map(|c| c.heap_size()).sum();
309 let names: usize = self.names.iter().map(|n| n.text().len()).sum();
310 data + names + self.system.heap_size()
311 }
312
313 pub fn len(&self) -> usize {
314 self.columns.len()
315 }
316
317 pub fn is_empty(&self) -> bool {
318 self.columns.is_empty()
319 }
320
321 pub fn iter(&self) -> impl Iterator<Item = ColumnRef<'_>> + '_ {
322 self.names.iter().zip(self.columns.iter()).map(|(n, d)| ColumnRef::new(n, d))
323 }
324
325 pub fn first(&self) -> Option<ColumnRef<'_>> {
326 self.get(0)
327 }
328
329 pub fn last(&self) -> Option<ColumnRef<'_>> {
330 let n = self.len();
331 if n == 0 {
332 None
333 } else {
334 self.get(n - 1)
335 }
336 }
337
338 pub fn get(&self, index: usize) -> Option<ColumnRef<'_>> {
339 if index < self.len() {
340 Some(ColumnRef::new(&self.names[index], &self.columns[index]))
341 } else {
342 None
343 }
344 }
345
346 pub fn name_at(&self, index: usize) -> &Fragment {
347 &self.names[index]
348 }
349
350 pub fn data_at(&self, index: usize) -> &ColumnBuffer {
351 &self.columns[index]
352 }
353
354 pub fn data_at_mut(&mut self, index: usize) -> &mut ColumnBuffer {
355 &mut self.columns[index]
356 }
357
358 pub fn row(&self, i: usize) -> Vec<Value> {
359 self.columns.iter().map(|c| c.get_value(i)).collect()
360 }
361
362 pub fn column(&self, name: &str) -> Option<ColumnRef<'_>> {
363 self.names.iter().position(|n| n.text() == name).and_then(|i| self.get(i))
364 }
365
366 pub fn row_count(&self) -> usize {
367 if !self.row_numbers().is_empty() {
368 self.row_numbers().len()
369 } else {
370 self.columns.first().map_or(0, |col| col.len())
371 }
372 }
373
374 pub fn has_rows(&self) -> bool {
375 self.row_count() > 0
376 }
377
378 pub fn is_scalar(&self) -> bool {
379 self.len() == 1 && self.row_count() == 1
380 }
381
382 pub fn get_row(&self, index: usize) -> Vec<Value> {
383 self.columns.iter().map(|col| col.get_value(index)).collect()
384 }
385
386 #[track_caller]
387 pub fn assert_invariants(&self, ctx: &str) {
388 let n = self.columns.first().map_or(0, |c| c.len());
389 for (i, col) in self.columns.iter().enumerate() {
390 assert_eq!(
391 col.len(),
392 n,
393 "{ctx}: Columns column[{i}] has length {} but columns[0] has length {n}",
394 col.len(),
395 );
396 }
397 self.system.assert_invariants(n, ctx);
398 }
399}
400
401impl Columns {
402 pub fn from_rows(names: &[&str], result_rows: &[Vec<Value>]) -> Self {
403 let column_count = names.len();
404
405 let mut name_vec: Vec<Fragment> = names.iter().map(Fragment::internal).collect();
406 let mut buffers: Vec<ColumnBuffer> =
407 (0..column_count).map(|_| ColumnBuffer::none_typed(ValueType::Boolean, 0)).collect();
408
409 for row in result_rows {
410 assert_eq!(row.len(), column_count, "row length does not match column count");
411 for (i, value) in row.iter().enumerate() {
412 buffers[i].push_value(value.clone());
413 }
414 }
415
416 let _ = &mut name_vec;
417 Self {
418 system: SystemColumns::empty(),
419 columns: buffers,
420 names: name_vec,
421 }
422 }
423
424 pub fn from_encoded_bytes(shape: &RowShape, ids: &[RowNumber], bytes_slice: &[EncodedBytes]) -> Self {
425 assert_eq!(ids.len(), bytes_slice.len(), "ids length must match rows length");
426 let fields = shape.fields();
427 let row_count = bytes_slice.len();
428
429 let mut columns_vec: Vec<ColumnWithName> = Vec::with_capacity(fields.len());
430 for field in fields.iter() {
431 let mut data = ColumnBuffer::with_capacity(field.constraint.get_type(), row_count);
432 if field.constraint.get_type() == ValueType::DictionaryId
433 && let ColumnBuffer::DictionaryId(container) = &mut data
434 && let Some(Constraint::Dictionary(dict_id, _)) = field.constraint.constraint()
435 {
436 container.set_dictionary_id(*dict_id);
437 }
438 columns_vec.push(ColumnWithName {
439 name: Fragment::internal(&field.name),
440 data,
441 });
442 }
443
444 for encoded in bytes_slice {
445 for (i, _) in fields.iter().enumerate() {
446 columns_vec[i].data.push_value(shape.get_value(encoded, i));
447 }
448 }
449
450 let row_numbers: Vec<RowNumber> = ids.to_vec();
451 let (created_at, updated_at): (Vec<DateTime>, Vec<DateTime>) = match shape.family() {
452 RowFamily::Pod => (Vec::new(), Vec::new()),
453 _ => (
454 bytes_slice.iter().map(|r| shape.created_at(r)).collect(),
455 bytes_slice.iter().map(|r| shape.updated_at(r)).collect(),
456 ),
457 };
458 let time: Vec<DateTime> = bytes_slice.iter().filter_map(|r| shape.time(r)).collect();
459
460 Self::with_system(
461 columns_vec,
462 SystemColumns::new(row_numbers, Vec::new(), created_at, updated_at, time),
463 )
464 }
465}
466
467impl Columns {
468 pub fn empty() -> Self {
469 Self {
470 system: SystemColumns::empty(),
471 columns: Vec::new(),
472 names: Vec::new(),
473 }
474 }
475}
476
477impl Default for Columns {
478 fn default() -> Self {
479 Self::empty()
480 }
481}
482
483impl Columns {
484 pub fn extract_by_indices(&self, indices: &[usize]) -> Columns {
485 if indices.is_empty() {
486 return Columns::empty();
487 }
488
489 let mut new_buffers: Vec<ColumnBuffer> = Vec::with_capacity(self.columns.len());
490 for col in self.columns.iter() {
491 let mut new_data = col.empty_like(indices.len());
492 for &idx in indices {
493 new_data.push_value(col.get_value(idx));
494 }
495 new_buffers.push(new_data);
496 }
497
498 Columns {
499 system: self.system.permute(indices),
500 columns: new_buffers,
501 names: self.names.clone(),
502 }
503 }
504
505 pub fn extract_row(&self, index: usize) -> Columns {
506 self.extract_by_indices(&[index])
507 }
508
509 pub fn append(&mut self, source: Columns) -> Result<()> {
510 if source.row_count() == 0 {
511 return Ok(());
512 }
513 if self.columns.is_empty() {
514 *self = source;
515 return Ok(());
516 }
517
518 self.validate_append_compatibility(&source)?;
519 self.system.extend(&source.system)?;
520 self.extend_data_columns(source.columns)?;
521 Ok(())
522 }
523
524 #[inline]
525 fn validate_append_compatibility(&self, source: &Columns) -> Result<()> {
526 if self.columns.len() != source.columns.len() {
527 return_internal_error!(
528 "Columns::append: column count mismatch (self={}, source={})",
529 self.columns.len(),
530 source.columns.len()
531 );
532 }
533 Ok(())
534 }
535
536 #[inline]
537 fn extend_data_columns(&mut self, source_columns: Vec<ColumnBuffer>) -> Result<()> {
538 let dest_cols = &mut self.columns;
539 reifydb_assertions! {
540 let dest_len = dest_cols.len();
541 let src_len = source_columns.len();
542 assert!(
543 dest_len == src_len,
544 "append extends destination columns by source index, so a source with more columns than \
545 the destination would index dest_cols out of bounds and panic mid-append, leaving self \
546 partially extended (dest_len={dest_len}, src_len={src_len})"
547 );
548 }
549 for (i, src_col) in source_columns.into_iter().enumerate() {
550 dest_cols[i].extend(src_col)?;
551 }
552 Ok(())
553 }
554
555 pub fn concat(batches: Vec<Columns>) -> Result<Option<Columns>> {
556 let mut iter = batches.into_iter();
557 let mut merged = match iter.next() {
558 Some(first) => first,
559 None => return Ok(None),
560 };
561 for cols in iter {
562 merged.append(cols)?;
563 }
564 if merged.row_count() == 0 {
565 return Ok(None);
566 }
567 Ok(Some(merged))
568 }
569
570 pub fn remove_row(&mut self, row_number: RowNumber) -> bool {
571 let pos = self.row_numbers().iter().position(|&r| r == row_number);
572 let Some(idx) = pos else {
573 return false;
574 };
575
576 let kept_indices: Vec<usize> = (0..self.row_count()).filter(|&i| i != idx).collect();
577 *self = self.extract_by_indices(&kept_indices);
578 true
579 }
580
581 pub fn project_by_names(&self, names: &[String]) -> Columns {
582 let mut new_names = Vec::new();
583 let mut new_buffers = Vec::new();
584
585 for name in names {
586 if let Some(pos) = self.names.iter().position(|n| n.text() == name.as_str()) {
587 new_names.push(self.names[pos].clone());
588 new_buffers.push(self.columns[pos].clone());
589 }
590 }
591
592 if new_buffers.is_empty() {
593 return Columns::empty();
594 }
595
596 Columns {
597 system: self.system.clone(),
598 columns: new_buffers,
599 names: new_names,
600 }
601 }
602
603 pub fn partition_by_keys<K: Hash + Eq + Clone>(&self, keys: &[K]) -> IndexMap<K, Columns> {
604 assert_eq!(keys.len(), self.row_count(), "keys length must match row count");
605
606 let mut key_to_indices: IndexMap<K, Vec<usize>> = IndexMap::new();
607 for (idx, key) in keys.iter().enumerate() {
608 key_to_indices.entry(key.clone()).or_default().push(idx);
609 }
610
611 key_to_indices.into_iter().map(|(key, indices)| (key, self.extract_by_indices(&indices))).collect()
612 }
613
614 pub fn from_row(row: &Row) -> Self {
615 let mut out = Columns::empty();
616 out.reset_from_row(row);
617 out
618 }
619
620 pub fn reset_from_row(&mut self, row: &Row) {
621 let field_count = row.shape.fields().len();
622
623 self.system.clear();
624 self.columns.clear();
625 self.names.clear();
626
627 self.columns.reserve(field_count);
628 self.names.reserve(field_count);
629
630 let (created_at, updated_at) = match row.shape.family() {
631 RowFamily::Pod | RowFamily::Operator => (None, None),
632 _ => (Some(row.shape.created_at(&row.encoded)), Some(row.shape.updated_at(&row.encoded))),
633 };
634
635 self.system.push(RowStamps {
636 row_number: Some(row.number),
637 partition: None,
638 created_at,
639 updated_at,
640 time: row.shape.time(&row.encoded),
641 });
642
643 for (idx, field) in row.shape.fields().iter().enumerate() {
644 let value = row.shape.get_value(&row.encoded, idx);
645
646 let column_type = if matches!(value, Value::None { .. }) {
647 field.constraint.get_type()
648 } else {
649 value.get_type()
650 };
651
652 let mut data = if column_type.is_option() {
653 ColumnBuffer::none_typed(column_type.clone(), 0)
654 } else {
655 ColumnBuffer::with_capacity(column_type.clone(), 1)
656 };
657 data.push_value(value);
658
659 if column_type == ValueType::DictionaryId
660 && let ColumnBuffer::DictionaryId(container) = &mut data
661 && let Some(Constraint::Dictionary(dict_id, _)) = field.constraint.constraint()
662 {
663 container.set_dictionary_id(*dict_id);
664 }
665
666 let name = row.shape.get_field_name(idx).expect("RowShape missing name for field");
667
668 self.names.push(Fragment::internal(name));
669 self.columns.push(data);
670 }
671 }
672}
673
674#[cfg(test)]
675pub mod tests {
676 use std::str::FromStr;
677
678 use reifydb_value::value::{
679 blob::Blob,
680 constraint::{bytes::MaxBytes, precision::Precision, scale::Scale},
681 date::Date,
682 datetime::DateTime,
683 decimal::Decimal,
684 dictionary::{DictionaryEntryId, DictionaryId},
685 duration::Duration,
686 identity::IdentityId,
687 int::Int,
688 time::Time,
689 uint::Uint,
690 uuid::{Uuid4, Uuid7},
691 };
692 use uuid::{Timestamp, Uuid};
693
694 use super::*;
695
696 fn uuid7_at(a: u64, b: u16) -> Uuid7 {
697 Uuid7::from(Uuid::new_v7(Timestamp::from_gregorian_time(a, b)))
698 }
699
700 fn assert_extract_preserves_values(buffer: ColumnBuffer, indices: &[usize]) {
703 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
704 let extracted = original.extract_by_indices(indices);
705
706 assert_eq!(extracted.len(), 1, "column count must be preserved");
707 assert_eq!(extracted.row_count(), indices.len(), "row count must equal number of indices");
708
709 let src = original.data_at(0);
710 let dst = extracted.data_at(0);
711 assert_eq!(dst.get_type(), src.get_type(), "value type must be preserved");
712 for (j, &idx) in indices.iter().enumerate() {
713 assert_eq!(
714 dst.get_value(j),
715 src.get_value(idx),
716 "value at extracted row {j} must equal source row {idx}"
717 );
718 }
719 }
720
721 #[test]
722 fn extract_by_indices_preserves_bool_values() {
723 assert_extract_preserves_values(ColumnBuffer::bool([true, false, true, false]), &[3, 1, 2]);
724 }
725
726 #[test]
727 fn extract_by_indices_preserves_float4_values() {
728 assert_extract_preserves_values(ColumnBuffer::float4([1.0f32, 2.5, -3.0, 4.25]), &[3, 1, 2]);
729 }
730
731 #[test]
732 fn extract_by_indices_preserves_float8_values() {
733 assert_extract_preserves_values(ColumnBuffer::float8([1.0f64, 2.5, -3.0, 4.25]), &[3, 1, 2]);
734 }
735
736 #[test]
737 fn extract_by_indices_preserves_int1_values() {
738 assert_extract_preserves_values(ColumnBuffer::int1([-1i8, 2, -3, 4]), &[3, 1, 2]);
739 }
740
741 #[test]
742 fn extract_by_indices_preserves_int2_values() {
743 assert_extract_preserves_values(ColumnBuffer::int2([-1i16, 2, -3, 4]), &[3, 1, 2]);
744 }
745
746 #[test]
747 fn extract_by_indices_preserves_int4_values() {
748 assert_extract_preserves_values(ColumnBuffer::int4([-1i32, 2, -3, 4]), &[3, 1, 2]);
749 }
750
751 #[test]
752 fn extract_by_indices_preserves_int8_values() {
753 assert_extract_preserves_values(ColumnBuffer::int8([-1i64, 2, -3, 4]), &[3, 1, 2]);
754 }
755
756 #[test]
757 fn extract_by_indices_preserves_int16_values() {
758 assert_extract_preserves_values(ColumnBuffer::int16([-1i128, 2, -3, 4]), &[3, 1, 2]);
759 }
760
761 #[test]
762 fn extract_by_indices_preserves_uint1_values() {
763 assert_extract_preserves_values(ColumnBuffer::uint1([1u8, 2, 3, 4]), &[3, 1, 2]);
764 }
765
766 #[test]
767 fn extract_by_indices_preserves_uint2_values() {
768 assert_extract_preserves_values(ColumnBuffer::uint2([1u16, 2, 3, 4]), &[3, 1, 2]);
769 }
770
771 #[test]
772 fn extract_by_indices_preserves_uint4_values() {
773 assert_extract_preserves_values(ColumnBuffer::uint4([1u32, 2, 3, 4]), &[3, 1, 2]);
774 }
775
776 #[test]
777 fn extract_by_indices_preserves_uint8_values() {
778 assert_extract_preserves_values(ColumnBuffer::uint8([1u64, 2, 3, 4]), &[3, 1, 2]);
779 }
780
781 #[test]
782 fn extract_by_indices_preserves_uint16_values() {
783 assert_extract_preserves_values(ColumnBuffer::uint16([1u128, 2, 3, 4]), &[3, 1, 2]);
784 }
785
786 #[test]
787 fn extract_by_indices_preserves_utf8_values() {
788 assert_extract_preserves_values(ColumnBuffer::utf8(["a", "bb", "ccc", "dddd"]), &[3, 1, 2]);
789 }
790
791 #[test]
792 fn extract_by_indices_preserves_date_values() {
793 let data = [
794 Date::from_ymd(2025, 1, 1).unwrap(),
795 Date::from_ymd(2025, 6, 15).unwrap(),
796 Date::from_ymd(2024, 12, 31).unwrap(),
797 Date::from_ymd(2000, 2, 29).unwrap(),
798 ];
799 assert_extract_preserves_values(ColumnBuffer::date(data), &[3, 1, 2]);
800 }
801
802 #[test]
803 fn extract_by_indices_preserves_datetime_values() {
804 let data = [
805 DateTime::from_epoch_secs(1000).unwrap(),
806 DateTime::from_epoch_secs(2000).unwrap(),
807 DateTime::from_epoch_secs(3000).unwrap(),
808 DateTime::from_epoch_secs(4000).unwrap(),
809 ];
810 assert_extract_preserves_values(ColumnBuffer::datetime(data), &[3, 1, 2]);
811 }
812
813 #[test]
814 fn extract_by_indices_preserves_time_values() {
815 let data = [
816 Time::from_hms(0, 0, 0).unwrap(),
817 Time::from_hms(12, 30, 45).unwrap(),
818 Time::from_hms(23, 59, 59).unwrap(),
819 Time::from_hms(6, 15, 0).unwrap(),
820 ];
821 assert_extract_preserves_values(ColumnBuffer::time(data), &[3, 1, 2]);
822 }
823
824 #[test]
825 fn extract_by_indices_preserves_duration_values() {
826 let data = [
827 Duration::from_days(1).unwrap(),
828 Duration::from_days(7).unwrap(),
829 Duration::from_days(30).unwrap(),
830 Duration::from_days(365).unwrap(),
831 ];
832 assert_extract_preserves_values(ColumnBuffer::duration(data), &[3, 1, 2]);
833 }
834
835 #[test]
836 fn extract_by_indices_preserves_identity_id_values() {
837 let data = [IdentityId::root(), IdentityId::system(), IdentityId::anonymous(), IdentityId::root()];
838 assert_extract_preserves_values(ColumnBuffer::identity_id(data), &[3, 1, 2]);
839 }
840
841 #[test]
842 fn extract_by_indices_preserves_uuid4_values() {
843 let data = [Uuid4::generate(), Uuid4::generate(), Uuid4::generate(), Uuid4::generate()];
844 assert_extract_preserves_values(ColumnBuffer::uuid4(data), &[3, 1, 2]);
845 }
846
847 #[test]
848 fn extract_by_indices_preserves_uuid7_values() {
849 let data = [uuid7_at(1, 1), uuid7_at(1, 2), uuid7_at(2, 1), uuid7_at(2, 2)];
850 assert_extract_preserves_values(ColumnBuffer::uuid7(data), &[3, 1, 2]);
851 }
852
853 #[test]
854 fn extract_by_indices_preserves_blob_values() {
855 let data = [
856 Blob::new(vec![1]),
857 Blob::new(vec![2, 3]),
858 Blob::new(vec![4, 5, 6]),
859 Blob::new(vec![7, 8, 9, 10]),
860 ];
861 assert_extract_preserves_values(ColumnBuffer::blob(data), &[3, 1, 2]);
862 }
863
864 #[test]
865 fn extract_by_indices_preserves_int_values() {
866 let data = [Int::from(-1i64), Int::from(2i64), Int::from(-3i64), Int::from(4i64)];
867 assert_extract_preserves_values(ColumnBuffer::int(data), &[3, 1, 2]);
868 }
869
870 #[test]
871 fn extract_by_indices_preserves_uint_values() {
872 let data = [Uint::from(1u64), Uint::from(2u64), Uint::from(3u64), Uint::from(4u64)];
873 assert_extract_preserves_values(ColumnBuffer::uint(data), &[3, 1, 2]);
874 }
875
876 #[test]
877 fn extract_by_indices_preserves_decimal_values() {
878 let data = [
879 Decimal::from_str("1.50").unwrap(),
880 Decimal::from_str("2.25").unwrap(),
881 Decimal::from_str("-3.75").unwrap(),
882 Decimal::from_str("4.00").unwrap(),
883 ];
884 assert_extract_preserves_values(ColumnBuffer::decimal(data), &[3, 1, 2]);
885 }
886
887 #[test]
888 fn extract_by_indices_preserves_any_values() {
889 let data = [Value::Int4(1), Value::Utf8("two".to_string()), Value::Boolean(true), Value::none()];
890 assert_extract_preserves_values(ColumnBuffer::any(data), &[3, 1, 2]);
891 }
892
893 #[test]
894 fn extract_by_indices_preserves_dictionary_id_values() {
895 let data = [
896 DictionaryEntryId::U2(10),
897 DictionaryEntryId::U2(20),
898 DictionaryEntryId::U2(30),
899 DictionaryEntryId::U2(40),
900 ];
901 assert_extract_preserves_values(ColumnBuffer::dictionary_id(data), &[3, 1, 2]);
902 }
903
904 #[test]
905 fn extract_by_indices_preserves_option_values_including_none() {
906 let mut buffer = ColumnBuffer::with_capacity(ValueType::Option(Box::new(ValueType::Int4)), 0);
907 buffer.push_value(Value::Int4(1));
908 buffer.push_value(Value::none());
909 buffer.push_value(Value::Int4(3));
910 buffer.push_value(Value::none());
911 assert_extract_preserves_values(buffer, &[3, 1, 2, 0]);
912 }
913
914 #[test]
915 fn extract_by_indices_empty_indices_yields_empty_columns() {
916 let original = Columns::new(vec![ColumnWithName::int4("c", [1, 2, 3])]);
917 let extracted = original.extract_by_indices(&[]);
918 assert_eq!(extracted.row_count(), 0);
919 assert!(extracted.is_empty());
920 }
921
922 #[test]
923 fn extract_by_indices_full_identity_reproduces_all_rows() {
924 assert_extract_preserves_values(ColumnBuffer::int4([10, 20, 30, 40]), &[0, 1, 2, 3]);
925 }
926
927 #[test]
928 fn heap_size_grows_with_row_count() {
929 let small = Columns::new(vec![ColumnWithName::int4("c", [1i32, 2, 3, 4])]);
930 let large = Columns::new(vec![ColumnWithName::int4("c", 0..4000i32)]);
931 assert!(
932 large.heap_size() > small.heap_size() + 4000,
933 "heap_size must scale with the number of buffered rows (small={}, large={})",
934 small.heap_size(),
935 large.heap_size()
936 );
937 }
938
939 #[test]
940 fn heap_size_counts_utf8_payload_not_just_row_count() {
941 let short = Columns::new(vec![ColumnWithName::new("c", ColumnBuffer::utf8(["a", "b", "c"]))]);
944 let long_value = "x".repeat(4096);
945 let long = Columns::new(vec![ColumnWithName::new(
946 "c",
947 ColumnBuffer::utf8([long_value.clone(), long_value.clone(), long_value.clone()]),
948 )]);
949 assert_eq!(short.row_count(), long.row_count(), "same row count is the point of the test");
950 assert!(
951 long.heap_size() >= short.heap_size() + 3 * 4096,
952 "heap_size must account for utf8 payload bytes (short={}, long={})",
953 short.heap_size(),
954 long.heap_size()
955 );
956 }
957
958 #[test]
959 fn extract_by_indices_duplicate_index_duplicates_row() {
960 let original = Columns::new(vec![ColumnWithName::int4("c", [10, 20, 30])]);
961 let extracted = original.extract_by_indices(&[1, 1, 1]);
962 assert_eq!(extracted.row_count(), 3);
963 assert_eq!(extracted.data_at(0).get_value(0), Value::Int4(20));
964 assert_eq!(extracted.data_at(0).get_value(1), Value::Int4(20));
965 assert_eq!(extracted.data_at(0).get_value(2), Value::Int4(20));
966 }
967
968 #[test]
969 fn extract_by_indices_extracts_multiple_columns_consistently() {
970 let original = Columns::new(vec![
971 ColumnWithName::int4("id", [1, 2, 3, 4]),
972 ColumnWithName::utf8(
973 "name",
974 ["a".to_string(), "b".to_string(), "c".to_string(), "d".to_string()],
975 ),
976 ColumnWithName::bool("flag", [true, false, true, false]),
977 ]);
978 let extracted = original.extract_by_indices(&[2, 0]);
979
980 assert_eq!(extracted.len(), 3);
981 assert_eq!(extracted.row_count(), 2);
982 assert_eq!(extracted.column("id").unwrap().data().get_value(0), Value::Int4(3));
983 assert_eq!(extracted.column("id").unwrap().data().get_value(1), Value::Int4(1));
984 assert_eq!(extracted.column("name").unwrap().data().get_value(0), Value::Utf8("c".to_string()));
985 assert_eq!(extracted.column("name").unwrap().data().get_value(1), Value::Utf8("a".to_string()));
986 assert_eq!(extracted.column("flag").unwrap().data().get_value(0), Value::Boolean(true));
987 assert_eq!(extracted.column("flag").unwrap().data().get_value(1), Value::Boolean(true));
988 }
989
990 #[test]
991 fn extract_by_indices_extracts_system_columns_in_order() {
992 let columns = vec![ColumnWithName::int4("id", [10, 20, 30, 40])];
993 let row_numbers = vec![RowNumber::from(1), RowNumber::from(2), RowNumber::from(3), RowNumber::from(4)];
994 let created_at = vec![
995 DateTime::from_epoch_secs(1000).unwrap(),
996 DateTime::from_epoch_secs(2000).unwrap(),
997 DateTime::from_epoch_secs(3000).unwrap(),
998 DateTime::from_epoch_secs(4000).unwrap(),
999 ];
1000 let updated_at = vec![
1001 DateTime::from_epoch_secs(1100).unwrap(),
1002 DateTime::from_epoch_secs(2200).unwrap(),
1003 DateTime::from_epoch_secs(3300).unwrap(),
1004 DateTime::from_epoch_secs(4400).unwrap(),
1005 ];
1006 let time = created_at.clone();
1007 let original = Columns::with_system(
1008 columns,
1009 SystemColumns::new(row_numbers, Vec::new(), created_at, updated_at, time),
1010 );
1011
1012 let extracted = original.extract_by_indices(&[3, 0]);
1013
1014 let rns: Vec<RowNumber> = extracted.row_numbers().iter().cloned().collect();
1015 assert_eq!(rns, vec![RowNumber::from(4), RowNumber::from(1)], "row_numbers must follow indices");
1016 assert_eq!(
1017 extracted.created_at().iter().cloned().collect::<Vec<_>>(),
1018 vec![DateTime::from_epoch_secs(4000).unwrap(), DateTime::from_epoch_secs(1000).unwrap()],
1019 "created_at must follow indices"
1020 );
1021 assert_eq!(
1022 extracted.updated_at().iter().cloned().collect::<Vec<_>>(),
1023 vec![DateTime::from_epoch_secs(4400).unwrap(), DateTime::from_epoch_secs(1100).unwrap()],
1024 "updated_at must follow indices"
1025 );
1026 }
1027
1028 #[test]
1034 fn extract_by_indices_preserves_dictionary_id_metadata() {
1035 let mut buffer = ColumnBuffer::dictionary_id([
1036 DictionaryEntryId::U2(10),
1037 DictionaryEntryId::U2(20),
1038 DictionaryEntryId::U2(30),
1039 ]);
1040 match &mut buffer {
1041 ColumnBuffer::DictionaryId(container) => container.set_dictionary_id(DictionaryId(42)),
1042 _ => unreachable!("dictionary_id factory must build a DictionaryId buffer"),
1043 }
1044
1045 let original = Columns::new(vec![ColumnWithName::new("token", buffer)]);
1046 let extracted = original.extract_by_indices(&[2, 0]);
1047
1048 match extracted.data_at(0) {
1049 ColumnBuffer::DictionaryId(container) => {
1050 assert_eq!(
1051 container.dictionary_id(),
1052 Some(DictionaryId(42)),
1053 "dictionary_id metadata must survive extraction"
1054 );
1055 }
1056 other => panic!("expected DictionaryId buffer, got {:?}", other.get_type()),
1057 }
1058 }
1059
1060 #[test]
1061 fn extract_by_indices_preserves_utf8_max_bytes_metadata() {
1062 let mut buffer = ColumnBuffer::utf8(["a", "bb", "ccc"]);
1063 match &mut buffer {
1064 ColumnBuffer::Utf8 {
1065 max_bytes,
1066 ..
1067 } => *max_bytes = MaxBytes::new(255),
1068 _ => unreachable!(),
1069 }
1070
1071 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
1072 let extracted = original.extract_by_indices(&[2, 0]);
1073
1074 match extracted.data_at(0) {
1075 ColumnBuffer::Utf8 {
1076 max_bytes,
1077 ..
1078 } => assert_eq!(*max_bytes, MaxBytes::new(255), "Utf8 max_bytes must survive extraction"),
1079 other => panic!("expected Utf8 buffer, got {:?}", other.get_type()),
1080 }
1081 }
1082
1083 #[test]
1084 fn extract_by_indices_preserves_blob_max_bytes_metadata() {
1085 let mut buffer = ColumnBuffer::blob([Blob::new(vec![1]), Blob::new(vec![2, 3]), Blob::new(vec![4])]);
1086 match &mut buffer {
1087 ColumnBuffer::Blob {
1088 max_bytes,
1089 ..
1090 } => *max_bytes = MaxBytes::new(1024),
1091 _ => unreachable!(),
1092 }
1093
1094 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
1095 let extracted = original.extract_by_indices(&[2, 0]);
1096
1097 match extracted.data_at(0) {
1098 ColumnBuffer::Blob {
1099 max_bytes,
1100 ..
1101 } => assert_eq!(*max_bytes, MaxBytes::new(1024), "Blob max_bytes must survive extraction"),
1102 other => panic!("expected Blob buffer, got {:?}", other.get_type()),
1103 }
1104 }
1105
1106 #[test]
1107 fn extract_by_indices_preserves_int_max_bytes_metadata() {
1108 let mut buffer = ColumnBuffer::int([Int::from(1i64), Int::from(2i64), Int::from(3i64)]);
1109 match &mut buffer {
1110 ColumnBuffer::Int {
1111 max_bytes,
1112 ..
1113 } => *max_bytes = MaxBytes::new(16),
1114 _ => unreachable!(),
1115 }
1116
1117 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
1118 let extracted = original.extract_by_indices(&[2, 0]);
1119
1120 match extracted.data_at(0) {
1121 ColumnBuffer::Int {
1122 max_bytes,
1123 ..
1124 } => assert_eq!(*max_bytes, MaxBytes::new(16), "Int max_bytes must survive extraction"),
1125 other => panic!("expected Int buffer, got {:?}", other.get_type()),
1126 }
1127 }
1128
1129 #[test]
1130 fn extract_by_indices_preserves_uint_max_bytes_metadata() {
1131 let mut buffer = ColumnBuffer::uint([Uint::from(1u64), Uint::from(2u64), Uint::from(3u64)]);
1132 match &mut buffer {
1133 ColumnBuffer::Uint {
1134 max_bytes,
1135 ..
1136 } => *max_bytes = MaxBytes::new(8),
1137 _ => unreachable!(),
1138 }
1139
1140 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
1141 let extracted = original.extract_by_indices(&[2, 0]);
1142
1143 match extracted.data_at(0) {
1144 ColumnBuffer::Uint {
1145 max_bytes,
1146 ..
1147 } => assert_eq!(*max_bytes, MaxBytes::new(8), "Uint max_bytes must survive extraction"),
1148 other => panic!("expected Uint buffer, got {:?}", other.get_type()),
1149 }
1150 }
1151
1152 #[test]
1153 fn extract_by_indices_preserves_decimal_precision_and_scale_metadata() {
1154 let mut buffer = ColumnBuffer::decimal([
1155 Decimal::from_str("1.50").unwrap(),
1156 Decimal::from_str("2.25").unwrap(),
1157 Decimal::from_str("3.75").unwrap(),
1158 ]);
1159 match &mut buffer {
1160 ColumnBuffer::Decimal {
1161 precision,
1162 scale,
1163 ..
1164 } => {
1165 *precision = Precision::new(10);
1166 *scale = Scale::new(2);
1167 }
1168 _ => unreachable!(),
1169 }
1170
1171 let original = Columns::new(vec![ColumnWithName::new("c", buffer)]);
1172 let extracted = original.extract_by_indices(&[2, 0]);
1173
1174 match extracted.data_at(0) {
1175 ColumnBuffer::Decimal {
1176 precision,
1177 scale,
1178 ..
1179 } => {
1180 assert_eq!(*precision, Precision::new(10), "Decimal precision must survive extraction");
1181 assert_eq!(*scale, Scale::new(2), "Decimal scale must survive extraction");
1182 }
1183 other => panic!("expected Decimal buffer, got {:?}", other.get_type()),
1184 }
1185 }
1186
1187 #[test]
1188 fn test_single_row_temporal_types() {
1189 let date = Date::from_ymd(2025, 1, 15).unwrap();
1190 let datetime = DateTime::from_epoch_secs(1642694400).unwrap();
1191 let time = Time::from_hms(14, 30, 45).unwrap();
1192 let duration = Duration::from_days(30).unwrap();
1193
1194 let columns = Columns::single_row([
1195 ("date_col", Value::Date(date.clone())),
1196 ("datetime_col", Value::DateTime(datetime.clone())),
1197 ("time_col", Value::Time(time.clone())),
1198 ("interval_col", Value::Duration(duration.clone())),
1199 ]);
1200
1201 assert_eq!(columns.len(), 4);
1202 assert_eq!(columns.shape(), (1, 4));
1203
1204 assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
1205 assert_eq!(columns.column("datetime_col").unwrap().data().get_value(0), Value::DateTime(datetime));
1206 assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
1207 assert_eq!(columns.column("interval_col").unwrap().data().get_value(0), Value::Duration(duration));
1208 }
1209
1210 #[test]
1211 fn test_single_row_mixed_types() {
1212 let date = Date::from_ymd(2025, 7, 15).unwrap();
1213 let time = Time::from_hms(9, 15, 30).unwrap();
1214
1215 let columns = Columns::single_row([
1216 ("bool_col", Value::Boolean(true)),
1217 ("int_col", Value::Int4(42)),
1218 ("str_col", Value::Utf8("hello".to_string())),
1219 ("date_col", Value::Date(date.clone())),
1220 ("time_col", Value::Time(time.clone())),
1221 ("none_col", Value::none()),
1222 ]);
1223
1224 assert_eq!(columns.len(), 6);
1225 assert_eq!(columns.shape(), (1, 6));
1226
1227 assert_eq!(columns.column("bool_col").unwrap().data().get_value(0), Value::Boolean(true));
1228 assert_eq!(columns.column("int_col").unwrap().data().get_value(0), Value::Int4(42));
1229 assert_eq!(columns.column("str_col").unwrap().data().get_value(0), Value::Utf8("hello".to_string()));
1230 assert_eq!(columns.column("date_col").unwrap().data().get_value(0), Value::Date(date));
1231 assert_eq!(columns.column("time_col").unwrap().data().get_value(0), Value::Time(time));
1232 assert_eq!(columns.column("none_col").unwrap().data().get_value(0), Value::none());
1233 }
1234
1235 #[test]
1236 fn test_single_row_none_of_int4_is_int4_typed() {
1237 let columns = Columns::single_row([("n", Value::none_of(ValueType::Int4))]);
1240 match columns.column("n").unwrap().data().get_value(0) {
1241 Value::None {
1242 inner,
1243 } => assert_eq!(inner, ValueType::Int4),
1244 other => panic!("expected Value::None, got {other:?}"),
1245 }
1246 }
1247
1248 #[test]
1249 fn test_single_row_none_of_utf8_is_utf8_typed() {
1250 let columns = Columns::single_row([("n", Value::none_of(ValueType::Utf8))]);
1251 match columns.column("n").unwrap().data().get_value(0) {
1252 Value::None {
1253 inner,
1254 } => assert_eq!(inner, ValueType::Utf8),
1255 other => panic!("expected Value::None, got {other:?}"),
1256 }
1257 }
1258
1259 #[test]
1260 fn test_single_row_bare_none_is_any_typed() {
1261 let columns = Columns::single_row([("n", Value::none())]);
1262 match columns.column("n").unwrap().data().get_value(0) {
1263 Value::None {
1264 inner,
1265 } => assert_eq!(inner, ValueType::Any),
1266 other => panic!("expected Value::None, got {other:?}"),
1267 }
1268 }
1269
1270 #[test]
1271 fn test_single_row_none_of_nested_option_collapses_to_base_type() {
1272 let inner_ty = ValueType::Option(Box::new(ValueType::Duration));
1275 let columns = Columns::single_row([("n", Value::none_of(inner_ty))]);
1276 match columns.column("n").unwrap().data().get_value(0) {
1277 Value::None {
1278 inner,
1279 } => assert_eq!(inner, ValueType::Duration),
1280 other => panic!("expected Value::None, got {other:?}"),
1281 }
1282 }
1283
1284 #[test]
1285 fn test_single_row_none_of_boolean_is_boolean_typed() {
1286 let columns = Columns::single_row([("n", Value::none_of(ValueType::Boolean))]);
1289 match columns.column("n").unwrap().data().get_value(0) {
1290 Value::None {
1291 inner,
1292 } => assert_eq!(inner, ValueType::Boolean),
1293 other => panic!("expected Value::None, got {other:?}"),
1294 }
1295 }
1296
1297 #[test]
1298 fn test_single_row_normal_column_names_work() {
1299 let columns = Columns::single_row([("normal_column", Value::Int4(42))]);
1300 assert_eq!(columns.len(), 1);
1301 assert_eq!(columns.column("normal_column").unwrap().data().get_value(0), Value::Int4(42));
1302 }
1303}