use std::collections::BTreeMap;
use derive_more::{Display, Error, From};
use itertools::Itertools;
use tempest_core::tempest_str::TempestStr;
use tempest_io::Io;
use tempest_tql::ast::{
BinaryOpKind, CreateDatabaseStmt, CreateTableStmt, CreateTypeStmt, DeleteFromStmt,
Fields, InsertIntoStmt, ProjectionKind, SelectFromStmt, Stmt,
};
use crate::{
catalog::{
Catalog, CatalogError,
flatten_schema,
schema::{
DatabaseSchema, EnumSchema, EnumVariantDef, FieldDef, FieldId, StructSchema, TableId,
TableSchema, TypeExpr, VariantId,
},
},
query::{
eval::{CompiledExpr, compile_expr, eval_row},
resolve::{
ResolveError, resolve_database, resolve_table, resolve_type, resolve_type_schema,
},
},
types::{TempestType, TempestValue},
};
#[derive(Debug, Display, From, Error)]
pub enum PlanError {
#[display("catalog error: {}", _0)]
CatalogError(CatalogError),
#[display("resolve error: {}", _0)]
ResolveError(ResolveError),
#[from(skip)]
#[display("field with the name '{}' not found on type", _0)]
FieldNotFound(#[error(not(source))] TempestStr<'static>),
#[from(skip)]
#[display("field supplied named '{}' not found on type", _0)]
UnknownField(#[error(not(source))] TempestStr<'static>),
#[from(skip)]
#[display("missing field named '{}' for creating this type", _0)]
MissingField(#[error(not(source))] TempestStr<'static>),
#[display("type mismatch: expected value of type {}, but got value of type {}", expected.name(), got.name())]
TypeMismatch {
expected: TempestType,
got: TempestType,
},
#[from(skip)]
#[display("feature not supported yet: {}", _0)]
UnsupportedFeature(#[error(not(source))] &'static str),
#[from(skip)]
#[display("received column ref '{}', but expected literal", _0)]
ColumnRefInLiteral(#[error(not(source))] TempestStr<'static>),
#[from(skip)]
#[display("unsupported operand {:?} for values of type {} and {}", op, left.name(), right.name())]
UnsupportedOperand {
left: TempestType,
right: TempestType,
op: BinaryOpKind,
},
#[from(skip)]
#[display("expected type {} for predicate, but got type {}", TempestType::Bool.name(), _0.name())]
InvalidTypeForPredicate(#[error(not(source))] TempestType),
}
pub(crate) enum LogicalPlanNode {
CreateDatabase(DatabaseSchema),
CreateType(StructSchema),
CreateEnum(EnumSchema),
CreateTable(TableSchema),
Insert {
table_id: TableId,
row: Vec<TempestValue<'static>>,
},
Select {
table_id: TableId,
columns: Vec<(usize, TempestStr<'static>)>,
predicate: Option<CompiledExpr>,
},
Delete {
table_id: TableId,
predicate: Option<CompiledExpr>,
},
}
pub(crate) struct LogicalPlanner<'a, I: Io> {
catalog: &'a Catalog<I>,
}
impl<'a, I: Io> LogicalPlanner<'a, I> {
pub(crate) fn new(catalog: &'a Catalog<I>) -> Self {
Self { catalog }
}
fn plan_create_database(&self, stmt: CreateDatabaseStmt) -> Result<LogicalPlanNode, PlanError> {
let name = stmt.name.name.into_owned();
Ok(LogicalPlanNode::CreateDatabase(DatabaseSchema::new(name)))
}
fn plan_create_type(&self, stmt: CreateTypeStmt) -> Result<LogicalPlanNode, PlanError> {
let (db_id, _) = resolve_database(&stmt.path, self.catalog)?;
let generic_params = stmt
.generic_params
.params
.into_iter()
.map(|ident| ident.name.into_owned())
.collect_vec();
let generic_param_refs = generic_params.iter().map(|s| s).collect_vec();
let type_name = stmt.path.name().unwrap().name.clone().into_owned();
match stmt.fields {
Fields::Enum(variants) => {
let variants: BTreeMap<VariantId, EnumVariantDef> = variants
.into_iter()
.enumerate()
.map(|(i, v)| {
let fields = match v.fields {
Fields::Unit => vec![],
Fields::Unnamed(tys) => tys
.iter()
.map(|ty| resolve_type(ty, self.catalog, &generic_param_refs))
.collect::<Result<Vec<_>, _>>()?,
Fields::Named(named) => named
.iter()
.map(|f| resolve_type(&f.ty, self.catalog, &generic_param_refs))
.collect::<Result<Vec<_>, _>>()?,
Fields::Enum(_) => {
return Err(PlanError::UnsupportedFeature("nested enum variant"))
}
};
Ok((VariantId(i as u32), EnumVariantDef { name: v.name.name.into_owned(), fields }))
})
.collect::<Result<BTreeMap<_, _>, PlanError>>()?;
Ok(LogicalPlanNode::CreateEnum(EnumSchema {
database_id: Some(db_id),
name: type_name,
generic_params,
variants,
}))
}
non_enum => {
let fields = match non_enum {
Fields::Unit => BTreeMap::new(),
Fields::Named(fields) => fields
.into_iter()
.enumerate()
.map(|(i, field)| {
let ty = resolve_type(&field.ty, self.catalog, &generic_param_refs)?;
Ok::<_, PlanError>((
FieldId(i as u32),
FieldDef {
name: field.name.name.into_owned(),
ty,
},
))
})
.try_collect()?,
Fields::Unnamed(fields) => fields
.into_iter()
.enumerate()
.map(|(i, ty_expr)| {
let ty = resolve_type(&ty_expr, self.catalog, &generic_param_refs)?;
Ok::<_, PlanError>((
FieldId(i as u32),
FieldDef {
name: (i as u32).into(),
ty,
},
))
})
.try_collect()?,
Fields::Enum(_) => unreachable!(),
};
Ok(LogicalPlanNode::CreateType(StructSchema {
database_id: Some(db_id),
name: type_name,
generic_params,
fields,
}))
}
}
}
fn plan_create_table(&self, stmt: CreateTableStmt) -> Result<LogicalPlanNode, PlanError> {
let (db_id, _) = resolve_database(&stmt.path, self.catalog)?;
let (type_id, type_schema) = resolve_type_schema(&stmt.ty.path, self.catalog)?;
let struct_schema = type_schema
.as_struct()
.ok_or(PlanError::UnsupportedFeature("enum type as table backing type"))?;
let generic_args: Vec<_> = stmt
.ty
.generic_args
.args
.into_iter()
.map(|ty| resolve_type(&ty, self.catalog, &[]))
.try_collect()?;
let flat = flatten_schema(&struct_schema.fields, &generic_args, self.catalog, "")?;
let primary_key: Vec<Vec<FieldId>> = stmt
.body
.primary_key
.columns
.iter()
.map(|path| -> Result<Vec<FieldId>, PlanError> {
let field_ids =
resolve_pk_path(path, &struct_schema.fields, &generic_args, self.catalog)?;
crate::catalog::pk_path_to_flat_idx(
&field_ids,
&struct_schema.fields,
&generic_args,
self.catalog,
&flat,
)
.ok_or_else(|| PlanError::FieldNotFound(path.render().into_owned()))?;
Ok(field_ids)
})
.try_collect()?;
Ok(LogicalPlanNode::CreateTable(TableSchema {
database_id: db_id,
name: stmt.path.name().unwrap().name.clone().into_owned(),
type_id,
generic_args,
primary_key,
}))
}
pub(crate) fn plan_insert_into<'b>(
&self,
stmt: InsertIntoStmt<'b>,
) -> Result<LogicalPlanNode, PlanError> {
let table = resolve_table(&stmt.table, self.catalog)?;
let mut row = Vec::new();
eval_row(stmt.values, table.fields, table.generic_args, self.catalog, &mut row)?;
Ok(LogicalPlanNode::Insert {
table_id: table.id,
row,
})
}
pub(crate) fn plan_select_from(
&self,
stmt: SelectFromStmt,
) -> Result<LogicalPlanNode, PlanError> {
let table = resolve_table(&stmt.table, self.catalog)?;
let flat = flatten_schema(table.fields, table.generic_args, self.catalog, "")?;
let alias_str = stmt.alias.as_ref().map(|a| a.name.as_ref());
let columns: Vec<(usize, TempestStr<'static>)> = match stmt.projection.kind {
ProjectionKind::All => flat
.iter()
.enumerate()
.map(|(i, ff)| (i, ff.name.clone()))
.collect(),
ProjectionKind::Columns(col_refs) => col_refs
.iter()
.map(|col_ref| {
let lookup = resolve_path_lookup(col_ref, alias_str);
let pos = flat
.iter()
.position(|f| f.name.as_ref() == lookup.as_str())
.ok_or_else(|| {
PlanError::FieldNotFound(
TempestStr::from_owned(lookup.clone())
.expect("field name is valid"),
)
})?;
Ok((pos, flat[pos].name.clone()))
})
.collect::<Result<Vec<_>, PlanError>>()?,
};
let predicate = stmt
.selection
.map(|expr| {
let (expr, ty) = compile_expr(expr, &flat, alias_str, self.catalog)?;
if ty != TempestType::Bool {
return Err(PlanError::InvalidTypeForPredicate(ty));
}
Ok(expr)
})
.transpose()?;
Ok(LogicalPlanNode::Select {
table_id: table.id,
columns,
predicate,
})
}
pub(crate) fn plan_delete_from(
&self,
stmt: DeleteFromStmt,
) -> Result<LogicalPlanNode, PlanError> {
let table = resolve_table(&stmt.table, self.catalog)?;
let flat = flatten_schema(table.fields, table.generic_args, self.catalog, "")?;
let alias_str = stmt.alias.as_ref().map(|a| a.name.as_ref());
let predicate = stmt
.selection
.map(|expr| {
let (expr, ty) = compile_expr(expr, &flat, alias_str, self.catalog)?;
if ty != TempestType::Bool {
return Err(PlanError::InvalidTypeForPredicate(ty));
}
Ok(expr)
})
.transpose()?;
Ok(LogicalPlanNode::Delete {
table_id: table.id,
predicate,
})
}
pub(crate) fn plan(&self, stmt: Stmt<'_>) -> Result<LogicalPlanNode, PlanError> {
match stmt {
Stmt::CreateDatabase(stmt) => self.plan_create_database(stmt),
Stmt::CreateType(stmt) => self.plan_create_type(stmt),
Stmt::CreateTable(stmt) => self.plan_create_table(stmt),
Stmt::InsertInto(stmt) => self.plan_insert_into(stmt),
Stmt::SelectFrom(stmt) => self.plan_select_from(stmt),
Stmt::DeleteFrom(stmt) => self.plan_delete_from(stmt),
}
}
}
fn resolve_pk_path(
path: &tempest_tql::ast::Path<'_>,
fields: &BTreeMap<FieldId, crate::catalog::schema::FieldDef>,
generic_args: &[TypeExpr],
catalog: &crate::catalog::CatalogState,
) -> Result<Vec<FieldId>, PlanError> {
use tempest_tql::ast::PathSegment;
use std::borrow::Cow;
use crate::catalog::schema::StructSchema;
let mut result = Vec::new();
let mut current_fields = fields;
let mut current_args: Cow<[TypeExpr]> = Cow::Borrowed(generic_args);
let mut _schema_store: Option<StructSchema> = None;
for (i, segment) in path.segments.iter().enumerate() {
let (fid, def) = match segment {
PathSegment::Named(ident) => current_fields
.iter()
.find(|(_, def)| def.name == ident.name)
.map(|(fid, def)| (*fid, def))
.ok_or_else(|| PlanError::FieldNotFound(ident.name.clone().into_owned()))?,
&PathSegment::Indexed(idx) => current_fields
.iter()
.nth(idx as usize)
.map(|(fid, def)| (*fid, def))
.ok_or_else(|| PlanError::FieldNotFound(idx.into()))?,
};
result.push(fid);
if i < path.segments.len() - 1 {
let resolved = match &def.ty {
TypeExpr::GenericParam(idx) => current_args
.get(*idx as usize)
.cloned()
.ok_or_else(|| PlanError::FieldNotFound(path.render().into_owned()))?,
other => other.clone(),
};
match resolved {
TypeExpr::Ref(type_id, ref_args) => {
_schema_store = catalog
.get_type(type_id)
.and_then(|ts| ts.as_struct().cloned());
current_fields = &_schema_store
.as_ref()
.ok_or_else(|| PlanError::FieldNotFound(path.render().into_owned()))?
.fields;
current_args = Cow::Owned(ref_args);
}
_ => return Err(PlanError::FieldNotFound(path.render().into_owned())),
}
}
}
Ok(result)
}
fn resolve_path_lookup(path: &tempest_tql::ast::Path<'_>, alias: Option<&str>) -> String {
if let Some(alias) = alias {
if let tempest_tql::ast::PathSegment::Named(first) = &path.segments[0] {
if first.name.as_ref() == alias && path.segments.len() > 1 {
return path.segments[1..]
.iter()
.map(|s| s.to_string())
.collect::<Vec<_>>()
.join(".");
}
}
}
path.render().to_string()
}