1use rudb_common::{Error, LogicalType, Result, Value};
15
16use crate::selection::Selection;
17use crate::vector::{VECTOR_SIZE, Vector};
18
19#[derive(Debug, Clone, PartialEq)]
21pub struct Chunk {
22 columns: Vec<Vector>,
23 rows: usize,
24}
25
26const _: () = {
38 const fn assert_send<T: Send>() {}
39 assert_send::<Chunk>();
40};
41
42impl Chunk {
43 pub fn new(columns: Vec<Vector>) -> Result<Self> {
49 let rows = columns.first().map_or(0, Vector::len);
50 Self::with_rows(columns, rows)
51 }
52
53 pub fn with_rows(columns: Vec<Vector>, rows: usize) -> Result<Self> {
60 if rows > VECTOR_SIZE {
61 return Err(Error::internal(format!(
62 "a chunk of {rows} rows is longer than the {VECTOR_SIZE} row vector"
63 )));
64 }
65 for (index, column) in columns.iter().enumerate() {
66 if column.len() != rows {
67 return Err(Error::internal(format!(
68 "column {index} of a chunk is {} rows and the chunk is {rows}",
69 column.len()
70 )));
71 }
72 }
73 Ok(Self { columns, rows })
74 }
75
76 #[must_use]
82 pub fn empty(types: &[LogicalType]) -> Self {
83 let columns =
84 types.iter().map(|ty| Vector::constant(ty.clone(), Value::Null, 0)).collect::<Vec<_>>();
85 Self { columns, rows: 0 }
86 }
87
88 #[must_use]
90 pub fn columns(&self) -> &[Vector] {
91 &self.columns
92 }
93
94 pub fn column(&self, index: usize) -> Result<&Vector> {
100 self.columns.get(index).ok_or_else(|| {
101 Error::internal(format!(
102 "column {index} of a chunk that has {} columns",
103 self.columns.len()
104 ))
105 })
106 }
107
108 #[must_use]
110 pub fn into_columns(self) -> Vec<Vector> {
111 self.columns
112 }
113
114 #[must_use]
116 pub fn width(&self) -> usize {
117 self.columns.len()
118 }
119
120 #[must_use]
122 pub fn len(&self) -> usize {
123 self.rows
124 }
125
126 #[must_use]
128 pub fn is_empty(&self) -> bool {
129 self.rows == 0
130 }
131
132 #[must_use]
134 pub fn types(&self) -> Vec<LogicalType> {
135 self.columns.iter().map(|column| column.logical_type().clone()).collect()
136 }
137
138 #[must_use]
143 pub fn value_at(&self, row: usize, column: usize) -> Value {
144 match self.columns.get(column) {
145 Some(held) => held.value_at(row),
146 None => Value::Null,
147 }
148 }
149
150 pub fn row(&self, row: usize) -> impl Iterator<Item = Value> + '_ {
152 self.columns.iter().map(move |column| column.value_at(row))
153 }
154
155 pub fn select(self, selection: &Selection) -> Result<Self> {
167 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
168 return Err(Error::internal(format!(
169 "a selection keeps row {bad} of a chunk that has {} rows",
170 self.rows
171 )));
172 }
173 let rows = selection.len();
174 let codes = selection.indices();
175 let mut columns = Vec::with_capacity(self.columns.len());
176 for column in self.columns {
177 columns.push(Vector::dictionary(codes.to_vec(), column)?);
178 }
179 Self::with_rows(columns, rows)
180 }
181
182 pub fn compact(self, selection: &Selection) -> Result<Self> {
218 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
219 return Err(Error::internal(format!(
220 "a selection keeps row {bad} of a chunk that has {} rows",
221 self.rows
222 )));
223 }
224 let rows = selection.len();
225 let indices = selection.indices();
226 let mut columns = Vec::with_capacity(self.columns.len());
227 for column in &self.columns {
228 columns.push(column.gather(indices)?);
229 }
230 Self::with_rows(columns, rows)
231 }
232
233 pub fn project(self, positions: &[usize]) -> Result<Self> {
242 let width = self.columns.len();
243 if let Some(&bad) = positions.iter().find(|&&position| position >= width) {
244 return Err(Error::internal(format!(
245 "column {bad} of a chunk that has {width} columns"
246 )));
247 }
248 let rows = self.rows;
249 let mut sources: Vec<Option<Vector>> = self.columns.into_iter().map(Some).collect();
250 let mut columns = Vec::with_capacity(positions.len());
251 for (at, &position) in positions.iter().enumerate() {
252 let last_use = !positions[at + 1..].contains(&position);
253 let taken = if last_use { sources[position].take() } else { sources[position].clone() };
254 match taken {
255 Some(column) => columns.push(column),
256 None => {
259 return Err(Error::internal(format!("column {position} was taken twice")));
260 }
261 }
262 }
263 Self::with_rows(columns, rows)
264 }
265
266 pub fn flatten(&self) -> Result<Self> {
276 let mut columns = Vec::with_capacity(self.columns.len());
277 for column in &self.columns {
278 columns.push(column.flatten()?);
279 }
280 Self::with_rows(columns, self.rows)
281 }
282}
283
284#[cfg(test)]
285mod tests {
286 use rudb_common::LogicalType;
287
288 use super::*;
289 use crate::vector::{Data, Form};
290
291 fn integers(values: &[i32]) -> Vector {
292 Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
293 .expect("integers are an i32 layout")
294 }
295
296 #[test]
297 fn a_chunk_takes_its_length_from_its_columns() {
298 let chunk = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4, 5, 6])])
299 .expect("two columns of three");
300 assert_eq!(chunk.len(), 3);
301 assert_eq!(chunk.width(), 2);
302 assert_eq!(chunk.value_at(2, 1), Value::Integer(6));
303 }
304
305 #[test]
306 fn a_ragged_chunk_is_caught() {
307 let error = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4])])
308 .expect_err("a chunk is not ragged");
309 assert!(error.message().contains("column 1"), "{error}");
310 }
311
312 #[test]
315 fn a_chunk_with_no_columns_can_still_have_rows() {
316 let chunk = Chunk::with_rows(Vec::new(), 900).expect("no columns and nine hundred rows");
317 assert_eq!(chunk.len(), 900);
318 assert_eq!(chunk.width(), 0);
319 assert!(!chunk.is_empty(), "nine hundred rows is not empty");
320 }
321
322 #[test]
323 fn a_chunk_longer_than_a_vector_is_caught() {
324 let error = Chunk::with_rows(Vec::new(), VECTOR_SIZE + 1).expect_err("too long");
325 assert!(error.message().contains("longer than"), "{error}");
326 }
327
328 #[test]
329 fn an_empty_chunk_keeps_its_types() {
330 let chunk = Chunk::empty(&[LogicalType::Integer, LogicalType::Varchar]);
331 assert_eq!(chunk.len(), 0);
332 assert_eq!(chunk.types(), vec![LogicalType::Integer, LogicalType::Varchar]);
333 }
334
335 #[test]
336 fn selecting_keeps_the_rows_it_selected_and_no_others() {
337 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
338 .expect("four rows");
339 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
340 let chunk = chunk.select(&kept).expect("rows one and three exist");
341 assert_eq!(chunk.len(), 2);
342 assert_eq!(chunk.row(0).collect::<Vec<_>>(), vec![Value::Integer(20), Value::Integer(2)]);
343 assert_eq!(chunk.row(1).collect::<Vec<_>>(), vec![Value::Integer(40), Value::Integer(4)]);
344 }
345
346 #[test]
349 fn selecting_leaves_the_values_where_they_were() {
350 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40])]).expect("four rows");
351 let kept = Selection::from_predicate(4, |index| index == 0);
352 let chunk = chunk.select(&kept).expect("row zero exists");
353 assert_eq!(chunk.column(0).expect("one column").form(), Form::Dictionary);
354 }
355
356 #[test]
357 fn a_selection_past_the_end_is_caught() {
358 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
359 let mut kept = Selection::empty();
360 kept.push(7);
361 let error = chunk.select(&kept).expect_err("row seven does not exist");
362 assert!(error.message().contains("row 7"), "{error}");
363 }
364
365 #[test]
368 fn compacting_keeps_the_same_rows_selecting_does_and_leaves_no_indirection() {
369 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
370 .expect("four rows");
371 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
372 let selected = chunk.clone().select(&kept).expect("rows one and three exist");
373 let compacted = chunk.compact(&kept).expect("rows one and three exist");
374 assert_eq!(compacted.len(), selected.len());
375 for row in 0..compacted.len() {
376 assert_eq!(
377 compacted.row(row).collect::<Vec<_>>(),
378 selected.row(row).collect::<Vec<_>>()
379 );
380 }
381 assert_eq!(compacted.column(0).expect("one column").form(), Form::Flat);
382 }
383
384 #[test]
385 fn a_selection_past_the_end_is_caught_by_compacting_too() {
386 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
387 let mut kept = Selection::empty();
388 kept.push(7);
389 let error = chunk.compact(&kept).expect_err("row seven does not exist");
390 assert!(error.message().contains("row 7"), "{error}");
391 }
392
393 #[test]
394 fn projecting_reorders_and_can_repeat_a_column() {
395 let chunk = Chunk::new(vec![integers(&[1, 2]), integers(&[3, 4])]).expect("two by two");
396 let chunk = chunk.project(&[1, 0, 1]).expect("both columns exist");
397 assert_eq!(chunk.width(), 3);
398 assert_eq!(
399 chunk.row(0).collect::<Vec<_>>(),
400 vec![Value::Integer(3), Value::Integer(1), Value::Integer(3)]
401 );
402 }
403
404 #[test]
405 fn projecting_a_column_that_is_not_there_is_caught() {
406 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("one column");
407 let error = chunk.project(&[0, 4]).expect_err("there is no column four");
408 assert!(error.message().contains("column 4"), "{error}");
409 }
410
411 #[test]
412 fn flattening_a_selected_chunk_gives_the_same_values() {
413 let chunk = Chunk::new(vec![integers(&[10, 20, 30])]).expect("three rows");
414 let kept = Selection::from_predicate(3, |index| index != 1);
415 let selected = chunk.select(&kept).expect("rows zero and two exist");
416 let flat = selected.flatten().expect("integers flatten");
417 assert_eq!(flat.column(0).expect("one column").form(), Form::Flat);
418 for row in 0..flat.len() {
419 assert_eq!(flat.value_at(row, 0), selected.value_at(row, 0), "row {row}");
420 }
421 }
422}