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
26impl Chunk {
27 pub fn new(columns: Vec<Vector>) -> Result<Self> {
33 let rows = columns.first().map_or(0, Vector::len);
34 Self::with_rows(columns, rows)
35 }
36
37 pub fn with_rows(columns: Vec<Vector>, rows: usize) -> Result<Self> {
44 if rows > VECTOR_SIZE {
45 return Err(Error::internal(format!(
46 "a chunk of {rows} rows is longer than the {VECTOR_SIZE} row vector"
47 )));
48 }
49 for (index, column) in columns.iter().enumerate() {
50 if column.len() != rows {
51 return Err(Error::internal(format!(
52 "column {index} of a chunk is {} rows and the chunk is {rows}",
53 column.len()
54 )));
55 }
56 }
57 Ok(Self { columns, rows })
58 }
59
60 #[must_use]
66 pub fn empty(types: &[LogicalType]) -> Self {
67 let columns =
68 types.iter().map(|ty| Vector::constant(ty.clone(), Value::Null, 0)).collect::<Vec<_>>();
69 Self { columns, rows: 0 }
70 }
71
72 #[must_use]
74 pub fn columns(&self) -> &[Vector] {
75 &self.columns
76 }
77
78 pub fn column(&self, index: usize) -> Result<&Vector> {
84 self.columns.get(index).ok_or_else(|| {
85 Error::internal(format!(
86 "column {index} of a chunk that has {} columns",
87 self.columns.len()
88 ))
89 })
90 }
91
92 #[must_use]
94 pub fn into_columns(self) -> Vec<Vector> {
95 self.columns
96 }
97
98 #[must_use]
100 pub fn width(&self) -> usize {
101 self.columns.len()
102 }
103
104 #[must_use]
106 pub fn len(&self) -> usize {
107 self.rows
108 }
109
110 #[must_use]
112 pub fn is_empty(&self) -> bool {
113 self.rows == 0
114 }
115
116 #[must_use]
118 pub fn types(&self) -> Vec<LogicalType> {
119 self.columns.iter().map(|column| column.logical_type().clone()).collect()
120 }
121
122 #[must_use]
127 pub fn value_at(&self, row: usize, column: usize) -> Value {
128 match self.columns.get(column) {
129 Some(held) => held.value_at(row),
130 None => Value::Null,
131 }
132 }
133
134 pub fn row(&self, row: usize) -> impl Iterator<Item = Value> + '_ {
136 self.columns.iter().map(move |column| column.value_at(row))
137 }
138
139 pub fn select(self, selection: &Selection) -> Result<Self> {
151 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
152 return Err(Error::internal(format!(
153 "a selection keeps row {bad} of a chunk that has {} rows",
154 self.rows
155 )));
156 }
157 let rows = selection.len();
158 let codes = selection.indices();
159 let mut columns = Vec::with_capacity(self.columns.len());
160 for column in self.columns {
161 columns.push(Vector::dictionary(codes.to_vec(), column)?);
162 }
163 Self::with_rows(columns, rows)
164 }
165
166 pub fn compact(self, selection: &Selection) -> Result<Self> {
202 if let Some(bad) = selection.iter().find(|&index| index >= self.rows) {
203 return Err(Error::internal(format!(
204 "a selection keeps row {bad} of a chunk that has {} rows",
205 self.rows
206 )));
207 }
208 let rows = selection.len();
209 let indices = selection.indices();
210 let mut columns = Vec::with_capacity(self.columns.len());
211 for column in &self.columns {
212 columns.push(column.gather(indices)?);
213 }
214 Self::with_rows(columns, rows)
215 }
216
217 pub fn project(self, positions: &[usize]) -> Result<Self> {
226 let width = self.columns.len();
227 if let Some(&bad) = positions.iter().find(|&&position| position >= width) {
228 return Err(Error::internal(format!(
229 "column {bad} of a chunk that has {width} columns"
230 )));
231 }
232 let rows = self.rows;
233 let mut sources: Vec<Option<Vector>> = self.columns.into_iter().map(Some).collect();
234 let mut columns = Vec::with_capacity(positions.len());
235 for (at, &position) in positions.iter().enumerate() {
236 let last_use = !positions[at + 1..].contains(&position);
237 let taken = if last_use { sources[position].take() } else { sources[position].clone() };
238 match taken {
239 Some(column) => columns.push(column),
240 None => {
243 return Err(Error::internal(format!("column {position} was taken twice")));
244 }
245 }
246 }
247 Self::with_rows(columns, rows)
248 }
249
250 pub fn flatten(&self) -> Result<Self> {
260 let mut columns = Vec::with_capacity(self.columns.len());
261 for column in &self.columns {
262 columns.push(column.flatten()?);
263 }
264 Self::with_rows(columns, self.rows)
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use rudb_common::LogicalType;
271
272 use super::*;
273 use crate::vector::{Data, Form};
274
275 fn integers(values: &[i32]) -> Vector {
276 Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec()))
277 .expect("integers are an i32 layout")
278 }
279
280 #[test]
281 fn a_chunk_takes_its_length_from_its_columns() {
282 let chunk = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4, 5, 6])])
283 .expect("two columns of three");
284 assert_eq!(chunk.len(), 3);
285 assert_eq!(chunk.width(), 2);
286 assert_eq!(chunk.value_at(2, 1), Value::Integer(6));
287 }
288
289 #[test]
290 fn a_ragged_chunk_is_caught() {
291 let error = Chunk::new(vec![integers(&[1, 2, 3]), integers(&[4])])
292 .expect_err("a chunk is not ragged");
293 assert!(error.message().contains("column 1"), "{error}");
294 }
295
296 #[test]
299 fn a_chunk_with_no_columns_can_still_have_rows() {
300 let chunk = Chunk::with_rows(Vec::new(), 900).expect("no columns and nine hundred rows");
301 assert_eq!(chunk.len(), 900);
302 assert_eq!(chunk.width(), 0);
303 assert!(!chunk.is_empty(), "nine hundred rows is not empty");
304 }
305
306 #[test]
307 fn a_chunk_longer_than_a_vector_is_caught() {
308 let error = Chunk::with_rows(Vec::new(), VECTOR_SIZE + 1).expect_err("too long");
309 assert!(error.message().contains("longer than"), "{error}");
310 }
311
312 #[test]
313 fn an_empty_chunk_keeps_its_types() {
314 let chunk = Chunk::empty(&[LogicalType::Integer, LogicalType::Varchar]);
315 assert_eq!(chunk.len(), 0);
316 assert_eq!(chunk.types(), vec![LogicalType::Integer, LogicalType::Varchar]);
317 }
318
319 #[test]
320 fn selecting_keeps_the_rows_it_selected_and_no_others() {
321 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
322 .expect("four rows");
323 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
324 let chunk = chunk.select(&kept).expect("rows one and three exist");
325 assert_eq!(chunk.len(), 2);
326 assert_eq!(chunk.row(0).collect::<Vec<_>>(), vec![Value::Integer(20), Value::Integer(2)]);
327 assert_eq!(chunk.row(1).collect::<Vec<_>>(), vec![Value::Integer(40), Value::Integer(4)]);
328 }
329
330 #[test]
333 fn selecting_leaves_the_values_where_they_were() {
334 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40])]).expect("four rows");
335 let kept = Selection::from_predicate(4, |index| index == 0);
336 let chunk = chunk.select(&kept).expect("row zero exists");
337 assert_eq!(chunk.column(0).expect("one column").form(), Form::Dictionary);
338 }
339
340 #[test]
341 fn a_selection_past_the_end_is_caught() {
342 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
343 let mut kept = Selection::empty();
344 kept.push(7);
345 let error = chunk.select(&kept).expect_err("row seven does not exist");
346 assert!(error.message().contains("row 7"), "{error}");
347 }
348
349 #[test]
352 fn compacting_keeps_the_same_rows_selecting_does_and_leaves_no_indirection() {
353 let chunk = Chunk::new(vec![integers(&[10, 20, 30, 40]), integers(&[1, 2, 3, 4])])
354 .expect("four rows");
355 let kept = Selection::from_predicate(4, |index| index % 2 == 1);
356 let selected = chunk.clone().select(&kept).expect("rows one and three exist");
357 let compacted = chunk.compact(&kept).expect("rows one and three exist");
358 assert_eq!(compacted.len(), selected.len());
359 for row in 0..compacted.len() {
360 assert_eq!(
361 compacted.row(row).collect::<Vec<_>>(),
362 selected.row(row).collect::<Vec<_>>()
363 );
364 }
365 assert_eq!(compacted.column(0).expect("one column").form(), Form::Flat);
366 }
367
368 #[test]
369 fn a_selection_past_the_end_is_caught_by_compacting_too() {
370 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("two rows");
371 let mut kept = Selection::empty();
372 kept.push(7);
373 let error = chunk.compact(&kept).expect_err("row seven does not exist");
374 assert!(error.message().contains("row 7"), "{error}");
375 }
376
377 #[test]
378 fn projecting_reorders_and_can_repeat_a_column() {
379 let chunk = Chunk::new(vec![integers(&[1, 2]), integers(&[3, 4])]).expect("two by two");
380 let chunk = chunk.project(&[1, 0, 1]).expect("both columns exist");
381 assert_eq!(chunk.width(), 3);
382 assert_eq!(
383 chunk.row(0).collect::<Vec<_>>(),
384 vec![Value::Integer(3), Value::Integer(1), Value::Integer(3)]
385 );
386 }
387
388 #[test]
389 fn projecting_a_column_that_is_not_there_is_caught() {
390 let chunk = Chunk::new(vec![integers(&[1, 2])]).expect("one column");
391 let error = chunk.project(&[0, 4]).expect_err("there is no column four");
392 assert!(error.message().contains("column 4"), "{error}");
393 }
394
395 #[test]
396 fn flattening_a_selected_chunk_gives_the_same_values() {
397 let chunk = Chunk::new(vec![integers(&[10, 20, 30])]).expect("three rows");
398 let kept = Selection::from_predicate(3, |index| index != 1);
399 let selected = chunk.select(&kept).expect("rows zero and two exist");
400 let flat = selected.flatten().expect("integers flatten");
401 assert_eq!(flat.column(0).expect("one column").form(), Form::Flat);
402 for row in 0..flat.len() {
403 assert_eq!(flat.value_at(row, 0), selected.value_at(row, 0), "row {row}");
404 }
405 }
406}