use std::collections::BTreeMap;
use std::sync::Arc;
use rudb_common::bounds::Zones;
use rudb_common::{Error, Field, LogicalType, Result, Span, Stat, Value};
use crate::expr::{Arm, ColumnBinding, Expr, SortKey};
use crate::node::{JoinKind, Node};
use crate::{ExprRef, NodeRef, Slice, StrRef, ValueRef};
#[derive(Debug, Clone)]
pub struct Plan {
nodes: Vec<Node>,
node_spans: Vec<Span>,
exprs: Vec<Expr>,
expr_spans: Vec<Span>,
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,
measured: BTreeMap<u32, Stat<u64>>,
distincts: BTreeMap<(u32, String), u64>,
zones: BTreeMap<u32, Arc<dyn Zones>>,
}
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(),
node_spans: Vec::new(),
exprs: Vec::new(),
expr_spans: 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,
measured: BTreeMap::new(),
distincts: BTreeMap::new(),
zones: BTreeMap::new(),
}
}
#[must_use]
pub fn root(&self) -> NodeRef {
self.root
}
pub fn set_root(&mut self, node: NodeRef) {
self.root = node;
}
pub fn measure(&mut self, index: u32, rows: Stat<u64>) {
self.measured.insert(index, rows);
}
#[must_use]
pub fn measured(&self, index: u32) -> Stat<u64> {
self.measured.get(&index).copied().unwrap_or(Stat::Unknown)
}
#[must_use]
pub fn measured_count(&self) -> usize {
self.measured.len()
}
pub fn measure_distinct(&mut self, index: u32, column: &str, distinct: u64) {
self.distincts.insert((index, column.to_owned()), distinct);
}
#[must_use]
pub fn distinct_measured(&self, index: u32, column: &str) -> Option<u64> {
self.distincts.get(&(index, column.to_owned())).copied()
}
#[must_use]
pub fn distinct_count(&self) -> usize {
self.distincts.len()
}
pub fn set_zones(&mut self, index: u32, zones: Arc<dyn Zones>) {
self.zones.insert(index, zones);
}
#[must_use]
pub fn zones(&self, index: u32) -> Option<&Arc<dyn Zones>> {
self.zones.get(&index)
}
#[must_use]
pub fn zones_count(&self) -> usize {
self.zones.len()
}
#[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 {
let span = node
.children()
.into_iter()
.flatten()
.map(|child| self.node_span(child))
.fold(Span::new(0, 0), merge_span);
self.add_node_at(node, span)
}
pub fn add_node_at(&mut self, node: Node, span: Span) -> NodeRef {
self.node_spans.push(span);
push(&mut self.nodes, node)
}
pub fn add_expr(&mut self, expr: Expr, ty: LogicalType) -> ExprRef {
let span = self.inferred_expr_span(&expr);
self.add_expr_at(expr, ty, span)
}
pub fn add_expr_at(&mut self, expr: Expr, ty: LogicalType, span: Span) -> ExprRef {
self.types.push(ty);
self.expr_spans.push(span);
push(&mut self.exprs, expr)
}
fn inferred_expr_span(&self, expr: &Expr) -> Span {
let mut span = Span::new(0, 0);
let mut include = |reference: ExprRef| {
span = merge_span(span, self.expr_span(reference));
};
match *expr {
Expr::Column(_) | Expr::Constant(_) => {}
Expr::Cast { input, .. } => include(input),
Expr::Compare { left, right, .. } => {
include(left);
include(right);
}
Expr::Conjunction { children, .. } => {
for &child in self.expr_list(children) {
include(child);
}
}
Expr::Function { args, .. } => {
for &arg in self.expr_list(args) {
include(arg);
}
}
Expr::Aggregate { args, filter, .. } | Expr::Window { args, filter, .. } => {
for &arg in self.expr_list(args) {
include(arg);
}
if let Some(filter) = filter {
include(filter);
}
}
Expr::Case { arms, otherwise } => {
for arm in self.arm_list(arms) {
include(arm.when);
include(arm.then);
}
if let Some(otherwise) = otherwise {
include(otherwise);
}
}
}
span
}
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 node_span(&self, reference: NodeRef) -> Span {
self.node_spans[reference as usize]
}
#[must_use]
pub fn expr_span(&self, reference: ExprRef) -> Span {
self.expr_spans[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 read_columns(&self, reference: ExprRef, found: &mut impl FnMut(ExprRef, ColumnBinding)) {
match *self.expr(reference) {
Expr::Column(binding) => found(reference, binding),
Expr::Constant(_) => {}
Expr::Cast { input, .. } => self.read_columns(input, found),
Expr::Compare { left, right, .. } => {
self.read_columns(left, found);
self.read_columns(right, found);
}
Expr::Conjunction { children, .. } | Expr::Function { args: children, .. } => {
for &child in self.expr_list(children) {
self.read_columns(child, found);
}
}
Expr::Aggregate { args, filter, .. } | Expr::Window { args, filter, .. } => {
for &arg in self.expr_list(args) {
self.read_columns(arg, found);
}
if let Some(inner) = filter {
self.read_columns(inner, found);
}
}
Expr::Case { arms, otherwise } => {
for arm in self.arm_list(arms) {
self.read_columns(arm.when, found);
self.read_columns(arm.then, found);
}
if let Some(inner) = otherwise {
self.read_columns(inner, found);
}
}
}
}
pub fn node_mut(&mut self, reference: NodeRef) -> &mut Node {
&mut self.nodes[reference as usize]
}
pub fn validate(&self) -> Result<()> {
if self.nodes.len() != self.node_spans.len() {
return Err(Error::internal(format!(
"the plan has {} nodes and {} node spans",
self.nodes.len(),
self.node_spans.len()
)));
}
if self.exprs.len() != self.expr_spans.len() {
return Err(Error::internal(format!(
"the plan has {} expressions and {} expression spans",
self.exprs.len(),
self.expr_spans.len()
)));
}
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, .. }
| Expr::Window { 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), .. }
| Expr::Window { 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, .. }
| Node::LateralFunction { 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::Fetch { args, columns, row, .. } => {
self.checked_field_list(columns, reference)?;
if self.checked_expr_list(args, reference)?.len() != 1 {
return fail("reads other than exactly one file, which no ordinal identifies");
}
self.checked_expr(row, reference)?;
if *self.expr_type(row) != LogicalType::BigInt {
return fail("takes its ordinals from an expression that is not BIGINT");
}
}
Node::TableFetch { catalog, schema, table, columns, row, .. } => {
for name in [catalog, schema, table] {
if name 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(row, reference)?;
if *self.expr_type(row) != LogicalType::BigInt {
return fail("takes its ordinals from an expression that is not BIGINT");
}
}
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::Window { partition, order, frame, expressions, .. } => {
self.checked_expr_list(partition, reference)?;
let end = order.start as usize + order.len as usize;
if end > self.sort_keys.len() {
return fail("names a window order run that is not in the pool");
}
for key in self.sort_key_list(order) {
self.checked_expr(key.expr, reference)?;
}
for bound in [frame.start, frame.end] {
match bound {
crate::WindowBound::Preceding(offset)
| crate::WindowBound::Following(offset) => {
self.checked_expr(offset, reference)?;
}
crate::WindowBound::UnboundedPreceding
| crate::WindowBound::CurrentRow
| crate::WindowBound::UnboundedFollowing => {}
}
}
if matches!(frame.start, crate::WindowBound::UnboundedFollowing)
|| matches!(frame.end, crate::WindowBound::UnboundedPreceding)
{
return fail("has an impossible frame boundary");
}
let expressions = self.checked_expr_list(expressions, reference)?;
if expressions.is_empty() {
return fail("has no window expressions");
}
for &expression in expressions {
if !matches!(self.expr(expression), Expr::Window { .. }) {
return fail(
"has something in its expression list that is not a window function",
);
}
}
}
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::DependentJoin { kind, conditions, .. } => {
if matches!(kind, JoinKind::Right | JoinKind::Full | JoinKind::Positional) {
return fail("has a join kind that cannot preserve an outer row dependency");
}
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::MaterializedCte { name, columns, .. } => {
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::CteScan { name, columns, .. } => {
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::SetOp { .. } => {}
}
for (expr, aggregate_allowed, window_allowed) in self.top_level_exprs(node) {
if window_allowed {
if let Expr::Window { args, filter, .. } = *self.expr(expr) {
let nested = self.expr_list(args).iter().chain(filter.iter()).any(|&child| {
self.reaches_a_window(child) || self.reaches_an_aggregate(child)
});
if nested {
return fail("has a window or aggregate inside a window function");
}
continue;
}
}
if self.reaches_a_window(expr) {
return fail("has a window function outside a window list");
}
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, bool)> {
let plain = |list: &[ExprRef]| -> Vec<(ExprRef, bool, bool)> {
list.iter().map(|&expr| (expr, false, false)).collect()
};
match *node {
Node::Get { .. }
| Node::Dummy
| Node::CrossProduct { .. }
| Node::MaterializedCte { .. }
| Node::CteScan { .. }
| 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, .. } | Node::LateralFunction { args, .. } => {
plain(self.expr_list(args))
}
Node::Fetch { args, row, .. } => {
let mut held = plain(self.expr_list(args));
held.push((row, false, false));
held
}
Node::TableFetch { row, .. } => vec![(row, false, false)],
Node::Filter { predicate, .. } => vec![(predicate, false, 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, false)));
all
}
Node::Window { partition, order, frame, expressions, .. } => {
let mut all = plain(self.expr_list(partition));
all.extend(self.sort_key_list(order).iter().map(|key| (key.expr, false, false)));
for bound in [frame.start, frame.end] {
if let crate::WindowBound::Preceding(offset)
| crate::WindowBound::Following(offset) = bound
{
all.push((offset, false, false));
}
}
all.extend(self.expr_list(expressions).iter().map(|&expr| (expr, false, true)));
all
}
Node::Sort { keys, .. } | Node::TopN { keys, .. } => {
self.sort_key_list(keys).iter().map(|key| (key.expr, false, false)).collect()
}
Node::Distinct { on, .. } => plain(self.expr_list(on)),
Node::Join { conditions, .. } | Node::DependentJoin { conditions, .. } => {
plain(self.expr_list(conditions))
}
}
}
fn reaches_an_aggregate(&self, reference: ExprRef) -> bool {
match *self.expr(reference) {
Expr::Aggregate { .. } => true,
Expr::Window { args, filter, .. } => {
self.expr_list(args).iter().any(|&child| self.reaches_an_aggregate(child))
|| filter.is_some_and(|child| self.reaches_an_aggregate(child))
}
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 reaches_a_window(&self, reference: ExprRef) -> bool {
match *self.expr(reference) {
Expr::Window { .. } => true,
Expr::Column(_) | Expr::Constant(_) => false,
Expr::Cast { input, .. } => self.reaches_a_window(input),
Expr::Compare { left, right, .. } => {
self.reaches_a_window(left) || self.reaches_a_window(right)
}
Expr::Conjunction { children, .. }
| Expr::Function { args: children, .. }
| Expr::Aggregate { args: children, filter: None, .. } => {
self.expr_list(children).iter().any(|&child| self.reaches_a_window(child))
}
Expr::Aggregate { args, filter: Some(filter), .. } => {
self.expr_list(args).iter().any(|&child| self.reaches_a_window(child))
|| self.reaches_a_window(filter)
}
Expr::Case { arms, otherwise } => {
self.arm_list(arms)
.iter()
.any(|arm| self.reaches_a_window(arm.when) || self.reaches_a_window(arm.then))
|| otherwise.is_some_and(|child| self.reaches_a_window(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 merge_span(left: Span, right: Span) -> Span {
match (left.is_empty(), right.is_empty()) {
(true, _) => right,
(_, true) => left,
(false, false) => Span::new(left.start.min(right.start), left.end.max(right.end)),
}
}
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 every_node_and_expression_carries_a_span_through_an_in_place_rewrite() {
let mut plan = Plan::new();
let expr = plan.add_expr_at(
Expr::Column(ColumnBinding::new(7, 0)),
LogicalType::Integer,
Span::new(7, 12),
);
let node = plan
.add_node_at(Node::Filter { input: plan.root(), predicate: expr }, Span::new(0, 18));
let root = plan.root();
let exprs = plan.add_expr_list(&[expr]);
*plan.node_mut(node) = Node::Project { input: root, index: 9, exprs, names: Slice::EMPTY };
assert_eq!(plan.expr_span(expr), Span::new(7, 12));
assert_eq!(plan.node_span(node), Span::new(0, 18));
assert_eq!(plan.node_spans.len(), plan.nodes.len());
assert_eq!(plan.expr_spans.len(), plan.exprs.len());
}
#[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_window_function_outside_a_window_list_is_caught() {
let mut plan = Plan::new();
let name = plan.intern("row_number");
let call = plan.add_expr(
Expr::Window {
name,
args: Slice::EMPTY,
distinct: false,
filter: None,
ignore_nulls: false,
},
LogicalType::BigInt,
);
let exprs = plan.add_expr_list(&[call]);
let label = plan.intern("n");
let names = plan.add_name_list(&[label]);
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("window function outside"), "unhelpful message: {message}");
}
#[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_table_nobody_measured_reads_back_as_unknown_rather_than_as_zero() {
let plan = Plan::new();
assert_eq!(plan.measured(0), Stat::Unknown);
assert_eq!(plan.measured_count(), 0);
}
#[test]
fn what_the_binder_measured_comes_back_the_way_it_went_in() {
let mut plan = Plan::new();
plan.measure(4, Stat::exact(4096, rudb_common::Provenance::RowCount));
assert_eq!(plan.measured(4), Stat::exact(4096, rudb_common::Provenance::RowCount));
assert_eq!(plan.measured(5), Stat::Unknown);
assert_eq!(plan.measured_count(), 1);
}
#[test]
fn a_counted_column_comes_back_by_name_and_an_uncounted_one_comes_back_as_nothing() {
let mut plan = Plan::new();
plan.measure_distinct(4, "n_nationkey", 25);
assert_eq!(plan.distinct_measured(4, "n_nationkey"), Some(25));
assert_eq!(plan.distinct_measured(4, "n_name"), None);
assert_eq!(plan.distinct_measured(5, "n_nationkey"), None);
assert_eq!(plan.distinct_count(), 1);
}
#[derive(Debug)]
struct Stub;
impl Zones for Stub {
fn column(&self, name: &str) -> Option<usize> {
(name == "d").then_some(0)
}
fn surviving(&self, _tests: &[rudb_common::bounds::Test]) -> Option<u64> {
Some(42)
}
fn spread(
&self,
tests: &[rudb_common::bounds::Test],
) -> Option<rudb_common::bounds::Spread> {
Some(rudb_common::bounds::Spread { fraction: 0.5, read: tests.len() })
}
}
#[test]
fn a_table_with_no_bounds_recorded_answers_nothing_and_is_not_an_empty_set_of_bounds() {
let mut plan = Plan::new();
assert!(plan.zones(0).is_none());
assert_eq!(plan.zones_count(), 0);
plan.set_zones(3, Arc::new(Stub));
assert!(plan.zones(3).is_some());
assert!(plan.zones(4).is_none(), "the index has to match, two scans are two stores");
assert_eq!(plan.zones_count(), 1);
}
#[test]
fn the_store_recorded_for_a_table_is_the_one_that_answers_for_it() {
let mut plan = Plan::new();
plan.set_zones(0, Arc::new(Stub));
let zones = plan.zones(0).expect("just recorded");
assert_eq!(zones.column("d"), Some(0));
assert_eq!(zones.column("e"), None);
assert_eq!(zones.surviving(&[]), Some(42));
}
#[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);
}
}