use rudb_common::{Error, Field, LogicalType, Result, Value};
use crate::expr::{Arm, ColumnBinding, Expr, SortKey};
use crate::node::Node;
use crate::{ExprRef, NodeRef, Slice, StrRef, ValueRef};
#[derive(Debug, Clone)]
pub struct Plan {
nodes: Vec<Node>,
exprs: Vec<Expr>,
types: Vec<LogicalType>,
values: Vec<Value>,
strings: Vec<String>,
expr_lists: Vec<ExprRef>,
name_lists: Vec<StrRef>,
fields: Vec<Field>,
sort_keys: Vec<SortKey>,
arms: Vec<Arm>,
rows: Vec<Slice>,
root: NodeRef,
}
impl Default for Plan {
fn default() -> Self {
Self::new()
}
}
impl Plan {
#[must_use]
pub fn new() -> Self {
let mut plan = Self::without_nodes();
plan.add_node(Node::Dummy);
plan
}
pub(crate) fn without_nodes() -> Self {
Self {
nodes: Vec::new(),
exprs: Vec::new(),
types: Vec::new(),
values: Vec::new(),
strings: Vec::new(),
expr_lists: Vec::new(),
name_lists: Vec::new(),
fields: Vec::new(),
sort_keys: Vec::new(),
arms: Vec::new(),
rows: Vec::new(),
root: 0,
}
}
#[must_use]
pub fn root(&self) -> NodeRef {
self.root
}
pub fn set_root(&mut self, node: NodeRef) {
self.root = node;
}
#[must_use]
pub fn node_count(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub fn expr_count(&self) -> usize {
self.exprs.len()
}
pub fn add_node(&mut self, node: Node) -> NodeRef {
push(&mut self.nodes, node)
}
pub fn add_expr(&mut self, expr: Expr, ty: LogicalType) -> ExprRef {
self.types.push(ty);
push(&mut self.exprs, expr)
}
pub fn add_value(&mut self, value: Value) -> ValueRef {
push(&mut self.values, value)
}
pub fn add_constant(&mut self, value: Value) -> ExprRef {
let ty = value.logical_type();
let reference = self.add_value(value);
self.add_expr(Expr::Constant(reference), ty)
}
pub fn intern(&mut self, text: &str) -> StrRef {
if let Some(found) = self.strings.iter().position(|held| held == text) {
return u32::try_from(found).expect("a string table this large cannot be built");
}
push(&mut self.strings, text.to_string())
}
pub fn add_expr_list(&mut self, exprs: &[ExprRef]) -> Slice {
extend(&mut self.expr_lists, exprs.iter().copied())
}
pub fn add_name_list(&mut self, names: &[StrRef]) -> Slice {
extend(&mut self.name_lists, names.iter().copied())
}
pub fn add_fields(&mut self, fields: &[Field]) -> Slice {
extend(&mut self.fields, fields.iter().cloned())
}
pub fn add_sort_keys(&mut self, keys: &[SortKey]) -> Slice {
extend(&mut self.sort_keys, keys.iter().copied())
}
pub fn add_arms(&mut self, arms: &[Arm]) -> Slice {
extend(&mut self.arms, arms.iter().copied())
}
pub fn add_rows(&mut self, rows: &[Slice]) -> Slice {
extend(&mut self.rows, rows.iter().copied())
}
#[must_use]
pub fn node(&self, reference: NodeRef) -> &Node {
&self.nodes[reference as usize]
}
#[must_use]
pub fn expr(&self, reference: ExprRef) -> &Expr {
&self.exprs[reference as usize]
}
#[must_use]
pub fn expr_type(&self, reference: ExprRef) -> &LogicalType {
&self.types[reference as usize]
}
#[must_use]
pub fn value(&self, reference: ValueRef) -> &Value {
&self.values[reference as usize]
}
#[must_use]
pub fn string(&self, reference: StrRef) -> &str {
&self.strings[reference as usize]
}
#[must_use]
pub fn expr_list(&self, slice: Slice) -> &[ExprRef] {
&self.expr_lists[slice.range()]
}
#[must_use]
pub fn name_list(&self, slice: Slice) -> &[StrRef] {
&self.name_lists[slice.range()]
}
#[must_use]
pub fn field_list(&self, slice: Slice) -> &[Field] {
&self.fields[slice.range()]
}
#[must_use]
pub fn sort_key_list(&self, slice: Slice) -> &[SortKey] {
&self.sort_keys[slice.range()]
}
#[must_use]
pub fn arm_list(&self, slice: Slice) -> &[Arm] {
&self.arms[slice.range()]
}
#[must_use]
pub fn row_list(&self, slice: Slice) -> &[Slice] {
&self.rows[slice.range()]
}
pub fn rebind(&mut self, reference: ExprRef, binding: ColumnBinding) {
match &mut self.exprs[reference as usize] {
Expr::Column(held) => *held = binding,
other => panic!("expression {reference} is {other:?}, not a column"),
}
}
pub fn node_mut(&mut self, reference: NodeRef) -> &mut Node {
&mut self.nodes[reference as usize]
}
pub fn validate(&self) -> Result<()> {
if self.exprs.len() != self.types.len() {
return Err(Error::internal(format!(
"the plan has {} expressions and {} types",
self.exprs.len(),
self.types.len()
)));
}
if self.root as usize >= self.nodes.len() {
return Err(Error::internal(format!(
"the plan is rooted at node {} and has {} nodes",
self.root,
self.nodes.len()
)));
}
for index in 0..self.exprs.len() {
self.validate_expr(u32::try_from(index).expect("index came from a length"))?;
}
for index in 0..self.nodes.len() {
self.validate_node(u32::try_from(index).expect("index came from a length"))?;
}
Ok(())
}
fn validate_expr(&self, reference: ExprRef) -> Result<()> {
let fail = |what: &str| Err(Error::internal(format!("expression {reference} {what}")));
let backwards = |operand: ExprRef| -> Result<()> {
if operand < reference {
Ok(())
} else {
Err(Error::internal(format!(
"expression {reference} refers to expression {operand}, which is not behind it"
)))
}
};
match *self.expr(reference) {
Expr::Column(_) => {}
Expr::Constant(value) => {
if value as usize >= self.values.len() {
return fail("names a constant that is not in the value table");
}
let held = self.value(value);
if !held.is_null() && held.logical_type() != *self.expr_type(reference) {
return fail("is a constant whose type disagrees with the value it holds");
}
}
Expr::Cast { input, .. } => backwards(input)?,
Expr::Compare { left, right, .. } => {
backwards(left)?;
backwards(right)?;
if *self.expr_type(reference) != LogicalType::Boolean {
return fail("is a comparison that does not produce BOOLEAN");
}
}
Expr::Conjunction { children, .. } => {
if children.len < 2 {
return fail("is a conjunction with fewer than two operands");
}
for &child in self.checked_expr_list(children, reference)? {
backwards(child)?;
}
if *self.expr_type(reference) != LogicalType::Boolean {
return fail("is a conjunction that does not produce BOOLEAN");
}
}
Expr::Function { name, args } | Expr::Aggregate { name, args, .. } => {
if name as usize >= self.strings.len() {
return fail("names a function that is not in the string table");
}
for &arg in self.checked_expr_list(args, reference)? {
backwards(arg)?;
}
if let Expr::Aggregate { filter: Some(filter), .. } = *self.expr(reference) {
backwards(filter)?;
if *self.expr_type(filter) != LogicalType::Boolean {
return fail("has a FILTER that is not BOOLEAN");
}
}
}
Expr::Case { arms, otherwise } => {
if arms.is_empty() {
return fail("is a CASE with no arms");
}
let end = arms.start as usize + arms.len as usize;
if end > self.arms.len() {
return fail("names an arm run that is not in the pool");
}
for arm in self.arm_list(arms) {
backwards(arm.when)?;
backwards(arm.then)?;
if *self.expr_type(arm.when) != LogicalType::Boolean {
return fail("has a WHEN that is not BOOLEAN");
}
}
if let Some(otherwise) = otherwise {
backwards(otherwise)?;
}
}
}
Ok(())
}
fn validate_node(&self, reference: NodeRef) -> Result<()> {
let node = self.node(reference);
let fail = |what: &str| {
Err(Error::internal(format!("node {reference}, which is a {}, {what}", node.keyword())))
};
for child in node.children().into_iter().flatten() {
if child >= reference {
return Err(Error::internal(format!(
"node {reference} has child {child}, which is not behind it"
)));
}
}
match *node {
Node::Dummy | Node::CrossProduct { .. } => {}
Node::Get { catalog, schema, table, alias, columns, .. } => {
for name in [catalog, schema, table, alias] {
if name as usize >= self.strings.len() {
return fail("names a string that is not in the table");
}
}
self.checked_field_list(columns, reference)?;
}
Node::Values { columns, rows, .. } => {
let width = self.checked_field_list(columns, reference)?.len();
let end = rows.start as usize + rows.len as usize;
if end > self.rows.len() {
return fail("names a row run that is not in the pool");
}
for row in self.row_list(rows) {
if self.checked_expr_list(*row, reference)?.len() != width {
return fail("has a row whose length is not the number of columns");
}
}
}
Node::TableFunction { function, args, options, settings, columns, .. } => {
if function as usize >= self.strings.len() {
return fail("names a string that is not in the table");
}
self.checked_field_list(columns, reference)?;
self.checked_expr_list(args, reference)?;
if self.checked_expr_list(settings, reference)?.len() != options.len as usize {
return fail("has a named parameter with no value or a value with no name");
}
let end = options.start as usize + options.len as usize;
if end > self.name_lists.len() {
return fail("names a name run that is not in the pool");
}
for &name in self.name_list(options) {
if name as usize >= self.strings.len() {
return fail("names a parameter that is not in the string table");
}
}
}
Node::Filter { predicate, .. } => {
self.checked_expr(predicate, reference)?;
if *self.expr_type(predicate) != LogicalType::Boolean {
return fail("filters on an expression that is not BOOLEAN");
}
}
Node::Project { exprs, names, .. } => {
let count = self.checked_expr_list(exprs, reference)?.len();
let end = names.start as usize + names.len as usize;
if end > self.name_lists.len() {
return fail("names a name run that is not in the pool");
}
if self.name_list(names).len() != count {
return fail("has a different number of names and expressions");
}
for &name in self.name_list(names) {
if name as usize >= self.strings.len() {
return fail("names an output name that is not in the string table");
}
}
}
Node::Aggregate { groups, aggregates, .. } => {
for &group in self.checked_expr_list(groups, reference)? {
if matches!(self.expr(group), Expr::Aggregate { .. }) {
return fail("groups by an aggregate");
}
}
for &aggregate in self.checked_expr_list(aggregates, reference)? {
if !matches!(self.expr(aggregate), Expr::Aggregate { .. }) {
return fail(
"has something in its aggregate list that is not an aggregate",
);
}
}
}
Node::Sort { keys, .. } | Node::TopN { keys, .. } => {
let end = keys.start as usize + keys.len as usize;
if end > self.sort_keys.len() {
return fail("names a sort key run that is not in the pool");
}
if keys.is_empty() {
return fail("sorts on nothing");
}
for key in self.sort_key_list(keys) {
self.checked_expr(key.expr, reference)?;
}
}
Node::Limit { .. } => {}
Node::Distinct { on, .. } => {
self.checked_expr_list(on, reference)?;
}
Node::Join { conditions, .. } => {
for &condition in self.checked_expr_list(conditions, reference)? {
if *self.expr_type(condition) != LogicalType::Boolean {
return fail("joins on a condition that is not BOOLEAN");
}
}
}
Node::SetOp { .. } => {}
}
for (expr, aggregate_allowed) in self.top_level_exprs(node) {
if aggregate_allowed {
if let Expr::Aggregate { args, filter, .. } = *self.expr(expr) {
let nested = self
.expr_list(args)
.iter()
.chain(filter.iter())
.any(|&child| self.reaches_an_aggregate(child));
if nested {
return fail("has an aggregate inside an aggregate");
}
continue;
}
}
if self.reaches_an_aggregate(expr) {
return fail("has an aggregate outside an aggregate list");
}
}
Ok(())
}
fn top_level_exprs(&self, node: &Node) -> Vec<(ExprRef, bool)> {
let plain = |list: &[ExprRef]| -> Vec<(ExprRef, bool)> {
list.iter().map(|&expr| (expr, false)).collect()
};
match *node {
Node::Get { .. }
| Node::Dummy
| Node::CrossProduct { .. }
| Node::SetOp { .. }
| Node::Limit { .. } => Vec::new(),
Node::Values { rows, .. } => {
self.row_list(rows).iter().flat_map(|row| plain(self.expr_list(*row))).collect()
}
Node::TableFunction { args, .. } => plain(self.expr_list(args)),
Node::Filter { predicate, .. } => vec![(predicate, false)],
Node::Project { exprs, .. } => plain(self.expr_list(exprs)),
Node::Aggregate { groups, aggregates, .. } => {
let mut all = plain(self.expr_list(groups));
all.extend(self.expr_list(aggregates).iter().map(|&expr| (expr, true)));
all
}
Node::Sort { keys, .. } | Node::TopN { keys, .. } => {
self.sort_key_list(keys).iter().map(|key| (key.expr, false)).collect()
}
Node::Distinct { on, .. } => plain(self.expr_list(on)),
Node::Join { conditions, .. } => plain(self.expr_list(conditions)),
}
}
fn reaches_an_aggregate(&self, reference: ExprRef) -> bool {
match *self.expr(reference) {
Expr::Aggregate { .. } => true,
Expr::Column(_) | Expr::Constant(_) => false,
Expr::Cast { input, .. } => self.reaches_an_aggregate(input),
Expr::Compare { left, right, .. } => {
self.reaches_an_aggregate(left) || self.reaches_an_aggregate(right)
}
Expr::Conjunction { children: list, .. } | Expr::Function { args: list, .. } => {
self.expr_list(list).iter().any(|&child| self.reaches_an_aggregate(child))
}
Expr::Case { arms, otherwise } => {
self.arm_list(arms).iter().any(|arm| {
self.reaches_an_aggregate(arm.when) || self.reaches_an_aggregate(arm.then)
}) || otherwise.is_some_and(|child| self.reaches_an_aggregate(child))
}
}
}
fn checked_expr(&self, reference: ExprRef, node: NodeRef) -> Result<()> {
if reference as usize >= self.exprs.len() {
return Err(Error::internal(format!(
"node {node} names expression {reference}, which is not in the arena"
)));
}
Ok(())
}
fn checked_expr_list(&self, slice: Slice, owner: u32) -> Result<&[ExprRef]> {
let end = slice.start as usize + slice.len as usize;
if end > self.expr_lists.len() {
return Err(Error::internal(format!(
"{owner} names an expression run that is not in the pool"
)));
}
let list = self.expr_list(slice);
for &reference in list {
if reference as usize >= self.exprs.len() {
return Err(Error::internal(format!(
"{owner} names expression {reference}, which is not in the arena"
)));
}
}
Ok(list)
}
fn checked_field_list(&self, slice: Slice, owner: u32) -> Result<&[Field]> {
let end = slice.start as usize + slice.len as usize;
if end > self.fields.len() {
return Err(Error::internal(format!(
"{owner} names a field run that is not in the pool"
)));
}
Ok(self.field_list(slice))
}
}
fn push<T>(pool: &mut Vec<T>, item: T) -> u32 {
let index = u32::try_from(pool.len()).expect("a plan arena cannot hold four billion entries");
pool.push(item);
index
}
fn extend<T>(pool: &mut Vec<T>, items: impl Iterator<Item = T>) -> Slice {
let start = u32::try_from(pool.len()).expect("a plan arena cannot hold four billion entries");
pool.extend(items);
let len =
u32::try_from(pool.len()).expect("a plan arena cannot hold four billion entries") - start;
Slice { start, len }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::expr::{ColumnBinding, CompareOp};
#[test]
fn a_fresh_plan_is_a_valid_plan() {
let plan = Plan::new();
assert_eq!(*plan.node(plan.root()), Node::Dummy);
plan.validate().expect("an empty plan is one row and no columns, which is legal");
}
#[test]
fn interning_the_same_string_twice_gives_the_same_reference() {
let mut plan = Plan::new();
let first = plan.intern("hits");
let second = plan.intern("hits");
let other = plan.intern("visits");
assert_eq!(first, second);
assert_ne!(first, other);
assert_eq!(plan.string(first), "hits");
}
#[test]
fn a_constant_takes_its_type_from_its_value() {
let mut plan = Plan::new();
let one = plan.add_constant(Value::Integer(1));
assert_eq!(*plan.expr_type(one), LogicalType::Integer);
plan.validate().expect("a constant that agrees with itself is valid");
}
#[test]
fn a_typed_null_is_allowed_to_disagree_with_its_value() {
let mut plan = Plan::new();
let null = plan.add_value(Value::Null);
plan.add_expr(Expr::Constant(null), LogicalType::Varchar);
plan.validate().expect("a typed null is the point of carrying types separately");
}
#[test]
fn a_constant_that_disagrees_with_its_value_is_caught() {
let mut plan = Plan::new();
let value = plan.add_value(Value::Integer(1));
plan.add_expr(Expr::Constant(value), LogicalType::Varchar);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("disagrees"), "unhelpful message: {message}");
}
#[test]
fn a_filter_on_something_that_is_not_boolean_is_caught() {
let mut plan = Plan::new();
let one = plan.add_constant(Value::Integer(1));
let filter = plan.add_node(Node::Filter { input: 0, predicate: one });
plan.set_root(filter);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("BOOLEAN"), "unhelpful message: {message}");
}
#[test]
fn a_projection_with_more_expressions_than_names_is_caught() {
let mut plan = Plan::new();
let one = plan.add_constant(Value::Integer(1));
let two = plan.add_constant(Value::Integer(2));
let exprs = plan.add_expr_list(&[one, two]);
let name = plan.intern("a");
let names = plan.add_name_list(&[name]);
let project = plan.add_node(Node::Project { input: 0, index: 1, exprs, names });
plan.set_root(project);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("names and expressions"), "unhelpful message: {message}");
}
#[test]
fn a_ragged_values_is_caught() {
let mut plan = Plan::new();
let one = plan.add_constant(Value::Integer(1));
let two = plan.add_constant(Value::Integer(2));
let wide = plan.add_expr_list(&[one, two]);
let narrow = plan.add_expr_list(&[one]);
let rows = plan.add_rows(&[wide, narrow]);
let columns = plan.add_fields(&[
Field::new("a", LogicalType::Integer),
Field::new("b", LogicalType::Integer),
]);
let values = plan.add_node(Node::Values { index: 0, columns, rows });
plan.set_root(values);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("number of columns"), "unhelpful message: {message}");
}
#[test]
fn an_aggregate_outside_an_aggregate_list_is_caught() {
let mut plan = Plan::new();
let name = plan.intern("count_star");
let count = plan.add_expr(
Expr::Aggregate { name, args: Slice::EMPTY, distinct: false, filter: None },
LogicalType::BigInt,
);
let zero = plan.add_constant(Value::BigInt(0));
let compare = plan.add_expr(
Expr::Compare { op: CompareOp::Greater, left: count, right: zero },
LogicalType::Boolean,
);
let filter = plan.add_node(Node::Filter { input: 0, predicate: compare });
plan.set_root(filter);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("aggregate outside"), "unhelpful message: {message}");
}
#[test]
fn an_aggregate_inside_an_aggregate_list_is_fine() {
let mut plan = Plan::new();
let name = plan.intern("count_star");
let count = plan.add_expr(
Expr::Aggregate { name, args: Slice::EMPTY, distinct: false, filter: None },
LogicalType::BigInt,
);
let aggregates = plan.add_expr_list(&[count]);
let aggregate =
plan.add_node(Node::Aggregate { input: 0, index: 1, groups: Slice::EMPTY, aggregates });
plan.set_root(aggregate);
plan.validate().expect("this is the one place an aggregate belongs");
}
#[test]
fn a_node_that_refers_to_itself_is_caught() {
let mut plan = Plan::new();
let filter = plan.add_node(Node::Filter { input: 0, predicate: 0 });
let one = plan.add_constant(Value::Boolean(true));
plan.nodes[filter as usize] = Node::Filter { input: filter, predicate: one };
plan.set_root(filter);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("not behind it"), "unhelpful message: {message}");
}
#[test]
fn an_expression_that_refers_forwards_is_caught() {
let mut plan = Plan::new();
let left = plan.add_constant(Value::Integer(1));
let compare = plan.add_expr(
Expr::Compare { op: CompareOp::Equal, left, right: left },
LogicalType::Boolean,
);
plan.exprs[compare as usize] =
Expr::Compare { op: CompareOp::Equal, left, right: compare + 1 };
plan.add_constant(Value::Integer(2));
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("not behind it"), "unhelpful message: {message}");
}
#[test]
fn a_root_that_is_not_in_the_arena_is_caught() {
let mut plan = Plan::new();
plan.set_root(17);
let message = plan.validate().unwrap_err().to_string();
assert!(message.contains("rooted at node 17"), "unhelpful message: {message}");
}
#[test]
fn a_column_binding_is_two_numbers_and_nothing_else() {
let binding = ColumnBinding::new(3, 7);
assert_eq!(binding.table, 3);
assert_eq!(binding.column, 7);
assert_eq!(size_of::<ColumnBinding>(), 8);
}
}