use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use rudb_catalog::Table;
use rudb_common::{Error, Field, LogicalType, Result};
use rudb_csv::Reader as CsvReader;
use rudb_functions::{Given, TableFunction, csv_given, open_csv, open_parquet, series_length};
use rudb_kernels::cast;
use rudb_parquet::Reader;
use rudb_pipeline::{Morsel, Progress, Source};
use rudb_plan::{ExprRef, Plan, Slice};
use rudb_vector::{Chunk, Data, VECTOR_SIZE, Vector};
use crate::expr::evaluate_all;
use crate::schema::Schema;
#[derive(Debug)]
pub(crate) struct Handout {
next: AtomicU64,
total: u64,
}
impl Handout {
pub(crate) fn new(total: usize) -> Self {
Self { next: AtomicU64::new(0), total: u64::try_from(total).unwrap_or(u64::MAX) }
}
pub(crate) fn take(&self) -> Option<Morsel> {
let at = self.next.fetch_add(1, Ordering::Relaxed);
(at < self.total).then(|| Morsel::new(at, at, at + 1))
}
}
pub(crate) fn position(morsel: &Morsel) -> usize {
usize::try_from(morsel.cursor()).unwrap_or(usize::MAX)
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while reading a file")
}
#[derive(Debug)]
pub(crate) struct Scan<'a> {
table: &'a Table,
columns: Vec<usize>,
schema: Schema,
chunks: Handout,
}
impl<'a> Scan<'a> {
pub(crate) fn new(
plan: &Plan,
table: &'a Table,
index: u32,
projection: Slice,
) -> Result<Self> {
let fields = plan.field_list(projection).to_vec();
let mut columns = Vec::with_capacity(fields.len());
for field in &fields {
let position = table.column_index(&field.name).ok_or_else(|| {
Error::catalog(format!(
"Table \"{}\" does not have a column named \"{}\"",
table.name().table,
field.name
))
})?;
columns.push(position);
}
let schema = Schema::numbered(fields, index);
let chunks = Handout::new(table.rows().chunk_count());
Ok(Self { table, columns, schema, chunks })
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
}
impl Source for Scan<'_> {
fn morsel(&self) -> Option<Morsel> {
self.chunks.take()
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let at = position(morsel);
if at >= self.table.rows().chunk_count() {
*out = Chunk::empty(&self.schema.types());
return Ok(Progress::Done);
}
*out = self.table.rows().read(at, &self.columns)?;
morsel.advance(1);
Ok(Progress::Done)
}
}
#[derive(Debug)]
pub(crate) struct Dummy {
schema: Schema,
one: Handout,
}
impl Dummy {
pub(crate) fn new() -> Self {
Self { schema: Schema::empty(), one: Handout::new(1) }
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
}
impl Source for Dummy {
fn morsel(&self) -> Option<Morsel> {
self.one.take()
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
*out = Chunk::with_rows(Vec::new(), 1)?;
morsel.advance(1);
Ok(Progress::Done)
}
}
#[derive(Debug)]
pub(crate) struct Values {
schema: Schema,
chunks: Vec<Chunk>,
handout: Handout,
}
impl Values {
pub(crate) fn new(plan: &Plan, index: u32, columns: Slice, rows: Slice) -> Result<Self> {
let fields = plan.field_list(columns).to_vec();
let schema = Schema::numbered(fields, index);
let types = schema.types();
let source = Schema::empty();
let one = Chunk::with_rows(Vec::new(), 1)?;
let mut down: Vec<Vec<rudb_common::Value>> = vec![Vec::new(); types.len()];
for row in plan.row_list(rows) {
let exprs: Vec<ExprRef> = plan.expr_list(*row).to_vec();
if exprs.len() != types.len() {
return Err(Error::internal(format!(
"a VALUES row of {} expressions in a {} column list",
exprs.len(),
types.len()
)));
}
let evaluated = evaluate_all(plan, &exprs, &source, &one)?;
for (position, vector) in evaluated.iter().enumerate() {
down[position].push(vector.value_at(0));
}
}
let total = down.first().map_or(0, Vec::len);
let mut chunks = Vec::new();
let mut start = 0;
while start < total {
let end = (start + VECTOR_SIZE).min(total);
let mut built = Vec::with_capacity(types.len());
for (position, ty) in types.iter().enumerate() {
built.push(Vector::from_values(ty.clone(), &down[position][start..end])?);
}
chunks.push(Chunk::with_rows(built, end - start)?);
start = end;
}
let handout = Handout::new(chunks.len());
Ok(Self { schema, chunks, handout })
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
}
#[derive(Debug)]
pub(crate) struct Series {
schema: Schema,
start: i64,
step: i64,
rows: u64,
morsels: AtomicU64,
}
const RUN: u64 = 16 * VECTOR_SIZE as u64;
impl Series {
pub(crate) fn new(plan: &Plan, index: u32, function: &str, args: Slice) -> Result<Self> {
let Some(function) = TableFunction::lookup(function) else {
return Err(Error::internal(format!("a plan with a table function called {function}")));
};
let fields = vec![Field::new(function.name(), LogicalType::BigInt)];
let schema = Schema::numbered(fields, index);
let exprs: Vec<ExprRef> = plan.expr_list(args).to_vec();
let source = Schema::empty();
let one = Chunk::with_rows(Vec::new(), 1)?;
let evaluated = evaluate_all(plan, &exprs, &source, &one)?;
let mut given = Vec::with_capacity(evaluated.len());
for vector in &evaluated {
match vector.value_at(0) {
rudb_common::Value::Null => return Ok(Self::empty(schema)),
rudb_common::Value::BigInt(n) => given.push(n),
other => {
return Err(Error::internal(format!(
"a table function argument bound as BIGINT arrived as {other}"
)));
}
}
}
let (start, stop, step) = match given.as_slice() {
[stop] => (0, *stop, 1),
[start, stop] => (*start, *stop, 1),
[start, stop, step] => (*start, *stop, *step),
_ => {
return Err(Error::internal(format!(
"{}() bound with {} arguments",
function.name(),
given.len()
)));
}
};
let rows = u64::try_from(series_length(function, start, stop, step)?).unwrap_or(u64::MAX);
Ok(Self { schema, start, step, rows, morsels: AtomicU64::new(0) })
}
fn empty(schema: Schema) -> Self {
Self { schema, start: 0, step: 1, rows: 0, morsels: AtomicU64::new(0) }
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn value_at(&self, position: u64) -> i64 {
let steps = i64::try_from(position).unwrap_or(i64::MAX);
self.start.saturating_add(self.step.saturating_mul(steps))
}
}
impl Source for Series {
fn morsel(&self) -> Option<Morsel> {
let index = self.morsels.fetch_add(1, Ordering::Relaxed);
let start = index.saturating_mul(RUN);
(start < self.rows)
.then(|| Morsel::new(index, start, self.rows.min(start.saturating_add(RUN))))
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let count = usize::try_from(morsel.remaining()).unwrap_or(usize::MAX).min(VECTOR_SIZE);
if count == 0 {
*out = Chunk::empty(&[LogicalType::BigInt]);
return Ok(Progress::Done);
}
let mut at = self.value_at(morsel.cursor());
let mut counted = Vec::with_capacity(count);
for _ in 0..count {
counted.push(at);
at = at.saturating_add(self.step);
}
morsel.advance(u64::try_from(count).unwrap_or(u64::MAX));
let vector = Vector::flat(LogicalType::BigInt, Data::Int64(counted.into()))?;
*out = Chunk::with_rows(vec![vector], count)?;
Ok(if morsel.is_drained() { Progress::Done } else { Progress::More })
}
}
#[derive(Debug)]
pub(crate) struct FileScan {
function: TableFunction,
paths: Vec<String>,
given: Given,
wanted: Vec<Field>,
schema: Schema,
reading: Mutex<Reading>,
one: Handout,
}
#[derive(Debug)]
struct Reading {
at: usize,
reader: Option<FileReader>,
}
impl FileScan {
pub(crate) fn new(
plan: &Plan,
index: u32,
function: TableFunction,
args: Slice,
options: Slice,
settings: Slice,
columns: Slice,
) -> Result<Self> {
let paths = file_arguments(plan, args, function)?;
let given = csv_options(plan, options, settings)?;
let wanted = plan.field_list(columns).to_vec();
let scan = Self {
function,
paths,
given,
wanted: wanted.clone(),
schema: Schema::numbered(wanted, index),
reading: Mutex::new(Reading { at: 0, reader: None }),
one: Handout::new(1),
};
{
let mut reading = scan.reading.lock().map_err(poisoned)?;
scan.advance(&mut reading)?;
}
Ok(scan)
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn advance(&self, reading: &mut Reading) -> Result<()> {
reading.reader = None;
let Some(path) = self.paths.get(reading.at) else { return Ok(()) };
let mut reader = FileReader::open(self.function, path, self.given)?;
let first = if reading.at == 0 { None } else { self.paths.first().map(String::as_str) };
reader.project(&positions(self.function, &self.wanted, &reader.fields(), path, first)?)?;
reader.settle(&self.wanted)?;
reading.reader = Some(reader);
reading.at += 1;
Ok(())
}
fn conform(&self, chunk: Chunk, file: usize) -> Result<Chunk> {
let rows = chunk.len();
let mut columns = Vec::with_capacity(self.wanted.len());
let mut changed = false;
for (at, field) in self.wanted.iter().enumerate() {
let column = chunk.column(at)?;
if column.logical_type() == &field.ty {
columns.push(column.clone());
continue;
}
changed = true;
columns.push(cast(column, &field.ty, false).map_err(|error| {
let path = self.paths.get(file.saturating_sub(1)).map_or("", String::as_str);
Error::conversion(format!(
"Error while reading file \"{path}\": failed to cast column \"{}\" from type \
{} to {}: {}",
field.name,
column.logical_type(),
field.ty,
error.message()
))
})?);
}
if !changed {
return Ok(chunk);
}
Chunk::with_rows(columns, rows)
}
}
impl Source for FileScan {
fn morsel(&self) -> Option<Morsel> {
self.one.take()
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let mut reading = self.reading.lock().map_err(poisoned)?;
loop {
let file = reading.at;
let Some(reader) = reading.reader.as_mut() else {
morsel.advance(1);
*out = Chunk::empty(&self.schema.types());
return Ok(Progress::Done);
};
if let Some(chunk) = reader.next_chunk()? {
*out = self.conform(chunk, file)?;
return Ok(Progress::More);
}
self.advance(&mut reading)?;
}
}
}
#[derive(Debug)]
enum FileReader {
Parquet(Reader),
Csv(CsvReader),
}
impl FileReader {
fn open(function: TableFunction, path: &str, given: Given) -> Result<Self> {
match function {
TableFunction::ReadCsv => Ok(Self::Csv(open_csv(path, given)?)),
_ => Ok(Self::Parquet(open_parquet(path)?)),
}
}
fn fields(&self) -> Vec<Field> {
match self {
Self::Parquet(reader) => reader.fields(),
Self::Csv(reader) => reader.fields(),
}
}
fn project(&mut self, columns: &[usize]) -> Result<()> {
match self {
Self::Parquet(reader) => reader.project(columns),
Self::Csv(reader) => reader.project(columns),
}
}
fn settle(&mut self, wanted: &[Field]) -> Result<()> {
match self {
Self::Parquet(reader) => {
let text: Vec<bool> =
wanted.iter().map(|field| field.ty == LogicalType::Varchar).collect();
reader.as_string(&text);
Ok(())
}
Self::Csv(reader) => {
let types: Vec<LogicalType> = wanted.iter().map(|field| field.ty.clone()).collect();
reader.retype(&types)
}
}
}
fn next_chunk(&mut self) -> Result<Option<Chunk>> {
match self {
Self::Parquet(reader) => reader.next_chunk(),
Self::Csv(reader) => reader.next_chunk(),
}
}
}
fn csv_options(plan: &Plan, options: Slice, settings: Slice) -> Result<Given> {
if options.len == 0 {
return Ok(Given::default());
}
let exprs: Vec<ExprRef> = plan.expr_list(settings).to_vec();
let source = Schema::empty();
let one = Chunk::with_rows(Vec::new(), 1)?;
let evaluated = evaluate_all(plan, &exprs, &source, &one)?;
let names: Vec<&str> = plan.name_list(options).iter().map(|name| plan.string(*name)).collect();
let written: Vec<(&str, rudb_common::Value)> =
names.into_iter().zip(evaluated.iter().map(|vector| vector.value_at(0))).collect();
csv_given(&written)
}
fn file_arguments(plan: &Plan, args: Slice, function: TableFunction) -> Result<Vec<String>> {
let exprs: Vec<ExprRef> = plan.expr_list(args).to_vec();
let source = Schema::empty();
let one = Chunk::with_rows(Vec::new(), 1)?;
let evaluated = evaluate_all(plan, &exprs, &source, &one)?;
let mut paths = Vec::with_capacity(evaluated.len());
for vector in &evaluated {
match vector.value_at(0) {
rudb_common::Value::Varchar(path) => paths.push(path),
other => {
return Err(Error::internal(format!(
"{}() bound with {other:?} rather than constant file names",
function.name()
)));
}
}
}
Ok(paths)
}
fn positions(
function: TableFunction,
wanted: &[Field],
held: &[Field],
path: &str,
first: Option<&str>,
) -> Result<Vec<usize>> {
let mut positions = Vec::with_capacity(wanted.len());
for field in wanted {
let at = held.iter().position(|column| column.name == field.name).ok_or_else(|| {
let Some(first) = first else {
return Error::io(format!(
"File \"{path}\" does not have a column named \"{}\"",
field.name
));
};
if matches!(function, TableFunction::ReadCsv) {
return rudb_csv::mismatch(first, path, &field.name);
}
let candidates: Vec<&str> = held.iter().map(|column| column.name.as_str()).collect();
Error::invalid_input(format!(
"Failed to read file \"{path}\": schema mismatch in glob: column \"{}\" was read \
from the original file \"{first}\", but could not be found in file \
\"{path}\".\nCandidate names: {}\nIf you are trying to read files with different \
schemas, try setting union_by_name=True",
field.name,
candidates.join(", ")
))
})?;
positions.push(at);
}
Ok(positions)
}
impl Source for Values {
fn morsel(&self) -> Option<Morsel> {
self.handout.take()
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
*out = match self.chunks.get(position(morsel)) {
Some(chunk) => chunk.clone(),
None => Chunk::empty(&self.schema.types()),
};
morsel.advance(1);
Ok(Progress::Done)
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicU64;
use rudb_common::{LogicalType, Value};
use rudb_pipeline::{Progress, Source};
use rudb_vector::Chunk;
use super::{Handout, RUN, Schema, Series, VECTOR_SIZE};
fn series(start: i64, step: i64, rows: u64) -> Series {
Series { schema: Schema::empty(), start, step, rows, morsels: AtomicU64::new(0) }
}
fn drained(series: &Series) -> (Vec<i64>, usize) {
let mut values = Vec::new();
let mut morsels = 0;
while let Some(mut morsel) = series.morsel() {
morsels += 1;
loop {
let mut chunk = Chunk::empty(&[LogicalType::BigInt]);
let progress = series.read(&mut morsel, &mut chunk).expect("a series reads");
for row in 0..chunk.len() {
match chunk.value_at(row, 0) {
Value::BigInt(value) => values.push(value),
other => panic!("a series produced {other}"),
}
}
if progress == Progress::Done {
break;
}
}
}
(values, morsels)
}
#[test]
fn a_morsel_of_a_series_is_read_a_chunk_at_a_time() {
let rows = VECTOR_SIZE as u64 * 2 + 5;
let (values, morsels) = drained(&series(0, 1, rows));
assert_eq!(morsels, 1);
assert_eq!(values.len(), rows as usize);
assert_eq!(values[0], 0);
assert_eq!(values[values.len() - 1], rows as i64 - 1);
}
#[test]
fn a_series_longer_than_a_morsel_carries_on_where_the_last_one_stopped() {
let rows = RUN + 3;
let (values, morsels) = drained(&series(10, 3, rows));
assert_eq!(morsels, 2);
assert_eq!(values.len(), rows as usize);
assert_eq!(values[0], 10);
assert_eq!(values[RUN as usize], 10 + 3 * RUN as i64);
assert_eq!(values[values.len() - 1], 10 + 3 * (rows as i64 - 1));
}
#[test]
fn a_series_of_nothing_hands_out_no_work() {
let (values, morsels) = drained(&series(0, 1, 0));
assert!(values.is_empty());
assert_eq!(morsels, 0);
}
#[test]
fn a_handout_gives_each_position_to_one_caller_and_then_stops() {
let handout = Handout::new(3);
let taken: Vec<u64> = (0..3).map(|_| handout.take().expect("a position").start()).collect();
assert_eq!(taken, [0, 1, 2]);
assert!(handout.take().is_none());
assert!(handout.take().is_none());
}
}