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]
142 pub fn footprint(&self) -> usize {
143 size_of::<Self>() + self.columns.iter().map(Vector::footprint).sum::<usize>()
144 }
145
146 #[must_use]
148 pub fn types(&self) -> Vec<LogicalType> {
149 self.columns.iter().map(|column| column.logical_type().clone()).collect()
150 }
151
152 #[must_use]
157 pub fn value_at(&self, row: usize, column: usize) -> Value {
158 match self.columns.get(column) {
159 Some(held) => held.value_at(row),
160 None => Value::Null,
161 }
162 }
163
164 pub fn row(&self, row: usize) -> impl Iterator<Item = Value> + '_ {
166 self.columns.iter().map(move |column| column.value_at(row))
167 }
168
169 pub fn select(self, selection: &Selection) -> Result<Self> {
181 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
182 return Err(Error::internal(format!(
183 "a selection keeps row {bad} of a chunk that has {} rows",
184 self.rows
185 )));
186 }
187 let rows = selection.len();
188 let codes = selection.indices();
189 let mut columns = Vec::with_capacity(self.columns.len());
190 for column in self.columns {
191 columns.push(Vector::dictionary(codes.to_vec(), column)?);
192 }
193 Self::with_rows(columns, rows)
194 }
195
196 pub fn compact(self, selection: &Selection) -> Result<Self> {
232 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
233 return Err(Error::internal(format!(
234 "a selection keeps row {bad} of a chunk that has {} rows",
235 self.rows
236 )));
237 }
238 let rows = selection.len();
239 let indices = selection.indices();
240 let mut columns = Vec::with_capacity(self.columns.len());
241 for column in &self.columns {
242 columns.push(column.gather(indices)?);
243 }
244 Self::with_rows(columns, rows)
245 }
246
247 pub fn project(self, positions: &[usize]) -> Result<Self> {
256 let width = self.columns.len();
257 if let Some(&bad) = positions.iter().find(|&&position| position >= width) {
258 return Err(Error::internal(format!(
259 "column {bad} of a chunk that has {width} columns"
260 )));
261 }
262 let rows = self.rows;
263 let mut sources: Vec<Option<Vector>> = self.columns.into_iter().map(Some).collect();
264 let mut columns = Vec::with_capacity(positions.len());
265 for (at, &position) in positions.iter().enumerate() {
266 let last_use = !positions[at + 1..].contains(&position);
267 let taken = if last_use { sources[position].take() } else { sources[position].clone() };
268 match taken {
269 Some(column) => columns.push(column),
270 None => {
273 return Err(Error::internal(format!("column {position} was taken twice")));
274 }
275 }
276 }
277 Self::with_rows(columns, rows)
278 }
279
280 pub fn flatten(&self) -> Result<Self> {
290 let mut columns = Vec::with_capacity(self.columns.len());
291 for column in &self.columns {
292 columns.push(column.flatten()?);
293 }
294 Self::with_rows(columns, self.rows)
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use rudb_common::LogicalType;
301
302 use super::*;
303 use crate::vector::{Data, Form};
304
305 fn integers(values: &[i32]) -> Vector {
306 Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
307 .expect("integers are an i32 layout")
308 }
309
310 #[test]
311 fn a_chunk_takes_its_length_from_its_columns() {
312 let chunk = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4, 5, 6])])
313 .expect("two columns of three");
314 assert_eq!(chunk.len(), 3);
315 assert_eq!(chunk.width(), 2);
316 assert_eq!(chunk.value_at(2, 1), Value::Integer(6));
317 }
318
319 #[test]
320 fn a_ragged_chunk_is_caught() {
321 let error = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4])])
322 .expect_err("a chunk is not ragged");
323 assert!(error.message().contains("column 1"), "{error}");
324 }
325
326 #[test]
329 fn a_chunk_with_no_columns_can_still_have_rows() {
330 let chunk = Chunk::with_rows(Vec::new(), 900).expect("no columns and nine hundred rows");
331 assert_eq!(chunk.len(), 900);
332 assert_eq!(chunk.width(), 0);
333 assert!(!chunk.is_empty(), "nine hundred rows is not empty");
334 }
335
336 #[test]
337 fn a_chunk_longer_than_a_vector_is_caught() {
338 let error = Chunk::with_rows(Vec::new(), VECTOR_SIZE + 1).expect_err("too long");
339 assert!(error.message().contains("longer than"), "{error}");
340 }
341
342 #[test]
343 fn an_empty_chunk_keeps_its_types() {
344 let chunk = Chunk::empty(&[LogicalType::Integer, LogicalType::Varchar]);
345 assert_eq!(chunk.len(), 0);
346 assert_eq!(chunk.types(), vec![LogicalType::Integer, LogicalType::Varchar]);
347 }
348
349 #[test]
350 fn selecting_keeps_the_rows_it_selected_and_no_others() {
351 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
352 .expect("four rows");
353 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
354 let chunk = chunk.select(&kept).expect("rows one and three exist");
355 assert_eq!(chunk.len(), 2);
356 assert_eq!(chunk.row(0).collect::<Vec<_>>(), vec![Value::Integer(20), Value::Integer(2)]);
357 assert_eq!(chunk.row(1).collect::<Vec<_>>(), vec![Value::Integer(40), Value::Integer(4)]);
358 }
359
360 #[test]
363 fn selecting_leaves_the_values_where_they_were() {
364 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40])]).expect("four rows");
365 let kept = Selection::from_predicate(4, |index| index == 0);
366 let chunk = chunk.select(&kept).expect("row zero exists");
367 assert_eq!(chunk.column(0).expect("one column").form(), Form::Dictionary);
368 }
369
370 #[test]
371 fn a_selection_past_the_end_is_caught() {
372 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
373 let mut kept = Selection::empty();
374 kept.push(7);
375 let error = chunk.select(&kept).expect_err("row seven does not exist");
376 assert!(error.message().contains("row 7"), "{error}");
377 }
378
379 #[test]
382 fn compacting_keeps_the_same_rows_selecting_does_and_leaves_no_indirection() {
383 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
384 .expect("four rows");
385 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
386 let selected = chunk.clone().select(&kept).expect("rows one and three exist");
387 let compacted = chunk.compact(&kept).expect("rows one and three exist");
388 assert_eq!(compacted.len(), selected.len());
389 for row in 0..compacted.len() {
390 assert_eq!(
391 compacted.row(row).collect::<Vec<_>>(),
392 selected.row(row).collect::<Vec<_>>()
393 );
394 }
395 assert_eq!(compacted.column(0).expect("one column").form(), Form::Flat);
396 }
397
398 #[test]
399 fn a_selection_past_the_end_is_caught_by_compacting_too() {
400 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
401 let mut kept = Selection::empty();
402 kept.push(7);
403 let error = chunk.compact(&kept).expect_err("row seven does not exist");
404 assert!(error.message().contains("row 7"), "{error}");
405 }
406
407 #[test]
408 fn projecting_reorders_and_can_repeat_a_column() {
409 let chunk = Chunk::new(vec![integers(&[1, 2]), integers(&[3, 4])]).expect("two by two");
410 let chunk = chunk.project(&[1, 0, 1]).expect("both columns exist");
411 assert_eq!(chunk.width(), 3);
412 assert_eq!(
413 chunk.row(0).collect::<Vec<_>>(),
414 vec![Value::Integer(3), Value::Integer(1), Value::Integer(3)]
415 );
416 }
417
418 #[test]
419 fn projecting_a_column_that_is_not_there_is_caught() {
420 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("one column");
421 let error = chunk.project(&[0, 4]).expect_err("there is no column four");
422 assert!(error.message().contains("column 4"), "{error}");
423 }
424
425 #[test]
426 fn flattening_a_selected_chunk_gives_the_same_values() {
427 let chunk = Chunk::new(vec![integers(&[10, 20, 30])]).expect("three rows");
428 let kept = Selection::from_predicate(3, |index| index != 1);
429 let selected = chunk.select(&kept).expect("rows zero and two exist");
430 let flat = selected.flatten().expect("integers flatten");
431 assert_eq!(flat.column(0).expect("one column").form(), Form::Flat);
432 for row in 0..flat.len() {
433 assert_eq!(flat.value_at(row, 0), selected.value_at(row, 0), "row {row}");
434 }
435 }
436
437 #[test]
438 fn a_chunk_costs_what_its_columns_cost() {
439 let chunk = Chunk::new(vec![integers(&[1; 1000]), integers(&[2; 1000])])
440 .expect("two columns of a thousand");
441 let columns: usize = chunk.columns().iter().map(Vector::footprint).sum();
442 assert_eq!(chunk.footprint(), size_of::<Chunk>() + columns);
443 assert!(chunk.footprint() >= 8000, "two thousand i32: {}", chunk.footprint());
444 }
445}