use std::cell::Cell;
use std::collections::HashMap;
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::Reader as CsvReader;
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_parquet::{Bound, Op, Reader, Test, skips};
use rudb_pipeline::{Compaction, Gauge, Morsel, Progress, Source, narrow};
use rudb_plan::{ConjunctionOp, Expr, ExprRef, 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::Sideways;
use crate::stream::later_passes;
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 })
}
pub(crate) fn new(
plan: &Plan,
input: &Schema,
schema: Schema,
groups: Slice,
column: usize,
entries: Vec<(Value, u64)>,
session: &Session,
) -> Result<Self> {
let output_types = schema.types();
let group_exprs = plan.expr_list(groups);
if output_types.len() != group_exprs.len() + 1
|| output_types.last() != Some(&LogicalType::BigInt)
{
return Err(Error::internal("a frequency source count is not BIGINT"));
}
let input_types = input.types();
let key_type = input_types
.get(column)
.ok_or_else(|| Error::internal("a frequency column is outside its scan"))?;
let mut chunks = Vec::with_capacity(entries.len().div_ceil(VECTOR_SIZE));
for entries in entries.chunks(VECTOR_SIZE) {
let mut keys = Vec::with_capacity(entries.len());
let mut counts = Vec::with_capacity(entries.len());
for (key, count) in entries {
keys.push(key.clone());
counts.push(
i64::try_from(*count)
.map_err(|_| Error::internal("a stored frequency exceeds BIGINT"))?,
);
}
let mut columns = input_types
.iter()
.map(|ty| Vector::constant(ty.clone(), Value::Null, keys.len()))
.collect::<Vec<_>>();
columns[column] = Vector::from_values(key_type.clone(), &keys)?;
let input_chunk = Chunk::with_rows(columns, keys.len())?;
let mut output = evaluate_all_in_time_zone(
plan,
group_exprs,
input,
&input_chunk,
session.session_time_zone(),
)?;
output.push(Vector::flat(LogicalType::BigInt, Data::Int64(counts.into()))?);
chunks.push(Chunk::with_rows(output, keys.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 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,
paying: Paying,
counters: Option<Arc<Counters>>,
}
const WARMUP: usize = 1 << 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 saw(&self, rows: usize, kept: usize) {
self.seen.fetch_add(rows, Ordering::Relaxed);
self.kept.fetch_add(kept, Ordering::Relaxed);
}
}
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: rudb_plan::NodeRef,
pub(crate) predicate: ExprRef,
pub(crate) tests: Vec<(usize, Op, Bound)>,
pub(crate) whole: 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>,
compaction: &'static dyn Compaction,
passes: u32,
probes: 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);
}
}
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 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
};
Ok(Self {
predicate: Prepared::one(plan, pushdown.predicate, schema)?.in_session(session),
late,
compaction,
passes: later_passes(plan, pushdown.node),
probes: (whole && probes.len() == wanted).then_some(probes),
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>>,
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 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,
paying: Paying::default(),
})
}
fn apply(&self, at: usize, chunk: &mut Chunk) -> 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 = pushed.predicate.evaluate_filter(chunk, &mut working.scratch)?;
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);
}
if kept.len() != chunk.len() {
narrow(pushed.compaction, chunk, &kept, &mut working.gauge)?;
}
pushed.give(slot, working);
Ok(())
}
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)?;
Ok(true)
}
fn sift(&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))?;
}
}
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 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)
}
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();
if probes.is_empty() || !worth_sifting(working, threads, rows, weight) {
return (live, rows);
}
let mut sifted = 0;
for stripe in &mut live {
stripe.parts.retain(|&at| !self.table.rows().skips(at, probes));
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() {
if !probes.is_empty() && self.table.rows().stripe_skips(stripe, probes) {
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
}
impl Source for Scan<'_> {
fn morsel(&self) -> Option<Morsel> {
let Some(spread) = self.spread.get() else {
return self.chunks.take();
};
let index = spread.handout.number()?;
let run = spread.runs.get(usize::try_from(index).unwrap_or(usize::MAX))?;
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(index, 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 _ = self.spread.set(Spread { handout: Handout::new(runs.len()), runs });
}
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_late(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.apply(at, out)?;
self.sift(out)?;
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.apply(at, out)?;
self.sift(out)?;
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,
}
#[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,
}),
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.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);
Ok(())
}
fn cut(&self, cutting: &mut Cutting) -> Result<Option<Piece>> {
let file = cutting.at.saturating_sub(1);
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);
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),
}
impl FileReader {
fn open(function: TableFunction, path: &str, given: Given) -> Result<Self> {
match function {
TableFunction::ReadCsv => Ok(Self::Csv(open_csv(path, given)?)),
_ => Ok(Self::Parquet(open_parquet(path)?)),
}
}
fn fields(&self) -> Vec<Field> {
match self {
Self::Parquet(reader) => reader.fields(),
Self::Csv(reader) => reader.fields(),
}
}
fn project(&mut self, columns: &[usize]) -> Result<()> {
match self {
Self::Parquet(reader) => reader.project(columns),
Self::Csv(reader) => reader.project(columns),
}
}
fn settle(&mut self, wanted: &[Field]) -> Result<()> {
match self {
Self::Parquet(reader) => {
let text: Vec<bool> =
wanted.iter().map(|field| field.ty == LogicalType::Varchar).collect();
reader.as_string(&text);
Ok(())
}
Self::Csv(reader) => {
let types: Vec<LogicalType> = wanted.iter().map(|field| field.ty.clone()).collect();
reader.retype(&types)
}
}
}
fn next_chunk(&mut self) -> Result<Option<Chunk>> {
match self {
Self::Parquet(reader) => reader.next_chunk(),
Self::Csv(reader) => reader.next_chunk(),
}
}
fn row_groups(&self) -> usize {
match self {
Self::Parquet(reader) => reader.metadata().row_groups.len(),
Self::Csv(_) => 0,
}
}
fn bytes_read(&self) -> u64 {
match self {
Self::Parquet(reader) => reader.bytes_read(),
Self::Csv(_) => 0,
}
}
}
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,
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));
}
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,
paying: 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 };
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 };
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 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_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_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);
}
}