alopex_sql/planner/
mod.rs

1//! Query planning module for the Alopex SQL dialect.
2//!
3//! This module provides:
4//! - [`PlannerError`]: Error types for planning phase
5//! - [`ResolvedType`]: Normalized type information for type checking
6//! - [`TypedExpr`]: Type-checked expressions with resolved types
7//! - [`LogicalPlan`]: Logical query plan representation
8//! - [`NameResolver`]: Table and column reference resolution
9//! - [`TypeChecker`]: Expression type inference and validation
10//! - [`Planner`]: Main entry point for converting AST to LogicalPlan
11
12mod error;
13pub mod knn_optimizer;
14pub mod logical_plan;
15pub mod name_resolver;
16pub mod type_checker;
17pub mod typed_expr;
18pub mod types;
19
20#[cfg(test)]
21mod planner_tests;
22
23pub use error::PlannerError;
24pub use knn_optimizer::{KnnPattern, SortDirection, detect_knn_pattern};
25pub use logical_plan::LogicalPlan;
26pub use name_resolver::{NameResolver, ResolvedColumn};
27pub use type_checker::TypeChecker;
28pub use typed_expr::{
29    ProjectedColumn, Projection, SortExpr, TypedAssignment, TypedExpr, TypedExprKind,
30};
31pub use types::ResolvedType;
32
33use crate::ast::ddl::{
34    ColumnConstraint, ColumnDef, CreateIndex, CreateTable, DropIndex, DropTable,
35};
36use crate::ast::dml::{Delete, Insert, OrderByExpr, Select, SelectItem, Update};
37use crate::ast::expr::Literal;
38use crate::ast::{Statement, StatementKind};
39use crate::catalog::{Catalog, ColumnMetadata, IndexMetadata, TableMetadata};
40use crate::{DataSourceFormat, TableType};
41use std::collections::HashMap;
42
43/// The SQL query planner.
44///
45/// The planner converts AST statements into logical plans. It performs:
46/// - Name resolution: Validates table and column references
47/// - Type checking: Infers and validates expression types
48/// - Plan construction: Builds the logical plan tree
49///
50/// # Design Notes
51///
52/// - The planner uses an immutable reference to the catalog (`&C`)
53/// - DDL statements produce plans but don't modify the catalog
54/// - The executor is responsible for applying catalog changes
55///
56/// # Examples
57///
58/// ```
59/// use alopex_sql::catalog::MemoryCatalog;
60/// use alopex_sql::planner::Planner;
61///
62/// let catalog = MemoryCatalog::new();
63/// let planner = Planner::new(&catalog);
64///
65/// // Parse and plan a statement
66/// // let stmt = parser.parse("SELECT * FROM users")?;
67/// // let plan = planner.plan(&stmt)?;
68/// ```
69pub struct Planner<'a, C: Catalog> {
70    catalog: &'a C,
71    name_resolver: NameResolver<'a, C>,
72    type_checker: TypeChecker<'a, C>,
73}
74
75impl<'a, C: Catalog> Planner<'a, C> {
76    /// Create a new planner with the given catalog.
77    pub fn new(catalog: &'a C) -> Self {
78        Self {
79            catalog,
80            name_resolver: NameResolver::new(catalog),
81            type_checker: TypeChecker::new(catalog),
82        }
83    }
84
85    /// Plan a SQL statement.
86    ///
87    /// This is the main entry point for converting an AST statement into a logical plan.
88    ///
89    /// # Errors
90    ///
91    /// Returns a `PlannerError` if:
92    /// - Referenced tables or columns don't exist
93    /// - Type checking fails
94    /// - DDL validation fails (e.g., table already exists for CREATE TABLE)
95    pub fn plan(&self, stmt: &Statement) -> Result<LogicalPlan, PlannerError> {
96        match &stmt.kind {
97            // DDL statements
98            StatementKind::CreateTable(ct) => self.plan_create_table(ct),
99            StatementKind::DropTable(dt) => self.plan_drop_table(dt),
100            StatementKind::CreateIndex(ci) => self.plan_create_index(ci),
101            StatementKind::DropIndex(di) => self.plan_drop_index(di),
102
103            // DML statements
104            StatementKind::Select(sel) => self.plan_select(sel),
105            StatementKind::Insert(ins) => self.plan_insert(ins),
106            StatementKind::Update(upd) => self.plan_update(upd),
107            StatementKind::Delete(del) => self.plan_delete(del),
108        }
109    }
110
111    // ============================================================
112    // DDL Planning Methods (Task 16)
113    // ============================================================
114
115    /// Plan a CREATE TABLE statement.
116    ///
117    /// Validates that the table doesn't already exist (unless IF NOT EXISTS is specified),
118    /// and converts the AST column definitions to catalog metadata.
119    fn plan_create_table(&self, stmt: &CreateTable) -> Result<LogicalPlan, PlannerError> {
120        // Check if table already exists
121        if !stmt.if_not_exists && self.catalog.table_exists(&stmt.name) {
122            return Err(PlannerError::table_already_exists(&stmt.name));
123        }
124
125        // Convert column definitions to metadata
126        let columns: Vec<ColumnMetadata> = stmt
127            .columns
128            .iter()
129            .map(|col| self.convert_column_def(col))
130            .collect();
131
132        // Collect primary key from table constraints
133        let primary_key = Self::extract_primary_key(stmt);
134
135        // Build table metadata
136        // Note: table_id defaults to 0 as placeholder; Executor assigns the actual ID
137        let mut table = TableMetadata::new(stmt.name.clone(), columns);
138        if let Some(pk) = primary_key {
139            table = table.with_primary_key(pk);
140        }
141        table.catalog_name = "default".to_string();
142        table.namespace_name = "default".to_string();
143        table.table_type = TableType::Managed;
144        table.data_source_format = DataSourceFormat::Alopex;
145        table.properties = HashMap::new();
146
147        Ok(LogicalPlan::CreateTable {
148            table,
149            if_not_exists: stmt.if_not_exists,
150            with_options: stmt.with_options.clone(),
151        })
152    }
153
154    /// Convert an AST column definition to catalog column metadata.
155    fn convert_column_def(&self, col: &ColumnDef) -> ColumnMetadata {
156        let data_type = ResolvedType::from_ast(&col.data_type);
157        let mut meta = ColumnMetadata::new(col.name.clone(), data_type);
158
159        // Process constraints
160        for constraint in &col.constraints {
161            meta = Self::apply_column_constraint(meta, constraint);
162        }
163
164        meta
165    }
166
167    /// Apply a column constraint to column metadata.
168    fn apply_column_constraint(
169        mut meta: ColumnMetadata,
170        constraint: &ColumnConstraint,
171    ) -> ColumnMetadata {
172        match constraint {
173            ColumnConstraint::NotNull => {
174                meta.not_null = true;
175            }
176            ColumnConstraint::Null => {
177                meta.not_null = false;
178            }
179            ColumnConstraint::PrimaryKey => {
180                meta.primary_key = true;
181                meta.not_null = true; // PRIMARY KEY implies NOT NULL
182            }
183            ColumnConstraint::Unique => {
184                meta.unique = true;
185            }
186            ColumnConstraint::Default(expr) => {
187                meta.default = Some(expr.clone());
188            }
189            ColumnConstraint::WithSpan { kind, .. } => {
190                meta = Self::apply_column_constraint(meta, kind);
191            }
192        }
193        meta
194    }
195
196    /// Extract primary key columns from table constraints.
197    fn extract_primary_key(stmt: &CreateTable) -> Option<Vec<String>> {
198        use crate::ast::ddl::TableConstraint;
199
200        // First check table-level constraints
201        // Note: Currently only PrimaryKey variant exists; when more variants are added,
202        // this should iterate to find the first PrimaryKey constraint
203        if let Some(TableConstraint::PrimaryKey { columns, .. }) = stmt.constraints.first() {
204            return Some(columns.clone());
205        }
206
207        // Then check column-level PRIMARY KEY constraints
208        let pk_columns: Vec<String> = stmt
209            .columns
210            .iter()
211            .filter(|col| col.constraints.iter().any(Self::is_primary_key_constraint))
212            .map(|col| col.name.clone())
213            .collect();
214
215        if pk_columns.is_empty() {
216            None
217        } else {
218            Some(pk_columns)
219        }
220    }
221
222    /// Check if a column constraint is a PRIMARY KEY constraint.
223    fn is_primary_key_constraint(constraint: &ColumnConstraint) -> bool {
224        match constraint {
225            ColumnConstraint::PrimaryKey => true,
226            ColumnConstraint::WithSpan { kind, .. } => Self::is_primary_key_constraint(kind),
227            _ => false,
228        }
229    }
230
231    /// Plan a DROP TABLE statement.
232    ///
233    /// Validates that the table exists (unless IF EXISTS is specified).
234    fn plan_drop_table(&self, stmt: &DropTable) -> Result<LogicalPlan, PlannerError> {
235        // Check if table exists
236        if !stmt.if_exists && !self.table_exists_in_default(&stmt.name) {
237            return Err(PlannerError::TableNotFound {
238                name: stmt.name.clone(),
239                line: stmt.span.start.line,
240                column: stmt.span.start.column,
241            });
242        }
243
244        Ok(LogicalPlan::DropTable {
245            name: stmt.name.clone(),
246            if_exists: stmt.if_exists,
247        })
248    }
249
250    fn table_exists_in_default(&self, name: &str) -> bool {
251        match self.catalog.get_table(name) {
252            Some(table) => table.catalog_name == "default" && table.namespace_name == "default",
253            None => false,
254        }
255    }
256
257    /// Plan a CREATE INDEX statement.
258    ///
259    /// Validates that:
260    /// - The index doesn't already exist (unless IF NOT EXISTS is specified)
261    /// - The target table exists
262    /// - The target column exists in the table
263    fn plan_create_index(&self, stmt: &CreateIndex) -> Result<LogicalPlan, PlannerError> {
264        // Check if index already exists
265        if !stmt.if_not_exists && self.catalog.index_exists(&stmt.name) {
266            return Err(PlannerError::index_already_exists(&stmt.name));
267        }
268
269        // Validate table exists
270        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
271
272        // Validate column exists
273        self.name_resolver
274            .resolve_column(table, &stmt.column, stmt.span)?;
275
276        // Build index metadata
277        // Note: index_id is set to 0 as placeholder; Executor assigns the actual ID
278        // Note: column_indices will be resolved by Executor when table schema is available
279        let mut index = IndexMetadata::new(
280            0,
281            stmt.name.clone(),
282            stmt.table.clone(),
283            vec![stmt.column.clone()],
284        );
285
286        if let Some(method) = stmt.method {
287            index = index.with_method(method);
288        }
289
290        let options: Vec<(String, String)> = stmt
291            .options
292            .iter()
293            .map(|opt| (opt.key.clone(), opt.value.clone()))
294            .collect();
295        if !options.is_empty() {
296            index = index.with_options(options);
297        }
298
299        Ok(LogicalPlan::CreateIndex {
300            index,
301            if_not_exists: stmt.if_not_exists,
302        })
303    }
304
305    /// Plan a DROP INDEX statement.
306    ///
307    /// Validates that the index exists (unless IF EXISTS is specified).
308    fn plan_drop_index(&self, stmt: &DropIndex) -> Result<LogicalPlan, PlannerError> {
309        // Check if index exists
310        if !stmt.if_exists && !self.index_exists_in_default(&stmt.name) {
311            return Err(PlannerError::index_not_found(&stmt.name));
312        }
313
314        Ok(LogicalPlan::DropIndex {
315            name: stmt.name.clone(),
316            if_exists: stmt.if_exists,
317        })
318    }
319
320    fn index_exists_in_default(&self, name: &str) -> bool {
321        match self.catalog.get_index(name) {
322            Some(index) => index.catalog_name == "default" && index.namespace_name == "default",
323            None => false,
324        }
325    }
326
327    // ============================================================
328    // DML Planning Methods (Task 17 & 18)
329    // ============================================================
330
331    /// Plan a SELECT statement.
332    ///
333    /// Builds a logical plan tree: Scan -> Filter -> Sort -> Limit
334    /// Each layer is optional and only added if the corresponding clause is present.
335    fn plan_select(&self, stmt: &Select) -> Result<LogicalPlan, PlannerError> {
336        // 1. Resolve the FROM table
337        let table = self
338            .name_resolver
339            .resolve_table(&stmt.from.name, stmt.from.span)?;
340
341        // 2. Build the projection
342        let projection = self.build_projection(&stmt.projection, table)?;
343
344        // 3. Create the base Scan plan
345        let mut plan = LogicalPlan::Scan {
346            table: table.name.clone(),
347            projection,
348        };
349
350        // 4. Add Filter if WHERE clause is present
351        if let Some(ref selection) = stmt.selection {
352            let predicate = self.type_checker.infer_type(selection, table)?;
353
354            // Verify predicate returns Boolean
355            if predicate.resolved_type != ResolvedType::Boolean {
356                return Err(PlannerError::type_mismatch(
357                    "Boolean",
358                    predicate.resolved_type.to_string(),
359                    selection.span,
360                ));
361            }
362
363            plan = LogicalPlan::Filter {
364                input: Box::new(plan),
365                predicate,
366            };
367        }
368
369        // 5. Add Sort if ORDER BY clause is present
370        if !stmt.order_by.is_empty() {
371            let order_by = self.build_sort_exprs(&stmt.order_by, table)?;
372            plan = LogicalPlan::Sort {
373                input: Box::new(plan),
374                order_by,
375            };
376        }
377
378        // 6. Add Limit if LIMIT/OFFSET is present
379        if stmt.limit.is_some() || stmt.offset.is_some() {
380            let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
381            let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
382            plan = LogicalPlan::Limit {
383                input: Box::new(plan),
384                limit,
385                offset,
386            };
387        }
388
389        Ok(plan)
390    }
391
392    /// Build the projection for a SELECT statement.
393    ///
394    /// Handles wildcard expansion and expression type checking.
395    fn build_projection(
396        &self,
397        items: &[SelectItem],
398        table: &TableMetadata,
399    ) -> Result<Projection, PlannerError> {
400        // Check for wildcard - if present, expand it
401        if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
402            let columns = self.name_resolver.expand_wildcard(table);
403            return Ok(Projection::All(columns));
404        }
405
406        // Process each select item
407        let mut projected_columns = Vec::new();
408        for item in items {
409            match item {
410                SelectItem::Wildcard { span } => {
411                    // Wildcard mixed with other items - expand inline
412                    for col in &table.columns {
413                        let column_index = table.get_column_index(&col.name).unwrap();
414                        let typed_expr = TypedExpr::column_ref(
415                            table.name.clone(),
416                            col.name.clone(),
417                            column_index,
418                            col.data_type.clone(),
419                            *span,
420                        );
421                        projected_columns.push(ProjectedColumn::new(typed_expr));
422                    }
423                }
424                SelectItem::Expr { expr, alias, .. } => {
425                    let typed_expr = self.type_checker.infer_type(expr, table)?;
426                    let projected = if let Some(alias) = alias {
427                        ProjectedColumn::with_alias(typed_expr, alias.clone())
428                    } else {
429                        ProjectedColumn::new(typed_expr)
430                    };
431                    projected_columns.push(projected);
432                }
433            }
434        }
435
436        Ok(Projection::Columns(projected_columns))
437    }
438
439    /// Build sort expressions from ORDER BY clause.
440    fn build_sort_exprs(
441        &self,
442        order_by: &[OrderByExpr],
443        table: &TableMetadata,
444    ) -> Result<Vec<SortExpr>, PlannerError> {
445        let mut sort_exprs = Vec::new();
446
447        for order_expr in order_by {
448            let typed_expr = self.type_checker.infer_type(&order_expr.expr, table)?;
449
450            // Determine sort direction (default: ASC)
451            let asc = order_expr.asc.unwrap_or(true);
452
453            // Determine NULLS ordering (default: NULLS LAST for both ASC and DESC)
454            let nulls_first = order_expr.nulls_first.unwrap_or(false);
455
456            sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
457        }
458
459        Ok(sort_exprs)
460    }
461
462    /// Extract a numeric value from a LIMIT or OFFSET expression.
463    ///
464    /// Currently only supports literal integer values.
465    fn extract_limit_value(
466        &self,
467        expr: &Option<crate::ast::expr::Expr>,
468        stmt_span: crate::ast::Span,
469    ) -> Result<Option<u64>, PlannerError> {
470        match expr {
471            None => Ok(None),
472            Some(e) => {
473                // For now, only support literal integers
474                if let crate::ast::expr::ExprKind::Literal(Literal::Number(s)) = &e.kind {
475                    s.parse::<u64>().map(Some).map_err(|_| {
476                        PlannerError::type_mismatch("unsigned integer", s.clone(), e.span)
477                    })
478                } else {
479                    Err(PlannerError::unsupported_feature(
480                        "non-literal LIMIT/OFFSET",
481                        "v0.3.0+",
482                        stmt_span,
483                    ))
484                }
485            }
486        }
487    }
488
489    /// Plan an INSERT statement.
490    ///
491    /// Handles column list specification or implicit column ordering.
492    /// When columns are omitted, uses table definition order from TableMetadata.
493    fn plan_insert(&self, stmt: &Insert) -> Result<LogicalPlan, PlannerError> {
494        // Resolve the target table
495        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
496
497        // Determine the column list
498        let columns: Vec<String> = if let Some(ref cols) = stmt.columns {
499            // Explicit column list - validate each column exists
500            for col in cols {
501                self.name_resolver.resolve_column(table, col, stmt.span)?;
502            }
503            cols.clone()
504        } else {
505            // Implicit - use all columns in table definition order
506            table.column_names().into_iter().map(String::from).collect()
507        };
508
509        // Validate and type-check each row of values
510        let mut typed_values: Vec<Vec<TypedExpr>> = Vec::new();
511
512        for row in &stmt.values {
513            // Check column count matches
514            if row.len() != columns.len() {
515                return Err(PlannerError::column_value_count_mismatch(
516                    columns.len(),
517                    row.len(),
518                    stmt.span,
519                ));
520            }
521
522            // Type-check each value
523            let typed_row = self.type_check_insert_values(row, &columns, table)?;
524            typed_values.push(typed_row);
525        }
526
527        Ok(LogicalPlan::Insert {
528            table: table.name.clone(),
529            columns,
530            values: typed_values,
531        })
532    }
533
534    /// Type-check INSERT values against column definitions.
535    fn type_check_insert_values(
536        &self,
537        values: &[crate::ast::expr::Expr],
538        columns: &[String],
539        table: &TableMetadata,
540    ) -> Result<Vec<TypedExpr>, PlannerError> {
541        let mut typed_values = Vec::new();
542
543        for (i, value) in values.iter().enumerate() {
544            let column_name = &columns[i];
545            let column_meta = table.get_column(column_name).ok_or_else(|| {
546                PlannerError::column_not_found(column_name, &table.name, value.span)
547            })?;
548
549            // Type-check the value expression
550            let typed_value = self.type_checker.infer_type(value, table)?;
551
552            // Check for NOT NULL constraint violation (except for NULL literal which is allowed if nullable)
553            if column_meta.not_null
554                && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
555            {
556                return Err(PlannerError::null_constraint_violation(
557                    column_name,
558                    value.span,
559                ));
560            }
561
562            // Validate type compatibility
563            self.validate_type_assignment(&typed_value, &column_meta.data_type, value.span)?;
564
565            typed_values.push(typed_value);
566        }
567
568        Ok(typed_values)
569    }
570
571    /// Validate that a value type can be assigned to a column type.
572    fn validate_type_assignment(
573        &self,
574        value: &TypedExpr,
575        target_type: &ResolvedType,
576        span: crate::ast::Span,
577    ) -> Result<(), PlannerError> {
578        // NULL can be assigned to any nullable column
579        if value.resolved_type == ResolvedType::Null {
580            return Ok(());
581        }
582
583        // Check for exact match or implicit conversion compatibility
584        if self.types_compatible(&value.resolved_type, target_type) {
585            return Ok(());
586        }
587
588        Err(PlannerError::type_mismatch(
589            target_type.to_string(),
590            value.resolved_type.to_string(),
591            span,
592        ))
593    }
594
595    /// Check if two types are compatible for assignment.
596    fn types_compatible(&self, source: &ResolvedType, target: &ResolvedType) -> bool {
597        use ResolvedType::*;
598
599        // Same type is always compatible
600        if source == target {
601            return true;
602        }
603
604        // Numeric promotions
605        match (source, target) {
606            // Integer can be assigned to BigInt, Float, Double
607            (Integer, BigInt) | (Integer, Float) | (Integer, Double) => true,
608            // BigInt can be assigned to Float, Double
609            (BigInt, Float) | (BigInt, Double) => true,
610            // Float can be assigned to Double
611            (Float, Double) => true,
612            // Vector dimensions must match
613            (Vector { dimension: d1, .. }, Vector { dimension: d2, .. }) => d1 == d2,
614            _ => false,
615        }
616    }
617
618    /// Plan an UPDATE statement.
619    ///
620    /// Validates assignments and optional WHERE clause.
621    fn plan_update(&self, stmt: &Update) -> Result<LogicalPlan, PlannerError> {
622        // Resolve the target table
623        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
624
625        // Process assignments
626        let mut typed_assignments = Vec::new();
627
628        for assignment in &stmt.assignments {
629            // Resolve the column
630            let column_meta =
631                self.name_resolver
632                    .resolve_column(table, &assignment.column, assignment.span)?;
633            let column_index = table.get_column_index(&assignment.column).unwrap();
634
635            // Type-check the value expression
636            let typed_value = self.type_checker.infer_type(&assignment.value, table)?;
637
638            // Check NOT NULL constraint
639            if column_meta.not_null
640                && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
641            {
642                return Err(PlannerError::null_constraint_violation(
643                    &assignment.column,
644                    assignment.value.span,
645                ));
646            }
647
648            // Validate type compatibility
649            self.validate_type_assignment(
650                &typed_value,
651                &column_meta.data_type,
652                assignment.value.span,
653            )?;
654
655            typed_assignments.push(TypedAssignment::new(
656                assignment.column.clone(),
657                column_index,
658                typed_value,
659            ));
660        }
661
662        // Process optional WHERE clause
663        let filter = if let Some(ref selection) = stmt.selection {
664            let predicate = self.type_checker.infer_type(selection, table)?;
665
666            // Verify predicate returns Boolean
667            if predicate.resolved_type != ResolvedType::Boolean {
668                return Err(PlannerError::type_mismatch(
669                    "Boolean",
670                    predicate.resolved_type.to_string(),
671                    selection.span,
672                ));
673            }
674
675            Some(predicate)
676        } else {
677            None
678        };
679
680        Ok(LogicalPlan::Update {
681            table: table.name.clone(),
682            assignments: typed_assignments,
683            filter,
684        })
685    }
686
687    /// Plan a DELETE statement.
688    ///
689    /// Validates optional WHERE clause.
690    fn plan_delete(&self, stmt: &Delete) -> Result<LogicalPlan, PlannerError> {
691        // Resolve the target table
692        let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
693
694        // Process optional WHERE clause
695        let filter = if let Some(ref selection) = stmt.selection {
696            let predicate = self.type_checker.infer_type(selection, table)?;
697
698            // Verify predicate returns Boolean
699            if predicate.resolved_type != ResolvedType::Boolean {
700                return Err(PlannerError::type_mismatch(
701                    "Boolean",
702                    predicate.resolved_type.to_string(),
703                    selection.span,
704                ));
705            }
706
707            Some(predicate)
708        } else {
709            None
710        };
711
712        Ok(LogicalPlan::Delete {
713            table: table.name.clone(),
714            filter,
715        })
716    }
717}