use std::error::Error; use std::fmt; use crate::catalog::{CatalogError, PredicateCatalog}; use crate::planner::logical::{LogicalExpr, LogicalPlan, NamedExpr}; use crate::relational::{DataType, Field, Schema, Value}; use crate::sql::ast::{BinaryOp, Expr, Literal, Select, SelectItem}; /// Errors returned when translating SQL AST into a logical plan. #[derive(Debug)] pub enum PlannerError { /// Catalog lookup failed. Catalog(CatalogError), /// A referenced column does not exist in the input schema. UnknownColumn(String), } impl fmt::Display for PlannerError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Catalog(err) => write!(f, "catalog error: {}", err), Self::UnknownColumn(column) => write!(f, "unknown column `{}`", column), } } } impl Error for PlannerError { fn source(&self) -> Option<&(dyn Error + 'static)> { match self { Self::Catalog(err) => Some(err), Self::UnknownColumn(_) => None, } } } impl From for PlannerError { fn from(value: CatalogError) -> Self { Self::Catalog(value) } } /// Plan a parsed `SELECT` statement into the current logical plan subset. pub fn plan_select( select: &Select, catalog: &PredicateCatalog, ) -> Result { let scan_schema = catalog.schema_for(&select.from)?.clone(); let mut plan = LogicalPlan::Scan { table: select.from.clone(), schema: scan_schema.clone(), }; if let Some(selection) = &select.selection { let predicate = plan_expr(selection, &scan_schema)?; plan = LogicalPlan::Filter { input: Box::new(plan), predicate, }; } if is_wildcard_projection(&select.projection) { return Ok(plan); } let mut expressions = Vec::new(); let mut fields = Vec::new(); for (index, item) in select.projection.iter().enumerate() { match item { SelectItem::Expr { expr, alias } => { let planned_expr = plan_expr(expr, &scan_schema)?; let output_name = alias .clone() .unwrap_or_else(|| default_projection_name(expr, index + 1)); let (data_type, nullable) = projection_metadata(expr, &scan_schema)?; expressions.push(NamedExpr { name: output_name.clone(), expr: planned_expr, }); fields.push(Field::new(output_name, data_type, nullable)); } SelectItem::Wildcard => unreachable!("wildcard projections are handled earlier"), } } Ok(LogicalPlan::Project { input: Box::new(plan), expressions, schema: Schema::new(fields), }) } fn is_wildcard_projection(items: &[SelectItem]) -> bool { matches!(items, [SelectItem::Wildcard]) } fn plan_expr(expr: &Expr, schema: &Schema) -> Result { match expr { Expr::Identifier(name) => { if schema.index_of(name).is_none() { return Err(PlannerError::UnknownColumn(name.clone())); } Ok(LogicalExpr::Column(name.clone())) } Expr::Literal(literal) => Ok(LogicalExpr::Literal(plan_literal(literal))), Expr::Binary { left, op, right } => match op { BinaryOp::Eq => Ok(LogicalExpr::Eq( Box::new(plan_expr(left, schema)?), Box::new(plan_expr(right, schema)?), )), }, } } fn plan_literal(literal: &Literal) -> Value { match literal { Literal::String(value) => Value::text(value.clone()), Literal::Null => Value::Null, } } fn projection_metadata(expr: &Expr, schema: &Schema) -> Result<(DataType, bool), PlannerError> { match expr { Expr::Identifier(name) => { let index = schema .index_of(name) .ok_or_else(|| PlannerError::UnknownColumn(name.clone()))?; let field = &schema.fields()[index]; Ok((field.data_type().clone(), field.nullable())) } Expr::Literal(Literal::String(_)) => Ok((DataType::Text, false)), Expr::Literal(Literal::Null) => Ok((DataType::Text, true)), Expr::Binary { .. } => Ok((DataType::Boolean, true)), } } fn default_projection_name(expr: &Expr, ordinal: usize) -> String { match expr { Expr::Identifier(name) => name.clone(), Expr::Literal(_) | Expr::Binary { .. } => format!("expr{}", ordinal), } } #[cfg(test)] mod tests { use super::*; use crate::catalog::PredicateCatalog; use crate::chase::{Atom, Instance, Term}; use crate::sql::parser::parse_select; #[test] fn plans_projection_and_filter() { let instance: Instance = vec![Atom::new( "Parent", vec![Term::constant("alice"), Term::constant("bob")], )] .into_iter() .collect(); let catalog = PredicateCatalog::from_instance(&instance).unwrap(); let select = parse_select("SELECT c0 FROM Parent WHERE c1 = 'bob'").unwrap(); let plan = plan_select(&select, &catalog).unwrap(); assert_eq!(plan.output_schema().len(), 1); } #[test] fn plans_aliases_and_literal_projection() { let instance: Instance = vec![Atom::new( "Parent", vec![Term::constant("alice"), Term::constant("bob")], )] .into_iter() .collect(); let catalog = PredicateCatalog::from_instance(&instance).unwrap(); let select = parse_select("SELECT c0 AS parent_name, 'seed' AS label, NULL FROM Parent").unwrap(); let plan = plan_select(&select, &catalog).unwrap(); let schema = plan.output_schema(); assert_eq!(schema.len(), 3); assert_eq!(schema.fields()[0].name(), "parent_name"); assert_eq!(schema.fields()[1].name(), "label"); assert_eq!(schema.fields()[2].name(), "expr3"); assert_eq!(schema.fields()[1].data_type(), &DataType::Text); } }