use std::borrow::Cow;
use rudb_common::{Error, LogicalType, Result, Value};
use rudb_pipeline::Lease;
use rudb_vector::{Assembly, Chunk, Form, Vector};
use crate::pairs::in_parallel;
pub(crate) const PAD: u32 = u32::MAX;
#[derive(Debug, Default)]
pub(crate) struct Build {
columns: Vec<Vector>,
rows: usize,
}
impl Build {
pub(crate) fn new(
types: &[LogicalType],
chunks: &[Chunk],
threads: &Lease<'_>,
) -> Result<Self> {
let rows: usize = chunks.iter().map(Chunk::len).sum();
if rows >= PAD as usize {
return Err(Error::out_of_range(format!(
"a join cannot gather {rows} rows, which is more than a position can name"
)));
}
let columns = laid_out(types, chunks, threads)?;
Ok(Self { columns, rows })
}
pub(crate) fn rows(&self) -> usize {
self.rows
}
pub(crate) fn column(&self, index: usize) -> Option<&Vector> {
self.columns.get(index)
}
pub(crate) fn footprint(&self) -> u64 {
self.columns
.iter()
.map(|column| u64::try_from(column.footprint()).unwrap_or(u64::MAX))
.sum()
}
pub(crate) fn gather(&self, at: &[u32]) -> Result<Vec<Vector>> {
self.columns.iter().map(|column| column.gather(at)).collect()
}
pub(crate) fn gather_wanted(&self, at: &[u32], wanted: &[bool]) -> Result<Vec<Vector>> {
self.columns
.iter()
.enumerate()
.map(|(index, column)| {
if wanted.get(index).copied().unwrap_or(true) {
column.gather(at)
} else {
Ok(Vector::constant(column.logical_type().clone(), Value::Null, at.len()))
}
})
.collect()
}
pub(crate) fn chunk(&self, at: &[u32]) -> Result<Chunk> {
Chunk::with_rows(self.gather(at)?, at.len())
}
pub(crate) fn row(&self, at: u32) -> Vec<Value> {
self.columns.iter().map(|column| column.value_at(at as usize)).collect()
}
}
pub(crate) fn laid_out(
types: &[LogicalType],
chunks: &[Chunk],
threads: &Lease<'_>,
) -> Result<Vec<Vector>> {
let rows: usize = chunks.iter().map(Chunk::len).sum();
let one = |index: usize| -> Result<Vector> {
if let Some(laid) = end_to_end(&types[index], chunks, index)? {
return Ok(laid);
}
let mut assembly = Assembly::new(types[index].clone(), rows)?;
let mut at: Vec<u32> = Vec::new();
let mut base: u32 = 0;
for chunk in chunks {
let len = u32::try_from(chunk.len()).unwrap_or(PAD);
at.clear();
at.extend(base..base + len);
assembly.place(&at, chunk.column(index)?)?;
base += len;
}
assembly.finish()
};
in_parallel(threads, types.len(), threads.degree(), "gathered column", one)
}
fn end_to_end(ty: &LogicalType, chunks: &[Chunk], index: usize) -> Result<Option<Vector>> {
if matches!(ty, LogicalType::List(_) | LogicalType::Struct(_) | LogicalType::Map(_, _)) {
return Ok(None);
}
let mut pieces: Vec<Cow<'_, Vector>> = Vec::with_capacity(chunks.len());
for chunk in chunks.iter().filter(|chunk| !chunk.is_empty()) {
let column = chunk.column(index)?;
let piece = match column.form() {
Form::Flat | Form::StringView => Cow::Borrowed(column),
_ => Cow::Owned(column.flatten()?),
};
if piece.form() == Form::Flat && short(&piece) {
return Ok(None);
}
pieces.push(piece);
}
if let Some(laid) = rudb_vector::concat(ty, &pieces)? {
return Ok(Some(laid));
}
for piece in &mut pieces {
if piece.form() == Form::StringView {
*piece = Cow::Owned(piece.flatten()?);
if short(piece) {
return Ok(None);
}
}
}
rudb_vector::concat(ty, &pieces)
}
fn short(piece: &Vector) -> bool {
piece.data().is_none_or(|data| data.len() != piece.len())
}
#[cfg(test)]
mod tests {
use rudb_common::{LogicalType, Value};
use rudb_pipeline::Lease;
use rudb_vector::{Chunk, Data, Vector};
use super::{Build, PAD};
fn chunk(values: &[i32], text: &[&str]) -> Chunk {
let numbers = Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
.expect("integers are an i32 layout");
let strings = Vector::from_values(
LogicalType::Varchar,
&text.iter().map(|&word| Value::Varchar(word.to_string())).collect::<Vec<_>>(),
)
.expect("strings are a varlen layout");
Chunk::new(vec![numbers, strings]).expect("two columns of the same length")
}
fn alone() -> Lease<'static> {
Lease::alone()
}
fn types() -> Vec<LogicalType> {
vec![LogicalType::Integer, LogicalType::Varchar]
}
#[test]
fn chunks_laid_end_to_end_read_back_in_the_order_they_were_given() {
let side =
Build::new(&types(), &[chunk(&[1, 2], &["a", "b"]), chunk(&[3], &["c"])], &alone())
.expect("two chunks of two columns");
assert_eq!(side.rows(), 3);
let gathered = side.gather(&[0, 1, 2]).expect("three positions in range");
assert_eq!(gathered[0].value_at(0), Value::Integer(1));
assert_eq!(gathered[0].value_at(2), Value::Integer(3));
assert_eq!(gathered[1].value_at(1), Value::Varchar("b".to_string()));
assert_eq!(gathered[1].value_at(2), Value::Varchar("c".to_string()));
}
#[test]
fn a_position_may_be_asked_for_more_than_once_and_in_any_order() {
let side = Build::new(&types(), &[chunk(&[10, 20], &["x", "y"])], &alone())
.expect("one chunk of two columns");
let gathered = side.gather(&[1, 1, 0]).expect("three positions in range");
assert_eq!(gathered[0].value_at(0), Value::Integer(20));
assert_eq!(gathered[0].value_at(1), Value::Integer(20));
assert_eq!(gathered[0].value_at(2), Value::Integer(10));
}
#[test]
fn the_padding_position_reads_as_null_in_every_column() {
let side = Build::new(&types(), &[chunk(&[7], &["z"])], &alone())
.expect("one chunk of two columns");
let gathered = side.gather(&[PAD, 0]).expect("a padded position and a real one");
assert_eq!(gathered[0].value_at(0), Value::Null);
assert_eq!(gathered[1].value_at(0), Value::Null);
assert_eq!(gathered[0].value_at(1), Value::Integer(7));
}
#[test]
fn a_side_with_no_chunks_still_has_its_columns_and_every_one_of_them_is_null() {
let side = Build::new(&types(), &[], &alone()).expect("no chunks at all");
assert_eq!(side.rows(), 0);
let gathered = side.gather(&[PAD, PAD]).expect("two padded positions");
assert_eq!(gathered.len(), 2);
assert_eq!(gathered[0].value_at(0), Value::Null);
assert_eq!(gathered[1].value_at(1), Value::Null);
}
#[test]
fn a_row_read_as_values_is_the_row_that_went_in() {
let side =
Build::new(&types(), &[chunk(&[4, 5], &["p", "q"])], &alone()).expect("one chunk");
assert_eq!(side.row(1), vec![Value::Integer(5), Value::Varchar("q".to_string())]);
}
#[test]
fn a_side_laid_out_of_pieces_of_every_form_reads_back_as_the_values_that_went_in() {
let int = LogicalType::Integer;
let text = LogicalType::Varchar;
let words = |list: &[&str]| {
let values: Vec<Value> =
list.iter().map(|&word| Value::Varchar(word.to_string())).collect();
Vector::from_values(LogicalType::Varchar, &values).expect("strings build")
};
let page = rudb_vector::concat(
&text,
&[words(&["one", "a string longer than twelve", "three", "four"])],
)
.expect("a flat piece lays")
.expect("and comes back as views");
let numbers = |list: &[Option<i32>]| {
let values: Vec<Value> =
list.iter().map(|value| value.map_or(Value::Null, Value::Integer)).collect();
Vector::from_values(LogicalType::Integer, &values).expect("integers build")
};
let coded =
Vector::dictionary(vec![1, 0], numbers(&[Some(5), Some(6)])).expect("codes in range");
let shared = [
Chunk::new(vec![numbers(&[Some(1), None]), page.gather(&[0, 1]).expect("in range")]),
Chunk::new(vec![coded, page.gather(&[3, 2]).expect("in range")]),
];
let mixed = [
Chunk::new(vec![
Vector::constant(int.clone(), Value::Integer(7), 2),
page.gather(&[1, 0]).expect("in range"),
]),
Chunk::new(vec![Vector::constant(int.clone(), Value::Null, 2), words(&["x", "y"])]),
];
let want_shared = (
vec![Value::Integer(1), Value::Null, Value::Integer(6), Value::Integer(5)],
["one", "a string longer than twelve", "four", "three"],
);
let want_mixed = (
vec![Value::Integer(7), Value::Integer(7), Value::Null, Value::Null],
["a string longer than twelve", "one", "x", "y"],
);
for (chunks, (ints, strings)) in [(shared, want_shared), (mixed, want_mixed)] {
let chunks: Vec<Chunk> =
chunks.into_iter().map(|chunk| chunk.expect("two columns of two rows")).collect();
let side = Build::new(&[int.clone(), text.clone()], &chunks, &alone()).expect("lays");
let got = side.gather(&[0, 1, 2, 3]).expect("four positions in range");
let read: Vec<Value> = (0..4).map(|at| got[0].value_at(at)).collect();
assert_eq!(read, ints);
let read: Vec<Value> = (0..4).map(|at| got[1].value_at(at)).collect();
let strings: Vec<Value> =
strings.iter().map(|&word| Value::Varchar(word.to_string())).collect();
assert_eq!(read, strings);
}
}
}