use std::sync::Arc;
use rudb_catalog::{Catalog, QualifiedName};
use rudb_common::{Cancel, Memory, Result};
use rudb_functions::TableFunction;
use rudb_metrics::{Counters, Driver, Report};
use rudb_parquet::{Bound, Op};
use rudb_pipeline::{
BufferId, DynSink, DynStream, Pipeline, PipelineId, Source, Watched, root, root_in_order,
};
use rudb_plan::{
CompareOp, ConjunctionOp, Expr, ExprRef, Node, NodeRef, PipelineRef, Plan, ROOT, Shape, Slice,
seams_of,
};
use rudb_seam::Settings;
use crate::fetch::Fetch;
use crate::gather::{Gather, Keep};
use crate::group::{Aggregate, Distinct};
use crate::join::{CrossProduct, Gathered, Join};
use crate::query::Query;
use crate::register::registries;
use crate::schema::Schema;
use crate::setop::SetOp;
use crate::sort::Sort;
use crate::source::{Dummy, FileScan, Scan, Series, Values};
use crate::strategies::Strategies;
use crate::stream::{Filter, Limit, Project};
use crate::topn::TopN;
pub fn build<'a>(plan: &'a Plan, catalog: &'a Catalog) -> Result<Query<'a>> {
build_with(plan, catalog, &Cancel::new(), &Memory::unlimited(), &Settings::new())
}
pub fn build_with<'a>(
plan: &'a Plan,
catalog: &'a Catalog,
cancel: &Cancel,
memory: &Memory,
seams: &Settings,
) -> Result<Query<'a>> {
build_measured(plan, catalog, cancel, memory, seams, &Report::new())
}
pub fn build_measured<'a>(
plan: &'a Plan,
catalog: &'a Catalog,
cancel: &Cancel,
memory: &Memory,
seams: &Settings,
report: &Report,
) -> Result<Query<'a>> {
let shape = Shape::of(plan);
for pipeline in shape.all() {
report.pipeline(pipeline);
for waits_for in shape.waits_for(pipeline) {
report.depends(pipeline, *waits_for);
}
}
let mut building = Building {
plan,
catalog,
cancel,
memory,
seams,
report,
shape,
done: Vec::new(),
drivers: Vec::new(),
pruning: Vec::new(),
};
let segment = building.node(plan.root())?;
let schema = segment.schema.clone();
let (sink, reader) = if ordered(plan, plan.root()) {
root(BufferId(0), None)
} else {
root_in_order(BufferId(0), None)
};
building.close(segment, ROOT, Arc::new(sink));
let Building { done, drivers, .. } = building;
Query::new(done, drivers, reader, schema)
}
fn ordered(plan: &Plan, node: NodeRef) -> bool {
match *plan.node(node) {
Node::Sort { .. } | Node::TopN { .. } => true,
Node::Project { input, .. }
| Node::Filter { input, .. }
| Node::Limit { input, .. }
| Node::Fetch { input, .. } => ordered(plan, input),
_ => false,
}
}
fn bounds(plan: &Plan, input: NodeRef, predicate: ExprRef) -> Vec<(usize, Op, Bound)> {
let Node::TableFunction { index, function, .. } = *plan.node(input) else { return Vec::new() };
if TableFunction::lookup(plan.string(function)) != Some(TableFunction::ReadParquet) {
return Vec::new();
}
let mut tests = Vec::new();
conjuncts(plan, predicate, index, &mut tests);
tests
}
fn conjuncts(plan: &Plan, predicate: ExprRef, index: u32, out: &mut Vec<(usize, Op, Bound)>) {
match *plan.expr(predicate) {
Expr::Conjunction { op: ConjunctionOp::And, children } => {
for child in plan.expr_list(children) {
conjuncts(plan, *child, index, out);
}
}
Expr::Compare { op, left, right } => {
if let Some(test) = comparison(plan, op, left, right, index) {
out.push(test);
}
}
_ => {}
}
}
fn comparison(
plan: &Plan,
op: CompareOp,
left: ExprRef,
right: ExprRef,
index: u32,
) -> Option<(usize, Op, Bound)> {
let op = match op {
CompareOp::Equal => Op::Equal,
CompareOp::Less => Op::Less,
CompareOp::LessOrEqual => Op::LessOrEqual,
CompareOp::Greater => Op::Greater,
CompareOp::GreaterOrEqual => Op::GreaterOrEqual,
CompareOp::NotEqual | CompareOp::DistinctFrom | CompareOp::NotDistinctFrom => return None,
};
let (op, binding, value) = match (plan.expr(left), plan.expr(right)) {
(Expr::Column(binding), Expr::Constant(value)) => (op, *binding, *value),
(Expr::Constant(value), Expr::Column(binding)) => (op.flipped(), *binding, *value),
_ => return None,
};
if binding.table != index {
return None;
}
Some((binding.column as usize, op, Bound::of_value(plan.value(value))?))
}
struct Segment<'a> {
source: Arc<dyn Source + 'a>,
streams: Vec<Arc<dyn DynStream + 'a>>,
schema: Schema,
after: Vec<PipelineRef>,
}
impl<'a> Segment<'a> {
fn new(source: Arc<dyn Source + 'a>, schema: Schema) -> Self {
Self { source, streams: Vec::new(), schema, after: Vec::new() }
}
fn reading(source: Arc<dyn Source + 'a>, schema: Schema, after: PipelineRef) -> Self {
Self { source, streams: Vec::new(), schema, after: vec![after] }
}
fn then(mut self, stream: Arc<dyn DynStream + 'a>, schema: Schema) -> Self {
self.streams.push(stream);
self.schema = schema;
self
}
}
struct Building<'a, 'b> {
plan: &'a Plan,
catalog: &'a Catalog,
cancel: &'b Cancel,
memory: &'b Memory,
seams: &'b Settings,
report: &'b Report,
shape: Shape,
done: Vec<Pipeline<'a>>,
drivers: Vec<Arc<Driver>>,
pruning: Vec<(usize, Op, Bound)>,
}
impl<'a> Building<'a, '_> {
fn gathered(&self, node: NodeRef) -> u32 {
self.shape.gathered(node).expect("a node with two inputs has a second operator")
}
fn close(&mut self, segment: Segment<'a>, id: PipelineRef, sink: Arc<dyn DynSink + 'a>) {
let mut pipeline = Pipeline::new(PipelineId(id), segment.source, sink);
for stream in segment.streams {
pipeline = pipeline.then(stream);
}
for after in segment.after {
pipeline = pipeline.after(PipelineId(after));
}
self.done.push(pipeline);
self.drivers.push(self.report.driving(id));
}
fn watch(
&self,
node: NodeRef,
id: u32,
pipeline: u32,
kind: &str,
detail: Option<&str>,
) -> Arc<Counters> {
let mut counters = Counters::new(id, pipeline, kind);
if let Some(detail) = detail {
counters = counters.detailed(detail);
}
for seam in seams_of(self.plan.node(node)) {
if let Some(running) = registries().running(*seam, self.seams) {
counters = counters.chose(seam.name(), &running.name, running.is_reference);
}
}
self.report.watch(counters)
}
fn aggregate(
&mut self,
reference: NodeRef,
input: NodeRef,
index: u32,
groups: Slice,
aggregates: Slice,
max_groups: Option<usize>,
) -> Result<Segment<'a>> {
let below = self.node(input)?;
let (aggregate, out) =
Aggregate::new(self.plan, &below.schema, index, groups, aggregates, self.memory)?;
let aggregate = match max_groups {
Some(limit) => aggregate.limit_groups(limit),
None => aggregate,
};
let schema = aggregate.schema().clone();
let id = self.shape.operator(reference);
let pipeline = self.shape.pipeline(reference);
let counters = self.watch(reference, id, pipeline, "Aggregate", None);
let reading = Arc::clone(&counters);
self.close(below, pipeline, Arc::new(Watched::new(aggregate, counters)));
Ok(Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline))
}
fn node(&mut self, reference: NodeRef) -> Result<Segment<'a>> {
let plan = self.plan;
let memory = self.memory;
let id = self.shape.operator(reference);
let pipeline = self.shape.pipeline(reference);
let segment = match *plan.node(reference) {
Node::Get { catalog: database, schema, table, index, columns, .. } => {
let name = QualifiedName::new(
plan.string(database),
plan.string(schema),
plan.string(table),
);
let scan = Scan::new(plan, self.catalog.table(&name)?, index, columns)?;
let schema = scan.schema().clone();
let counters =
self.watch(reference, id, pipeline, "Scan", Some(plan.string(table)));
Segment::new(Arc::new(Watched::new(scan, counters)), schema)
}
Node::Dummy => {
let dummy = Dummy::new();
let schema = dummy.schema().clone();
let counters = self.watch(reference, id, pipeline, "Dummy", None);
Segment::new(Arc::new(Watched::new(dummy, counters)), schema)
}
Node::Values { index, columns, rows } => {
let values = Values::new(plan, index, columns, rows)?;
let schema = values.schema().clone();
let counters = self.watch(reference, id, pipeline, "Values", None);
Segment::new(Arc::new(Watched::new(values, counters)), schema)
}
Node::TableFunction { index, function, args, options, settings, columns } => {
let name = plan.string(function);
match TableFunction::lookup(name) {
Some(function @ (TableFunction::ReadParquet | TableFunction::ReadCsv)) => {
let counters = self.watch(reference, id, pipeline, "FileScan", Some(name));
let tests = std::mem::take(&mut self.pruning);
let scan = FileScan::new(
plan, index, function, args, options, settings, columns, tests,
)?
.watched(counters.clone());
let schema = scan.schema().clone();
Segment::new(Arc::new(Watched::new(scan, counters)), schema)
}
Some(TableFunction::RudbStrategies) => {
let table = Strategies::new(plan, index, columns)?;
let schema = table.schema().clone();
let counters = self.watch(reference, id, pipeline, "Strategies", None);
Segment::new(Arc::new(Watched::new(table, counters)), schema)
}
_ => {
let series = Series::new(plan, index, name, args)?;
let schema = series.schema().clone();
let counters = self.watch(reference, id, pipeline, "Series", Some(name));
Segment::new(Arc::new(Watched::new(series, counters)), schema)
}
}
}
Node::Fetch { input, index, args, columns, row } => {
let below = self.node(input)?;
let counters = self.watch(reference, id, pipeline, "Fetch", None);
let fetch = Fetch::new(plan, &below.schema, index, args, columns, row)?
.watched(counters.clone());
let schema = fetch.schema().clone();
below.then(Arc::new(Watched::new(fetch, counters)), schema)
}
Node::Filter { input, predicate } => {
self.pruning = bounds(plan, input, predicate);
let below = self.node(input)?;
self.pruning = Vec::new();
let schema = below.schema.clone();
let filter = Filter::new(plan, reference, predicate, &schema, self.seams)?;
let counters = self.watch(reference, id, pipeline, "Filter", None);
below.then(Arc::new(Watched::new(filter, counters)), schema)
}
Node::Project { input, index, exprs, names } => {
let below = self.node(input)?;
let project = Project::new(plan, &below.schema, index, exprs, names)?;
let schema = project.schema().clone();
let counters = self.watch(reference, id, pipeline, "Project", None);
below.then(Arc::new(Watched::new(project, counters)), schema)
}
Node::Aggregate { input, index, groups, aggregates } => {
self.aggregate(reference, input, index, groups, aggregates, None)?
}
Node::Sort { input, keys } => {
let below = self.node(input)?;
let schema = below.schema.clone();
let (sort, out) = Sort::new(plan, &schema, keys, memory)?;
let counters = self.watch(reference, id, pipeline, "Sort", None);
let reading = Arc::clone(&counters);
self.close(below, pipeline, Arc::new(Watched::new(sort, counters)));
Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
}
Node::Limit { input, count, offset } => {
let max_groups = count
.and_then(|count| count.checked_add(offset))
.and_then(|count| usize::try_from(count).ok());
let below = match (plan.node(input).clone(), max_groups) {
(
Node::Aggregate { input: under, index, groups, aggregates },
Some(max_groups),
) => {
self.aggregate(input, under, index, groups, aggregates, Some(max_groups))?
}
_ => self.node(input)?,
};
let schema = below.schema.clone();
let limit = Limit::new(count, offset);
let counters = self.watch(reference, id, pipeline, "Limit", None);
below.then(Arc::new(Watched::new(limit, counters)), schema)
}
Node::TopN { input, keys, count, offset } => {
let below = self.node(input)?;
let schema = below.schema.clone();
let (top, out) = TopN::new(plan, &schema, keys, count, offset, memory)?;
let counters = self.watch(reference, id, pipeline, "TopN", None);
let reading = Arc::clone(&counters);
self.close(below, pipeline, Arc::new(Watched::new(top, counters)));
Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
}
Node::Distinct { input, on } => {
let below = self.node(input)?;
let schema = below.schema.clone();
let (distinct, out) = Distinct::new(plan, &schema, on, memory)?;
let counters = self.watch(reference, id, pipeline, "Distinct", None);
let reading = Arc::clone(&counters);
self.close(below, pipeline, Arc::new(Watched::new(distinct, counters)));
Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
}
Node::Join { left, right, kind, conditions } => {
let gather_id = self.gathered(reference);
let gathering = self.shape.pipeline(right);
let right = self.node(right)?;
let right_schema = right.schema.clone();
let (gather, gathered) = Gather::new(memory);
let kept = self.watch(reference, gather_id, gathering, "Gather", None);
self.close(right, gathering, Arc::new(Watched::new(gather, kept)));
let mut left = self.node(left)?;
let side = Gathered { schema: &right_schema, rows: gathered };
let (join, out) =
Join::new(plan, &left.schema, side, kind, conditions, self.cancel, memory);
let schema = join.schema().clone();
let counters = self.watch(reference, id, pipeline, "Join", None);
let reading = Arc::clone(&counters);
left.after.push(gathering);
self.close(left, pipeline, Arc::new(Watched::new(join, counters)));
Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
}
Node::CrossProduct { left, right } => {
let keep_id = self.gathered(reference);
let aside = self.shape.pipeline(right);
let right = self.node(right)?;
let right_schema = right.schema.clone();
let (keep, kept) = Keep::new(memory);
let held = self.watch(reference, keep_id, aside, "Keep", None);
self.close(right, aside, Arc::new(Watched::new(keep, held)));
let mut left = self.node(left)?;
let cross = CrossProduct::new(&left.schema, &right_schema, kept);
let schema = cross.schema().clone();
let counters = self.watch(reference, id, pipeline, "CrossProduct", None);
left.after.push(aside);
left.then(Arc::new(Watched::new(cross, counters)), schema)
}
Node::SetOp { left, right, kind, all, index } => {
let gather_id = self.gathered(reference);
let counting = self.shape.pipeline(right);
let right = self.node(right)?;
let (gather, gathered) = Gather::new(memory);
let kept = self.watch(reference, gather_id, counting, "Gather", None);
self.close(right, counting, Arc::new(Watched::new(gather, kept)));
let mut left = self.node(left)?;
let (setop, out) = SetOp::new(&left.schema, gathered, kind, all, index, memory);
let schema = setop.schema().clone();
let counters = self.watch(reference, id, pipeline, "SetOp", None);
let reading = Arc::clone(&counters);
left.after.push(counting);
self.close(left, pipeline, Arc::new(Watched::new(setop, counters)));
Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
}
};
Ok(segment)
}
}