use rudb_common::{Error, LogicalType, Result, Value};
use rudb_kernels::{
Comparison, Connective, cast, combine, compare, is_true, refine, refine_flags, selection,
};
use rudb_plan::{CompareOp, ConjunctionOp, Expr, ExprRef, Plan};
use rudb_vector::{Chunk, Selection, Vector};
use crate::schema::Schema;
const _: () = {
const fn assert_shareable<T: Send + Sync>() {}
assert_shareable::<Prepared>();
};
#[derive(Debug)]
pub struct Prepared {
steps: Vec<Step>,
types: Vec<LogicalType>,
operands: Vec<usize>,
last_use: Vec<usize>,
roots: Vec<usize>,
}
#[derive(Debug)]
enum Step {
Column(usize),
Constant(Value),
Cast {
input: usize,
try_cast: bool,
},
Compare {
op: Comparison,
left: usize,
right: usize,
},
Conjunction {
op: Connective,
start: usize,
len: usize,
},
Function {
name: String,
start: usize,
len: usize,
},
Case {
arms: Vec<PreparedArm>,
otherwise: Option<Prepared>,
},
}
#[derive(Debug)]
struct PreparedArm {
when: Prepared,
then: Prepared,
}
#[derive(Debug)]
pub struct Scratch {
slots: Vec<Option<Vector>>,
}
impl Prepared {
pub fn new(plan: &Plan, exprs: &[ExprRef], schema: &Schema) -> Result<Self> {
let mut prepared = Self {
steps: Vec::new(),
types: Vec::new(),
operands: Vec::new(),
last_use: Vec::new(),
roots: Vec::new(),
};
for &expr in exprs {
let root = prepared.push(plan, expr, schema)?;
prepared.roots.push(root);
}
prepared.last_use = prepared.last_uses();
Ok(prepared)
}
fn last_uses(&self) -> Vec<usize> {
let mut last = vec![usize::MAX; self.steps.len()];
for index in 0..self.steps.len() {
self.for_each_operand(index, |operand| last[operand] = index);
}
for &root in &self.roots {
last[root] = usize::MAX;
}
last
}
fn for_each_operand(&self, index: usize, mut visit: impl FnMut(usize)) {
match &self.steps[index] {
Step::Column(_) | Step::Constant(_) | Step::Case { .. } => {}
Step::Cast { input, .. } => visit(*input),
Step::Compare { left, right, .. } => {
visit(*left);
visit(*right);
}
Step::Conjunction { start, len, .. } | Step::Function { start, len, .. } => {
for &operand in &self.operands[*start..*start + *len] {
visit(operand);
}
}
}
}
pub fn one(plan: &Plan, expr: ExprRef, schema: &Schema) -> Result<Self> {
Self::new(plan, &[expr], schema)
}
#[must_use]
pub fn scratch(&self) -> Scratch {
Scratch { slots: (0..self.steps.len()).map(|_| None).collect() }
}
#[must_use]
pub fn len(&self) -> usize {
self.roots.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.roots.is_empty()
}
pub fn evaluate(
&self,
chunk: &Chunk,
scratch: &mut Scratch,
out: &mut Vec<Vector>,
) -> Result<()> {
self.run(chunk, scratch)?;
for &root in &self.roots {
match self.steps[root] {
Step::Column(position) => out.push(chunk.column(position)?.clone()),
_ => out.push(scratch.slots[root].take().ok_or_else(|| missing(root))?),
}
}
Ok(())
}
pub fn evaluate_one<'s>(
&'s self,
chunk: &'s Chunk,
scratch: &'s mut Scratch,
) -> Result<&'s Vector> {
let [root] = self.roots[..] else {
return Err(Error::internal(format!(
"evaluate_one over a prepared expression of {} roots",
self.roots.len()
)));
};
self.run(chunk, scratch)?;
self.operand(root, chunk, &scratch.slots)
}
pub fn evaluate_filter(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<Selection> {
let [root] = self.roots[..] else {
return Err(Error::internal(format!(
"evaluate_filter over a prepared expression of {} roots",
self.roots.len()
)));
};
let Step::Conjunction { op: Connective::And, start, len } = self.steps[root] else {
let flags = self.evaluate_one(chunk, scratch)?;
return Ok(selection(flags, chunk.len()));
};
scratch.slots.clear();
scratch.slots.resize_with(self.steps.len(), || None);
let mut begin = 0;
let mut kept: Option<Selection> = None;
for at in 0..len {
let conjunct = self.operands[start + at];
if kept.as_ref().is_some_and(Selection::is_empty) {
break;
}
for index in begin..conjunct {
self.run_step(index, chunk, scratch)?;
}
let next = self.thread(conjunct, chunk, scratch, kept.as_ref())?;
kept = Some(next);
for index in begin..=conjunct {
scratch.slots[index] = None;
}
begin = conjunct + 1;
}
Ok(kept.unwrap_or_else(|| Selection::identity(chunk.len())))
}
fn thread(
&self,
index: usize,
chunk: &Chunk,
scratch: &mut Scratch,
kept: Option<&Selection>,
) -> Result<Selection> {
if let Step::Compare { op, left, right } = self.steps[index] {
let left = self.operand(left, chunk, &scratch.slots)?;
let right = self.operand(right, chunk, &scratch.slots)?;
return match kept {
None => Ok(selection(&compare(op, left, right)?, chunk.len())),
Some(kept) => refine(op, left, right, kept),
};
}
self.run_step(index, chunk, scratch)?;
let flags = self.operand(index, chunk, &scratch.slots)?;
match kept {
None => Ok(selection(flags, chunk.len())),
Some(kept) => refine_flags(flags, kept),
}
}
fn run(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<()> {
scratch.slots.clear();
scratch.slots.resize_with(self.steps.len(), || None);
for index in 0..self.steps.len() {
self.run_step(index, chunk, scratch)?;
}
Ok(())
}
fn run_step(&self, index: usize, chunk: &Chunk, scratch: &mut Scratch) -> Result<()> {
let produced = self.step(index, chunk, &scratch.slots)?;
scratch.slots[index] = produced;
let slots = &mut scratch.slots;
self.for_each_operand(index, |operand| {
if self.last_use[operand] == index {
slots[operand] = None;
}
});
Ok(())
}
fn step(
&self,
index: usize,
chunk: &Chunk,
slots: &[Option<Vector>],
) -> Result<Option<Vector>> {
let ty = &self.types[index];
let produced = match &self.steps[index] {
Step::Column(_) => None,
Step::Constant(value) => Some(Vector::constant(ty.clone(), value.clone(), chunk.len())),
Step::Cast { input, try_cast } => {
Some(cast(self.operand(*input, chunk, slots)?, ty, *try_cast)?)
}
Step::Compare { op, left, right } => Some(compare(
*op,
self.operand(*left, chunk, slots)?,
self.operand(*right, chunk, slots)?,
)?),
Step::Conjunction { op, start, len } => {
Some(
self.with_operands(*start, *len, chunk, slots, |children| {
combine(*op, children)
})?,
)
}
Step::Function { name, start, len } => {
Some(self.with_operands(*start, *len, chunk, slots, |args| {
rudb_kernels::call(name, args, ty)
})?)
}
Step::Case { arms, otherwise } => {
Some(self.case(chunk, arms, otherwise.as_ref(), ty)?)
}
};
Ok(produced)
}
fn operand<'v>(
&self,
index: usize,
chunk: &'v Chunk,
slots: &'v [Option<Vector>],
) -> Result<&'v Vector> {
if let Step::Column(position) = self.steps[index] {
return chunk.column(position);
}
slots[index].as_ref().ok_or_else(|| missing(index))
}
fn with_operands<'v, T>(
&self,
start: usize,
len: usize,
chunk: &'v Chunk,
slots: &'v [Option<Vector>],
run: impl FnOnce(&[&'v Vector]) -> Result<T>,
) -> Result<T> {
match self.operands[start..start + len] {
[a] => run(&[self.operand(a, chunk, slots)?]),
[a, b] => run(&[self.operand(a, chunk, slots)?, self.operand(b, chunk, slots)?]),
[a, b, c] => run(&[
self.operand(a, chunk, slots)?,
self.operand(b, chunk, slots)?,
self.operand(c, chunk, slots)?,
]),
_ => {
let gathered = self.gather(start, len, chunk, slots)?;
run(&gathered)
}
}
}
fn gather<'v>(
&self,
start: usize,
len: usize,
chunk: &'v Chunk,
slots: &'v [Option<Vector>],
) -> Result<Vec<&'v Vector>> {
let mut gathered = Vec::with_capacity(len);
for &operand in &self.operands[start..start + len] {
gathered.push(self.operand(operand, chunk, slots)?);
}
Ok(gathered)
}
fn case(
&self,
chunk: &Chunk,
arms: &[PreparedArm],
otherwise: Option<&Prepared>,
ty: &LogicalType,
) -> Result<Vector> {
let mut answers = vec![Value::Null; chunk.len()];
let mut pending: Vec<usize> = (0..chunk.len()).collect();
for arm in arms {
if pending.is_empty() {
break;
}
let narrowed = narrow(chunk, &pending)?;
let mut scratch = arm.when.scratch();
let flags = arm.when.evaluate_one(&narrowed, &mut scratch)?;
let mut taken = Vec::new();
let mut still = Vec::new();
for (at, &row) in pending.iter().enumerate() {
if is_true(&flags.value_at(at)) {
taken.push((at, row));
} else {
still.push(row);
}
}
if !taken.is_empty() {
let positions: Vec<usize> = taken.iter().map(|&(at, _)| at).collect();
let matched = narrow(&narrowed, &positions)?;
let mut scratch = arm.then.scratch();
let results = arm.then.evaluate_one(&matched, &mut scratch)?;
for (slot, &(_, row)) in taken.iter().enumerate() {
answers[row] = results.value_at(slot);
}
}
pending = still;
}
if let Some(otherwise) = otherwise {
if !pending.is_empty() {
let narrowed = narrow(chunk, &pending)?;
let mut scratch = otherwise.scratch();
let results = otherwise.evaluate_one(&narrowed, &mut scratch)?;
for (slot, &row) in pending.iter().enumerate() {
answers[row] = results.value_at(slot);
}
}
}
Vector::from_values(ty.clone(), &answers)
}
fn push(&mut self, plan: &Plan, expr: ExprRef, schema: &Schema) -> Result<usize> {
let ty = plan.expr_type(expr).clone();
let step = match *plan.expr(expr) {
Expr::Column(binding) => {
let position = schema.position_of(binding).ok_or_else(|| {
Error::internal(format!(
"column #{}.{} is not in the schema this operator was given",
binding.table, binding.column
))
})?;
Step::Column(position)
}
Expr::Constant(reference) => Step::Constant(plan.value(reference).clone()),
Expr::Cast { input, try_cast } => {
Step::Cast { input: self.push(plan, input, schema)?, try_cast }
}
Expr::Compare { op, left, right } => Step::Compare {
op: comparison(op),
left: self.push(plan, left, schema)?,
right: self.push(plan, right, schema)?,
},
Expr::Conjunction { op, children } => {
let (start, len) = self.push_list(plan, plan.expr_list(children), schema)?;
Step::Conjunction { op: connective(op), start, len }
}
Expr::Function { name, args } => {
let (start, len) = self.push_list(plan, plan.expr_list(args), schema)?;
Step::Function { name: plan.string(name).to_string(), start, len }
}
Expr::Aggregate { name, .. } => {
return Err(Error::internal(format!(
"the {} aggregate was evaluated as an ordinary expression",
plan.string(name)
)));
}
Expr::Case { arms, otherwise } => {
let mut prepared = Vec::new();
for &arm in plan.arm_list(arms) {
prepared.push(PreparedArm {
when: Self::one(plan, arm.when, schema)?,
then: Self::one(plan, arm.then, schema)?,
});
}
let otherwise = match otherwise {
Some(otherwise) => Some(Self::one(plan, otherwise, schema)?),
None => None,
};
Step::Case { arms: prepared, otherwise }
}
};
self.steps.push(step);
self.types.push(ty);
Ok(self.steps.len() - 1)
}
fn push_list(
&mut self,
plan: &Plan,
exprs: &[ExprRef],
schema: &Schema,
) -> Result<(usize, usize)> {
let mut indices = Vec::with_capacity(exprs.len());
for &expr in exprs {
indices.push(self.push(plan, expr, schema)?);
}
let start = self.operands.len();
let len = indices.len();
self.operands.extend(indices);
Ok((start, len))
}
}
fn missing(index: usize) -> Error {
Error::internal(format!("step {index} was used as an operand before it produced anything"))
}
pub(crate) fn narrow(chunk: &Chunk, rows: &[usize]) -> Result<Chunk> {
let mut selection = Selection::with_capacity(rows.len());
for &row in rows {
selection.push(row);
}
chunk.clone().select(&selection)
}
pub(crate) fn comparison(op: CompareOp) -> Comparison {
match op {
CompareOp::Equal => Comparison::Equal,
CompareOp::NotEqual => Comparison::NotEqual,
CompareOp::Less => Comparison::Less,
CompareOp::LessOrEqual => Comparison::LessOrEqual,
CompareOp::Greater => Comparison::Greater,
CompareOp::GreaterOrEqual => Comparison::GreaterOrEqual,
CompareOp::DistinctFrom => Comparison::DistinctFrom,
CompareOp::NotDistinctFrom => Comparison::NotDistinctFrom,
}
}
pub(crate) fn connective(op: ConjunctionOp) -> Connective {
match op {
ConjunctionOp::And => Connective::And,
ConjunctionOp::Or => Connective::Or,
}
}
#[cfg(test)]
mod tests {
use rudb_common::{Field, LogicalType, Value};
use rudb_kernels::is_true;
use rudb_plan::{ExprRef, Node, Plan};
use rudb_vector::{Chunk, Selection, Vector};
use super::{Prepared, narrow};
use crate::expr::evaluate;
use crate::schema::Schema;
fn input() -> (Schema, Chunk) {
let schema = Schema::numbered(
vec![Field::new("x", LogicalType::Integer), Field::new("s", LogicalType::Varchar)],
0,
);
let x = Vector::from_values(
LogicalType::Integer,
&[Value::Integer(3), Value::Integer(1), Value::Null, Value::Integer(2)],
)
.expect("four integers");
let s = Vector::from_values(
LogicalType::Varchar,
&[
Value::Varchar("a".to_string()),
Value::Null,
Value::Varchar("c".to_string()),
Value::Varchar("a".to_string()),
],
)
.expect("four strings");
(schema, Chunk::new(vec![x, s]).expect("two columns of four rows"))
}
fn projection(exprs: &str) -> (Plan, Vec<ExprRef>) {
let text =
format!("Project #1 [{exprs}]\n Get memory.main.t AS t #0 [x::INTEGER, s::VARCHAR]");
let plan = Plan::parse(&text).expect("a well formed plan");
let Node::Project { exprs, .. } = *plan.node(plan.root()) else {
panic!("the root of that text is a projection");
};
let list = plan.expr_list(exprs).to_vec();
(plan, list)
}
fn agrees(exprs: &str) {
let (schema, chunk) = input();
let (plan, list) = projection(exprs);
let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
let mut scratch = prepared.scratch();
let mut fast = Vec::new();
prepared.evaluate(&chunk, &mut scratch, &mut fast).expect("the prepared form runs");
for (at, &expr) in list.iter().enumerate() {
let slow = evaluate(&plan, expr, &schema, &chunk).expect("the tree walk runs");
for row in 0..chunk.len() {
assert_eq!(
fast[at].value_at(row),
slow.value_at(row),
"expression {at} of `{exprs}` at row {row}"
);
}
}
}
#[test]
fn a_column_reference_agrees() {
agrees("#0.0::INTEGER AS a, #0.1::VARCHAR AS b");
}
#[test]
fn a_constant_agrees() {
agrees("7::INTEGER AS a, NULL::INTEGER AS b");
}
#[test]
fn a_cast_agrees() {
agrees("CAST(#0.0::INTEGER)::BIGINT AS a, CAST(#0.0::INTEGER)::VARCHAR AS b");
}
#[test]
fn a_comparison_agrees() {
agrees("(#0.0::INTEGER > 1::INTEGER)::BOOLEAN AS a");
}
#[test]
fn a_conjunction_agrees() {
agrees(
"((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 3::INTEGER)::BOOLEAN)\
::BOOLEAN AS a",
);
}
#[test]
fn a_function_agrees() {
agrees("\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER AS a");
}
#[test]
fn a_case_agrees() {
agrees(
"CASE WHEN (#0.0::INTEGER > 1::INTEGER)::BOOLEAN THEN 10::INTEGER \
ELSE 20::INTEGER END::INTEGER AS a",
);
}
#[test]
fn a_column_mentioned_three_times_agrees() {
agrees("\"+\"(\"+\"(#0.0::INTEGER, #0.0::INTEGER)::INTEGER, #0.0::INTEGER)::INTEGER AS a");
}
#[test]
fn a_chain_holds_one_intermediate_at_a_time() {
let (schema, chunk) = input();
let mut expr = "#0.0::INTEGER".to_string();
for _ in 0..8 {
expr = format!("\"+\"({expr}, 1::INTEGER)::INTEGER");
}
let (plan, list) = projection(&format!("{expr} AS a"));
let prepared = Prepared::new(&plan, &list, &schema).expect("the chain resolves");
let mut scratch = prepared.scratch();
prepared.run(&chunk, &mut scratch).expect("the chain runs");
let live = scratch.slots.iter().filter(|slot| slot.is_some()).count();
assert_eq!(live, 1, "a chain that has run should be holding its answer and nothing else");
}
fn filters(predicate: &str) {
let (schema, chunk) = input();
let (plan, list) = projection(&format!("{predicate} AS p"));
let prepared = Prepared::new(&plan, &list, &schema).expect("the predicate resolves");
let mut scratch = prepared.scratch();
let threaded = prepared.evaluate_filter(&chunk, &mut scratch).expect("the filter runs");
let flags = evaluate(&plan, list[0], &schema, &chunk).expect("the tree walk runs");
let expected = Selection::from_predicate(chunk.len(), |row| is_true(&flags.value_at(row)));
assert_eq!(threaded, expected, "`{predicate}`");
let again = prepared.evaluate_filter(&chunk, &mut scratch).expect("the filter runs again");
assert_eq!(again, expected, "`{predicate}` a second time");
}
#[test]
fn a_single_comparison_filters_the_same_rows() {
filters("(#0.0::INTEGER > 1::INTEGER)::BOOLEAN");
filters("(#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN");
filters("(#0.0::INTEGER IS NOT DISTINCT FROM NULL::INTEGER)::BOOLEAN");
}
#[test]
fn a_chain_of_conjuncts_keeps_what_all_of_them_keep() {
filters(
"((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 3::INTEGER)::BOOLEAN)\
::BOOLEAN",
);
filters(
"((#0.0::INTEGER >= 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER <= 3::INTEGER)::BOOLEAN \
AND (#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN AND (#0.0::INTEGER <> 2::INTEGER)\
::BOOLEAN)::BOOLEAN",
);
}
#[test]
fn a_conjunct_that_keeps_nothing_ends_the_predicate() {
filters(
"((#0.0::INTEGER > 9::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 9::INTEGER)::BOOLEAN)\
::BOOLEAN",
);
}
#[test]
fn a_conjunct_over_a_computed_operand_keeps_the_same_rows() {
filters(
"((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND \
(\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER < 4::INTEGER)::BOOLEAN)::BOOLEAN",
);
}
#[test]
fn a_conjunct_that_is_not_a_comparison_is_threaded_too() {
filters(
"((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND ((#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN \
OR (#0.0::INTEGER = 1::INTEGER)::BOOLEAN)::BOOLEAN)::BOOLEAN",
);
filters(
"(((#0.1::VARCHAR = 'c'::VARCHAR)::BOOLEAN OR (#0.0::INTEGER = 3::INTEGER)::BOOLEAN)\
::BOOLEAN AND (#0.0::INTEGER <> 1::INTEGER)::BOOLEAN)::BOOLEAN",
);
}
#[test]
fn an_or_at_the_top_is_not_threaded() {
filters(
"((#0.0::INTEGER > 2::INTEGER)::BOOLEAN OR (#0.1::VARCHAR = 'c'::VARCHAR)::BOOLEAN)\
::BOOLEAN",
);
}
#[test]
fn a_filter_over_a_selected_chunk_keeps_the_same_rows() {
let (schema, chunk) = input();
let predicate = "((#0.0::INTEGER >= 1::INTEGER)::BOOLEAN AND \
(#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN)::BOOLEAN";
let (plan, list) = projection(&format!("{predicate} AS p"));
let prepared = Prepared::new(&plan, &list, &schema).expect("the predicate resolves");
let mut scratch = prepared.scratch();
let narrowed = narrow(&chunk, &[0, 3]).expect("two of the four rows");
let threaded = prepared.evaluate_filter(&narrowed, &mut scratch).expect("the filter runs");
let flags = evaluate(&plan, list[0], &schema, &narrowed).expect("the tree walk runs");
let expected =
Selection::from_predicate(narrowed.len(), |row| is_true(&flags.value_at(row)));
assert_eq!(threaded, expected);
}
#[test]
fn a_scratch_used_twice_gives_the_same_answer_twice() {
let (schema, chunk) = input();
let (plan, list) = projection("\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER AS a");
let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
let mut scratch = prepared.scratch();
let mut once = Vec::new();
prepared.evaluate(&chunk, &mut scratch, &mut once).expect("the first chunk runs");
let mut twice = Vec::new();
prepared.evaluate(&chunk, &mut scratch, &mut twice).expect("the second chunk runs");
assert_eq!(once, twice);
}
#[test]
fn a_shorter_chunk_after_a_longer_one_is_evaluated_at_its_own_length() {
let (schema, chunk) = input();
let (plan, list) = projection("7::INTEGER AS a");
let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
let mut scratch = prepared.scratch();
let mut full = Vec::new();
prepared.evaluate(&chunk, &mut scratch, &mut full).expect("the full chunk runs");
assert_eq!(full[0].len(), 4);
let short = chunk
.clone()
.select(&{
let mut selection = Selection::with_capacity(2);
selection.push(0);
selection.push(2);
selection
})
.expect("two of the four rows");
let mut cut = Vec::new();
prepared.evaluate(&short, &mut scratch, &mut cut).expect("the short chunk runs");
assert_eq!(cut[0].len(), 2);
}
#[test]
fn an_aggregate_is_refused_when_it_is_prepared() {
let (schema, _) = input();
let text = "Aggregate #1 groups=[] aggregates=[sum(#0.0::INTEGER)::HUGEINT]\n \
Get memory.main.t AS t #0 [x::INTEGER, s::VARCHAR]";
let plan = Plan::parse(text).expect("a well formed plan");
let Node::Aggregate { aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root of that text is an aggregate");
};
let list = plan.expr_list(aggregates).to_vec();
let error = Prepared::new(&plan, &list, &schema).expect_err("sum is not a scalar");
assert!(error.message().contains("sum"), "{error}");
}
}