query-engine/src/planner/sql.rs

use std::error::Error;
use std::fmt;

use crate::catalog::{CatalogError, PredicateCatalog};
use crate::planner::logical::{LogicalExpr, LogicalPlan, NamedExpr};
use crate::relational::{Field, Schema, Value};
use crate::sql::ast::{BinaryOp, Expr, Literal, Select, SelectItem};

#[derive(Debug)]
pub enum PlannerError {
    Catalog(CatalogError),
    UnknownColumn(String),
    UnsupportedProjection,
}

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),
            Self::UnsupportedProjection => {
                write!(f, "only wildcard and column projections are supported")
            }
        }
    }
}

impl Error for PlannerError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Catalog(err) => Some(err),
            Self::UnknownColumn(_) | Self::UnsupportedProjection => None,
        }
    }
}

impl From<CatalogError> for PlannerError {
    fn from(value: CatalogError) -> Self {
        Self::Catalog(value)
    }
}

pub fn plan_select(
    select: &Select,
    catalog: &PredicateCatalog,
) -> Result<LogicalPlan, PlannerError> {
    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 item in &select.projection {
        match item {
            SelectItem::Expr { expr, alias } => match expr {
                Expr::Identifier(name) => {
                    let index = scan_schema
                        .index_of(name)
                        .ok_or_else(|| PlannerError::UnknownColumn(name.clone()))?;
                    let input_field = &scan_schema.fields()[index];
                    let output_name = alias.clone().unwrap_or_else(|| name.clone());
                    expressions.push(NamedExpr {
                        name: output_name.clone(),
                        expr: LogicalExpr::Column(name.clone()),
                    });
                    fields.push(Field::new(
                        output_name,
                        input_field.data_type().clone(),
                        input_field.nullable(),
                    ));
                }
                _ => return Err(PlannerError::UnsupportedProjection),
            },
            SelectItem::Wildcard => return Err(PlannerError::UnsupportedProjection),
        }
    }

    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<LogicalExpr, PlannerError> {
    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,
    }
}

#[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);
    }
}