use std::cell::Cell;
use std::collections::{HashMap, HashSet};
use std::ops::Range;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use rudb_catalog::Table;
use rudb_common::{Error, Field, LogicalType, Result, Session, Value};
use rudb_csv::{Part, Reader as CsvReader, Split};
use rudb_encoding::sequence::Sequence;
use rudb_functions::{
FILE_ROW_NUMBER, Given, TableFunction, csv_given, open_csv, open_parquet, series_length,
};
use rudb_graph::Rids;
use rudb_kernels::{Stepping, cast, moment_steps};
use rudb_metrics::Counters;
use rudb_native::{Reader as NativeReader, RunProjectionPart, RunProjectionScan};
use rudb_parquet::{Bound, Op, Reader, Test, skips};
use rudb_pipeline::{Compaction, Gauge, Morsel, Progress, Source, narrow};
use rudb_plan::{ConjunctionOp, Expr, ExprRef, Node, NodeRef, Plan, Slice};
use rudb_seam::{Context, SeamId, Settings};
use rudb_storage::Probe;
use rudb_vector::{Chunk, Data, Selection, VECTOR_SIZE, Vector};
use crate::cutoff::Cutoff;
use crate::expr::{evaluate_all, evaluate_all_in_time_zone};
use crate::prepared::{Prepared, Scratch};
use crate::register::compaction;
use crate::schema::Schema;
use crate::sideways::{Keys, Sideways};
use crate::stream::{later_passes, marking_pays};
use crate::table::{Across, hash};
#[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 total(&self) -> usize {
usize::try_from(self.total).unwrap_or(usize::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 number(&self) -> Option<u64> {
let at = self.next.fetch_add(1, Ordering::Relaxed);
(at < self.total).then_some(at)
}
}
pub(crate) fn position(morsel: &Morsel) -> usize {
usize::try_from(morsel.cursor()).unwrap_or(usize::MAX)
}
#[derive(Debug)]
pub(crate) struct Frequencies {
chunks: Vec<Chunk>,
handout: Handout,
}
impl Frequencies {
pub(crate) fn records(schema: Schema, records: Vec<Vec<Value>>) -> Result<Self> {
let types = schema.types();
let mut chunks = Vec::with_capacity(records.len().div_ceil(VECTOR_SIZE));
for records in records.chunks(VECTOR_SIZE) {
let mut columns = vec![Vec::with_capacity(records.len()); types.len()];
for record in records {
if record.len() != types.len() {
return Err(Error::internal("a stored aggregate row has the wrong width"));
}
for (values, value) in columns.iter_mut().zip(record) {
values.push(value.clone());
}
}
let vectors = columns
.into_iter()
.zip(&types)
.map(|(values, ty)| Vector::from_values(ty.clone(), &values))
.collect::<Result<Vec<_>>>()?;
chunks.push(Chunk::with_rows(vectors, records.len())?);
}
let handout = Handout::new(chunks.len());
Ok(Self { chunks, handout })
}
pub(crate) fn grouped(schema: Schema, entries: Vec<(Vec<Value>, u64)>) -> Result<Self> {
let output_types = schema.types();
let Some((&LogicalType::BigInt, group_types)) = output_types.split_last() else {
return Err(Error::internal("a grouped frequency source count is not BIGINT"));
};
let mut chunks = Vec::with_capacity(entries.len().div_ceil(VECTOR_SIZE));
for entries in entries.chunks(VECTOR_SIZE) {
let mut columns = vec![Vec::with_capacity(entries.len()); group_types.len()];
let mut counts = Vec::with_capacity(entries.len());
for (keys, count) in entries {
if keys.len() != group_types.len() {
return Err(Error::internal("a grouped frequency key has the wrong width"));
}
for (column, key) in columns.iter_mut().zip(keys) {
column.push(key.clone());
}
counts.push(
i64::try_from(*count)
.map_err(|_| Error::internal("a stored frequency exceeds BIGINT"))?,
);
}
let mut output = columns
.into_iter()
.zip(group_types)
.map(|(values, ty)| Vector::from_values(ty.clone(), &values))
.collect::<Result<Vec<_>>>()?;
output.push(Vector::flat(LogicalType::BigInt, Data::Int64(counts.into()))?);
chunks.push(Chunk::with_rows(output, entries.len())?);
}
let handout = Handout::new(chunks.len());
Ok(Self { chunks, handout })
}
}
impl Source for Frequencies {
fn morsel(&self) -> Option<Morsel> {
self.handout.take()
}
fn morsels(&self, _threads: usize, _weight: usize) -> Option<usize> {
Some(self.handout.total())
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let at = position(morsel);
*out = self
.chunks
.get(at)
.ok_or_else(|| Error::internal("a frequency morsel is out of range"))?
.clone();
morsel.advance(1);
Ok(Progress::Done)
}
}
#[derive(Debug)]
pub(crate) struct ProjectionDistinct<'a> {
reader: &'a NativeReader,
order: usize,
covered: usize,
schema: Schema,
handout: Handout,
partitions: AtomicUsize,
scanner: Mutex<Option<Arc<RunProjectionScan<'a>>>>,
output: Mutex<ProjectionOutput>,
}
#[derive(Debug, Default)]
struct ProjectionOutput {
parts: Vec<Option<RunProjectionPart>>,
finished: usize,
chunks: Option<Vec<Chunk>>,
next: usize,
}
impl<'a> ProjectionDistinct<'a> {
pub(crate) fn new(
reader: &'a NativeReader,
order: usize,
covered: usize,
schema: Schema,
) -> Self {
Self {
reader,
order,
covered,
schema,
handout: Handout::new(8),
partitions: AtomicUsize::new(1),
scanner: Mutex::new(None),
output: Mutex::new(ProjectionOutput::default()),
}
}
fn chunks(&self, rows: Vec<(i32, u64)>) -> Result<Vec<Chunk>> {
let group_type = self.schema.types()[0].clone();
let entries = rows
.into_iter()
.map(|(group, count)| {
let value = match group_type {
LogicalType::TinyInt => Value::TinyInt(
i8::try_from(group)
.map_err(|_| Error::internal("a projection group exceeds TINYINT"))?,
),
LogicalType::SmallInt => Value::SmallInt(
i16::try_from(group)
.map_err(|_| Error::internal("a projection group exceeds SMALLINT"))?,
),
LogicalType::Integer => Value::Integer(group),
_ => return Err(Error::internal("a projection group has an unsupported type")),
};
Ok((vec![value], count))
})
.collect::<Result<Vec<_>>>()?;
Ok(Frequencies::grouped(self.schema.clone(), entries)?.chunks)
}
}
impl Source for ProjectionDistinct<'_> {
fn morsel(&self) -> Option<Morsel> {
self.handout.take().filter(|morsel| {
usize::try_from(morsel.index()).unwrap_or(usize::MAX)
< self.partitions.load(Ordering::Relaxed)
})
}
fn morsels(&self, threads: usize, _weight: usize) -> Option<usize> {
let rows = self.reader.table().rows();
let partitions = threads.clamp(1, 8).min(rows.div_ceil(1 << 19).max(1));
self.partitions.store(partitions, Ordering::Relaxed);
Some(partitions)
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let mut held =
self.output.lock().map_err(|_| Error::internal("projection output lock poisoned"))?;
if held.chunks.is_none() && self.partitions.load(Ordering::Relaxed) == 1 {
let rows = self
.reader
.grouped_distinct_run_projection_with_workers(
self.order,
self.covered,
usize::MAX,
1,
)?
.ok_or_else(|| Error::internal("a selected covering projection disappeared"))?;
held.chunks = Some(self.chunks(rows)?);
}
if held.chunks.is_none() && held.parts.get(position(morsel)).is_none_or(Option::is_none) {
drop(held);
let scanner = {
let mut slot = self
.scanner
.lock()
.map_err(|_| Error::internal("projection scanner lock poisoned"))?;
if slot.is_none() {
*slot = Some(Arc::new(
self.reader.run_projection_scan(self.order, self.covered)?.ok_or_else(
|| Error::internal("a selected covering projection disappeared"),
)?,
));
}
Arc::clone(slot.as_ref().expect("projection scanner initialized"))
};
let partitions = self.partitions.load(Ordering::Relaxed).min(scanner.pages());
let part = position(morsel);
if part >= partitions {
*out = Chunk::empty(&self.schema.types());
morsel.advance(1);
return Ok(Progress::Done);
}
let scanned = scanner.partition(part, partitions)?;
held = self
.output
.lock()
.map_err(|_| Error::internal("projection output lock poisoned"))?;
if held.parts.is_empty() {
held.parts.resize_with(partitions, || None);
}
held.parts[part] = Some(scanned);
held.finished += 1;
if held.finished < partitions {
*out = Chunk::empty(&self.schema.types());
morsel.advance(1);
return Ok(Progress::Done);
}
let parts = held
.parts
.iter_mut()
.map(|part| part.take().expect("all parts finished"))
.collect::<Vec<_>>();
let rows = scanner.finish(parts, usize::MAX)?;
held.chunks = Some(self.chunks(rows)?);
}
let chunks = held.chunks.as_ref().expect("projection output was initialized");
let total = chunks.len();
let next = chunks.get(held.next).cloned();
if let Some(chunk) = next {
*out = chunk;
held.next += 1;
if held.next < total {
return Ok(Progress::More);
}
} else {
*out = Chunk::empty(&self.schema.types());
}
morsel.advance(1);
Ok(Progress::Done)
}
}
#[derive(Debug)]
pub(crate) struct Summary {
row: Chunk,
handout: Handout,
}
impl Summary {
pub(crate) fn new(schema: &Schema, values: &[Value]) -> Result<Self> {
let types = schema.types();
if types.len() != values.len() {
return Err(Error::internal("a summary has a different number of values and columns"));
}
let columns = types
.iter()
.zip(values)
.map(|(ty, value)| Vector::from_values(ty.clone(), std::slice::from_ref(value)))
.collect::<Result<Vec<_>>>()?;
Ok(Self { row: Chunk::with_rows(columns, 1)?, handout: Handout::new(1) })
}
}
impl Source for Summary {
fn morsel(&self) -> Option<Morsel> {
self.handout.take()
}
fn morsels(&self, _threads: usize, _weight: usize) -> Option<usize> {
Some(self.handout.total())
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
*out = self.row.clone();
morsel.advance(1);
Ok(Progress::Done)
}
}
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<Option<usize>>,
offsets: Vec<i64>,
probes: Vec<Probe>,
sideways: Option<Arc<Sideways<'a>>>,
also: Vec<(Arc<Sideways<'a>>, Paying)>,
cutoff: Option<Arc<Cutoff>>,
index: u32,
testing: OnceLock<Vec<Probe>>,
schema: Schema,
chunks: Handout,
stripes: Vec<Range<usize>>,
pushed: Option<Pushed>,
waved: AtomicUsize,
spread: OnceLock<Spread>,
skipped: AtomicUsize,
deferred: Vec<usize>,
deferrable: Vec<usize>,
deferring: Paying,
unread: Vec<usize>,
paying: Paying,
passed: Paying,
counters: Option<Arc<Counters>>,
}
const WARMUP: usize = 1 << 16;
const TIGHT: usize = 16;
#[derive(Debug, Default)]
struct Paying {
seen: AtomicUsize,
kept: AtomicUsize,
}
impl Paying {
fn worth(&self) -> bool {
let seen = self.seen.load(Ordering::Relaxed);
seen < WARMUP
|| self.kept.load(Ordering::Relaxed).saturating_mul(4) < seen.saturating_mul(3)
}
fn loose(&self) -> bool {
let seen = self.seen.load(Ordering::Relaxed);
seen >= WARMUP && self.kept.load(Ordering::Relaxed).saturating_mul(2) > seen
}
fn tight(&self) -> bool {
let seen = self.seen.load(Ordering::Relaxed);
seen >= WARMUP && self.kept.load(Ordering::Relaxed).saturating_mul(TIGHT) < seen
}
fn saw(&self, rows: usize, kept: usize) {
self.seen.fetch_add(rows, Ordering::Relaxed);
self.kept.fetch_add(kept, Ordering::Relaxed);
}
}
fn read_above(plan: &Plan, filter: NodeRef, schema: &Schema) -> Option<Vec<bool>> {
let mut parents: HashMap<NodeRef, Vec<NodeRef>> = HashMap::new();
let mut seen = HashSet::new();
let mut stack = vec![plan.root()];
while let Some(at) = stack.pop() {
if !seen.insert(at) {
continue;
}
for child in plan.node(at).children().into_iter().flatten() {
parents.entry(child).or_default().push(at);
stack.push(child);
}
}
let mut above = vec![false; schema.types().len()];
let mut read = |expr: ExprRef| {
crate::join::columns(plan, expr, &mut |binding| {
if let Some(place) = schema.position_of(binding) {
above[place] = true;
}
});
};
let mut here = filter;
loop {
let [parent] = parents.get(&here)?.as_slice() else { return None };
match *plan.node(*parent) {
Node::Aggregate { groups, aggregates, .. } => {
for slice in [groups, aggregates] {
plan.expr_list(slice).iter().for_each(|&expr| read(expr));
}
break;
}
Node::Project { exprs, .. } => {
plan.expr_list(exprs).iter().for_each(|&expr| read(expr));
break;
}
Node::Filter { predicate, .. } => read(predicate),
Node::Join { conditions, .. } => {
plan.expr_list(conditions).iter().for_each(|&expr| read(expr));
}
Node::CrossProduct { .. } => {}
Node::Sort { keys, .. } | Node::TopN { keys, .. } => {
plan.sort_key_list(keys).iter().for_each(|key| read(key.expr));
}
Node::Limit { count, offset, .. }
if count.read().is_none() && offset.read().is_none() => {}
_ => return None,
}
here = *parent;
}
Some(above)
}
fn onto(columns: &[Option<usize>], tests: Vec<(usize, Op, Bound)>) -> Vec<Probe> {
tests
.into_iter()
.filter_map(|(at, op, value)| {
Some(Probe { column: columns.get(at)?.as_ref().copied()?, op, value })
})
.collect()
}
#[derive(Debug)]
pub(crate) struct Pushdown {
pub(crate) node: NodeRef,
pub(crate) predicate: ExprRef,
pub(crate) tests: Vec<(usize, Op, Bound)>,
pub(crate) whole: bool,
pub(crate) conjuncts: Vec<Option<Vec<(usize, Op, Bound)>>>,
pub(crate) marks: bool,
}
#[derive(Debug, Default)]
pub(crate) struct Filters<'a> {
pub(crate) pruning: Vec<(usize, Op, Bound)>,
pub(crate) pushed: Option<Pushdown>,
pub(crate) sideways: Option<Arc<Sideways<'a>>>,
pub(crate) also: Vec<Arc<Sideways<'a>>>,
pub(crate) cutoff: Option<Arc<Cutoff>>,
}
const SLOTS: usize = 64;
static READERS: AtomicUsize = AtomicUsize::new(0);
thread_local! {
static READER: Cell<usize> = const { Cell::new(usize::MAX) };
}
fn reduce(reduced: Option<(&Rids, u64)>, chunk: &mut Chunk) -> Result<()> {
let Some((rows, first)) = reduced else { return Ok(()) };
let len = chunk.len();
let kept = Selection::from_predicate(len, |row| rows.contains(first + row as u64));
if kept.len() == len {
return Ok(());
}
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.select(&kept)?;
Ok(())
}
fn reader() -> usize {
READER.with(|reader| {
let mut at = reader.get();
if at == usize::MAX {
at = READERS.fetch_add(1, Ordering::Relaxed);
reader.set(at);
}
at % SLOTS
})
}
#[derive(Debug)]
struct Pushed {
predicate: Prepared,
late: Option<Late>,
stored: Option<Stored>,
compaction: &'static dyn Compaction,
passes: u32,
marks: bool,
probes: Option<Vec<Probe>>,
settles: Option<Vec<Option<Vec<Probe>>>>,
spare: Vec<Mutex<Option<Working>>>,
}
#[derive(Debug)]
struct Working {
scratch: Scratch,
late_scratch: Option<Scratch>,
gauge: Gauge,
}
#[derive(Debug)]
struct Late {
input: usize,
predicate: Prepared,
seen: AtomicUsize,
kept: AtomicUsize,
}
const LATE_WARMUP: usize = 1 << 14;
impl Late {
fn worth(&self) -> bool {
let seen = self.seen.load(Ordering::Relaxed);
seen < LATE_WARMUP || self.kept.load(Ordering::Relaxed).saturating_mul(4) < seen
}
fn saw(&self, rows: usize, kept: usize) {
self.seen.fetch_add(rows, Ordering::Relaxed);
self.kept.fetch_add(kept, Ordering::Relaxed);
}
}
#[derive(Debug)]
struct Stored {
input: usize,
sequence: Sequence,
negated: bool,
}
fn stored_like(plan: &Plan, schema: &Schema, predicate: ExprRef) -> Option<Stored> {
let Expr::Function { name, args } = *plan.expr(predicate) else { return None };
let negated = match plan.string(name) {
"~~" => false,
"!~~" => true,
_ => return None,
};
let [column, constant] = plan.expr_list(args) else { return None };
let Expr::Column(binding) = *plan.expr(*column) else { return None };
let Expr::Constant(value) = *plan.expr(*constant) else { return None };
let Value::Varchar(spelling) = plan.value(value) else { return None };
let inner = spelling.strip_prefix('%')?.strip_suffix('%')?;
if inner.contains('_') {
return None;
}
let input = schema.position_of(binding)?;
if schema.types().get(input) != Some(&LogicalType::Varchar) {
return None;
}
let pieces: Vec<&[u8]> = inner.split('%').map(str::as_bytes).collect();
Some(Stored { input, sequence: Sequence::new(&pieces)?, negated })
}
fn late_like(plan: &Plan, schema: &Schema, predicate: ExprRef) -> Option<(usize, ExprRef)> {
if schema.width() != 2 {
return None;
}
let Expr::Conjunction { op: ConjunctionOp::And, children } = *plan.expr(predicate) else {
return None;
};
let [left, right] = plan.expr_list(children) else { return None };
for (candidate, other) in [(*left, *right), (*right, *left)] {
let Expr::Function { name, args } = *plan.expr(candidate) else { continue };
if plan.string(name) != "~~" {
continue;
}
let [column, constant] = plan.expr_list(args) else { continue };
let Expr::Column(binding) = *plan.expr(*column) else { continue };
if !matches!(*plan.expr(*constant), Expr::Constant(_)) {
continue;
}
let Some(input) = schema.position_of(binding) else { continue };
let Expr::Compare { left, right, .. } = *plan.expr(other) else { continue };
let compared = match (plan.expr(left), plan.expr(right)) {
(Expr::Column(binding), Expr::Constant(_))
| (Expr::Constant(_), Expr::Column(binding)) => *binding,
_ => continue,
};
if schema.position_of(compared) == Some(1 - input) {
return Some((input, candidate));
}
}
None
}
impl Pushed {
fn new(
plan: &Plan,
schema: &Schema,
columns: &[Option<usize>],
pushdown: Pushdown,
seams: &Settings,
session: &Session,
) -> Result<Self> {
let types = schema.types();
let context = Context::new(SeamId::ChunkCompaction, seams).with_types(&types);
let compaction = compaction().choose(&context)?.strategy();
let whole = pushdown.whole;
let wanted = pushdown.tests.len();
let probes = onto(columns, pushdown.tests);
let settles: Vec<Option<Vec<Probe>>> = pushdown
.conjuncts
.into_iter()
.map(|tests| {
let tests = tests?;
let wanted = tests.len();
Some(onto(columns, tests)).filter(|probes| probes.len() == wanted)
})
.collect();
let predicate = Prepared::one(plan, pushdown.predicate, schema)?.in_session(session);
let settles = (predicate.conjuncts() == Some(settles.len())
&& settles.iter().any(Option::is_some))
.then_some(settles);
let late = if columns.len() == 2 && columns.iter().all(Option::is_some) {
late_like(plan, schema, pushdown.predicate)
.map(|(input, expr)| {
Ok::<_, Error>(Late {
input,
predicate: Prepared::one(plan, expr, schema)?.in_session(session),
seen: AtomicUsize::new(0),
kept: AtomicUsize::new(0),
})
})
.transpose()?
} else {
None
};
let stored = stored_like(plan, schema, pushdown.predicate);
Ok(Self {
predicate,
late,
stored,
compaction,
passes: later_passes(plan, pushdown.node),
marks: pushdown.marks,
probes: (whole && probes.len() == wanted).then_some(probes),
settles,
spare: (0..SLOTS).map(|_| Mutex::new(None)).collect(),
})
}
fn take(&self, slot: usize) -> Working {
let waiting = self.spare[slot].lock().ok().and_then(|mut spare| spare.take());
waiting.unwrap_or_else(|| Working {
scratch: self.predicate.scratch(),
late_scratch: self.late.as_ref().map(|late| late.predicate.scratch()),
gauge: Gauge::new(self.passes),
})
}
fn give(&self, slot: usize, working: Working) {
if let Ok(mut spare) = self.spare[slot].lock() {
*spare = Some(working);
}
}
}
#[derive(Debug)]
struct Spread {
runs: Vec<Range<usize>>,
order: Vec<usize>,
handout: Handout,
}
impl<'a> Scan<'a> {
pub(crate) fn new(
plan: &Plan,
table: &'a Table,
index: u32,
projection: Slice,
filters: Filters<'a>,
seams: &Settings,
session: &Session,
) -> Result<Self> {
let Filters { pruning, pushed: pushdown, sideways, also, cutoff } = filters;
let fields = plan.field_list(projection).to_vec();
let mut columns = Vec::with_capacity(fields.len());
for field in &fields {
if field.name == FILE_ROW_NUMBER {
columns.push(None);
continue;
}
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(Some(position));
}
let probes = onto(&columns, pruning);
let schema = Schema::numbered(fields, index);
let chunks = Handout::new(table.rows().chunk_count());
let mut next = 0_i64;
let mut offsets = Vec::with_capacity(table.rows().chunk_count());
for at in 0..table.rows().chunk_count() {
offsets.push(next);
next =
next.saturating_add(i64::try_from(table.rows().chunk_len(at)?).unwrap_or(i64::MAX));
}
let stripes = table.rows().stripe_parts();
let mut read = vec![false; columns.len()];
if let Some(pushdown) = &pushdown {
crate::join::columns(plan, pushdown.predicate, &mut |binding| {
if let Some(at) = schema.position_of(binding) {
read[at] = true;
}
});
}
let types = schema.types();
let deferrable: Vec<usize> =
(0..columns.len()).filter(|&at| columns[at].is_some() && !read[at]).collect();
let deferred =
deferrable.iter().copied().filter(|&at| types[at] == LogicalType::Varchar).collect();
let unread = pushdown
.as_ref()
.and_then(|pushdown| read_above(plan, pushdown.node, &schema))
.map(|above| {
(0..columns.len())
.filter(|&at| columns[at].is_some() && read[at] && !above[at])
.collect()
})
.unwrap_or_default();
let pushed = pushdown
.map(|pushdown| Pushed::new(plan, &schema, &columns, pushdown, seams, session))
.transpose()?;
Ok(Self {
table,
columns,
offsets,
probes,
sideways,
also: also.into_iter().map(|sideways| (sideways, Paying::default())).collect(),
cutoff,
index,
testing: OnceLock::new(),
schema,
chunks,
stripes,
pushed,
waved: AtomicUsize::new(0),
spread: OnceLock::new(),
skipped: AtomicUsize::new(0),
counters: None,
deferred,
deferrable,
deferring: Paying::default(),
unread,
paying: Paying::default(),
passed: Paying::default(),
})
}
fn apply(&self, at: usize, chunk: &mut Chunk, mark: bool) -> Result<()> {
let reduced = self.reduced(at).filter(|(rows, _)| !rows.is_full());
let Some(pushed) = self.pushed.as_ref() else { return reduce(reduced, chunk) };
let whole =
pushed.probes.as_deref().is_some_and(|probes| self.table.rows().certain(at, probes));
if whole {
self.waved.fetch_add(1, Ordering::Relaxed);
return reduce(reduced, chunk);
}
let slot = reader();
let mut working = pushed.take(slot);
let mut kept = match &pushed.settles {
Some(settles) => {
let rows = self.table.rows();
let settled: Vec<bool> = settles
.iter()
.map(|probes| probes.as_deref().is_some_and(|probes| rows.certain(at, probes)))
.collect();
pushed.predicate.evaluate_settled(chunk, &mut working.scratch, &settled)?
}
None => pushed.predicate.evaluate_filter(chunk, &mut working.scratch)?,
};
self.passed.saw(chunk.len(), kept.len());
if let Some((rows, first)) = reduced {
let held: Vec<u32> = kept
.iter()
.filter(|&row| rows.contains(first + row as u64))
.filter_map(|row| u32::try_from(row).ok())
.collect();
kept = Selection::from_indices(held);
}
let alone = reduced.is_none() && self.sideways.is_none() && self.also.is_empty();
if kept.len() != chunk.len() {
if mark && alone && pushed.marks && marking_pays(kept.len(), chunk.len()) {
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.marked(kept);
} else {
self.narrow_read_columns(pushed.compaction, chunk, &kept, &mut working.gauge)?;
}
}
pushed.give(slot, working);
Ok(())
}
fn narrow_read_columns(
&self,
how: &dyn Compaction,
chunk: &mut Chunk,
kept: &Selection,
gauge: &mut Gauge,
) -> Result<()> {
let mut unread = self.unread.clone();
let joins = self.sideways.iter().chain(self.also.iter().map(|(sideways, _)| sideways));
for sideways in joins {
let domain = sideways.domain(self.index).map(|(key, _)| key);
let sifting = sideways.sifting(self.index).map(|(key, _)| key);
unread.retain(|&at| domain != Some(at) && sifting != Some(at));
}
if unread.is_empty() {
return narrow(how, chunk, kept, gauge);
}
let rows = chunk.len();
let types = chunk.types();
let mut columns: Vec<Option<Vector>> = std::mem::replace(chunk, Chunk::empty(&[]))
.into_columns()
.into_iter()
.map(Some)
.collect();
let live: Vec<Vector> = columns
.iter_mut()
.enumerate()
.filter(|(at, _)| !unread.contains(at))
.filter_map(|(_, column)| column.take())
.collect();
let mut part = Chunk::with_rows(live, rows)?;
narrow(how, &mut part, kept, gauge)?;
let len = part.len();
let mut live = part.into_columns().into_iter();
let mut out = Vec::with_capacity(columns.len());
for (at, ty) in types.into_iter().enumerate() {
if unread.contains(&at) {
out.push(Vector::constant(ty, Value::Null, len));
} else {
out.push(
live.next().ok_or_else(|| Error::internal("a narrowed column went missing"))?,
);
}
}
*chunk = Chunk::with_rows(out, len)?;
Ok(())
}
fn narrow_read(&self, at: usize, out: &mut Chunk, mark: bool) -> Result<()> {
let placed = self.reduced(at).is_some_and(|(rows, _)| !rows.is_full());
if placed {
self.apply(at, out, false)?;
return self.sift(out);
}
self.sift_exact(out)?;
if out.is_empty() {
return Ok(());
}
if self.passed.loose() {
self.sift_hashed(out)?;
if out.is_empty() {
return Ok(());
}
return self.apply(at, out, mark);
}
self.apply(at, out, mark)?;
self.sift_hashed(out)
}
fn read_stored(&self, at: usize, out: &mut Chunk) -> Result<bool> {
let Some(pushed) = &self.pushed else { return Ok(false) };
let Some(stored) = &pushed.stored else { return Ok(false) };
if !self.unread.contains(&stored.input) || self.reduced(at).is_some() {
return Ok(false);
}
let joins = self.sideways.iter().chain(self.also.iter().map(|(sideways, _)| sideways));
for sideways in joins {
let domain = sideways.domain(self.index).map(|(key, _)| key);
let sifting = sideways.sifting(self.index).map(|(key, _)| key);
if domain == Some(stored.input) || sifting == Some(stored.input) {
return Ok(false);
}
}
let Some(column) = self.columns[stored.input] else { return Ok(false) };
let rows = self.table.rows();
let Some(kept) = rows.rows_holding(at, column, &stored.sequence, stored.negated)? else {
return Ok(false);
};
let len = rows.chunk_len(at)?;
let others: Vec<usize> = (0..self.columns.len())
.filter(|&place| place != stored.input)
.filter_map(|place| self.columns[place])
.collect();
let read = if others.is_empty() { None } else { Some(rows.read(at, &others)?) };
let types = self.schema.types();
let mut held = Vec::with_capacity(self.columns.len());
let mut real = 0;
for (place, column) in self.columns.iter().enumerate() {
if place == stored.input {
held.push(Vector::constant(types[place].clone(), Value::Null, len));
} else if column.is_some() {
let read =
read.as_ref().ok_or_else(|| Error::internal("a read column is missing"))?;
held.push(read.column(real)?.clone());
real += 1;
} else {
held.push(Vector::sequence(self.offsets[at], 1, len));
}
}
*out = Chunk::with_rows(held, len)?;
self.passed.saw(len, kept.len());
if kept.len() != len {
let kept = Selection::from_indices(kept);
let alone = self.sideways.is_none() && self.also.is_empty();
if alone && pushed.marks && marking_pays(kept.len(), len) {
let whole = std::mem::replace(out, Chunk::empty(&[]));
*out = whole.marked(kept);
} else {
let slot = reader();
let mut working = pushed.take(slot);
narrow(pushed.compaction, out, &kept, &mut working.gauge)?;
pushed.give(slot, working);
}
}
self.sift_exact(out)?;
if !out.is_empty() {
self.sift_hashed(out)?;
}
Ok(true)
}
fn read_late(&self, at: usize, out: &mut Chunk) -> Result<bool> {
let Some(pushed) = &self.pushed else { return Ok(false) };
let Some(late) = &pushed.late else { return Ok(false) };
if self.sideways.is_some() || !self.also.is_empty() || !late.worth() {
return Ok(false);
}
let Some(primary) = self.columns[late.input] else { return Ok(false) };
let Some(secondary) = self.columns[1 - late.input] else { return Ok(false) };
let read = self.table.rows().read(at, &[primary])?;
let len = read.len();
let mut columns = self
.schema
.types()
.into_iter()
.map(|ty| Vector::constant(ty, Value::Null, len))
.collect::<Vec<_>>();
columns[late.input] = read.column(0)?.clone();
let first = Chunk::with_rows(columns, len)?;
let slot = reader();
let mut working = pushed.take(slot);
let selected = late.predicate.evaluate_filter(
&first,
working
.late_scratch
.as_mut()
.ok_or_else(|| Error::internal("a late filter has no scratch"))?,
)?;
pushed.give(slot, working);
late.saw(len, selected.len());
if selected.len().saturating_mul(4) > len {
return Ok(false);
}
if selected.is_empty() {
*out = Chunk::empty(&self.schema.types());
return Ok(true);
}
let fetched = self.table.rows().read_selected(at, &[secondary], selected.indices())?;
let first = read.column(0)?.gather(selected.indices())?;
let second = fetched.column(0)?.clone();
let columns = if late.input == 0 { vec![first, second] } else { vec![second, first] };
*out = Chunk::with_rows(columns, selected.len())?;
self.apply(at, out, false)?;
Ok(true)
}
fn read_deferring(&self, at: usize, out: &mut Chunk) -> Result<bool> {
if !self.deferring.worth() || self.reduced(at).is_some() {
return Ok(false);
}
let own = self.sideways.iter().map(|sideways| (sideways, &self.paying));
let joins = own.chain(self.also.iter().map(|(sideways, paying)| (sideways, paying)));
let mut keys = Vec::new();
let mut tight = false;
for (sideways, paying) in joins {
let before = keys.len();
keys.extend(sideways.domain(self.index).map(|(key, _)| key));
keys.extend(sideways.sifting(self.index).map(|(key, _)| key));
tight |= keys.len() > before && paying.tight();
}
let wide = if tight { &self.deferrable } else { &self.deferred };
if wide.is_empty() || (self.pushed.is_none() && keys.is_empty()) {
return Ok(false);
}
let deferred: Vec<usize> = wide.iter().copied().filter(|at| !keys.contains(at)).collect();
let first: Vec<usize> = (0..self.columns.len())
.filter(|at| !deferred.contains(at))
.filter_map(|at| self.columns[at])
.collect();
if deferred.is_empty() || first.is_empty() {
return Ok(false);
}
let read = self.table.rows().read(at, &first)?;
let len = read.len();
let types = self.schema.types();
let mut held = Vec::with_capacity(self.columns.len() + 1);
let mut real = 0;
for (place, column) in self.columns.iter().enumerate() {
if deferred.contains(&place) {
held.push(Vector::constant(types[place].clone(), Value::Null, len));
} else if column.is_some() {
held.push(read.column(real)?.clone());
real += 1;
} else {
held.push(Vector::sequence(self.offsets[at], 1, len));
}
}
held.push(Vector::sequence(0, 1, len));
*out = Chunk::with_rows(held, len)?;
self.narrow_read(at, out, false)?;
let kept = out.len();
self.deferring.saw(len, kept);
let mut columns = std::mem::replace(out, Chunk::empty(&[])).into_columns();
let numbers =
columns.pop().ok_or_else(|| Error::internal("the row numbers went missing"))?;
if kept > 0 {
let wanted: Vec<usize> = deferred.iter().filter_map(|&at| self.columns[at]).collect();
let fetched = if kept == len {
self.table.rows().read(at, &wanted)?
} else {
let mut block = Vec::new();
if !numbers.signed_block(&mut block) || block.len() < kept {
block = (0..kept)
.map(|row| numbers.signed_at(row).and_then(|at| i64::try_from(at).ok()))
.collect::<Option<Vec<i64>>>()
.ok_or_else(|| Error::internal("a row number is not a row"))?;
}
let positions = block[..kept]
.iter()
.map(|&number| u32::try_from(number))
.collect::<std::result::Result<Vec<u32>, _>>()
.map_err(|_| Error::internal("a row number is not a row"))?;
self.table.rows().read_rows(at, &wanted, &positions)?
};
for (from, &place) in deferred.iter().enumerate() {
columns[place] = fetched.column(from)?.clone();
}
}
*out = Chunk::with_rows(columns, kept)?;
Ok(true)
}
fn sift(&self, chunk: &mut Chunk) -> Result<()> {
self.sift_exact(chunk)?;
self.sift_hashed(chunk)
}
fn sift_exact(&self, chunk: &mut Chunk) -> Result<()> {
let handoffs = || {
let own = self.sideways.iter().map(|sideways| (sideways, &self.paying));
own.chain(self.also.iter().map(|(sideways, paying)| (sideways, paying)))
};
for (sideways, paying) in handoffs() {
let Some((at, domain)) = sideways.domain(self.index) else { continue };
if !paying.worth() {
continue;
}
let Ok(column) = chunk.column(at) else { continue };
let rows = chunk.len();
let kept = domain.keep(column, rows, &mut Vec::new());
paying.saw(rows, kept.len());
if kept.len() < rows {
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.select(&Selection::from_indices(kept))?;
}
}
Ok(())
}
fn sift_hashed(&self, chunk: &mut Chunk) -> Result<()> {
let handoffs = || {
let own = self.sideways.iter().map(|sideways| (sideways, &self.paying));
own.chain(self.also.iter().map(|(sideways, paying)| (sideways, paying)))
};
for (sideways, paying) in handoffs() {
if chunk.is_empty() {
return Ok(());
}
if sideways.domain(self.index).is_some() {
continue;
}
let Some((at, filter)) = sideways.sifting(self.index) else { continue };
if !paying.worth() {
continue;
}
let Ok(column) = chunk.column(at) else { continue };
let rows = chunk.len();
let mut hashes = Vec::new();
hash(std::slice::from_ref(column), rows, &mut hashes, Across::TwoInputs);
let mut held = Vec::new();
filter.holds_run(&hashes, &mut held);
let kept = Selection::from_predicate(rows, |row| held[row]);
paying.saw(rows, kept.len());
if kept.len() < rows {
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.select(&kept)?;
}
}
Ok(())
}
fn testing(&self) -> &[Probe] {
self.testing.get_or_init(|| {
let mut probes = self.probes.clone();
if let Some(sideways) = self.sideways.as_ref() {
if let (Some(counters), Some(reduced)) =
(&self.counters, sideways.reduction(self.index))
{
counters.reducing(reduced);
}
probes.extend(onto(&self.columns, sideways.tests(self.index)));
}
for (sideways, _) in &self.also {
probes.extend(onto(&self.columns, sideways.tests(self.index)));
}
probes
})
}
fn cutoff(&self) -> Vec<Probe> {
let Some(cutoff) = self.cutoff.as_ref() else { return Vec::new() };
let Some(test) = cutoff.probe(self.index) else { return Vec::new() };
onto(&self.columns, vec![test])
}
fn ordered(&self) -> Option<(usize, bool)> {
let (column, op) = self.cutoff.as_ref()?.ordered(self.index)?;
let column = (*self.columns.get(column)?)?;
Some((column, op == Op::GreaterOrEqual))
}
fn ruled(&self, at: usize, probes: &[Probe], cutoff: &[Probe]) -> bool {
let rows = self.table.rows();
(!probes.is_empty() && rows.skips(at, probes))
|| (!cutoff.is_empty() && rows.skips(at, cutoff))
|| self.reduced_away(at)
|| self.keyed_away(at)
}
fn keyed(&self) -> impl Iterator<Item = (usize, &Keys)> {
let joins = self.sideways.iter().chain(self.also.iter().map(|(sideways, _)| sideways));
joins.filter_map(|sideways| {
let (at, keys) = sideways.keys(self.index)?;
Some((self.columns.get(at).copied().flatten()?, keys))
})
}
fn keyed_away(&self, at: usize) -> bool {
self.keyed().any(|(column, keys)| {
self.table.rows().ruled_by(at, column, &|range| keys.misses(range))
})
}
fn reduced(&self, at: usize) -> Option<(&Rids, u64)> {
let rows = self.sideways.as_ref()?.rows(self.index)?;
let first = u64::try_from(*self.offsets.get(at)?).ok()?;
Some((rows, first))
}
fn reduced_away(&self, at: usize) -> bool {
let Some((rows, first)) = self.reduced(at) else { return false };
let Some(len) =
self.table.rows().chunk_len(at).ok().and_then(|len| u64::try_from(len).ok())
else {
return false;
};
len > 0 && !rows.any_between(first, first + len - 1)
}
fn living(&self, threads: usize, weight: usize) -> (Vec<Live>, usize) {
let (mut live, rows) = self.bounded();
let working = live.iter().filter(|stripe| !stripe.parts.is_empty()).count();
let probes = self.testing();
let keyed = self.keyed().next().is_some();
let sifting = !probes.is_empty() && worth_sifting(working, threads, rows, weight);
if !keyed && !sifting {
return (live, rows);
}
let mut sifted = 0;
for stripe in &mut live {
stripe.parts.retain(|&at| {
!(sifting && self.table.rows().skips(at, probes)) && !self.keyed_away(at)
});
stripe.rows =
stripe.parts.iter().map(|&at| self.table.rows().chunk_len(at).unwrap_or(0)).sum();
sifted += stripe.rows;
}
(live, sifted)
}
fn bounded(&self) -> (Vec<Live>, usize) {
let probes = self.testing();
let mut rows = 0;
let mut live = Vec::with_capacity(self.stripes.len());
for (stripe, parts) in self.stripes.iter().enumerate() {
let keyed_away = || {
self.keyed().any(|(column, keys)| {
self.table.rows().stripe_ruled_by(stripe, column, &|range| keys.misses(range))
})
};
if (!probes.is_empty() && self.table.rows().stripe_skips(stripe, probes))
|| keyed_away()
{
live.push(Live::default());
continue;
}
let held = self.table.rows().stripe_rows(stripe);
rows += held;
live.push(Live { parts: parts.clone().collect(), rows: held });
}
(live, rows)
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
pub(crate) fn watched(mut self, counters: Arc<Counters>) -> Self {
self.counters = Some(counters);
self
}
fn unreached(&self, runs: &[Range<usize>]) {
let Some(counters) = self.counters.as_ref() else { return };
let covered: usize = runs.iter().map(|run| run.end.saturating_sub(run.start)).sum();
for _ in 0..self.table.rows().chunk_count().saturating_sub(covered) {
counters.part_pruned();
}
}
}
#[derive(Debug, Default)]
struct Live {
parts: Vec<usize>,
rows: usize,
}
fn worth_sifting(working: usize, threads: usize, rows: usize, weight: usize) -> bool {
working < threads.min(instances_for(rows, weight))
}
fn runs_of(live: &[Live], instances: usize, rows: impl Fn(usize) -> usize) -> Vec<Range<usize>> {
let total: usize = live.iter().map(|stripe| stripe.rows).sum();
let working = live.iter().filter(|stripe| !stripe.parts.is_empty()).count();
let share = total.div_ceil(instances.max(1)).max(1);
let (whole, piece) = if working >= instances {
(share, share.div_ceil(2).max(1))
} else {
let piece = share.div_ceil(4).max(1);
(piece, piece)
};
let mut runs: Vec<Range<usize>> = Vec::with_capacity(instances.saturating_mul(2));
for stripe in live {
let parts = &stripe.parts;
let Some((&first, &last)) = parts.first().zip(parts.last()) else { continue };
if stripe.rows <= whole {
runs.push(first..last + 1);
continue;
}
let mut taken = 0;
let mut open = false;
for &at in parts {
let size = rows(at);
if open && taken + size <= piece {
if let Some(run) = runs.last_mut() {
run.end = at + 1;
}
taken += size;
} else {
runs.push(at..at + 1);
taken = size;
open = true;
}
}
}
runs
}
const ORDERED_RUN: usize = 8;
fn best_first(
runs: Vec<Range<usize>>,
front: usize,
descending: bool,
range: impl Fn(usize) -> Option<rudb_storage::Range>,
) -> (Vec<Range<usize>>, Vec<usize>) {
let mut pieces = Vec::with_capacity(runs.len());
for run in runs.iter().cloned() {
let mut start = run.start;
while start < run.end {
let end = (start + ORDERED_RUN).min(run.end);
pieces.push(start..end);
start = end;
}
}
let reach: Vec<Option<Bound>> = pieces
.iter()
.map(|piece| {
piece
.clone()
.filter_map(|at| {
let range = range(at)?;
if descending { range.high } else { range.low }
})
.reduce(
|one, other| if descending { one.larger(other) } else { one.smaller(other) },
)
})
.collect();
let better = if descending { std::cmp::Ordering::Greater } else { std::cmp::Ordering::Less };
let mut first = vec![false; pieces.len()];
let mut chosen = Vec::with_capacity(front);
for _ in 0..front.min(pieces.len()) {
let mut best: Option<usize> = None;
for (at, bound) in reach.iter().enumerate() {
let Some(bound) = bound else { continue };
if first[at] {
continue;
}
let wins = best
.and_then(|best| reach[best].as_ref())
.is_none_or(|held| bound.order(held) == Some(better));
if wins {
best = Some(at);
}
}
let Some(best) = best else { break };
first[best] = true;
chosen.push(best);
}
if chosen.is_empty() {
return (runs, Vec::new());
}
let mut out = Vec::with_capacity(runs.len() + 2 * chosen.len());
let mut landed = vec![0; pieces.len()];
let mut early = Vec::with_capacity(runs.len() + 2 * chosen.len());
let mut next = 0;
for run in &runs {
let mut rest: Option<Range<usize>> = None;
while let Some(piece) = pieces.get(next).filter(|piece| piece.start < run.end) {
if first[next] {
if let Some(rest) = rest.take() {
out.push(rest);
early.push(false);
}
landed[next] = out.len();
out.push(piece.clone());
early.push(true);
} else {
match rest.as_mut() {
Some(rest) => rest.end = piece.end,
None => rest = Some(piece.clone()),
}
}
next += 1;
}
if let Some(rest) = rest {
out.push(rest);
early.push(false);
}
}
let mut order: Vec<usize> = chosen.iter().map(|&piece| landed[piece]).collect();
order.extend((0..out.len()).filter(|&at| !early[at]));
(out, order)
}
impl Source for Scan<'_> {
fn morsel(&self) -> Option<Morsel> {
let Some(spread) = self.spread.get() else {
return self.chunks.take();
};
let taken = usize::try_from(spread.handout.number()?).unwrap_or(usize::MAX);
let index = spread.order.get(taken).copied().unwrap_or(taken);
let run = spread.runs.get(index)?;
let start = u64::try_from(run.start).unwrap_or(u64::MAX);
let end = u64::try_from(run.end).unwrap_or(u64::MAX);
Some(Morsel::new(u64::try_from(index).unwrap_or(u64::MAX), start, end))
}
fn morsels(&self, threads: usize, weight: usize) -> Option<usize> {
let chunks = self.chunks.total();
if self.stripes.is_empty() {
return Some(chunks);
}
let (live, rows) = self.living(threads, weight);
let instances = chunks.min(threads).min(instances_for(rows, weight));
if instances > 0 && self.stripes.len() >= instances {
self.table.rows().keep_stripes(instances.saturating_mul(2));
let runs = runs_of(&live, instances, |at| self.table.rows().chunk_len(at).unwrap_or(0));
self.unreached(&runs);
let (runs, order) = match self.ordered() {
Some((column, descending)) => best_first(runs, instances * 2, descending, |at| {
self.table.rows().range_of(at, column)
}),
None => (runs, Vec::new()),
};
let _ = self.spread.set(Spread { handout: Handout::new(runs.len()), runs, order });
}
Some(instances)
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let probes = self.testing();
let cutoff = self.cutoff();
let at = loop {
let at = position(morsel);
if at >= self.table.rows().chunk_count() || morsel.is_drained() {
*out = Chunk::empty(&self.schema.types());
return Ok(Progress::Done);
}
morsel.advance(1);
if !self.ruled(at, probes, &cutoff) {
break at;
}
self.skipped.fetch_add(1, Ordering::Relaxed);
if let Some(counters) = &self.counters {
counters.part_pruned();
}
};
if let Some(counters) = &self.counters {
counters.part_read();
}
if self.read_stored(at, out)? || self.read_late(at, out)? || self.read_deferring(at, out)? {
return Ok(more(morsel));
}
let projected: Vec<usize> = self.columns.iter().flatten().copied().collect();
let read = self.table.rows().read(at, &projected)?;
if self.columns.iter().all(Option::is_some) {
*out = read;
self.narrow_read(at, out, true)?;
return Ok(more(morsel));
}
let mut held = Vec::with_capacity(self.columns.len());
let mut real = 0;
for column in &self.columns {
if column.is_some() {
held.push(read.column(real)?.clone());
real += 1;
} else {
held.push(Vector::sequence(self.offsets[at], 1, read.len()));
}
}
*out = Chunk::with_rows(held, read.len())?;
self.narrow_read(at, out, true)?;
Ok(more(morsel))
}
fn weight(&self) -> usize {
self.pushed.as_ref().map_or(0, |pushed| pushed.predicate.passes())
}
}
fn more(morsel: &Morsel) -> Progress {
if morsel.is_drained() { Progress::Done } else { Progress::More }
}
fn instances_for(rows: usize, weight: usize) -> usize {
let small = rows.saturating_mul(weight.max(1)).div_ceil(25_000).min(8);
let large = rows.div_ceil(62_500);
small.max(large).max(1)
}
#[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 morsels(&self, _threads: usize, _weight: usize) -> Option<usize> {
Some(self.one.total())
}
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,
session: &Session,
) -> Result<Self> {
let time_zone = session.session_time_zone();
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<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_in_time_zone(plan, &exprs, &source, &one, time_zone)?;
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,
ty: LogicalType,
listed: Option<Arc<[i64]>>,
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 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)?;
if let [start, stop, step] = evaluated.as_slice() {
if *step.logical_type() == LogicalType::Interval {
let ty = start.logical_type().clone();
let schema = Schema::numbered(vec![Field::new(function.name(), ty.clone())], index);
let empty = Self { ty, ..Self::empty(schema.clone()) };
let values = (start.value_at(0), stop.value_at(0), step.value_at(0));
let Some(stepping) = moments(function, &values.0, &values.1, &values.2)? else {
return Ok(empty);
};
let rows = u64::try_from(stepping.len()).unwrap_or(u64::MAX);
return Ok(match stepping {
Stepping::Even { start, step, .. } => Self { start, step, rows, ..empty },
Stepping::Listed(stamps) => Self { rows, listed: Some(stamps.into()), ..empty },
});
}
}
let fields = vec![Field::new(function.name(), LogicalType::BigInt)];
let schema = Schema::numbered(fields, index);
let mut given = Vec::with_capacity(evaluated.len());
for vector in &evaluated {
match vector.value_at(0) {
Value::Null => return Ok(Self::empty(schema)),
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 { start, step, rows, ..Self::empty(schema) })
}
fn empty(schema: Schema) -> Self {
Self {
schema,
start: 0,
step: 1,
rows: 0,
ty: LogicalType::BigInt,
listed: None,
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 morsels(&self, _threads: usize, _weight: usize) -> Option<usize> {
Some(usize::try_from(self.rows.div_ceil(RUN)).unwrap_or(usize::MAX))
}
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(std::slice::from_ref(&self.ty));
return Ok(Progress::Done);
}
let counted = if let Some(listed) = &self.listed {
let from = usize::try_from(morsel.cursor()).unwrap_or(usize::MAX);
let Some(run) = listed.get(from..from.saturating_add(count)) else {
return Err(Error::internal("a morsel past the end of a series of moments"));
};
run.to_vec()
} else {
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);
}
counted
};
morsel.advance(u64::try_from(count).unwrap_or(u64::MAX));
let vector = Vector::flat(self.ty.clone(), Data::Int64(counted.into()))?;
*out = Chunk::with_rows(vec![vector], count)?;
Ok(if morsel.is_drained() { Progress::Done } else { Progress::More })
}
}
pub(crate) fn moments(
function: TableFunction,
start: &Value,
stop: &Value,
step: &Value,
) -> Result<Option<Stepping>> {
let moment = |value: &Value| match value {
Value::Timestamp(stamp) | Value::TimestampTz(stamp) => Ok(Some(*stamp)),
Value::Null => Ok(None),
other => Err(Error::internal(format!("a range of moments from a {other}"))),
};
let (Some(start), Some(stop), Value::Interval { months, days, micros }) =
(moment(start)?, moment(stop)?, step)
else {
return Ok(None);
};
if [start, stop].iter().any(|&stamp| stamp == i64::MAX || stamp == -i64::MAX) {
return Err(Error::binder("RANGE with infinite bounds is not supported"));
}
let forward = *months > 0 || *days > 0 || *micros > 0;
let backward = *months < 0 || *days < 0 || *micros < 0;
if !forward && !backward {
return Err(Error::binder("interval cannot be 0!"));
}
if forward && backward {
return Err(Error::binder(
"RANGE with composite interval that has mixed signs is not supported",
));
}
moment_steps(function.inclusive(), start, stop, (*months, *days, *micros)).map(Some)
}
#[derive(Debug)]
pub(crate) struct FileScan<'a> {
function: TableFunction,
paths: Vec<String>,
given: Given,
wanted: Vec<Field>,
numbered: bool,
schema: Schema,
tests: Vec<(usize, Op, Bound)>,
cutting: Mutex<Cutting>,
open: Mutex<HashMap<u64, Arc<Mutex<Piece>>>>,
counters: Option<Arc<Counters>>,
sideways: Option<Arc<Sideways<'a>>>,
index: u32,
paying: Paying,
}
const MORSEL_ROWS: usize = 32_768;
fn gather_rows(reader: &FileReader, cut: usize, threads: usize, asked: usize) -> usize {
let FileReader::Parquet(parquet) = reader else { return 0 };
if cut != 0 || asked == 0 {
return 0;
}
let total = parquet
.metadata()
.row_groups
.iter()
.map(|group| usize::try_from(group.rows).unwrap_or(usize::MAX))
.fold(0_usize, usize::saturating_add);
gather_target(total, threads, asked)
}
fn gather_target(total: usize, threads: usize, asked: usize) -> usize {
(total / threads.max(1)).min(asked)
}
fn gathered(first: usize, later: impl Iterator<Item = usize>, gather: usize) -> usize {
let mut total = first;
let mut taken = 1;
for rows in later {
match total.checked_add(rows) {
Some(next) if next <= gather => {
total = next;
taken += 1;
}
_ => break,
}
}
taken
}
const FINE: u64 = 4;
fn morsel_rows(reader: &FileReader, groups: usize, threads: usize) -> usize {
let FileReader::Parquet(parquet) = reader else { return 0 };
if groups == 0 || threads <= groups {
return 0;
}
let page = parquet.page_bytes().unwrap_or(u64::MAX);
cut_rows(page, parquet.chunk_bytes(), group_rows(parquet, 0), groups, threads)
}
fn cut_rows(page: u64, whole: u64, rows: usize, groups: usize, threads: usize) -> usize {
if groups == 0 || threads <= groups || rows == 0 || page == 0 {
return 0;
}
let cut = rows.div_ceil(threads.div_ceil(groups)).max(MORSEL_ROWS);
let taking = u64::try_from(cut.min(rows)).unwrap_or(u64::MAX);
let each = whole / u64::try_from(rows).unwrap_or(u64::MAX) * taking;
if page.saturating_mul(FINE) > each { 0 } else { cut }
}
fn next_piece(rows: usize, part: usize, target: usize) -> Range<usize> {
let each = rows.div_ceil(parts(rows, target));
let upto = part.saturating_add(each).min(rows);
part.min(upto)..upto
}
fn parts(rows: usize, target: usize) -> usize {
if target == 0 {
return 1;
}
rows.div_ceil(target).max(1)
}
fn group_rows(reader: &Reader, at: usize) -> usize {
reader
.metadata()
.row_groups
.get(at)
.map_or(0, |group| usize::try_from(group.rows).unwrap_or(usize::MAX))
}
fn aim(cutting: &mut Cutting) {
let Some(reader) = cutting.reader.as_ref() else { return };
cutting.cut = morsel_rows(reader, cutting.groups, cutting.threads);
cutting.gather = gather_rows(reader, cutting.cut, cutting.threads, cutting.asked);
cutting.pieces = pieces(reader, cutting.cut, cutting.gather);
}
fn pieces(reader: &FileReader, target: usize, gather: usize) -> usize {
let FileReader::Parquet(reader) = reader else { return 1 };
let rows: Vec<usize> = reader
.metadata()
.row_groups
.iter()
.map(|group| usize::try_from(group.rows).unwrap_or(usize::MAX))
.collect();
morsels_of(&rows, target, gather)
}
fn morsels_of(rows: &[usize], target: usize, gather: usize) -> usize {
if gather == 0 {
return rows.iter().map(|&group| parts(group, target)).sum::<usize>().max(1);
}
let mut at = 0;
let mut morsels = 0;
while let Some(&first) = rows.get(at) {
at += gathered(first, rows[at + 1..].iter().copied(), gather);
morsels += 1;
}
morsels.max(1)
}
#[derive(Debug)]
struct Cutting {
at: usize,
reader: Option<FileReader>,
group: usize,
groups: usize,
part: usize,
cut: usize,
gather: usize,
asked: usize,
threads: usize,
pieces: usize,
skipping: Vec<Test>,
skipped: usize,
row: i64,
given: u64,
split: Option<(Arc<Split>, usize)>,
}
#[derive(Debug)]
struct Piece {
file: usize,
reader: Option<FileReader>,
failure: Option<Error>,
row: i64,
}
impl<'a> FileScan<'a> {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
plan: &Plan,
index: u32,
function: TableFunction,
args: Slice,
options: Slice,
settings: Slice,
columns: Slice,
tests: Vec<(usize, Op, Bound)>,
sideways: Option<Arc<Sideways<'a>>>,
) -> Result<Self> {
let paths = file_arguments(plan, args, function)?;
let given = csv_options(plan, options, settings)?;
let produced = plan.field_list(columns).to_vec();
let numbered = produced.last().is_some_and(|field| field.name == FILE_ROW_NUMBER);
let wanted =
if numbered { produced[..produced.len() - 1].to_vec() } else { produced.clone() };
let scan = Self {
function,
paths,
given,
wanted,
numbered,
schema: Schema::numbered(produced, index),
tests,
cutting: Mutex::new(Cutting {
at: 0,
reader: None,
group: 0,
groups: 0,
part: 0,
cut: 0,
gather: 0,
asked: 0,
threads: 1,
pieces: 1,
skipping: Vec::new(),
skipped: 0,
row: 0,
given: 0,
split: None,
}),
open: Mutex::new(HashMap::new()),
counters: None,
sideways,
index,
paying: Paying::default(),
};
{
let mut cutting = scan.cutting.lock().map_err(poisoned)?;
scan.advance(&mut cutting)?;
}
Ok(scan)
}
fn sift(&self, chunk: &mut Chunk) -> Result<()> {
let Some(sideways) = self.sideways.as_ref() else { return Ok(()) };
if !self.paying.worth() {
return Ok(());
}
if let Some((at, domain)) = sideways.domain(self.index) {
let Ok(column) = chunk.column(at) else { return Ok(()) };
let rows = chunk.len();
let kept = domain.keep(column, rows, &mut Vec::new());
self.paying.saw(rows, kept.len());
if kept.len() < rows {
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.select(&Selection::from_indices(kept))?;
}
return Ok(());
}
let Some((at, filter)) = sideways.sifting(self.index) else { return Ok(()) };
let Ok(column) = chunk.column(at) else { return Ok(()) };
let rows = chunk.len();
let mut hashes = Vec::new();
hash(std::slice::from_ref(column), rows, &mut hashes, Across::TwoInputs);
let mut held = Vec::new();
filter.holds_run(&hashes, &mut held);
let kept = Selection::from_predicate(rows, |row| held[row]);
self.paying.saw(rows, kept.len());
if kept.len() == rows {
return Ok(());
}
let whole = std::mem::replace(chunk, Chunk::empty(&[]));
*chunk = whole.select(&kept)?;
Ok(())
}
pub(crate) fn watched(mut self, counters: Arc<Counters>) -> Self {
self.counters = Some(counters);
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn advance(&self, cutting: &mut Cutting) -> Result<()> {
cutting.reader = None;
cutting.split = None;
cutting.row = 0;
cutting.group = 0;
cutting.groups = 0;
cutting.part = 0;
cutting.cut = 0;
cutting.gather = 0;
cutting.pieces = 1;
cutting.skipping = Vec::new();
let Some(path) = self.paths.get(cutting.at) else { return Ok(()) };
let mut reader = FileReader::open(self.function, path, self.given)?;
let first = if cutting.at == 0 { None } else { self.paths.first().map(String::as_str) };
let held = positions(self.function, &self.wanted, &reader.fields(), path, first)?;
reader.project(&held)?;
reader.settle(&self.wanted)?;
cutting.groups = reader.row_groups();
cutting.skipping = self
.tests
.iter()
.filter_map(|(at, op, value)| {
Some(Test { column: *held.get(*at)?, op: *op, value: value.clone() })
})
.collect();
cutting.reader = Some(reader);
cutting.at += 1;
aim(cutting);
self.divide(cutting);
Ok(())
}
fn divide(&self, cutting: &mut Cutting) {
if cutting.threads <= 1 || self.numbered {
return;
}
let size = rudb_csv::split::size();
match cutting.reader.take() {
Some(FileReader::Csv(reader)) if reader.ranges(size) > 1 => {
let ranges = reader.ranges(size);
let split = Split::new(reader, ranges);
cutting.pieces = split.ranges();
cutting.split = Some((Arc::new(split), 0));
}
other => cutting.reader = other,
}
}
fn cut(&self, cutting: &mut Cutting) -> Result<Option<Piece>> {
let file = cutting.at.saturating_sub(1);
if let Some((split, next)) = cutting.split.as_mut() {
if *next == split.ranges() {
return Ok(None);
}
let part = split.part(*next);
*next += 1;
return Ok(Some(Piece {
file,
reader: Some(FileReader::Part(part)),
failure: None,
row: 0,
}));
}
if let Some(FileReader::Parquet(reader)) = cutting.reader.as_ref() {
let metadata = reader.metadata();
while cutting.group < cutting.groups {
let Some(group) = metadata.row_groups.get(cutting.group) else { break };
if !skips(&cutting.skipping, group, &metadata.schema) {
break;
}
cutting.group += 1;
cutting.row = cutting.row.saturating_add(group.rows);
cutting.skipped = cutting.skipped.saturating_add(1);
if let Some(counters) = &self.counters {
counters.part_pruned();
}
}
}
let piece = match cutting.reader.as_ref() {
Some(FileReader::Parquet(reader)) if cutting.group < cutting.groups => {
let at = cutting.group;
let rows = group_rows(reader, at);
let taken = if cutting.part == 0 && cutting.gather > rows {
let metadata = reader.metadata();
let later = metadata.row_groups.get(at + 1..cutting.groups).unwrap_or(&[]);
let later = later
.iter()
.take_while(|group| !skips(&cutting.skipping, group, &metadata.schema))
.map(|group| usize::try_from(group.rows).unwrap_or(usize::MAX));
gathered(rows, later, cutting.gather)
} else {
1
};
if taken > 1 {
let covered = (at..at + taken)
.map(|group| group_rows(reader, group))
.fold(0_usize, usize::saturating_add);
if let Some(counters) = &self.counters {
for _ in 0..taken {
counters.part_read();
}
}
let split = reader.split(at..at + taken)?;
cutting.group = at + taken;
let row = cutting.row;
cutting.row =
cutting.row.saturating_add(i64::try_from(covered).unwrap_or(i64::MAX));
return Ok(Some(Piece {
file,
reader: Some(FileReader::Parquet(split)),
failure: None,
row,
}));
}
let piece = next_piece(rows, cutting.part, cutting.cut);
if piece.start == 0 {
if let Some(counters) = &self.counters {
counters.part_read();
}
}
let split = reader.split_rows(at, piece.clone())?;
if piece.end >= rows {
cutting.group += 1;
cutting.part = 0;
} else {
cutting.part = piece.end;
}
let row = cutting.row;
cutting.row = cutting
.row
.saturating_add(i64::try_from(piece.end - piece.start).unwrap_or(i64::MAX));
Piece { file, reader: Some(FileReader::Parquet(split)), failure: None, row }
}
Some(FileReader::Csv(_)) => {
Piece { file, reader: cutting.reader.take(), failure: None, row: 0 }
}
_ => return Ok(None),
};
Ok(Some(piece))
}
fn hand(&self, cutting: &mut Cutting, piece: Piece) -> Option<Morsel> {
let index = cutting.given;
cutting.given += 1;
let covers = u64::from(piece.failure.is_none());
self.open.lock().ok()?.insert(index, Arc::new(Mutex::new(piece)));
Some(Morsel::new(index, 0, covers))
}
fn piece(&self, index: u64) -> Result<Arc<Mutex<Piece>>> {
let open = self.open.lock().map_err(poisoned)?;
open.get(&index)
.map(Arc::clone)
.ok_or_else(|| Error::internal(format!("a file scan was read at morsel {index}")))
}
fn number(&self, chunk: Chunk, piece: &mut Piece) -> Result<Chunk> {
let rows = chunk.len();
let mut columns = Vec::with_capacity(chunk.width() + 1);
for at in 0..chunk.width() {
columns.push(chunk.column(at)?.clone());
}
let first = piece.row;
piece.row = piece.row.saturating_add(i64::try_from(rows).unwrap_or(i64::MAX));
let mut data = Vec::with_capacity(rows);
for at in 0..rows {
data.push(first.saturating_add(i64::try_from(at).unwrap_or(i64::MAX)));
}
columns.push(Vector::flat(LogicalType::BigInt, Data::Int64(data.into()))?);
Chunk::with_rows(columns, rows)
}
fn conform(&self, chunk: Chunk, file: usize) -> Result<Chunk> {
let rows = chunk.len();
let settled = self.wanted.iter().enumerate().all(|(at, field)| {
chunk.column(at).is_ok_and(|column| column.logical_type() == &field.ty)
});
if settled {
return Ok(chunk);
}
let mut columns = Vec::with_capacity(self.wanted.len());
for (at, field) in self.wanted.iter().enumerate() {
let column = chunk.column(at)?;
if column.logical_type() == &field.ty {
columns.push(column.clone());
continue;
}
columns.push(cast(column, &field.ty, false).map_err(|error| {
let path = self.paths.get(file).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()
))
})?);
}
Chunk::with_rows(columns, rows)
}
}
impl Source for FileScan<'_> {
fn morsel(&self) -> Option<Morsel> {
let mut cutting = self.cutting.lock().ok()?;
loop {
match self.cut(&mut cutting) {
Ok(Some(piece)) => return self.hand(&mut cutting, piece),
Ok(None) => {}
Err(error) => {
let file = cutting.at.saturating_sub(1);
let piece = Piece { file, reader: None, failure: Some(error), row: 0 };
return self.hand(&mut cutting, piece);
}
}
if cutting.at >= self.paths.len() {
return None;
}
if let Err(error) = self.advance(&mut cutting) {
let file = cutting.at.saturating_sub(1);
let piece = Piece { file, reader: None, failure: Some(error), row: 0 };
return self.hand(&mut cutting, piece);
}
}
}
fn morsels(&self, threads: usize, _weight: usize) -> Option<usize> {
let mut cutting = self.cutting.lock().ok()?;
cutting.threads = threads;
aim(&mut cutting);
self.divide(&mut cutting);
Some(cutting.pieces.max(1).saturating_mul(self.paths.len().max(1)))
}
fn gather(&self, rows: usize) {
if let Ok(mut cutting) = self.cutting.lock() {
cutting.asked = rows;
}
}
fn read(&self, morsel: &mut Morsel, out: &mut Chunk) -> Result<Progress> {
let piece = self.piece(morsel.index())?;
let mut piece = piece.lock().map_err(poisoned)?;
if let Some(error) = piece.failure.take() {
return Err(error);
}
loop {
let file = piece.file;
let Some(reader) = piece.reader.as_mut() else {
morsel.advance(1);
*out = Chunk::empty(&self.schema.types());
return Ok(Progress::Done);
};
let before = reader.bytes_read();
let next = reader.next_chunk()?;
if let Some(counters) = &self.counters {
counters.read(reader.bytes_read().saturating_sub(before));
}
let Some(chunk) = next else {
piece.reader = None;
continue;
};
let mut chunk = self.conform(chunk, file)?;
if self.numbered {
chunk = self.number(chunk, &mut piece)?;
}
self.sift(&mut chunk)?;
*out = chunk;
return Ok(Progress::More);
}
}
}
#[derive(Debug)]
enum FileReader {
Parquet(Reader),
Csv(CsvReader),
Part(Part),
}
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(),
Self::Part(_) => Vec::new(),
}
}
fn project(&mut self, columns: &[usize]) -> Result<()> {
match self {
Self::Parquet(reader) => reader.project(columns),
Self::Csv(reader) => reader.project(columns),
Self::Part(_) => Err(cut_already()),
}
}
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)
}
Self::Part(_) => Err(cut_already()),
}
}
fn next_chunk(&mut self) -> Result<Option<Chunk>> {
match self {
Self::Parquet(reader) => reader.next_chunk(),
Self::Csv(reader) => reader.next_chunk(),
Self::Part(part) => part.next_chunk(),
}
}
fn row_groups(&self) -> usize {
match self {
Self::Parquet(reader) => reader.metadata().row_groups.len(),
Self::Csv(_) | Self::Part(_) => 0,
}
}
fn bytes_read(&self) -> u64 {
match self {
Self::Parquet(reader) => reader.bytes_read(),
Self::Csv(reader) => reader.bytes_read(),
Self::Part(part) => part.bytes_read(),
}
}
}
fn cut_already() -> Error {
Error::internal("a range of a CSV file was reshaped after the file was cut")
}
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, Value)> =
names.into_iter().zip(evaluated.iter().map(|vector| vector.value_at(0))).collect();
csv_given(&written)
}
pub(crate) 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) {
Value::Varchar(path) => paths.push(path),
other => {
return Err(Error::internal(format!(
"{}() bound with {other:?} rather than constant file names",
function.name()
)));
}
}
}
Ok(paths)
}
pub(crate) 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 morsels(&self, _threads: usize, _weight: usize) -> Option<usize> {
Some(self.handout.total())
}
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::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use rudb_catalog::{QualifiedName, Table};
use rudb_common::{Field, LogicalType, Value};
use rudb_functions::TableFunction;
use rudb_metrics::Counters;
use rudb_pipeline::{Progress, Source};
use rudb_plan::{ColumnBinding, Node, Plan};
use rudb_storage::Blocked;
use rudb_vector::{Chunk, Data, Vector};
use super::{
Across, Bound, Cutoff, FileScan, Filters, Handout, Live, OnceLock, Op, Paying, Probe,
Pushdown, RUN, Scan, Schema, Series, Session, Settings, Sideways, VECTOR_SIZE, WARMUP,
best_first, cut_rows, gather_target, gathered, hash, instances_for, morsels_of, next_piece,
parts, runs_of, worth_sifting,
};
use crate::sideways::Found;
fn cutting(rows: usize, target: usize) -> Vec<std::ops::Range<usize>> {
let mut out = Vec::new();
let mut part = 0;
loop {
let piece = next_piece(rows, part, target);
if piece.is_empty() {
return out;
}
part = piece.end;
out.push(piece);
assert!(out.len() <= rows + 1, "cutting {rows} rows into {target} did not terminate");
}
}
fn series(start: i64, step: i64, rows: u64) -> Series {
Series { start, step, rows, ..Series::empty(Schema::empty()) }
}
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());
}
fn native(label: &str, parts: usize, rows: usize) -> (Table, std::path::PathBuf) {
let stamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time advances")
.as_nanos();
let path = std::env::temp_dir()
.join(format!("rudb-exec-{label}-{}-{stamp}.rdb", std::process::id()));
let mut writer = rudb_native::Writer::create(
&path,
"items",
vec![Field::required("id", LogicalType::Integer)],
)
.expect("a new file");
for part in 0..parts {
let values: Vec<Value> =
(0..rows).map(|row| Value::Integer((part * rows + row) as i32)).collect();
let chunk = Chunk::new(vec![
Vector::from_values(LogicalType::Integer, &values).expect("integers"),
])
.expect("matching rows");
writer.append(&chunk).expect("one part");
}
writer.finish().expect("commit");
let reader = rudb_native::Reader::open(&path).expect("reopen from disk");
let table = Table::native(QualifiedName::new("memory", "main", "items"), reader)
.expect("a native table");
(table, path)
}
fn scan_of(table: &Table) -> Scan<'_> {
let plan = Plan::parse("Get memory.main.items AS items #0 [id::INTEGER]")
.expect("the plan text round trips");
let Node::Get { index, columns, .. } = *plan.node(plan.root()) else {
panic!("the plan is a get");
};
Scan::new(
&plan,
table,
index,
columns,
Filters::default(),
&Settings::default(),
&Session::default(),
)
.expect("the column is there")
}
fn taken(scan: &Scan<'_>) -> Vec<(std::ops::Range<u64>, usize)> {
let mut out = Vec::new();
while let Some(mut morsel) = scan.morsel() {
let mut rows = 0;
loop {
let mut chunk = Chunk::empty(&[]);
let progress = scan.read(&mut morsel, &mut chunk).expect("a part reads");
rows += chunk.len();
if progress == Progress::Done {
break;
}
}
out.push((morsel.start()..morsel.end(), rows));
}
out
}
#[test]
fn a_native_scan_with_a_stripe_for_every_worker_hands_out_stripes() {
let (table, path) = native("stripes", 128, 200);
let scan = scan_of(&table);
let stripes = table.rows().stripe_parts();
assert_eq!(stripes.len(), 2, "two full stripes of sixty four parts");
assert_eq!(scan.morsels(2, 1), Some(2), "one instance per stripe");
let taken = taken(&scan);
assert_eq!(taken.len(), 2, "one morsel per stripe");
assert_eq!(taken[0].0, 0..64);
assert_eq!(taken[1].0, 64..128);
assert_eq!(taken.iter().map(|(_, rows)| rows).sum::<usize>(), 128 * 200);
std::fs::remove_file(path).expect("remove scratch file");
}
#[test]
fn a_native_scan_with_one_stripe_hands_out_parts() {
let (table, path) = native("one-stripe", 64, 400);
let scan = scan_of(&table);
assert_eq!(table.rows().stripe_parts().len(), 1, "one stripe and more workers than that");
assert_eq!(scan.morsels(4, 1), Some(2), "the rows behind it pay for two instances");
let taken = taken(&scan);
assert_eq!(taken.len(), 64, "one morsel per part");
assert_eq!(taken[0].0, 0..1);
assert_eq!(taken.iter().map(|(_, rows)| rows).sum::<usize>(), 64 * 400);
std::fs::remove_file(path).expect("remove scratch file");
}
fn pruned_scan(table: &Table, tests: Vec<(usize, Op, Bound)>) -> Scan<'_> {
let plan = Plan::parse("Get memory.main.items AS items #0 [id::INTEGER]")
.expect("the plan text round trips");
let Node::Get { index, columns, .. } = *plan.node(plan.root()) else {
panic!("the plan is a get");
};
Scan::new(
&plan,
table,
index,
columns,
Filters { pruning: tests, ..Filters::default() },
&Settings::default(),
&Session::default(),
)
.expect("the column is there")
}
#[test]
fn a_native_scan_divides_the_stripe_that_survives_pruning() {
let (table, path) = native("pruned", 128, 400);
assert_eq!(table.rows().stripe_parts().len(), 2, "two full stripes of sixty four parts");
let scan = pruned_scan(&table, vec![(0, Op::GreaterOrEqual, Bound::Int(25_600))]);
assert_eq!(scan.morsels(2, 1), Some(2), "the surviving half pays for two instances");
let taken = taken(&scan);
assert!(taken.len() > 2, "the one stripe with work is cut up, not handed out whole");
assert!(
taken.iter().all(|(run, _)| run.start >= 64),
"no morsel covers the ruled out stripe"
);
assert_eq!(taken.iter().map(|(_, rows)| rows).sum::<usize>(), 64 * 400);
std::fs::remove_file(path).expect("remove scratch file");
}
#[test]
fn the_sieves_are_read_early_only_when_the_bounds_have_left_the_work_in_a_heap() {
assert!(worth_sifting(2, 32, 1_000_000, 1), "two stripes of work and a machine to fill");
assert!(
!worth_sifting(16, 32, 1_000_000, 1),
"sixteen is as many instances as the rows buy"
);
assert!(!worth_sifting(1, 1, 1_000_000, 1), "one worker has nowhere to spread it anyway");
assert!(!worth_sifting(1, 32, 1_000, 1), "a thousand rows pay for one worker either way");
assert!(worth_sifting(1, 32, 1_000, 64), "a thousand rows of heavy work pay for more");
}
fn living(stripes: &[&[usize]], rows: impl Fn(usize) -> usize) -> Vec<Live> {
stripes
.iter()
.map(|parts| Live {
parts: parts.to_vec(),
rows: parts.iter().map(|&at| rows(at)).sum(),
})
.collect()
}
#[test]
fn a_stripe_holding_no_more_than_a_share_stays_one_run() {
let live = living(&[&[0, 1, 2, 3], &[4, 5, 6, 7], &[8, 9, 10, 11]], |_| 1);
assert_eq!(runs_of(&live, 1, |_| 1), [0..4, 4..8, 8..12], "one worker, one run a stripe");
assert_eq!(runs_of(&live, 3, |_| 1), [0..4, 4..8, 8..12], "a stripe is a share exactly");
let four = runs_of(&live, 4, |_| 1);
assert_eq!(four.len(), 12, "three stripes for four workers, so a part apiece");
assert_eq!(four.first(), Some(&(0..1)));
assert_eq!(four.last(), Some(&(11..12)));
}
#[test]
fn the_stripe_holding_the_rows_is_the_one_that_gets_cut() {
let rows = |at: usize| if (1..9).contains(&at) { 1_000 } else { 10 };
let live = living(&[&[0], &[1, 2, 3, 4, 5, 6, 7, 8], &[9]], rows);
let runs = runs_of(&live, 2, rows);
assert_eq!(runs, [0..1, 1..3, 3..5, 5..7, 7..9, 9..10], "the big stripe in four, not one");
let held: Vec<usize> = runs.iter().map(|run| run.clone().map(rows).sum()).collect();
assert_eq!(held, [10, 2_000, 2_000, 2_000, 2_000, 10], "and the four are the same size");
}
#[test]
fn a_stripe_that_holds_all_the_work_is_cut_up() {
let live = living(&[&[1, 3, 5, 7], &[]], |_| 1);
let runs = runs_of(&live, 2, |_| 1);
assert_eq!(runs, [1..2, 3..4, 5..6, 7..8], "a run apiece, two workers, four to go round");
assert!(runs.iter().all(|run| run.start >= 1 && run.end <= 8));
}
#[test]
fn the_pieces_a_top_n_wants_most_go_out_first() {
let range = |at: usize| {
let low = i128::try_from(at % 13).expect("small");
(!(30..40).contains(&at)).then(|| rudb_storage::Range {
low: Some(Bound::Int(low)),
high: Some(Bound::Int(low + 5)),
..rudb_storage::Range::default()
})
};
let runs = vec![0..20, 20..45];
let (out, order) = best_first(runs.clone(), 2, false, range);
assert_eq!(out, [0..8, 8..16, 16..20, 20..45]);
assert_eq!(order, [0, 1, 2, 3]);
let (out, order) = best_first(runs.clone(), 3, false, range);
assert_eq!(out, [0..8, 8..16, 16..20, 20..28, 28..45]);
assert_eq!(order, [0, 1, 3, 2, 4], "the third best is the first to reach 0 after 13");
let (out, order) = best_first(runs.clone(), 4, false, range);
assert_eq!(out, [0..8, 8..16, 16..20, 20..28, 28..36, 36..44, 44..45]);
assert_eq!(order, [0, 1, 3, 5, 2, 4, 6]);
let (out, order) = best_first(runs.clone(), 1, true, range);
assert_eq!(out, [0..8, 8..16, 16..20, 20..45]);
assert_eq!(order, [1, 0, 2, 3], "part 12 reaches 17 and part 25 only ties it");
let (out, order) = best_first(runs.clone(), 99, false, |_| None);
assert_eq!((out, order), (runs, Vec::new()), "nothing stored leaves the runs alone");
let (out, _) = best_first(vec![0..20, 20..45], 3, false, range);
let seen: Vec<usize> = out.iter().flat_map(Clone::clone).collect();
assert_eq!(seen, (0..45).collect::<Vec<_>>());
}
fn fixture() -> FileScan<'static> {
pruned(Vec::new())
}
fn pruned(tests: Vec<(usize, Op, Bound)>) -> FileScan<'static> {
let scan = untold(tests);
assert_eq!(scan.morsels(2, 0), Some(2), "a group for each of the two threads");
scan
}
fn untold(tests: Vec<(usize, Op, Bound)>) -> FileScan<'static> {
let path = format!("{}/../rudb-parquet/testdata/mixed.parquet", env!("CARGO_MANIFEST_DIR"));
let path = path.replace('\\', "\\\\").replace('\'', "''");
let text = format!(
"TableFunction read_parquet args=['{path}'::VARCHAR] #0 [a::INTEGER, b::BIGINT]"
);
let plan = Plan::parse(&text).expect("the plan text round trips");
let Node::TableFunction { index, args, options, settings, columns, .. } =
*plan.node(plan.root())
else {
panic!("the plan is a table function");
};
FileScan::new(
&plan,
index,
TableFunction::ReadParquet,
args,
options,
settings,
columns,
tests,
None,
)
.expect("the fixture is there")
}
fn morsels(scan: &FileScan<'_>) -> Vec<usize> {
let mut rows = Vec::new();
while let Some(mut morsel) = scan.morsel() {
let mut chunk = Chunk::empty(&[]);
let mut read = 0;
while let Progress::More = scan.read(&mut morsel, &mut chunk).expect("decodes") {
read += chunk.len();
}
rows.push(read);
}
rows
}
#[test]
fn a_parquet_file_is_one_morsel_per_row_group() {
let scan = fixture();
let rows = morsels(&scan);
assert_eq!(rows, [2048, 2048], "two row groups of 2048");
}
#[test]
fn a_row_group_whose_bounds_rule_out_the_filter_is_never_handed_out() {
let scan = pruned(vec![(0, Op::Greater, Bound::Int(1_000))]);
let rows = morsels(&scan);
assert_eq!(rows, Vec::<usize>::new(), "both row groups are ruled out");
let cutting = scan.cutting.lock().expect("the lock holds");
assert_eq!(cutting.skipped, 2, "and both were skipped rather than read");
}
#[test]
fn a_row_group_whose_bounds_overlap_the_filter_is_handed_out_as_usual() {
let scan = pruned(vec![(0, Op::Greater, Bound::Int(50))]);
let rows = morsels(&scan);
assert_eq!(rows, [2048, 2048], "nothing is ruled out");
let cutting = scan.cutting.lock().expect("the lock holds");
assert_eq!(cutting.skipped, 0);
}
#[test]
fn a_test_against_a_position_the_scan_does_not_produce_rules_nothing_out() {
let scan = pruned(vec![(7, Op::Greater, Bound::Int(1_000))]);
assert_eq!(morsels(&scan), [2048, 2048]);
}
#[test]
fn a_scan_that_has_handed_out_every_morsel_hands_out_no_more() {
let scan = fixture();
let _ = morsels(&scan);
assert!(scan.morsel().is_none());
assert!(scan.morsel().is_none());
}
#[test]
fn every_morsel_can_be_taken_before_any_of_them_is_read() {
let scan = fixture();
let mut taken = Vec::new();
while let Some(morsel) = scan.morsel() {
taken.push(morsel);
}
assert_eq!(taken.len(), 2);
let mut rows = Vec::new();
for morsel in &mut taken {
let mut chunk = Chunk::empty(&[]);
let mut read = 0;
while let Progress::More = scan.read(morsel, &mut chunk).expect("decodes") {
read += chunk.len();
}
rows.push(read);
}
assert_eq!(rows, [2048, 2048]);
}
fn counted(rows: usize) -> Table {
let mut table = Table::new(
QualifiedName::new("memory", "main", "t"),
vec![Field::new("n", LogicalType::Integer)],
)
.expect("one column");
let values: Vec<Vec<Value>> = (0..rows).map(|n| vec![Value::Integer(n as i32)]).collect();
table.append_rows(&values).expect("integers");
table
}
fn scanning(table: &Table, tests: Vec<(usize, Op, Bound)>) -> Scan<'_> {
let fields = vec![Field::new("n", LogicalType::Integer)];
let probes =
tests.into_iter().map(|(column, op, value)| Probe { column, op, value }).collect();
Scan {
table,
columns: vec![Some(0)],
offsets: (0..table.rows().chunk_count())
.map(|at| i64::try_from(at * VECTOR_SIZE).expect("a small table"))
.collect(),
probes,
also: Vec::new(),
sideways: None,
cutoff: None,
index: 0,
testing: OnceLock::new(),
schema: Schema::numbered(fields, 0),
chunks: Handout::new(table.rows().chunk_count()),
stripes: Vec::new(),
pushed: None,
waved: AtomicUsize::new(0),
spread: OnceLock::new(),
skipped: AtomicUsize::new(0),
counters: None,
deferred: Vec::new(),
deferrable: Vec::new(),
deferring: Paying::default(),
unread: Vec::new(),
paying: Paying::default(),
passed: Paying::default(),
}
}
fn beaten(table: &Table, op: Op, bound: Option<Bound>) -> Scan<'_> {
let cutoff = Cutoff::new();
cutoff.about(ColumnBinding::new(0, 0), op);
if let Some(bound) = bound {
cutoff.reached(bound);
}
let mut scan = scanning(table, Vec::new());
scan.cutoff = Some(cutoff);
scan
}
fn applying<'a>(table: &'a Table, predicate: &str) -> (Plan, Scan<'a>) {
let plan =
Plan::parse(&format!("Filter {predicate}\n Get memory.main.t AS t #0 [n::INTEGER]"))
.expect("the plan text round trips");
let Node::Filter { input, predicate } = *plan.node(plan.root()) else {
panic!("the plan is a filter");
};
let Node::Get { index, columns, .. } = *plan.node(input) else {
panic!("under a get");
};
let moved =
rudb_opt::bounds::into_scan(&plan, plan.root()).expect("a filter over a stored table");
let pruning = rudb_opt::bounds::of(&plan, input, predicate);
let pushdown = Pushdown {
node: plan.root(),
predicate,
tests: moved.tests,
whole: moved.whole,
conjuncts: moved.conjuncts,
marks: false,
};
let filters = Filters { pruning, pushed: Some(pushdown), ..Filters::default() };
let scan = Scan::new(
&plan,
table,
index,
columns,
filters,
&Settings::default(),
&Session::default(),
)
.expect("the column is there");
(plan, scan)
}
#[test]
fn a_sparse_like_fetches_the_second_column_after_selection() {
let mut table = Table::new(
QualifiedName::new("memory", "main", "t"),
vec![
Field::new("URL", LogicalType::Varchar),
Field::new("SearchPhrase", LogicalType::Varchar),
],
)
.expect("two columns");
let rows = (0..VECTOR_SIZE * 3)
.map(|row| {
let matching =
row == 3 || row == 5 || (VECTOR_SIZE..VECTOR_SIZE + 400).contains(&row);
let phrase = if row == 5 || row == VECTOR_SIZE + 10 { "" } else { "phrase" };
vec![
Value::Varchar(if matching { "google.test" } else { "example.test" }.into()),
Value::Varchar(phrase.into()),
]
})
.collect::<Vec<_>>();
table.append_rows(&rows).expect("rows");
let plan = Plan::parse(
"Filter (\"~~\"(#0.0::VARCHAR, '%google%'::VARCHAR)::BOOLEAN AND (#0.1::VARCHAR <> ''::VARCHAR)::BOOLEAN)::BOOLEAN\n Get memory.main.t AS t #0 [URL::VARCHAR, SearchPhrase::VARCHAR]",
)
.expect("the plan text round trips");
let Node::Filter { input, predicate } = *plan.node(plan.root()) else {
panic!("the plan is a filter");
};
let Node::Get { index, columns, .. } = *plan.node(input) else { panic!("under a get") };
let moved =
rudb_opt::bounds::into_scan(&plan, plan.root()).expect("a filter over a stored table");
let pushdown = Pushdown {
node: plan.root(),
predicate,
tests: moved.tests,
whole: moved.whole,
conjuncts: moved.conjuncts,
marks: false,
};
let filters = Filters { pushed: Some(pushdown), ..Filters::default() };
let scan = Scan::new(
&plan,
&table,
index,
columns,
filters,
&Settings::default(),
&Session::default(),
)
.expect("two projected columns");
assert!(scan.pushed.as_ref().and_then(|pushed| pushed.late.as_ref()).is_some());
let mut found = Vec::new();
while let Some(mut morsel) = scan.morsel() {
loop {
let mut chunk = Chunk::empty(&[]);
let progress = scan.read(&mut morsel, &mut chunk).expect("a part reads");
found.extend(
(0..chunk.len()).map(|row| (chunk.value_at(row, 0), chunk.value_at(row, 1))),
);
if progress == Progress::Done {
break;
}
}
}
let expected = rows
.iter()
.filter(|row| {
matches!(&row[0], Value::Varchar(url) if url.contains("google"))
&& matches!(&row[1], Value::Varchar(phrase) if !phrase.is_empty())
})
.map(|row| (row[0].clone(), row[1].clone()))
.collect::<Vec<_>>();
assert_eq!(found, expected, "sparse, dense, and empty parts keep the same rows");
}
#[test]
fn a_filter_the_zone_maps_prove_is_applied_to_no_rows_at_all() {
let table = counted(VECTOR_SIZE * 5);
let (_plan, scan) = applying(&table, "(#0.0::INTEGER >= 0::INTEGER)::BOOLEAN");
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 5, "every row came through");
assert_eq!(scan.waved.load(Ordering::Relaxed), 5, "and no chunk was compared");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0, "nor ruled out");
}
#[test]
fn a_scan_skips_waves_through_and_compares_in_the_one_pass() {
let table = counted(VECTOR_SIZE * 5);
let cutoff = VECTOR_SIZE * 2 + VECTOR_SIZE / 2;
let (_plan, scan) =
applying(&table, &format!("(#0.0::INTEGER >= {cutoff}::INTEGER)::BOOLEAN"));
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 5 - cutoff);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 2, "the first two hold nothing wanted");
assert_eq!(scan.waved.load(Ordering::Relaxed), 2, "the last two are all of them wanted");
}
#[test]
fn a_conjunct_a_chunk_proves_is_left_out_and_the_rest_still_runs() {
let table = counted(VECTOR_SIZE * 5);
let cutoff = VECTOR_SIZE * 2 + VECTOR_SIZE / 2;
let odd = VECTOR_SIZE * 4 + 3;
let (_plan, scan) = applying(
&table,
&format!(
"((#0.0::INTEGER >= {cutoff}::INTEGER)::BOOLEAN AND \
(#0.0::INTEGER <> {odd}::INTEGER)::BOOLEAN)::BOOLEAN"
),
);
let pushed = scan.pushed.as_ref().expect("the filter moved into the scan");
assert!(pushed.probes.is_none(), "the `<>` is not a test, so no chunk is proved whole");
let settles = pushed.settles.as_ref().expect("the other conjunct is");
assert_eq!(settles.iter().map(Option::is_some).collect::<Vec<_>>(), [true, false]);
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 5 - cutoff - 1);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 2, "the first two hold nothing wanted");
assert_eq!(scan.waved.load(Ordering::Relaxed), 0, "and none is wanted whole");
}
#[test]
fn the_two_part_counts_add_up_to_the_table() {
let table = counted(VECTOR_SIZE * 5);
let cutoff = VECTOR_SIZE * 2 + VECTOR_SIZE / 2;
let (_plan, mut scan) =
applying(&table, &format!("(#0.0::INTEGER >= {cutoff}::INTEGER)::BOOLEAN"));
let counters = Arc::new(Counters::new(0, 0, "Scan"));
scan = scan.watched(Arc::clone(&counters));
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 5 - cutoff);
let operator = counters.snapshot();
assert_eq!(operator.parts_pruned, 2, "the first two hold nothing wanted");
assert_eq!(
operator.parts_read + operator.parts_pruned,
table.rows().chunk_count() as u64,
"every part is counted once and only once"
);
}
#[test]
fn a_chunk_the_zone_cannot_decide_comes_back_narrowed_to_the_rows_that_pass() {
let table = counted(VECTOR_SIZE);
let (_plan, scan) = applying(&table, "(#0.0::INTEGER < 10::INTEGER)::BOOLEAN");
assert_eq!(counted_rows(&scan), 10);
assert_eq!(scan.waved.load(Ordering::Relaxed), 0);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0);
}
#[test]
fn a_filter_with_an_or_in_it_moves_into_the_scan_and_no_chunk_is_waved_through() {
let table = counted(VECTOR_SIZE * 5);
let high = VECTOR_SIZE * 5 - 10;
let (_plan, scan) = applying(
&table,
&format!(
"((#0.0::INTEGER < 10::INTEGER)::BOOLEAN \
OR (#0.0::INTEGER >= {high}::INTEGER)::BOOLEAN)::BOOLEAN"
),
);
assert_eq!(counted_rows(&scan), 20, "the ten at each end");
assert_eq!(scan.waved.load(Ordering::Relaxed), 0, "and nothing was proved about a chunk");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0);
}
fn counted_rows(scan: &Scan<'_>) -> usize {
let mut rows = 0;
while let Some(mut morsel) = scan.morsel() {
let mut chunk = Chunk::empty(&[LogicalType::Integer]);
loop {
let progress = scan.read(&mut morsel, &mut chunk).expect("a table scan reads");
rows += chunk.len();
if progress == Progress::Done {
break;
}
}
}
rows
}
#[test]
fn a_filter_on_a_table_reads_only_the_chunks_that_can_hold_a_match() {
let table = counted(VECTOR_SIZE * 5);
assert_eq!(table.rows().chunk_count(), 5);
let scan = scanning(&table, vec![(0, Op::Equal, Bound::Int(5_000))]);
assert_eq!(counted_rows(&scan), VECTOR_SIZE, "one chunk's worth");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 4);
}
#[test]
fn a_scan_skips_the_parts_the_top_n_above_it_has_already_beaten() {
let table = counted(VECTOR_SIZE * 5);
let scan = beaten(&table, Op::LessOrEqual, Some(Bound::Int(100)));
assert_eq!(counted_rows(&scan), VECTOR_SIZE, "only the chunk holding 0 to 1023");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 4);
}
#[test]
fn a_descending_top_n_skips_the_parts_below_its_cutoff() {
let table = counted(VECTOR_SIZE * 5);
let cutoff = i128::try_from(VECTOR_SIZE * 4 + VECTOR_SIZE / 2).expect("a small number");
let scan = beaten(&table, Op::GreaterOrEqual, Some(Bound::Int(cutoff)));
assert_eq!(counted_rows(&scan), VECTOR_SIZE, "only the last chunk");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 4);
}
#[test]
fn a_part_that_only_ties_the_cutoff_is_still_read() {
let table = counted(VECTOR_SIZE * 5);
let cutoff = i128::try_from(VECTOR_SIZE * 2).expect("a small number");
let scan = beaten(&table, Op::LessOrEqual, Some(Bound::Int(cutoff)));
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 3);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 2);
}
#[test]
fn a_top_n_that_has_not_filled_its_candidates_rules_nothing_out() {
let table = counted(VECTOR_SIZE * 3);
let scan = beaten(&table, Op::LessOrEqual, None);
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 3);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0);
}
#[test]
fn a_scan_with_no_tests_reads_every_chunk() {
let table = counted(VECTOR_SIZE * 3);
let scan = scanning(&table, Vec::new());
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 3);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0);
}
#[test]
fn a_scan_skips_the_chunks_a_joins_build_side_ruled_out() {
let table = counted(VECTOR_SIZE * 5);
let sideways = Sideways::new();
sideways.about(ColumnBinding::new(0, 0));
sideways.found(Found::of(Some((Bound::Int(5_000), Bound::Int(5_100))), None));
let mut scan = scanning(&table, Vec::new());
scan.sideways = Some(sideways);
assert_eq!(counted_rows(&scan), VECTOR_SIZE, "one chunk's worth");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 4);
}
#[test]
fn a_scan_skips_the_chunks_between_a_joins_keys() {
let table = counted(VECTOR_SIZE * 5);
let sideways = Sideways::new();
sideways.about(ColumnBinding::new(0, 0));
let range = Some((Bound::Int(100), Bound::Int(9_000)));
sideways.found(Found::listing(range, vec![9_000, 100]));
let mut scan = scanning(&table, Vec::new());
scan.sideways = Some(sideways);
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 2, "the first chunk and the last");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 3);
}
#[test]
fn a_scan_keeps_exactly_the_rows_a_reduction_handed_down() {
let table = counted(VECTOR_SIZE * 5);
let size = u64::try_from(VECTOR_SIZE).expect("a small number");
let kept: Vec<u64> = (2 * size + 5..2 * size + 15).chain([4 * size + 7]).collect();
let sideways = Sideways::new();
sideways.about(ColumnBinding::new(0, 0));
let rows = rudb_graph::Rids::from_sorted(5 * size, kept).expect("in order");
sideways.found(Found::exactly(None, rows));
let mut scan = scanning(&table, Vec::new());
scan.sideways = Some(sideways);
assert_eq!(counted_rows(&scan), 11);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 3);
}
fn holding(values: &[i32]) -> Blocked {
let mut filter = Blocked::sized(values.len(), 1 << 20).expect("a filter over three keys");
let column = Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
.expect("integers are an i32 layout");
let mut hashes = Vec::new();
hash(std::slice::from_ref(&column), values.len(), &mut hashes, Across::TwoInputs);
for word in hashes {
filter.add(word);
}
filter
}
#[test]
fn a_scan_drops_the_rows_a_joins_build_side_cannot_match() {
let table = counted(VECTOR_SIZE * 3);
let sideways = Sideways::new();
sideways.about(ColumnBinding::new(0, 0));
sideways.found(Found::of(
Some((Bound::Int(2_100), Bound::Int(2_300))),
Some(holding(&[2_100, 2_200, 2_300])),
));
let mut scan = scanning(&table, Vec::new());
scan.sideways = Some(sideways);
assert_eq!(counted_rows(&scan), 3, "the three rows the build side holds a key for");
assert_eq!(scan.skipped.load(Ordering::Relaxed), 2, "and two chunks were never read");
}
#[test]
fn a_runtime_filter_that_turns_nothing_away_is_given_up_on() {
let idle = Paying::default();
let busy = Paying::default();
for _ in 0..(WARMUP / VECTOR_SIZE) {
assert!(idle.worth(), "nothing is decided under the warmup");
idle.saw(VECTOR_SIZE, VECTOR_SIZE);
busy.saw(VECTOR_SIZE, VECTOR_SIZE / 16);
}
assert!(!idle.worth(), "a filter that kept every row is not worth hashing for");
assert!(busy.worth(), "one that kept a sixteenth of them is");
assert!(!idle.worth());
}
#[test]
fn a_filter_is_worth_hashing_for_once_it_drops_a_quarter_of_the_rows() {
for (kept, worth) in [(0, true), (740, true), (749, true), (750, false), (1_000, false)] {
let paying = Paying::default();
for _ in 0..(WARMUP / 1_000 + 1) {
paying.saw(1_000, kept);
}
assert_eq!(paying.worth(), worth, "{kept} of every thousand rows kept");
}
}
#[test]
fn a_pushed_filter_is_loose_once_it_keeps_more_than_half() {
for (kept, loose) in [(0, false), (250, false), (500, false), (501, true), (1_000, true)] {
let passed = Paying::default();
assert!(!passed.loose(), "nothing is decided under the warmup");
for _ in 0..(WARMUP / 1_000 + 1) {
passed.saw(1_000, kept);
}
assert_eq!(passed.loose(), loose, "{kept} of every thousand rows kept");
}
}
#[test]
fn a_scan_under_a_join_that_offered_nothing_reads_every_chunk() {
let table = counted(VECTOR_SIZE * 3);
let mut scan = scanning(&table, Vec::new());
scan.sideways = Some(Sideways::new());
assert_eq!(counted_rows(&scan), VECTOR_SIZE * 3);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 0);
}
#[test]
fn conjuncts_that_rule_out_every_chunk_read_nothing() {
let table = counted(VECTOR_SIZE * 4);
let split = i128::try_from(VECTOR_SIZE * 2).expect("a small number");
let scan = scanning(
&table,
vec![(0, Op::GreaterOrEqual, Bound::Int(split)), (0, Op::Less, Bound::Int(split))],
);
assert_eq!(counted_rows(&scan), 0);
assert_eq!(scan.skipped.load(Ordering::Relaxed), 4);
}
#[test]
fn the_morsels_of_a_row_group_tile_it_once_each() {
for rows in [1, 2, 7, 9, 10, 100, 4_095, 4_096, 4_097, 122_880, 123_554] {
for target in [1, 2, 3, 300, 4_096, 32_768] {
let pieces = cutting(rows, target);
assert_eq!(pieces.len(), parts(rows, target), "{rows} rows in morsels of {target}");
let mut at = 0;
for piece in &pieces {
assert_eq!(piece.start, at, "{rows} rows in morsels of {target}");
assert!(piece.len() <= target, "{rows} rows in morsels of {target}");
at = piece.end;
}
assert_eq!(at, rows, "{rows} rows in morsels of {target}");
}
}
}
#[test]
fn a_row_group_is_cut_into_even_morsels() {
let pieces = cutting(123_554, 32_768);
assert_eq!(pieces.len(), 4);
let longest = pieces.iter().map(std::ops::Range::len).max().expect("four morsels");
let shortest = pieces.iter().map(std::ops::Range::len).min().expect("four morsels");
assert_eq!(longest, 30_889);
assert_eq!(shortest, 30_887);
assert!(longest - shortest < pieces.len(), "morsels of {longest} and {shortest} rows");
}
#[test]
fn a_row_group_no_larger_than_a_morsel_is_one_morsel() {
assert_eq!(cutting(2_048, 32_768), vec![0..2_048]);
assert_eq!(cutting(32_768, 32_768), vec![0..32_768]);
assert_eq!(parts(2_048, 32_768), 1);
}
#[test]
fn an_empty_row_group_is_no_morsels() {
assert!(next_piece(0, 0, 32_768).is_empty());
assert_eq!(cutting(0, 32_768), Vec::new());
assert_eq!(cutting(0, 0), Vec::new());
}
const LOAD: usize = 131_072;
#[test]
fn a_scan_a_load_asked_to_gather_takes_small_row_groups_as_one_morsel() {
let scan = untold(Vec::new());
scan.gather(LOAD);
assert_eq!(scan.morsels(1, 0), Some(1));
assert_eq!(morsels(&scan), [4096], "both groups in one morsel");
assert_eq!(morsels(&fixture()), [2048, 2048], "and apart when there are two threads");
}
#[test]
fn a_scan_nobody_asked_to_gather_hands_out_a_group_at_a_time() {
let scan = untold(Vec::new());
scan.gather(0);
assert_eq!(scan.morsels(1, 0), Some(2));
assert_eq!(morsels(&scan), [2048, 2048]);
}
#[test]
fn a_gathered_run_stops_at_a_group_the_filter_rules_out() {
let scan = untold(vec![(0, Op::Greater, Bound::Int(1_000))]);
scan.gather(LOAD);
assert_eq!(scan.morsels(1, 0), Some(1));
assert_eq!(morsels(&scan), Vec::<usize>::new());
assert_eq!(scan.cutting.lock().expect("the lock holds").skipped, 2);
}
#[test]
fn small_row_groups_are_gathered_up_to_the_target_and_large_ones_are_left_alone() {
assert_eq!(gathered(8_000, [8_000, 8_000, 8_000].into_iter(), 20_000), 2);
assert_eq!(gathered(8_000, [8_000; 40].into_iter(), LOAD), 16);
assert_eq!(gathered(122_880, [122_880].into_iter(), LOAD), 1, "a DuckDB file");
assert_eq!(gathered(8_000, [200_000, 8_000].into_iter(), LOAD), 1);
assert_eq!(gathered(8_000, std::iter::empty(), LOAD), 1);
assert_eq!(gathered(0, [0, 0].into_iter(), 0), 3, "empty groups cost nothing to join");
}
#[test]
fn the_clickbench_sample_is_gathered_into_as_many_morsels_as_duckdb_writes_groups() {
let rows = vec![8_312; 1_203];
let total: usize = rows.iter().sum();
assert_eq!(morsels_of(&rows, 0, 0), 1_203, "one a group without gathering");
assert_eq!(gather_target(total, 32, 0), 0, "a query asks for nothing");
let gather = gather_target(total, 32, LOAD);
assert_eq!(gather, LOAD);
assert_eq!(morsels_of(&rows, 0, gather), 1_203_usize.div_ceil(15));
let many = gather_target(total, 1_000, LOAD);
assert!(morsels_of(&rows, 0, many) >= 1_000.min(rows.len()) / 2, "threads are kept busy");
assert_eq!(gather_target(total, 2_000, LOAD), 4_999, "a share each, below a group");
assert_eq!(morsels_of(&rows, 0, gather_target(total, 2_000, LOAD)), 1_203);
}
#[test]
fn a_row_group_that_is_not_worth_cutting_is_one_morsel() {
assert_eq!(cutting(123_554, 0), vec![0..123_554]);
assert_eq!(parts(123_554, 0), 1);
assert_eq!(parts(0, 0), 1);
}
#[test]
fn a_file_is_cut_only_as_finely_as_there_are_threads_to_want_it() {
let rows = 123_554;
let whole = 12_355_400;
for (threads, wanted) in [(1, 1), (8, 1), (9, 1), (16, 2), (32, 4), (64, 4)] {
let cut = cut_rows(1024, whole, rows, 9, threads);
assert_eq!(parts(rows, cut), wanted, "{threads} threads over nine row groups");
}
}
#[test]
fn native_workers_grow_with_the_rows_there_are_to_divide() {
for (rows, workers) in [
(0, 1),
(1_000, 1),
(25_000, 1),
(50_000, 2),
(100_000, 4),
(200_000, 8),
(500_000, 8),
(1_000_000, 16),
(2_500_000, 40),
(9_999_750, 160),
] {
assert_eq!(instances_for(rows, 1), workers, "{rows} rows");
}
}
#[test]
fn native_workers_grow_with_the_work_behind_each_row() {
for (rows, weight, workers) in [
(24_576, 1, 1),
(24_576, 2, 2),
(24_576, 7, 7),
(24_576, 12, 8),
(1_000, 25, 1),
(3_000, 9, 2),
(0, 100, 1),
] {
assert_eq!(instances_for(rows, weight), workers, "{rows} rows at weight {weight}");
}
}
#[test]
fn a_weight_never_moves_the_slope_that_divides_a_large_query() {
for rows in [200_000, 500_000, 1_000_000, 2_500_000, 9_999_750] {
let plain = instances_for(rows, 1);
for weight in [2, 4, 12, 64] {
assert_eq!(instances_for(rows, weight), plain, "{rows} rows at weight {weight}");
}
}
}
#[test]
fn a_file_of_coarse_pages_is_not_cut_however_many_threads_there_are() {
let rows = 123_554;
let whole = 12_355_400;
assert_eq!(cut_rows(whole, whole, rows, 9, 64), 0, "one page per chunk, as DuckDB writes");
assert_eq!(cut_rows(819_201, whole, rows, 9, 64), 0);
assert_eq!(cut_rows(0, whole, rows, 9, 64), 0);
assert_eq!(cut_rows(1024, 0, rows, 9, 64), 0, "a group of no bytes is never worth cutting");
assert!(cut_rows(819_200, whole, rows, 9, 64) > 0);
}
}