1pub mod aggregate_expr;
13mod error;
14pub mod knn_optimizer;
15pub mod logical_plan;
16pub mod name_resolver;
17pub mod type_checker;
18pub mod typed_expr;
19pub mod types;
20
21#[cfg(test)]
22mod planner_tests;
23
24pub use aggregate_expr::{AggregateExpr, AggregateFunction};
25pub use error::PlannerError;
26pub use knn_optimizer::{KnnPattern, SortDirection, detect_knn_pattern};
27pub use logical_plan::{JoinType, LogicalPlan};
28pub use name_resolver::{NameResolver, ResolvedColumn};
29pub use type_checker::{ScopedTable, TypeChecker};
30pub use typed_expr::{
31 ProjectedColumn, Projection, SortExpr, TypedAssignment, TypedExpr, TypedExprKind,
32};
33pub use types::ResolvedType;
34
35use crate::ast::ddl::{
36 ColumnConstraint, ColumnDef, CreateIndex, CreateTable, DropIndex, DropTable,
37};
38use crate::ast::dml::{
39 Delete, FromItem, Insert, InsertSource, LITERAL_TABLE, OrderByExpr, Select, SelectItem, Update,
40};
41use crate::ast::expr::Literal;
42use crate::ast::{PragmaValue, Spanned, Statement, StatementKind};
43use crate::catalog::{Catalog, ColumnMetadata, IndexMetadata, TableMetadata};
44use crate::{AlopexDialect, DataSourceFormat, Parser, SqlError, TableType};
45use std::collections::HashMap;
46
47struct PlannedRelation {
48 plan: LogicalPlan,
49 schema: Vec<ColumnMetadata>,
50 scope: Vec<ScopedTable>,
51}
52
53#[derive(Debug, Clone)]
59pub struct PlannedStatement {
60 pub plan: LogicalPlan,
62 pub routing_input: RoutingInput,
64}
65
66impl PlannedStatement {
67 pub fn statement_kind(&self) -> &StatementKind {
69 &self.routing_input.statement_kind
70 }
71
72 pub fn table_references(&self) -> &[TableReference] {
74 &self.routing_input.table_references
75 }
76
77 pub fn diagnostics(&self) -> &[PlanningDiagnostic] {
80 &self.routing_input.diagnostics
81 }
82}
83
84#[derive(Debug, Clone)]
86pub struct RoutingInput {
87 pub statement_kind: StatementKind,
90 pub table_references: Vec<TableReference>,
92 pub diagnostics: Vec<PlanningDiagnostic>,
94}
95
96#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct TableReference {
99 pub table_name: String,
101 pub access: TableReferenceAccess,
103 pub source: TableReferenceSource,
105}
106
107impl TableReference {
108 pub fn new(
109 table_name: impl Into<String>,
110 access: TableReferenceAccess,
111 source: TableReferenceSource,
112 ) -> Self {
113 Self {
114 table_name: table_name.into(),
115 access,
116 source,
117 }
118 }
119}
120
121#[derive(Debug, Clone, Copy, PartialEq, Eq)]
123pub enum TableReferenceAccess {
124 Read,
126 Write,
128 Create,
130 Drop,
132 Metadata,
134}
135
136#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138pub enum TableReferenceSource {
139 TopLevelPlanTableName,
141 LogicalPlanScan,
143 LogicalPlanMutationTarget,
145 LogicalPlanDdlTarget,
147 LogicalPlanIndexTarget,
149 TypedExprSubquery,
151}
152
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
155pub enum PlanningDiagnosticSeverity {
156 Info,
157 Warning,
158}
159
160#[derive(Debug, Clone, PartialEq, Eq)]
162pub struct PlanningDiagnostic {
163 pub code: &'static str,
165 pub severity: PlanningDiagnosticSeverity,
167 pub message: String,
169}
170
171impl PlanningDiagnostic {
172 pub fn info(code: &'static str, message: impl Into<String>) -> Self {
173 Self {
174 code,
175 severity: PlanningDiagnosticSeverity::Info,
176 message: message.into(),
177 }
178 }
179
180 pub fn warning(code: &'static str, message: impl Into<String>) -> Self {
181 Self {
182 code,
183 severity: PlanningDiagnosticSeverity::Warning,
184 message: message.into(),
185 }
186 }
187}
188
189pub fn plan_sql_for_routing<C: Catalog + ?Sized>(
191 catalog: &C,
192 sql: &str,
193) -> Result<Vec<PlannedStatement>, SqlError> {
194 let statements = Parser::parse_sql(&AlopexDialect, sql).map_err(SqlError::from)?;
195 statements
196 .iter()
197 .map(|statement| plan_statement_for_routing(catalog, statement).map_err(SqlError::from))
198 .collect()
199}
200
201pub fn plan_statement_for_routing<C: Catalog + ?Sized>(
203 catalog: &C,
204 statement: &Statement,
205) -> Result<PlannedStatement, PlannerError> {
206 let planner = Planner::new(catalog);
207 let plan = planner.plan(statement)?;
208 let routing_input = routing_input_for_plan(&statement.kind, &plan);
209 Ok(PlannedStatement {
210 plan,
211 routing_input,
212 })
213}
214
215fn routing_input_for_plan(statement_kind: &StatementKind, plan: &LogicalPlan) -> RoutingInput {
216 let mut diagnostics = Vec::new();
217 let extractor = TableReferenceExtractor::new();
218 let table_references = extractor.extract_from_logical_plan(
219 plan,
220 table_reference_access(statement_kind),
221 &mut diagnostics,
222 );
223
224 RoutingInput {
225 statement_kind: statement_kind.clone(),
226 table_references,
227 diagnostics,
228 }
229}
230
231#[derive(Debug, Default, Clone, Copy)]
233pub struct TableReferenceExtractor;
234
235impl TableReferenceExtractor {
236 pub fn new() -> Self {
237 Self
238 }
239
240 pub fn extract_from_logical_plan(
243 &self,
244 plan: &LogicalPlan,
245 root_access: TableReferenceAccess,
246 diagnostics: &mut Vec<PlanningDiagnostic>,
247 ) -> Vec<TableReference> {
248 let mut references = Vec::new();
249 self.extract_plan(
250 plan,
251 root_access,
252 TableReferenceSource::LogicalPlanScan,
253 diagnostics,
254 &mut references,
255 );
256 if references.is_empty() {
257 diagnostics.push(PlanningDiagnostic::info(
258 "ALOPEX-PLAN-ROUTE-001",
259 "statement has no physical table reference",
260 ));
261 }
262 references
263 }
264
265 pub fn extract_from_subquery_context(
267 &self,
268 plan: &LogicalPlan,
269 diagnostics: &mut Vec<PlanningDiagnostic>,
270 ) -> Vec<TableReference> {
271 let mut references = Vec::new();
272 self.extract_plan(
273 plan,
274 TableReferenceAccess::Read,
275 TableReferenceSource::TypedExprSubquery,
276 diagnostics,
277 &mut references,
278 );
279 references
280 }
281
282 fn extract_plan(
283 &self,
284 plan: &LogicalPlan,
285 root_access: TableReferenceAccess,
286 scan_source: TableReferenceSource,
287 diagnostics: &mut Vec<PlanningDiagnostic>,
288 references: &mut Vec<TableReference>,
289 ) {
290 match plan {
291 LogicalPlan::Scan { table, projection } => {
292 if table != LITERAL_TABLE {
293 push_table_reference(
294 references,
295 table,
296 TableReferenceAccess::Read,
297 scan_source,
298 );
299 }
300 self.extract_projection(projection, diagnostics, references);
301 }
302 LogicalPlan::Filter { input, predicate } => {
303 self.extract_plan(input, root_access, scan_source, diagnostics, references);
304 self.extract_typed_expr(predicate, diagnostics, references);
305 }
306 LogicalPlan::Project { input, projection } => {
307 self.extract_plan(input, root_access, scan_source, diagnostics, references);
308 self.extract_projection(projection, diagnostics, references);
309 }
310 LogicalPlan::Join {
311 left,
312 right,
313 condition,
314 ..
315 } => {
316 self.extract_plan(
317 left,
318 TableReferenceAccess::Read,
319 scan_source,
320 diagnostics,
321 references,
322 );
323 self.extract_plan(
324 right,
325 TableReferenceAccess::Read,
326 scan_source,
327 diagnostics,
328 references,
329 );
330 if let Some(condition) = condition {
331 self.extract_typed_expr(condition, diagnostics, references);
332 }
333 }
334 LogicalPlan::Aggregate {
335 input,
336 group_keys,
337 aggregates,
338 having,
339 projection,
340 } => {
341 self.extract_plan(input, root_access, scan_source, diagnostics, references);
342 for expr in group_keys {
343 self.extract_typed_expr(expr, diagnostics, references);
344 }
345 for aggregate in aggregates {
346 if let Some(arg) = &aggregate.arg {
347 self.extract_typed_expr(arg, diagnostics, references);
348 }
349 }
350 if let Some(having) = having {
351 self.extract_typed_expr(having, diagnostics, references);
352 }
353 self.extract_projection(projection, diagnostics, references);
354 }
355 LogicalPlan::Sort { input, order_by } => {
356 self.extract_plan(input, root_access, scan_source, diagnostics, references);
357 for sort_expr in order_by {
358 self.extract_typed_expr(&sort_expr.expr, diagnostics, references);
359 }
360 }
361 LogicalPlan::Limit { input, .. } => {
362 self.extract_plan(input, root_access, scan_source, diagnostics, references);
363 }
364 LogicalPlan::Insert { table, values, .. } => {
365 push_table_reference(
366 references,
367 table,
368 root_access,
369 TableReferenceSource::LogicalPlanMutationTarget,
370 );
371 for row in values {
372 for value in row {
373 self.extract_typed_expr(value, diagnostics, references);
374 }
375 }
376 }
377 LogicalPlan::InsertSelect { table, source, .. } => {
378 push_table_reference(
379 references,
380 table,
381 root_access,
382 TableReferenceSource::LogicalPlanMutationTarget,
383 );
384 self.extract_plan(
385 source,
386 TableReferenceAccess::Read,
387 scan_source,
388 diagnostics,
389 references,
390 );
391 }
392 LogicalPlan::Update {
393 table,
394 assignments,
395 filter,
396 } => {
397 push_table_reference(
398 references,
399 table,
400 root_access,
401 TableReferenceSource::LogicalPlanMutationTarget,
402 );
403 for assignment in assignments {
404 self.extract_typed_expr(&assignment.value, diagnostics, references);
405 }
406 if let Some(filter) = filter {
407 self.extract_typed_expr(filter, diagnostics, references);
408 }
409 }
410 LogicalPlan::Delete { table, filter } => {
411 push_table_reference(
412 references,
413 table,
414 root_access,
415 TableReferenceSource::LogicalPlanMutationTarget,
416 );
417 if let Some(filter) = filter {
418 self.extract_typed_expr(filter, diagnostics, references);
419 }
420 }
421 LogicalPlan::CreateTable { table, .. } => push_table_reference(
422 references,
423 &table.name,
424 root_access,
425 TableReferenceSource::LogicalPlanDdlTarget,
426 ),
427 LogicalPlan::DropTable { name, .. } => push_table_reference(
428 references,
429 name,
430 root_access,
431 TableReferenceSource::LogicalPlanDdlTarget,
432 ),
433 LogicalPlan::CreateIndex { index, .. } => push_table_reference(
434 references,
435 &index.table,
436 root_access,
437 TableReferenceSource::LogicalPlanIndexTarget,
438 ),
439 LogicalPlan::DropIndex { name, .. } => diagnostics.push(PlanningDiagnostic::warning(
440 "ALOPEX-PLAN-ROUTE-003",
441 format!(
442 "DROP INDEX {name} does not expose a target table in the current logical plan"
443 ),
444 )),
445 LogicalPlan::Pragma { .. } => {}
446 }
447 }
448
449 fn extract_projection(
450 &self,
451 projection: &Projection,
452 diagnostics: &mut Vec<PlanningDiagnostic>,
453 references: &mut Vec<TableReference>,
454 ) {
455 if let Projection::Columns(columns) = projection {
456 for column in columns {
457 self.extract_typed_expr(&column.expr, diagnostics, references);
458 }
459 }
460 }
461
462 fn extract_typed_expr(
463 &self,
464 expr: &TypedExpr,
465 diagnostics: &mut Vec<PlanningDiagnostic>,
466 references: &mut Vec<TableReference>,
467 ) {
468 match &expr.kind {
469 TypedExprKind::Literal(_)
470 | TypedExprKind::ColumnRef { .. }
471 | TypedExprKind::VectorLiteral(_) => {}
472 TypedExprKind::BinaryOp { left, right, .. } => {
473 self.extract_typed_expr(left, diagnostics, references);
474 self.extract_typed_expr(right, diagnostics, references);
475 }
476 TypedExprKind::UnaryOp { operand, .. }
477 | TypedExprKind::Cast { expr: operand, .. }
478 | TypedExprKind::IsNull { expr: operand, .. } => {
479 self.extract_typed_expr(operand, diagnostics, references);
480 }
481 TypedExprKind::FunctionCall { args, .. } => {
482 for arg in args {
483 self.extract_typed_expr(arg, diagnostics, references);
484 }
485 }
486 TypedExprKind::Between {
487 expr, low, high, ..
488 } => {
489 self.extract_typed_expr(expr, diagnostics, references);
490 self.extract_typed_expr(low, diagnostics, references);
491 self.extract_typed_expr(high, diagnostics, references);
492 }
493 TypedExprKind::Like {
494 expr,
495 pattern,
496 escape,
497 ..
498 } => {
499 self.extract_typed_expr(expr, diagnostics, references);
500 self.extract_typed_expr(pattern, diagnostics, references);
501 if let Some(escape) = escape {
502 self.extract_typed_expr(escape, diagnostics, references);
503 }
504 }
505 TypedExprKind::InList { expr, list, .. } => {
506 self.extract_typed_expr(expr, diagnostics, references);
507 for item in list {
508 self.extract_typed_expr(item, diagnostics, references);
509 }
510 }
511 TypedExprKind::ScalarSubquery(subquery) => self.extract_plan(
512 subquery,
513 TableReferenceAccess::Read,
514 TableReferenceSource::TypedExprSubquery,
515 diagnostics,
516 references,
517 ),
518 TypedExprKind::InSubquery { expr, subquery, .. } => {
519 self.extract_typed_expr(expr, diagnostics, references);
520 self.extract_plan(
521 subquery,
522 TableReferenceAccess::Read,
523 TableReferenceSource::TypedExprSubquery,
524 diagnostics,
525 references,
526 );
527 }
528 TypedExprKind::Exists { subquery, .. } => self.extract_plan(
529 subquery,
530 TableReferenceAccess::Read,
531 TableReferenceSource::TypedExprSubquery,
532 diagnostics,
533 references,
534 ),
535 TypedExprKind::Quantified { expr, subquery, .. } => {
536 self.extract_typed_expr(expr, diagnostics, references);
537 self.extract_plan(
538 subquery,
539 TableReferenceAccess::Read,
540 TableReferenceSource::TypedExprSubquery,
541 diagnostics,
542 references,
543 );
544 }
545 }
546 }
547}
548
549fn push_table_reference(
550 references: &mut Vec<TableReference>,
551 table_name: &str,
552 access: TableReferenceAccess,
553 source: TableReferenceSource,
554) {
555 if !references.iter().any(|reference| {
556 reference.table_name == table_name
557 && reference.access == access
558 && reference.source == source
559 }) {
560 references.push(TableReference::new(table_name, access, source));
561 }
562}
563
564fn table_reference_access(statement_kind: &StatementKind) -> TableReferenceAccess {
565 match statement_kind {
566 StatementKind::Select(_) => TableReferenceAccess::Read,
567 StatementKind::Insert(_) | StatementKind::Update(_) | StatementKind::Delete(_) => {
568 TableReferenceAccess::Write
569 }
570 StatementKind::CreateTable(_) => TableReferenceAccess::Create,
571 StatementKind::DropTable(_) => TableReferenceAccess::Drop,
572 StatementKind::CreateIndex(_) | StatementKind::DropIndex(_) => {
573 TableReferenceAccess::Metadata
574 }
575 StatementKind::Pragma { .. } => TableReferenceAccess::Metadata,
576 }
577}
578
579pub struct Planner<'a, C: Catalog + ?Sized> {
606 catalog: &'a C,
607 name_resolver: NameResolver<'a, C>,
608 type_checker: TypeChecker<'a, C>,
609}
610
611impl<'a, C: Catalog + ?Sized> Planner<'a, C> {
612 pub fn new(catalog: &'a C) -> Self {
614 Self {
615 catalog,
616 name_resolver: NameResolver::new(catalog),
617 type_checker: TypeChecker::new(catalog),
618 }
619 }
620
621 pub fn plan(&self, stmt: &Statement) -> Result<LogicalPlan, PlannerError> {
632 match &stmt.kind {
633 StatementKind::CreateTable(ct) => self.plan_create_table(ct),
635 StatementKind::DropTable(dt) => self.plan_drop_table(dt),
636 StatementKind::CreateIndex(ci) => self.plan_create_index(ci),
637 StatementKind::DropIndex(di) => self.plan_drop_index(di),
638 StatementKind::Pragma { name, value } => self.plan_pragma(name, value),
639
640 StatementKind::Select(sel) => self.plan_select(sel),
642 StatementKind::Insert(ins) => self.plan_insert(ins),
643 StatementKind::Update(upd) => self.plan_update(upd),
644 StatementKind::Delete(del) => self.plan_delete(del),
645 }
646 }
647
648 fn plan_pragma(
649 &self,
650 raw_name: &str,
651 value: &Option<PragmaValue>,
652 ) -> Result<LogicalPlan, PlannerError> {
653 let name = raw_name.to_ascii_lowercase();
654 if !matches!(name.as_str(), "cache_size" | "memory_limit" | "io_stats") {
655 return Err(PlannerError::InvalidPragma {
656 name,
657 reason: "supported names are cache_size, memory_limit, and io_stats".to_string(),
658 });
659 }
660 match name.as_str() {
661 "cache_size" => match value {
662 Some(PragmaValue::Int(v)) if *v > 0 => {}
663 Some(PragmaValue::Int(_)) => {
664 return Err(PlannerError::InvalidPragma {
665 name,
666 reason: "cache_size must be a positive page count".to_string(),
667 });
668 }
669 Some(PragmaValue::Text(_)) => {
670 return Err(PlannerError::InvalidPragma {
671 name,
672 reason: "cache_size requires an integer page count".to_string(),
673 });
674 }
675 None => {}
676 },
677 "memory_limit" => {
678 if let Some(PragmaValue::Int(v)) = value
679 && *v < 0
680 {
681 return Err(PlannerError::InvalidPragma {
682 name,
683 reason: "memory_limit cannot be negative".to_string(),
684 });
685 }
686 }
687 "io_stats" => {
688 if value.is_some() {
689 return Err(PlannerError::InvalidPragma {
690 name,
691 reason: "io_stats does not accept a value".to_string(),
692 });
693 }
694 }
695 _ => unreachable!(),
696 }
697 Ok(LogicalPlan::Pragma {
698 name,
699 value: value.clone(),
700 })
701 }
702
703 fn plan_create_table(&self, stmt: &CreateTable) -> Result<LogicalPlan, PlannerError> {
712 if !stmt.if_not_exists && self.catalog.table_exists(&stmt.name) {
714 return Err(PlannerError::table_already_exists(&stmt.name));
715 }
716
717 let columns: Vec<ColumnMetadata> = stmt
719 .columns
720 .iter()
721 .map(|col| self.convert_column_def(col))
722 .collect();
723
724 let primary_key = Self::extract_primary_key(stmt);
726
727 let mut table = TableMetadata::new(stmt.name.clone(), columns);
730 if let Some(pk) = primary_key {
731 table = table.with_primary_key(pk);
732 }
733 table.catalog_name = "default".to_string();
734 table.namespace_name = "default".to_string();
735 table.table_type = TableType::Managed;
736 table.data_source_format = DataSourceFormat::Alopex;
737 table.properties = HashMap::new();
738
739 Ok(LogicalPlan::CreateTable {
740 table,
741 if_not_exists: stmt.if_not_exists,
742 with_options: stmt
743 .with_options
744 .iter()
745 .map(|opt| (opt.key.clone(), opt.value.clone()))
746 .collect(),
747 })
748 }
749
750 fn convert_column_def(&self, col: &ColumnDef) -> ColumnMetadata {
752 let data_type = ResolvedType::from_ast(&col.data_type);
753 let mut meta = ColumnMetadata::new(col.name.clone(), data_type);
754
755 for constraint in &col.constraints {
757 meta = Self::apply_column_constraint(meta, constraint);
758 }
759
760 meta
761 }
762
763 fn apply_column_constraint(
765 mut meta: ColumnMetadata,
766 constraint: &ColumnConstraint,
767 ) -> ColumnMetadata {
768 match constraint {
769 ColumnConstraint::NotNull { .. } => {
770 meta.not_null = true;
771 }
772 ColumnConstraint::PrimaryKey { .. } => {
773 meta.primary_key = true;
774 meta.not_null = true; }
776 ColumnConstraint::Unique { .. } => {
777 meta.unique = true;
778 }
779 ColumnConstraint::Default { value: expr, .. } => {
780 meta.default = Some(expr.clone());
781 }
782 }
783 meta
784 }
785
786 fn extract_primary_key(stmt: &CreateTable) -> Option<Vec<String>> {
788 use crate::ast::ddl::TableConstraint;
789
790 if let Some(TableConstraint::PrimaryKey { columns, .. }) = stmt.constraints.first() {
794 return Some(columns.clone());
795 }
796
797 let pk_columns: Vec<String> = stmt
799 .columns
800 .iter()
801 .filter(|col| col.constraints.iter().any(Self::is_primary_key_constraint))
802 .map(|col| col.name.clone())
803 .collect();
804
805 if pk_columns.is_empty() {
806 None
807 } else {
808 Some(pk_columns)
809 }
810 }
811
812 fn is_primary_key_constraint(constraint: &ColumnConstraint) -> bool {
814 matches!(constraint, ColumnConstraint::PrimaryKey { .. })
815 }
816
817 fn plan_drop_table(&self, stmt: &DropTable) -> Result<LogicalPlan, PlannerError> {
821 if !stmt.if_exists && !self.table_exists_in_default(&stmt.name) {
823 return Err(PlannerError::TableNotFound {
824 name: stmt.name.clone(),
825 line: stmt.span.start.line,
826 column: stmt.span.start.column,
827 });
828 }
829
830 Ok(LogicalPlan::DropTable {
831 name: stmt.name.clone(),
832 if_exists: stmt.if_exists,
833 })
834 }
835
836 fn table_exists_in_default(&self, name: &str) -> bool {
837 match self.catalog.get_table(name) {
838 Some(table) => table.catalog_name == "default" && table.namespace_name == "default",
839 None => false,
840 }
841 }
842
843 fn plan_create_index(&self, stmt: &CreateIndex) -> Result<LogicalPlan, PlannerError> {
850 if !stmt.if_not_exists && self.catalog.index_exists(&stmt.name) {
852 return Err(PlannerError::index_already_exists(&stmt.name));
853 }
854
855 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
857
858 self.name_resolver
860 .resolve_column(table, &stmt.column, stmt.span)?;
861
862 let mut index = IndexMetadata::new(
866 0,
867 stmt.name.clone(),
868 stmt.table.clone(),
869 vec![stmt.column.clone()],
870 );
871
872 if let Some(method) = stmt.method {
873 index = index.with_method(method);
874 }
875
876 let options: Vec<(String, String)> = stmt
877 .options
878 .iter()
879 .map(|opt| (opt.key.clone(), opt.value.clone()))
880 .collect();
881 if !options.is_empty() {
882 index = index.with_options(options);
883 }
884
885 Ok(LogicalPlan::CreateIndex {
886 index,
887 if_not_exists: stmt.if_not_exists,
888 })
889 }
890
891 fn plan_drop_index(&self, stmt: &DropIndex) -> Result<LogicalPlan, PlannerError> {
895 if !stmt.if_exists && !self.index_exists_in_default(&stmt.name) {
897 return Err(PlannerError::index_not_found(&stmt.name));
898 }
899
900 Ok(LogicalPlan::DropIndex {
901 name: stmt.name.clone(),
902 if_exists: stmt.if_exists,
903 })
904 }
905
906 fn index_exists_in_default(&self, name: &str) -> bool {
907 match self.catalog.get_index(name) {
908 Some(index) => index.catalog_name == "default" && index.namespace_name == "default",
909 None => false,
910 }
911 }
912
913 fn plan_select(&self, stmt: &Select) -> Result<LogicalPlan, PlannerError> {
922 self.plan_select_relation(stmt, &[])
923 .map(|relation| relation.plan)
924 }
925
926 fn plan_select_relation(
927 &self,
928 stmt: &Select,
929 outer_scope: &[ScopedTable],
930 ) -> Result<PlannedRelation, PlannerError> {
931 let mut relation = self.plan_from_items(&stmt.from, stmt.span, outer_scope)?;
932 let expr_scope = relation
933 .scope
934 .iter()
935 .cloned()
936 .chain(offset_scope(outer_scope, relation.schema.len()))
937 .collect::<Vec<_>>();
938
939 let has_group_by = stmt
940 .group_by
941 .as_ref()
942 .is_some_and(|items| !items.is_empty());
943 let has_aggregate = self.select_contains_aggregate(stmt);
944 let distinct_only =
945 stmt.distinct && !has_group_by && !has_aggregate && stmt.having.is_none();
946
947 let final_projection =
948 self.build_projection_with_scope(&stmt.projection, &relation.schema, &expr_scope)?;
949 install_base_projection(&mut relation.plan, &final_projection);
950 let needs_project_boundary = !matches!(relation.plan, LogicalPlan::Scan { .. });
951 let mut plan = relation.plan;
952
953 if let Some(ref selection) = stmt.selection {
955 let predicate = self.infer_expr_with_scope(selection, &expr_scope)?;
956
957 if predicate.resolved_type != ResolvedType::Boolean {
959 return Err(PlannerError::type_mismatch(
960 "Boolean",
961 predicate.resolved_type.to_string(),
962 selection.span,
963 ));
964 }
965
966 plan = LogicalPlan::Filter {
967 input: Box::new(plan),
968 predicate,
969 };
970 }
971
972 if has_group_by || has_aggregate || stmt.having.is_some() || stmt.distinct {
973 if !has_group_by && !has_aggregate && stmt.having.is_some() {
974 return Err(PlannerError::invalid_expression(
975 "HAVING requires GROUP BY or aggregate functions".to_string(),
976 ));
977 }
978
979 let (group_keys, projected) = if distinct_only {
980 let projected = self.build_projected_columns_for_distinct_with_scope(
981 &stmt.projection,
982 &relation.schema,
983 &expr_scope,
984 )?;
985 let group_keys = projected.iter().map(|col| col.expr.clone()).collect();
986 (group_keys, projected)
987 } else {
988 let group_keys = self.build_group_keys_with_scope(stmt, &expr_scope)?;
989 let projected = self.build_projected_columns_for_aggregate_with_scope(
990 &stmt.projection,
991 &expr_scope,
992 )?;
993 (group_keys, projected)
994 };
995 let mut aggregates = Vec::new();
996 let mut agg_map = HashMap::new();
997
998 for col in &projected {
999 self.collect_aggregates_from_typed_expr(&col.expr, &mut aggregates, &mut agg_map)?;
1000 }
1001
1002 let having_typed = if let Some(having) = &stmt.having {
1003 let typed = self.infer_expr_with_scope(having, &expr_scope)?;
1004 if typed.resolved_type != ResolvedType::Boolean {
1005 return Err(PlannerError::type_mismatch(
1006 "Boolean",
1007 typed.resolved_type.type_name().to_string(),
1008 typed.span,
1009 ));
1010 }
1011 self.collect_aggregates_from_typed_expr(&typed, &mut aggregates, &mut agg_map)?;
1012 Some(typed)
1013 } else {
1014 None
1015 };
1016
1017 let mut order_by = Vec::new();
1018 if !stmt.order_by.is_empty() {
1019 for order_expr in &stmt.order_by {
1020 let typed = self.infer_expr_with_scope(&order_expr.expr, &expr_scope)?;
1021 self.collect_aggregates_from_typed_expr(&typed, &mut aggregates, &mut agg_map)?;
1022 let asc = order_expr.asc.unwrap_or(true);
1023 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1024 order_by.push(SortExpr::new(typed, asc, nulls_first));
1025 }
1026 }
1027
1028 if let Some(ref having) = having_typed {
1029 self.type_checker
1030 .validate_having_expr(having, &group_keys, &aggregates)?;
1031 }
1032
1033 let output_schema = build_aggregate_schema(&group_keys, &aggregates);
1034 let output_names: Vec<String> = output_schema.iter().map(|c| c.name.clone()).collect();
1035
1036 let projection = self.build_aggregate_projection(
1037 projected,
1038 &group_keys,
1039 &aggregates,
1040 &output_names,
1041 )?;
1042
1043 let having = if let Some(having) = having_typed {
1044 Some(self.rewrite_expr_for_aggregate(
1045 &having,
1046 &group_keys,
1047 &aggregates,
1048 &output_names,
1049 )?)
1050 } else {
1051 None
1052 };
1053
1054 let order_by = order_by
1055 .into_iter()
1056 .map(|expr| {
1057 let rewritten = self.rewrite_expr_for_aggregate(
1058 &expr.expr,
1059 &group_keys,
1060 &aggregates,
1061 &output_names,
1062 )?;
1063 Ok(SortExpr::new(rewritten, expr.asc, expr.nulls_first))
1064 })
1065 .collect::<Result<Vec<_>, PlannerError>>()?;
1066
1067 let schema = projection_schema(&projection, &output_schema);
1068 plan = LogicalPlan::Aggregate {
1069 input: Box::new(plan),
1070 group_keys,
1071 aggregates,
1072 having,
1073 projection,
1074 };
1075
1076 if !order_by.is_empty() {
1077 plan = LogicalPlan::Sort {
1078 input: Box::new(plan),
1079 order_by,
1080 };
1081 }
1082
1083 if stmt.limit.is_some() || stmt.offset.is_some() {
1084 let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
1085 let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
1086 plan = LogicalPlan::Limit {
1087 input: Box::new(plan),
1088 limit,
1089 offset,
1090 };
1091 }
1092
1093 return Ok(PlannedRelation {
1094 plan,
1095 schema: schema.clone(),
1096 scope: vec![ScopedTable::new(
1097 TableMetadata::new(LITERAL_TABLE, schema),
1098 0,
1099 )],
1100 });
1101 }
1102
1103 if !stmt.order_by.is_empty() {
1105 let order_by = self.build_sort_exprs_with_scope(&stmt.order_by, &expr_scope)?;
1106 plan = LogicalPlan::Sort {
1107 input: Box::new(plan),
1108 order_by,
1109 };
1110 }
1111
1112 if stmt.limit.is_some() || stmt.offset.is_some() {
1113 let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
1114 let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
1115 plan = LogicalPlan::Limit {
1116 input: Box::new(plan),
1117 limit,
1118 offset,
1119 };
1120 }
1121
1122 let output_schema = projection_schema(&final_projection, &relation.schema);
1123 if needs_project_boundary {
1124 plan = LogicalPlan::Project {
1125 input: Box::new(plan),
1126 projection: final_projection,
1127 };
1128 }
1129 Ok(PlannedRelation {
1130 plan,
1131 schema: output_schema.clone(),
1132 scope: vec![ScopedTable::new(
1133 TableMetadata::new(LITERAL_TABLE, output_schema),
1134 0,
1135 )],
1136 })
1137 }
1138
1139 fn plan_from_items(
1143 &self,
1144 items: &[FromItem],
1145 select_span: crate::ast::Span,
1146 outer_scope: &[ScopedTable],
1147 ) -> Result<PlannedRelation, PlannerError> {
1148 match items {
1149 [] => {
1150 let schema = Vec::new();
1151 Ok(PlannedRelation {
1152 plan: LogicalPlan::Scan {
1153 table: LITERAL_TABLE.to_string(),
1154 projection: Projection::All(Vec::new()),
1155 },
1156 schema: schema.clone(),
1157 scope: vec![ScopedTable::new(
1158 TableMetadata::new(LITERAL_TABLE, schema),
1159 0,
1160 )],
1161 })
1162 }
1163 [single] => self.plan_from_item(single, 0, outer_scope),
1164 [first, rest @ ..] => {
1165 let mut relation = self.plan_from_item(first, 0, outer_scope)?;
1166 for item in rest {
1167 let right = self.plan_from_item(item, relation.schema.len(), outer_scope)?;
1168 relation = self.combine_join_relation(
1169 relation,
1170 right,
1171 JoinType::Cross,
1172 None,
1173 None,
1174 select_span,
1175 )?;
1176 }
1177 Ok(relation)
1178 }
1179 }
1180 }
1181
1182 fn plan_from_item(
1183 &self,
1184 item: &FromItem,
1185 start_index: usize,
1186 outer_scope: &[ScopedTable],
1187 ) -> Result<PlannedRelation, PlannerError> {
1188 match item {
1189 FromItem::Table { name, alias, span } => {
1190 let table = self.name_resolver.resolve_table(name, *span)?.clone();
1191 let mut scope_table = table.clone();
1192 if let Some(alias) = alias {
1193 scope_table.name = alias.clone();
1194 }
1195 let schema = table.columns.clone();
1196 Ok(PlannedRelation {
1197 plan: LogicalPlan::Scan {
1198 table: name.clone(),
1199 projection: Projection::All(
1200 schema.iter().map(|col| col.name.clone()).collect(),
1201 ),
1202 },
1203 schema,
1204 scope: vec![ScopedTable::new(scope_table, start_index)],
1205 })
1206 }
1207 FromItem::Join {
1208 left,
1209 right,
1210 join_type,
1211 condition,
1212 using,
1213 natural,
1214 span,
1215 } => {
1216 let left_relation = self.plan_from_item(left, start_index, outer_scope)?;
1217 let right_relation = self.plan_from_item(
1218 right,
1219 start_index + left_relation.schema.len(),
1220 outer_scope,
1221 )?;
1222 let expr_scope = left_relation
1223 .scope
1224 .iter()
1225 .cloned()
1226 .chain(right_relation.scope.iter().cloned())
1227 .chain(offset_scope(
1228 outer_scope,
1229 left_relation.schema.len() + right_relation.schema.len(),
1230 ))
1231 .collect::<Vec<_>>();
1232 let using = if *natural {
1233 Some(natural_join_columns(
1234 &left_relation.schema,
1235 &right_relation.schema,
1236 ))
1237 } else {
1238 using.clone()
1239 };
1240 let typed_condition = if let Some(expr) = condition {
1241 let typed = self.infer_expr_with_scope(expr, &expr_scope)?;
1242 if typed.resolved_type != ResolvedType::Boolean {
1243 return Err(PlannerError::type_mismatch(
1244 "Boolean",
1245 typed.resolved_type.to_string(),
1246 expr.span,
1247 ));
1248 }
1249 Some(typed)
1250 } else {
1251 self.build_using_condition(
1252 using.as_deref(),
1253 &left_relation,
1254 &right_relation,
1255 *span,
1256 )?
1257 };
1258 self.combine_join_relation(
1259 left_relation,
1260 right_relation,
1261 map_join_type(*join_type),
1262 typed_condition,
1263 using,
1264 *span,
1265 )
1266 }
1267 FromItem::Derived {
1268 subquery,
1269 alias,
1270 span,
1271 } => {
1272 let crate::ast::StatementKind::Select(select) = &subquery.kind else {
1273 return Err(PlannerError::unsupported_feature(
1274 "non-SELECT derived table",
1275 "v0.6.0-subquery Phase 6",
1276 *span,
1277 ));
1278 };
1279 let mut relation = self.plan_select_relation(select, outer_scope)?;
1280 let alias = alias.clone().ok_or_else(|| {
1281 PlannerError::invalid_expression("derived table requires an alias".to_string())
1282 })?;
1283 relation.plan = LogicalPlan::Project {
1284 input: Box::new(relation.plan),
1285 projection: Projection::All(
1286 relation.schema.iter().map(|col| col.name.clone()).collect(),
1287 ),
1288 };
1289 relation.scope = vec![ScopedTable::new(
1290 TableMetadata::new(alias, relation.schema.clone()),
1291 start_index,
1292 )];
1293 Ok(relation)
1294 }
1295 }
1296 }
1297
1298 fn combine_join_relation(
1299 &self,
1300 left: PlannedRelation,
1301 right: PlannedRelation,
1302 join_type: JoinType,
1303 condition: Option<TypedExpr>,
1304 using: Option<Vec<String>>,
1305 _span: crate::ast::Span,
1306 ) -> Result<PlannedRelation, PlannerError> {
1307 let mut schema = left.schema.clone();
1308 schema.extend(right.schema.clone());
1309 let mut scope = left.scope.clone();
1310 let mut right_scope = right.scope.clone();
1311 if let Some(columns) = &using {
1312 for column in columns {
1316 let right_index = right_scope.iter().find_map(|table| {
1317 table
1318 .table
1319 .get_column_index(column)
1320 .map(|index| table.start_index + index)
1321 });
1322 let Some(right_index) = right_index else {
1323 continue;
1324 };
1325 for table in &mut scope {
1326 if table.table.get_column_index(column).is_some() {
1327 table.merge_column_with(column, right_index);
1328 }
1329 }
1330 }
1331 for table in &mut right_scope {
1332 table.hide_unqualified_columns(columns);
1333 }
1334 }
1335 scope.extend(right_scope);
1336 Ok(PlannedRelation {
1337 plan: LogicalPlan::Join {
1338 left: Box::new(left.plan),
1339 right: Box::new(right.plan),
1340 join_type,
1341 condition,
1342 using,
1343 },
1344 schema,
1345 scope,
1346 })
1347 }
1348
1349 fn build_using_condition(
1350 &self,
1351 using: Option<&[String]>,
1352 left: &PlannedRelation,
1353 right: &PlannedRelation,
1354 span: crate::ast::Span,
1355 ) -> Result<Option<TypedExpr>, PlannerError> {
1356 let Some(columns) = using else {
1357 return Ok(None);
1358 };
1359 let mut condition = None;
1360 for column in columns {
1361 let left_col = find_scoped_column(&left.scope, column, span)?;
1362 let right_col = find_scoped_column(&right.scope, column, span)?;
1363 let left_expr = TypedExpr::column_ref(
1364 left_col.table,
1365 column.clone(),
1366 left_col.index,
1367 left_col.ty.clone(),
1368 span,
1369 );
1370 let right_expr = TypedExpr::column_ref(
1371 right_col.table,
1372 column.clone(),
1373 right_col.index,
1374 right_col.ty.clone(),
1375 span,
1376 );
1377 self.type_checker
1378 .check_comparison_op(&left_col.ty, &right_col.ty, span)?;
1379 let eq = TypedExpr::binary_op(
1380 left_expr,
1381 crate::ast::expr::BinaryOp::Eq,
1382 right_expr,
1383 ResolvedType::Boolean,
1384 span,
1385 );
1386 condition = Some(match condition {
1387 Some(prev) => TypedExpr::binary_op(
1388 prev,
1389 crate::ast::expr::BinaryOp::And,
1390 eq,
1391 ResolvedType::Boolean,
1392 span,
1393 ),
1394 None => eq,
1395 });
1396 }
1397 Ok(condition)
1398 }
1399
1400 fn infer_expr_with_scope(
1401 &self,
1402 expr: &crate::ast::expr::Expr,
1403 scope: &[ScopedTable],
1404 ) -> Result<TypedExpr, PlannerError> {
1405 self.type_checker
1406 .infer_type_with_scope(expr, scope, &|stmt, outer_scope| {
1407 let crate::ast::StatementKind::Select(select) = &stmt.kind else {
1408 return Err(PlannerError::unsupported_feature(
1409 "non-SELECT subquery",
1410 "v0.6.0-subquery Phase 6",
1411 stmt.span(),
1412 ));
1413 };
1414 let relation = self.plan_select_relation(select, outer_scope)?;
1415 Ok((relation.plan, relation.schema))
1416 })
1417 }
1418
1419 #[allow(dead_code)]
1420 fn build_projection(
1421 &self,
1422 items: &[SelectItem],
1423 table: &TableMetadata,
1424 ) -> Result<Projection, PlannerError> {
1425 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1427 let columns = self.name_resolver.expand_wildcard(table);
1428 return Ok(Projection::All(columns));
1429 }
1430
1431 let mut projected_columns = Vec::new();
1433 for item in items {
1434 match item {
1435 SelectItem::Wildcard { span } => {
1436 for col in &table.columns {
1438 let column_index = table.get_column_index(&col.name).unwrap();
1439 let typed_expr = TypedExpr::column_ref(
1440 table.name.clone(),
1441 col.name.clone(),
1442 column_index,
1443 col.data_type.clone(),
1444 *span,
1445 );
1446 projected_columns.push(ProjectedColumn::new(typed_expr));
1447 }
1448 }
1449 SelectItem::QualifiedWildcard {
1450 table: qualifier,
1451 span,
1452 } => {
1453 if qualifier != &table.name {
1454 return Err(PlannerError::invalid_expression(format!(
1455 "table '{qualifier}' is not available for wildcard projection"
1456 )));
1457 }
1458 for col in &table.columns {
1459 let column_index = table.get_column_index(&col.name).unwrap();
1460 projected_columns.push(ProjectedColumn::new(TypedExpr::column_ref(
1461 table.name.clone(),
1462 col.name.clone(),
1463 column_index,
1464 col.data_type.clone(),
1465 *span,
1466 )));
1467 }
1468 }
1469 SelectItem::Expr { expr, alias, .. } => {
1470 let typed_expr = self.type_checker.infer_type(expr, table)?;
1471 let projected = if let Some(alias) = alias {
1472 ProjectedColumn::with_alias(typed_expr, alias.clone())
1473 } else {
1474 ProjectedColumn::new(typed_expr)
1475 };
1476 projected_columns.push(projected);
1477 }
1478 }
1479 }
1480
1481 Ok(Projection::Columns(projected_columns))
1482 }
1483
1484 fn build_projection_with_scope(
1485 &self,
1486 items: &[SelectItem],
1487 schema: &[ColumnMetadata],
1488 scope: &[ScopedTable],
1489 ) -> Result<Projection, PlannerError> {
1490 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1491 return Ok(Projection::All(visible_wildcard_columns(schema, scope)));
1492 }
1493
1494 let mut projected_columns = Vec::new();
1495 for item in items {
1496 match item {
1497 SelectItem::Wildcard { span } => {
1498 for scoped in scope {
1499 for (local_idx, col) in scoped.table.columns.iter().enumerate() {
1500 projected_columns.push(ProjectedColumn::new(TypedExpr::column_ref(
1501 scoped.table.name.clone(),
1502 col.name.clone(),
1503 scoped.start_index + local_idx,
1504 col.data_type.clone(),
1505 *span,
1506 )));
1507 }
1508 }
1509 }
1510 SelectItem::QualifiedWildcard { table, span } => {
1511 let scoped = scope
1512 .iter()
1513 .filter(|scoped| scoped.table.name == *table)
1514 .collect::<Vec<_>>();
1515 match scoped.as_slice() {
1516 [] => {
1517 return Err(PlannerError::invalid_expression(format!(
1518 "table '{table}' is not available for wildcard projection"
1519 )));
1520 }
1521 [scoped] => {
1522 for (local_idx, col) in scoped.table.columns.iter().enumerate() {
1523 projected_columns.push(ProjectedColumn::new(
1524 TypedExpr::column_ref(
1525 scoped.table.name.clone(),
1526 col.name.clone(),
1527 scoped.start_index + local_idx,
1528 col.data_type.clone(),
1529 *span,
1530 ),
1531 ));
1532 }
1533 }
1534 _ => {
1535 return Err(PlannerError::ambiguous_column(
1536 table,
1537 scoped
1538 .iter()
1539 .map(|scoped| scoped.table.name.clone())
1540 .collect(),
1541 *span,
1542 ));
1543 }
1544 }
1545 }
1546 SelectItem::Expr { expr, alias, .. } => {
1547 let typed_expr = self.infer_expr_with_scope(expr, scope)?;
1548 let projected = if let Some(alias) = alias {
1549 ProjectedColumn::with_alias(typed_expr, alias.clone())
1550 } else {
1551 ProjectedColumn::new(typed_expr)
1552 };
1553 projected_columns.push(projected);
1554 }
1555 }
1556 }
1557
1558 Ok(Projection::Columns(projected_columns))
1559 }
1560
1561 #[allow(dead_code)]
1563 fn build_sort_exprs(
1564 &self,
1565 order_by: &[OrderByExpr],
1566 table: &TableMetadata,
1567 ) -> Result<Vec<SortExpr>, PlannerError> {
1568 let mut sort_exprs = Vec::new();
1569
1570 for order_expr in order_by {
1571 let typed_expr = self.type_checker.infer_type(&order_expr.expr, table)?;
1572
1573 let asc = order_expr.asc.unwrap_or(true);
1575
1576 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1578
1579 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1580 }
1581
1582 Ok(sort_exprs)
1583 }
1584
1585 fn build_sort_exprs_with_scope(
1586 &self,
1587 order_by: &[OrderByExpr],
1588 scope: &[ScopedTable],
1589 ) -> Result<Vec<SortExpr>, PlannerError> {
1590 let mut sort_exprs = Vec::new();
1591 for order_expr in order_by {
1592 let typed_expr = self.infer_expr_with_scope(&order_expr.expr, scope)?;
1593 let asc = order_expr.asc.unwrap_or(true);
1594 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1595 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1596 }
1597 Ok(sort_exprs)
1598 }
1599
1600 fn select_contains_aggregate(&self, stmt: &Select) -> bool {
1601 stmt.projection.iter().any(|item| match item {
1602 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => false,
1603 SelectItem::Expr { expr, .. } => expr_contains_aggregate(expr),
1604 }) || stmt
1605 .group_by
1606 .as_ref()
1607 .map(|items| items.iter().any(expr_contains_aggregate))
1608 .unwrap_or(false)
1609 || stmt
1610 .having
1611 .as_ref()
1612 .map(expr_contains_aggregate)
1613 .unwrap_or(false)
1614 || stmt
1615 .order_by
1616 .iter()
1617 .any(|order| expr_contains_aggregate(&order.expr))
1618 }
1619
1620 #[allow(dead_code)]
1621 fn build_group_keys(
1622 &self,
1623 stmt: &Select,
1624 table: &TableMetadata,
1625 ) -> Result<Vec<TypedExpr>, PlannerError> {
1626 let mut keys = Vec::new();
1627 if let Some(items) = &stmt.group_by {
1628 for expr in items {
1629 let typed = self.type_checker.infer_type(expr, table)?;
1630 if typed_expr_contains_aggregate(&typed) {
1631 return Err(PlannerError::invalid_expression(
1632 "GROUP BY cannot contain aggregate functions".to_string(),
1633 ));
1634 }
1635 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1636 return Err(PlannerError::invalid_expression(
1637 "GROUP BY expressions must be column references".to_string(),
1638 ));
1639 }
1640 keys.push(typed);
1641 }
1642 }
1643 Ok(keys)
1644 }
1645
1646 fn build_group_keys_with_scope(
1647 &self,
1648 stmt: &Select,
1649 scope: &[ScopedTable],
1650 ) -> Result<Vec<TypedExpr>, PlannerError> {
1651 let mut keys = Vec::new();
1652 if let Some(items) = &stmt.group_by {
1653 for expr in items {
1654 let typed = self.infer_expr_with_scope(expr, scope)?;
1655 if typed_expr_contains_aggregate(&typed) {
1656 return Err(PlannerError::invalid_expression(
1657 "GROUP BY cannot contain aggregate functions".to_string(),
1658 ));
1659 }
1660 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1661 return Err(PlannerError::invalid_expression(
1662 "GROUP BY expressions must be column references".to_string(),
1663 ));
1664 }
1665 keys.push(typed);
1666 }
1667 }
1668 Ok(keys)
1669 }
1670
1671 #[allow(dead_code)]
1672 fn build_projected_columns_for_aggregate(
1673 &self,
1674 items: &[SelectItem],
1675 table: &TableMetadata,
1676 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1677 let mut projected = Vec::new();
1678 for item in items {
1679 match item {
1680 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => {
1681 return Err(PlannerError::invalid_expression(
1682 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1683 ));
1684 }
1685 SelectItem::Expr { expr, alias, .. } => {
1686 let typed = self.type_checker.infer_type(expr, table)?;
1687 projected.push(ProjectedColumn {
1688 expr: typed,
1689 alias: alias.clone(),
1690 });
1691 }
1692 }
1693 }
1694 Ok(projected)
1695 }
1696
1697 fn build_projected_columns_for_aggregate_with_scope(
1698 &self,
1699 items: &[SelectItem],
1700 scope: &[ScopedTable],
1701 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1702 let mut projected = Vec::new();
1703 for item in items {
1704 match item {
1705 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => {
1706 return Err(PlannerError::invalid_expression(
1707 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1708 ));
1709 }
1710 SelectItem::Expr { expr, alias, .. } => {
1711 let typed = self.infer_expr_with_scope(expr, scope)?;
1712 projected.push(ProjectedColumn {
1713 expr: typed,
1714 alias: alias.clone(),
1715 });
1716 }
1717 }
1718 }
1719 Ok(projected)
1720 }
1721
1722 #[allow(dead_code)]
1723 fn build_projected_columns_for_distinct(
1724 &self,
1725 items: &[SelectItem],
1726 table: &TableMetadata,
1727 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1728 let projection = self.build_projection(items, table)?;
1729 match projection {
1730 Projection::All(columns) => {
1731 let mut projected = Vec::with_capacity(columns.len());
1732 for column in columns {
1733 let column_index = table.get_column_index(&column).ok_or_else(|| {
1734 PlannerError::invalid_expression(format!(
1735 "column '{column}' not found for DISTINCT projection"
1736 ))
1737 })?;
1738 let column_meta = table.get_column(&column).ok_or_else(|| {
1739 PlannerError::invalid_expression(format!(
1740 "column '{column}' not found for DISTINCT projection"
1741 ))
1742 })?;
1743 let typed_expr = TypedExpr::column_ref(
1744 table.name.clone(),
1745 column.clone(),
1746 column_index,
1747 column_meta.data_type.clone(),
1748 crate::ast::Span::default(),
1749 );
1750 projected.push(ProjectedColumn::new(typed_expr));
1751 }
1752 Ok(projected)
1753 }
1754 Projection::Columns(columns) => Ok(columns),
1755 }
1756 }
1757
1758 fn build_projected_columns_for_distinct_with_scope(
1759 &self,
1760 items: &[SelectItem],
1761 schema: &[ColumnMetadata],
1762 scope: &[ScopedTable],
1763 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1764 let projection = self.build_projection_with_scope(items, schema, scope)?;
1765 match projection {
1766 Projection::All(columns) => {
1767 let mut projected = Vec::with_capacity(columns.len());
1768 for (idx, column) in columns.into_iter().enumerate() {
1769 let column_meta = schema.get(idx).ok_or_else(|| {
1770 PlannerError::invalid_expression(format!(
1771 "column '{column}' not found for DISTINCT projection"
1772 ))
1773 })?;
1774 projected.push(ProjectedColumn::new(TypedExpr::column_ref(
1775 LITERAL_TABLE.to_string(),
1776 column,
1777 idx,
1778 column_meta.data_type.clone(),
1779 crate::ast::Span::default(),
1780 )));
1781 }
1782 Ok(projected)
1783 }
1784 Projection::Columns(columns) => Ok(columns),
1785 }
1786 }
1787
1788 fn collect_aggregates_from_typed_expr(
1789 &self,
1790 expr: &TypedExpr,
1791 aggregates: &mut Vec<AggregateExpr>,
1792 aggregate_map: &mut HashMap<AggregateSignature, usize>,
1793 ) -> Result<(), PlannerError> {
1794 match &expr.kind {
1795 TypedExprKind::FunctionCall {
1796 name,
1797 args,
1798 distinct,
1799 star,
1800 } if is_aggregate_function(name) => {
1801 for arg in args {
1802 if typed_expr_contains_aggregate(arg) {
1803 return Err(PlannerError::invalid_expression(
1804 "nested aggregate functions are not supported".to_string(),
1805 ));
1806 }
1807 }
1808 let (agg, signature) =
1809 self.build_aggregate_expr_from_typed(expr, name, args, *distinct, *star)?;
1810 aggregate_map.entry(signature).or_insert_with(|| {
1811 aggregates.push(agg);
1812 aggregates.len() - 1
1813 });
1814 Ok(())
1815 }
1816 TypedExprKind::BinaryOp { left, right, .. } => {
1817 self.collect_aggregates_from_typed_expr(left, aggregates, aggregate_map)?;
1818 self.collect_aggregates_from_typed_expr(right, aggregates, aggregate_map)?;
1819 Ok(())
1820 }
1821 TypedExprKind::UnaryOp { operand, .. } => {
1822 self.collect_aggregates_from_typed_expr(operand, aggregates, aggregate_map)
1823 }
1824 TypedExprKind::FunctionCall { args, .. } => {
1825 for arg in args {
1826 self.collect_aggregates_from_typed_expr(arg, aggregates, aggregate_map)?;
1827 }
1828 Ok(())
1829 }
1830 TypedExprKind::Between {
1831 expr, low, high, ..
1832 } => {
1833 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1834 self.collect_aggregates_from_typed_expr(low, aggregates, aggregate_map)?;
1835 self.collect_aggregates_from_typed_expr(high, aggregates, aggregate_map)?;
1836 Ok(())
1837 }
1838 TypedExprKind::Like {
1839 expr,
1840 pattern,
1841 escape,
1842 ..
1843 } => {
1844 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1845 self.collect_aggregates_from_typed_expr(pattern, aggregates, aggregate_map)?;
1846 if let Some(esc) = escape {
1847 self.collect_aggregates_from_typed_expr(esc, aggregates, aggregate_map)?;
1848 }
1849 Ok(())
1850 }
1851 TypedExprKind::InList { expr, list, .. } => {
1852 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1853 for item in list {
1854 self.collect_aggregates_from_typed_expr(item, aggregates, aggregate_map)?;
1855 }
1856 Ok(())
1857 }
1858 TypedExprKind::IsNull { expr, .. } => {
1859 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)
1860 }
1861 _ => Ok(()),
1862 }
1863 }
1864
1865 fn build_aggregate_expr_from_typed(
1866 &self,
1867 expr: &TypedExpr,
1868 name: &str,
1869 args: &[TypedExpr],
1870 distinct: bool,
1871 star: bool,
1872 ) -> Result<(AggregateExpr, AggregateSignature), PlannerError> {
1873 let lower = name.to_lowercase();
1874 match lower.as_str() {
1875 "count" => {
1876 if star {
1877 let agg = AggregateExpr::count_star();
1878 let signature = aggregate_signature(name, distinct, star, None, None, expr);
1879 return Ok((agg, signature));
1880 }
1881 if args.len() != 1 {
1882 return Err(PlannerError::type_mismatch(
1883 "1 argument",
1884 format!("{} arguments", args.len()),
1885 expr.span,
1886 ));
1887 }
1888 let agg = AggregateExpr {
1889 function: AggregateFunction::Count,
1890 arg: Some(args[0].clone()),
1891 distinct,
1892 result_type: ResolvedType::BigInt,
1893 };
1894 let signature =
1895 aggregate_signature(name, distinct, star, Some(&args[0]), None, expr);
1896 Ok((agg, signature))
1897 }
1898 "sum" => {
1899 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1900 let agg = AggregateExpr {
1901 function: AggregateFunction::Sum,
1902 arg: Some(arg.clone()),
1903 distinct,
1904 result_type: crate::planner::aggregate_expr::sum_result_type(
1905 &arg.resolved_type,
1906 ),
1907 };
1908 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1909 Ok((agg, signature))
1910 }
1911 "total" => {
1912 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1913 let agg = AggregateExpr {
1914 function: AggregateFunction::Total,
1915 arg: Some(arg.clone()),
1916 distinct: false,
1917 result_type: ResolvedType::Double,
1918 };
1919 let signature = aggregate_signature(name, false, star, Some(arg), None, expr);
1920 Ok((agg, signature))
1921 }
1922 "avg" => {
1923 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1924 let agg = AggregateExpr {
1925 function: AggregateFunction::Avg,
1926 arg: Some(arg.clone()),
1927 distinct,
1928 result_type: ResolvedType::Double,
1929 };
1930 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1931 Ok((agg, signature))
1932 }
1933 "min" => {
1934 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1935 let agg = AggregateExpr {
1936 function: AggregateFunction::Min,
1937 arg: Some(arg.clone()),
1938 distinct,
1939 result_type: arg.resolved_type.clone(),
1940 };
1941 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1942 Ok((agg, signature))
1943 }
1944 "max" => {
1945 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1946 let agg = AggregateExpr {
1947 function: AggregateFunction::Max,
1948 arg: Some(arg.clone()),
1949 distinct,
1950 result_type: arg.resolved_type.clone(),
1951 };
1952 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1953 Ok((agg, signature))
1954 }
1955 "group_concat" => {
1956 if args.is_empty() || args.len() > 2 {
1957 return Err(PlannerError::type_mismatch(
1958 "1 or 2 arguments",
1959 format!("{} arguments", args.len()),
1960 expr.span,
1961 ));
1962 }
1963 let arg = &args[0];
1964 let mut separator = None;
1965 if args.len() == 2 {
1966 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
1967 separator = Some(value.clone());
1968 } else {
1969 return Err(PlannerError::invalid_expression(
1970 "GROUP_CONCAT separator must be a string literal".to_string(),
1971 ));
1972 }
1973 }
1974 let agg = AggregateExpr {
1975 function: AggregateFunction::GroupConcat { separator },
1976 arg: Some(arg.clone()),
1977 distinct,
1978 result_type: ResolvedType::Text,
1979 };
1980 let signature = aggregate_signature(
1981 name,
1982 distinct,
1983 star,
1984 Some(arg),
1985 match &agg.function {
1986 AggregateFunction::GroupConcat { separator } => separator.as_ref(),
1987 _ => None,
1988 },
1989 expr,
1990 );
1991 Ok((agg, signature))
1992 }
1993 "string_agg" => {
1994 if args.len() != 2 {
1995 return Err(PlannerError::type_mismatch(
1996 "2 arguments",
1997 format!("{} arguments", args.len()),
1998 expr.span,
1999 ));
2000 }
2001 let arg = &args[0];
2002 let separator =
2003 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2004 Some(value.clone())
2005 } else {
2006 return Err(PlannerError::invalid_expression(
2007 "STRING_AGG separator must be a string literal".to_string(),
2008 ));
2009 };
2010 let agg = AggregateExpr {
2011 function: AggregateFunction::StringAgg { separator },
2012 arg: Some(arg.clone()),
2013 distinct,
2014 result_type: ResolvedType::Text,
2015 };
2016 let signature = aggregate_signature(
2017 name,
2018 distinct,
2019 star,
2020 Some(arg),
2021 match &agg.function {
2022 AggregateFunction::StringAgg { separator } => separator.as_ref(),
2023 _ => None,
2024 },
2025 expr,
2026 );
2027 Ok((agg, signature))
2028 }
2029 _ => Err(PlannerError::unsupported_feature(
2030 format!("function '{}'", name),
2031 "future",
2032 expr.span,
2033 )),
2034 }
2035 }
2036
2037 fn require_single_aggregate_arg<'b>(
2038 &self,
2039 args: &'b [TypedExpr],
2040 span: crate::ast::Span,
2041 ) -> Result<&'b TypedExpr, PlannerError> {
2042 if args.len() != 1 {
2043 return Err(PlannerError::type_mismatch(
2044 "1 argument",
2045 format!("{} arguments", args.len()),
2046 span,
2047 ));
2048 }
2049 Ok(&args[0])
2050 }
2051
2052 fn build_aggregate_projection(
2053 &self,
2054 projected: Vec<ProjectedColumn>,
2055 group_keys: &[TypedExpr],
2056 aggregates: &[AggregateExpr],
2057 output_names: &[String],
2058 ) -> Result<Projection, PlannerError> {
2059 let mut columns = Vec::new();
2060 for col in projected {
2061 let rewritten =
2062 self.rewrite_expr_for_aggregate(&col.expr, group_keys, aggregates, output_names)?;
2063 columns.push(ProjectedColumn {
2064 expr: rewritten,
2065 alias: col.alias,
2066 });
2067 }
2068 Ok(Projection::Columns(columns))
2069 }
2070
2071 fn rewrite_expr_for_aggregate(
2072 &self,
2073 expr: &TypedExpr,
2074 group_keys: &[TypedExpr],
2075 aggregates: &[AggregateExpr],
2076 output_names: &[String],
2077 ) -> Result<TypedExpr, PlannerError> {
2078 let group_key_map = build_group_key_map(group_keys);
2079 let aggregate_map = build_aggregate_map(aggregates);
2080
2081 rewrite_expr_with_maps(expr, &group_key_map, &aggregate_map, output_names)
2082 }
2083
2084 fn extract_limit_value(
2088 &self,
2089 expr: &Option<crate::ast::expr::Expr>,
2090 stmt_span: crate::ast::Span,
2091 ) -> Result<Option<u64>, PlannerError> {
2092 match expr {
2093 None => Ok(None),
2094 Some(e) => {
2095 if let crate::ast::expr::ExprKind::Literal {
2097 literal: Literal::Number(s),
2098 } = &e.kind
2099 {
2100 s.parse::<u64>().map(Some).map_err(|_| {
2101 PlannerError::type_mismatch("unsigned integer", s.clone(), e.span)
2102 })
2103 } else {
2104 Err(PlannerError::unsupported_feature(
2105 "non-literal LIMIT/OFFSET",
2106 "v0.3.0+",
2107 stmt_span,
2108 ))
2109 }
2110 }
2111 }
2112 }
2113
2114 fn plan_insert(&self, stmt: &Insert) -> Result<LogicalPlan, PlannerError> {
2119 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2121
2122 let columns: Vec<String> = if let Some(ref cols) = stmt.columns {
2124 for col in cols {
2126 self.name_resolver.resolve_column(table, col, stmt.span)?;
2127 }
2128 cols.clone()
2129 } else {
2130 table.column_names().into_iter().map(String::from).collect()
2132 };
2133
2134 match &stmt.source {
2135 InsertSource::Values { values } => {
2136 let mut typed_values: Vec<Vec<TypedExpr>> = Vec::new();
2137
2138 for row in values {
2139 if row.len() != columns.len() {
2140 return Err(PlannerError::column_value_count_mismatch(
2141 columns.len(),
2142 row.len(),
2143 stmt.span,
2144 ));
2145 }
2146
2147 typed_values.push(self.type_check_insert_values(row, &columns, table)?);
2148 }
2149
2150 Ok(LogicalPlan::Insert {
2151 table: table.name.clone(),
2152 columns,
2153 values: typed_values,
2154 })
2155 }
2156 InsertSource::Select { select } => {
2157 let source = self.plan_select_relation(select, &[])?;
2158 if source.schema.len() != columns.len() {
2159 return Err(PlannerError::column_value_count_mismatch(
2160 columns.len(),
2161 source.schema.len(),
2162 stmt.span,
2163 ));
2164 }
2165
2166 for (source_column, target_column) in source.schema.iter().zip(&columns) {
2167 let target = table
2168 .get_column(target_column)
2169 .expect("validated target column");
2170 if target.not_null && source_column.data_type == ResolvedType::Null {
2171 return Err(PlannerError::null_constraint_violation(
2172 target_column,
2173 stmt.span,
2174 ));
2175 }
2176 self.validate_resolved_type_assignment(
2177 &source_column.data_type,
2178 &target.data_type,
2179 stmt.span,
2180 )?;
2181 }
2182
2183 Ok(LogicalPlan::InsertSelect {
2184 table: table.name.clone(),
2185 columns,
2186 source: Box::new(source.plan),
2187 })
2188 }
2189 }
2190 }
2191
2192 fn type_check_insert_values(
2194 &self,
2195 values: &[crate::ast::expr::Expr],
2196 columns: &[String],
2197 table: &TableMetadata,
2198 ) -> Result<Vec<TypedExpr>, PlannerError> {
2199 let mut typed_values = Vec::new();
2200
2201 for (i, value) in values.iter().enumerate() {
2202 let column_name = &columns[i];
2203 let column_meta = table.get_column(column_name).ok_or_else(|| {
2204 PlannerError::column_not_found(column_name, &table.name, value.span)
2205 })?;
2206
2207 let typed_value = self.type_checker.infer_type(value, table)?;
2209
2210 if column_meta.not_null
2212 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2213 {
2214 return Err(PlannerError::null_constraint_violation(
2215 column_name,
2216 value.span,
2217 ));
2218 }
2219
2220 self.validate_type_assignment(&typed_value, &column_meta.data_type, value.span)?;
2222
2223 let typed_value =
2224 self.coerce_assignment_value(typed_value, &column_meta.data_type, value.span);
2225
2226 typed_values.push(typed_value);
2227 }
2228
2229 Ok(typed_values)
2230 }
2231
2232 fn validate_type_assignment(
2234 &self,
2235 value: &TypedExpr,
2236 target_type: &ResolvedType,
2237 span: crate::ast::Span,
2238 ) -> Result<(), PlannerError> {
2239 self.validate_resolved_type_assignment(&value.resolved_type, target_type, span)
2240 }
2241
2242 fn validate_resolved_type_assignment(
2243 &self,
2244 source_type: &ResolvedType,
2245 target_type: &ResolvedType,
2246 span: crate::ast::Span,
2247 ) -> Result<(), PlannerError> {
2248 if *source_type == ResolvedType::Null {
2250 return Ok(());
2251 }
2252
2253 if self.types_compatible(source_type, target_type) {
2255 return Ok(());
2256 }
2257
2258 Err(PlannerError::type_mismatch(
2259 target_type.to_string(),
2260 source_type.to_string(),
2261 span,
2262 ))
2263 }
2264
2265 fn types_compatible(&self, source: &ResolvedType, target: &ResolvedType) -> bool {
2267 use ResolvedType::*;
2268
2269 if source == target {
2271 return true;
2272 }
2273
2274 match (source, target) {
2276 (Integer, BigInt) | (Integer, Float) | (Integer, Double) => true,
2278 (BigInt, Float) | (BigInt, Double) => true,
2280 (Float, Double) => true,
2282 (Double, Float) => true,
2285 (Text | Integer | BigInt | Float | Double, Timestamp) => true,
2288 (Vector { dimension: d1, .. }, Vector { dimension: d2, .. }) => d1 == d2,
2290 _ => false,
2291 }
2292 }
2293
2294 fn coerce_assignment_value(
2295 &self,
2296 value: TypedExpr,
2297 target_type: &ResolvedType,
2298 span: crate::ast::Span,
2299 ) -> TypedExpr {
2300 if matches!(target_type, ResolvedType::Timestamp)
2301 && !matches!(
2302 value.resolved_type,
2303 ResolvedType::Timestamp | ResolvedType::Null
2304 )
2305 {
2306 TypedExpr::cast(value, ResolvedType::Timestamp, span)
2307 } else {
2308 value
2309 }
2310 }
2311
2312 fn plan_update(&self, stmt: &Update) -> Result<LogicalPlan, PlannerError> {
2316 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2318
2319 let mut typed_assignments = Vec::new();
2321
2322 for assignment in &stmt.assignments {
2323 let column_meta =
2325 self.name_resolver
2326 .resolve_column(table, &assignment.column, assignment.span)?;
2327 let column_index = table.get_column_index(&assignment.column).unwrap();
2328
2329 let typed_value = self.type_checker.infer_type(&assignment.value, table)?;
2331
2332 if column_meta.not_null
2334 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2335 {
2336 return Err(PlannerError::null_constraint_violation(
2337 &assignment.column,
2338 assignment.value.span,
2339 ));
2340 }
2341
2342 self.validate_type_assignment(
2344 &typed_value,
2345 &column_meta.data_type,
2346 assignment.value.span,
2347 )?;
2348
2349 let typed_value = self.coerce_assignment_value(
2350 typed_value,
2351 &column_meta.data_type,
2352 assignment.value.span,
2353 );
2354
2355 typed_assignments.push(TypedAssignment::new(
2356 assignment.column.clone(),
2357 column_index,
2358 typed_value,
2359 ));
2360 }
2361
2362 let filter = if let Some(ref selection) = stmt.selection {
2364 let predicate = self.type_checker.infer_type(selection, table)?;
2365
2366 if predicate.resolved_type != ResolvedType::Boolean {
2368 return Err(PlannerError::type_mismatch(
2369 "Boolean",
2370 predicate.resolved_type.to_string(),
2371 selection.span,
2372 ));
2373 }
2374
2375 Some(predicate)
2376 } else {
2377 None
2378 };
2379
2380 Ok(LogicalPlan::Update {
2381 table: table.name.clone(),
2382 assignments: typed_assignments,
2383 filter,
2384 })
2385 }
2386
2387 fn plan_delete(&self, stmt: &Delete) -> Result<LogicalPlan, PlannerError> {
2391 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2393
2394 let filter = if let Some(ref selection) = stmt.selection {
2396 let predicate = self.type_checker.infer_type(selection, table)?;
2397
2398 if predicate.resolved_type != ResolvedType::Boolean {
2400 return Err(PlannerError::type_mismatch(
2401 "Boolean",
2402 predicate.resolved_type.to_string(),
2403 selection.span,
2404 ));
2405 }
2406
2407 Some(predicate)
2408 } else {
2409 None
2410 };
2411
2412 Ok(LogicalPlan::Delete {
2413 table: table.name.clone(),
2414 filter,
2415 })
2416 }
2417}
2418
2419#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2420struct AggregateSignature {
2421 name: String,
2422 distinct: bool,
2423 star: bool,
2424 arg_key: Option<String>,
2425 separator: Option<String>,
2426}
2427
2428fn expr_contains_aggregate(expr: &crate::ast::expr::Expr) -> bool {
2429 use crate::ast::expr::ExprKind;
2430
2431 match &expr.kind {
2432 ExprKind::FunctionCall { name, args, .. } => {
2433 if is_aggregate_function(name) {
2434 return true;
2435 }
2436 args.iter().any(expr_contains_aggregate)
2437 }
2438 ExprKind::BinaryOp { left, right, .. } => {
2439 expr_contains_aggregate(left) || expr_contains_aggregate(right)
2440 }
2441 ExprKind::UnaryOp { operand, .. } => expr_contains_aggregate(operand),
2442 ExprKind::Cast { expr, .. } => expr_contains_aggregate(expr),
2443 ExprKind::Between {
2444 expr, low, high, ..
2445 } => {
2446 expr_contains_aggregate(expr)
2447 || expr_contains_aggregate(low)
2448 || expr_contains_aggregate(high)
2449 }
2450 ExprKind::Like {
2451 expr,
2452 pattern,
2453 escape,
2454 ..
2455 } => {
2456 expr_contains_aggregate(expr)
2457 || expr_contains_aggregate(pattern)
2458 || escape.as_deref().is_some_and(expr_contains_aggregate)
2459 }
2460 ExprKind::InList { expr, list, .. } => {
2461 expr_contains_aggregate(expr) || list.iter().any(expr_contains_aggregate)
2462 }
2463 ExprKind::IsNull { expr, .. } => expr_contains_aggregate(expr),
2464 ExprKind::ScalarSubquery { .. }
2465 | ExprKind::InSubquery { .. }
2466 | ExprKind::Exists { .. }
2467 | ExprKind::Quantified { .. }
2468 | ExprKind::Literal { .. }
2469 | ExprKind::VectorLiteral { .. }
2470 | ExprKind::ColumnRef { .. } => false,
2471 }
2472}
2473
2474fn typed_expr_contains_aggregate(expr: &TypedExpr) -> bool {
2475 match &expr.kind {
2476 TypedExprKind::FunctionCall { name, args, .. } => {
2477 if is_aggregate_function(name) {
2478 return true;
2479 }
2480 args.iter().any(typed_expr_contains_aggregate)
2481 }
2482 TypedExprKind::BinaryOp { left, right, .. } => {
2483 typed_expr_contains_aggregate(left) || typed_expr_contains_aggregate(right)
2484 }
2485 TypedExprKind::UnaryOp { operand, .. } => typed_expr_contains_aggregate(operand),
2486 TypedExprKind::Between {
2487 expr, low, high, ..
2488 } => {
2489 typed_expr_contains_aggregate(expr)
2490 || typed_expr_contains_aggregate(low)
2491 || typed_expr_contains_aggregate(high)
2492 }
2493 TypedExprKind::Like {
2494 expr,
2495 pattern,
2496 escape,
2497 ..
2498 } => {
2499 typed_expr_contains_aggregate(expr)
2500 || typed_expr_contains_aggregate(pattern)
2501 || escape
2502 .as_ref()
2503 .is_some_and(|inner| typed_expr_contains_aggregate(inner))
2504 }
2505 TypedExprKind::InList { expr, list, .. } => {
2506 typed_expr_contains_aggregate(expr) || list.iter().any(typed_expr_contains_aggregate)
2507 }
2508 TypedExprKind::IsNull { expr, .. } => typed_expr_contains_aggregate(expr),
2509 TypedExprKind::InSubquery { expr, .. } => typed_expr_contains_aggregate(expr),
2510 TypedExprKind::Quantified { expr, .. } => typed_expr_contains_aggregate(expr),
2511 TypedExprKind::ScalarSubquery(_) | TypedExprKind::Exists { .. } => false,
2512 _ => false,
2513 }
2514}
2515
2516fn map_join_type(join_type: crate::ast::dml::JoinType) -> JoinType {
2517 match join_type {
2518 crate::ast::dml::JoinType::Inner => JoinType::Inner,
2519 crate::ast::dml::JoinType::Left => JoinType::Left,
2520 crate::ast::dml::JoinType::Right => JoinType::Right,
2521 crate::ast::dml::JoinType::Full => JoinType::Full,
2522 crate::ast::dml::JoinType::Cross => JoinType::Cross,
2523 }
2524}
2525
2526struct FoundScopedColumn {
2527 table: String,
2528 index: usize,
2529 ty: ResolvedType,
2530}
2531
2532fn find_scoped_column(
2533 scope: &[ScopedTable],
2534 column: &str,
2535 span: crate::ast::Span,
2536) -> Result<FoundScopedColumn, PlannerError> {
2537 let mut matches = Vec::new();
2538 for table in scope {
2539 if let Some(local_idx) = table.table.get_column_index(column) {
2540 let meta = &table.table.columns[local_idx];
2541 matches.push(FoundScopedColumn {
2542 table: table.table.name.clone(),
2543 index: table.start_index + local_idx,
2544 ty: meta.data_type.clone(),
2545 });
2546 }
2547 }
2548 match matches.len() {
2549 0 => Err(PlannerError::column_not_found(column, "JOIN input", span)),
2550 1 => Ok(matches.remove(0)),
2551 _ => Err(PlannerError::ambiguous_column(
2552 column,
2553 scope.iter().map(|s| s.table.name.clone()).collect(),
2554 span,
2555 )),
2556 }
2557}
2558
2559fn projection_schema(
2560 projection: &Projection,
2561 input_schema: &[ColumnMetadata],
2562) -> Vec<ColumnMetadata> {
2563 match projection {
2564 Projection::All(names) => names
2565 .iter()
2566 .enumerate()
2567 .map(|(idx, name)| {
2568 let ty = (names.len() == input_schema.len())
2569 .then(|| input_schema.get(idx))
2570 .flatten()
2571 .or_else(|| input_schema.iter().find(|col| &col.name == name))
2572 .map(|col| col.data_type.clone())
2573 .unwrap_or(ResolvedType::Null);
2574 ColumnMetadata::new(name.clone(), ty)
2575 })
2576 .collect(),
2577 Projection::Columns(columns) => columns
2578 .iter()
2579 .enumerate()
2580 .map(|(idx, col)| {
2581 let name = col
2582 .alias
2583 .clone()
2584 .or_else(|| match &col.expr.kind {
2585 TypedExprKind::ColumnRef { column, .. } => Some(column.clone()),
2586 TypedExprKind::FunctionCall { name, args, .. }
2590 if name == "coalesce" && args.len() == 2 =>
2591 {
2592 match (&args[0].kind, &args[1].kind) {
2593 (
2594 TypedExprKind::ColumnRef { column: left, .. },
2595 TypedExprKind::ColumnRef { column: right, .. },
2596 ) if left == right => Some(left.clone()),
2597 _ => None,
2598 }
2599 }
2600 _ => None,
2601 })
2602 .unwrap_or_else(|| format!("col_{idx}"));
2603 ColumnMetadata::new(name, col.expr.resolved_type.clone())
2604 })
2605 .collect(),
2606 }
2607}
2608
2609fn visible_wildcard_columns(schema: &[ColumnMetadata], scope: &[ScopedTable]) -> Vec<String> {
2610 schema
2611 .iter()
2612 .enumerate()
2613 .filter(|(index, column)| {
2614 !scope.iter().any(|table| {
2615 *index >= table.start_index
2616 && *index < table.start_index + table.table.columns.len()
2617 && table.hidden_unqualified_columns.contains(&column.name)
2618 })
2619 })
2620 .map(|(_, column)| column.name.clone())
2621 .collect()
2622}
2623
2624fn offset_scope(scope: &[ScopedTable], offset: usize) -> Vec<ScopedTable> {
2625 scope
2626 .iter()
2627 .cloned()
2628 .map(|mut table| {
2629 table.start_index += offset;
2630 table.scope_level += 1;
2631 table
2632 })
2633 .collect()
2634}
2635
2636fn natural_join_columns(
2637 left_schema: &[ColumnMetadata],
2638 right_schema: &[ColumnMetadata],
2639) -> Vec<String> {
2640 left_schema
2641 .iter()
2642 .filter(|left| right_schema.iter().any(|right| right.name == left.name))
2643 .map(|column| column.name.clone())
2644 .collect()
2645}
2646
2647fn install_base_projection(plan: &mut LogicalPlan, projection: &Projection) {
2648 match plan {
2649 LogicalPlan::Scan {
2650 projection: scan_projection,
2651 ..
2652 } => *scan_projection = projection.clone(),
2653 LogicalPlan::Filter { input, .. } => install_base_projection(input, projection),
2654 _ => {}
2655 }
2656}
2657
2658fn is_aggregate_function(name: &str) -> bool {
2659 matches!(
2660 name.to_ascii_lowercase().as_str(),
2661 "count" | "sum" | "total" | "avg" | "min" | "max" | "group_concat" | "string_agg"
2662 )
2663}
2664
2665fn expr_key(expr: &TypedExpr) -> String {
2666 format!("{:?}", expr.kind)
2667}
2668
2669fn aggregate_signature(
2670 name: &str,
2671 distinct: bool,
2672 star: bool,
2673 arg: Option<&TypedExpr>,
2674 separator: Option<&String>,
2675 _expr: &TypedExpr,
2676) -> AggregateSignature {
2677 AggregateSignature {
2678 name: name.to_ascii_lowercase(),
2679 distinct,
2680 star,
2681 arg_key: arg.map(expr_key),
2682 separator: separator.cloned(),
2683 }
2684}
2685
2686fn build_group_key_map(group_keys: &[TypedExpr]) -> HashMap<String, usize> {
2687 let mut map = HashMap::new();
2688 for (idx, key) in group_keys.iter().enumerate() {
2689 map.insert(expr_key(key), idx);
2690 }
2691 map
2692}
2693
2694fn build_aggregate_map(aggregates: &[AggregateExpr]) -> HashMap<AggregateSignature, usize> {
2695 let mut map = HashMap::new();
2696 for (idx, agg) in aggregates.iter().enumerate() {
2697 let (name, separator, star, arg) = match &agg.function {
2698 AggregateFunction::Count => (
2699 "count".to_string(),
2700 None,
2701 agg.arg.is_none(),
2702 agg.arg.as_ref(),
2703 ),
2704 AggregateFunction::Sum => ("sum".to_string(), None, false, agg.arg.as_ref()),
2705 AggregateFunction::Total => ("total".to_string(), None, false, agg.arg.as_ref()),
2706 AggregateFunction::Avg => ("avg".to_string(), None, false, agg.arg.as_ref()),
2707 AggregateFunction::Min => ("min".to_string(), None, false, agg.arg.as_ref()),
2708 AggregateFunction::Max => ("max".to_string(), None, false, agg.arg.as_ref()),
2709 AggregateFunction::GroupConcat { separator } => (
2710 "group_concat".to_string(),
2711 separator.clone(),
2712 false,
2713 agg.arg.as_ref(),
2714 ),
2715 AggregateFunction::StringAgg { separator } => (
2716 "string_agg".to_string(),
2717 separator.clone(),
2718 false,
2719 agg.arg.as_ref(),
2720 ),
2721 };
2722 let signature = AggregateSignature {
2723 name,
2724 distinct: agg.distinct,
2725 star,
2726 arg_key: arg.map(expr_key),
2727 separator,
2728 };
2729 map.insert(signature, idx);
2730 }
2731 map
2732}
2733
2734fn build_aggregate_schema(
2735 group_keys: &[TypedExpr],
2736 aggregates: &[AggregateExpr],
2737) -> Vec<ColumnMetadata> {
2738 let mut schema = Vec::new();
2739 for (idx, key) in group_keys.iter().enumerate() {
2740 let name = match &key.kind {
2741 TypedExprKind::ColumnRef { column, .. } => column.clone(),
2742 _ => format!("group_{idx}"),
2743 };
2744 schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
2745 }
2746 for (idx, agg) in aggregates.iter().enumerate() {
2747 let name = match &agg.function {
2748 AggregateFunction::Count => format!("count_{idx}"),
2749 AggregateFunction::Sum => format!("sum_{idx}"),
2750 AggregateFunction::Total => format!("total_{idx}"),
2751 AggregateFunction::Avg => format!("avg_{idx}"),
2752 AggregateFunction::Min => format!("min_{idx}"),
2753 AggregateFunction::Max => format!("max_{idx}"),
2754 AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
2755 AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
2756 };
2757 schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
2758 }
2759 schema
2760}
2761
2762fn rewrite_expr_with_maps(
2763 expr: &TypedExpr,
2764 group_key_map: &HashMap<String, usize>,
2765 aggregate_map: &HashMap<AggregateSignature, usize>,
2766 output_names: &[String],
2767) -> Result<TypedExpr, PlannerError> {
2768 let group_key_count = output_names.len().saturating_sub(aggregate_map.len());
2769 let key = expr_key(expr);
2770 if let Some(idx) = group_key_map.get(&key) {
2771 return Ok(make_output_column_ref(
2772 *idx,
2773 output_names,
2774 expr.resolved_type.clone(),
2775 expr.span,
2776 ));
2777 }
2778
2779 match &expr.kind {
2780 TypedExprKind::FunctionCall {
2781 name,
2782 args,
2783 distinct,
2784 star,
2785 } if is_aggregate_function(name) => {
2786 let separator = if name.eq_ignore_ascii_case("group_concat") && args.len() == 2 {
2787 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2788 Some(value.clone())
2789 } else {
2790 return Err(PlannerError::invalid_expression(
2791 "GROUP_CONCAT separator must be a string literal".to_string(),
2792 ));
2793 }
2794 } else if name.eq_ignore_ascii_case("string_agg") && args.len() == 2 {
2795 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2796 Some(value.clone())
2797 } else {
2798 return Err(PlannerError::invalid_expression(
2799 "STRING_AGG separator must be a string literal".to_string(),
2800 ));
2801 }
2802 } else {
2803 None
2804 };
2805 let signature = AggregateSignature {
2806 name: name.to_ascii_lowercase(),
2807 distinct: *distinct,
2808 star: *star,
2809 arg_key: args.first().map(expr_key),
2810 separator,
2811 };
2812 let idx = aggregate_map.get(&signature).ok_or_else(|| {
2813 PlannerError::invalid_expression(
2814 "aggregate in expression is not part of plan".to_string(),
2815 )
2816 })?;
2817 let output_index = group_key_count + idx;
2818 Ok(make_output_column_ref(
2819 output_index,
2820 output_names,
2821 expr.resolved_type.clone(),
2822 expr.span,
2823 ))
2824 }
2825 TypedExprKind::FunctionCall {
2826 name,
2827 args,
2828 distinct,
2829 star,
2830 } => {
2831 if *distinct || *star {
2832 return Err(PlannerError::invalid_expression(
2833 "DISTINCT/STAR modifiers are only supported for aggregates".to_string(),
2834 ));
2835 }
2836 let mut rewritten_args = Vec::with_capacity(args.len());
2837 for arg in args {
2838 rewritten_args.push(rewrite_expr_with_maps(
2839 arg,
2840 group_key_map,
2841 aggregate_map,
2842 output_names,
2843 )?);
2844 }
2845 Ok(TypedExpr {
2846 kind: TypedExprKind::FunctionCall {
2847 name: name.clone(),
2848 args: rewritten_args,
2849 distinct: false,
2850 star: false,
2851 },
2852 resolved_type: expr.resolved_type.clone(),
2853 span: expr.span,
2854 })
2855 }
2856 TypedExprKind::BinaryOp { left, op, right } => {
2857 let left = rewrite_expr_with_maps(left, group_key_map, aggregate_map, output_names)?;
2858 let right = rewrite_expr_with_maps(right, group_key_map, aggregate_map, output_names)?;
2859 Ok(TypedExpr {
2860 kind: TypedExprKind::BinaryOp {
2861 left: Box::new(left),
2862 op: *op,
2863 right: Box::new(right),
2864 },
2865 resolved_type: expr.resolved_type.clone(),
2866 span: expr.span,
2867 })
2868 }
2869 TypedExprKind::UnaryOp { op, operand } => {
2870 let operand =
2871 rewrite_expr_with_maps(operand, group_key_map, aggregate_map, output_names)?;
2872 Ok(TypedExpr {
2873 kind: TypedExprKind::UnaryOp {
2874 op: *op,
2875 operand: Box::new(operand),
2876 },
2877 resolved_type: expr.resolved_type.clone(),
2878 span: expr.span,
2879 })
2880 }
2881 TypedExprKind::Between {
2882 expr: inner,
2883 low,
2884 high,
2885 negated,
2886 } => {
2887 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2888 let low = rewrite_expr_with_maps(low, group_key_map, aggregate_map, output_names)?;
2889 let high = rewrite_expr_with_maps(high, group_key_map, aggregate_map, output_names)?;
2890 Ok(TypedExpr {
2891 kind: TypedExprKind::Between {
2892 expr: Box::new(inner),
2893 low: Box::new(low),
2894 high: Box::new(high),
2895 negated: *negated,
2896 },
2897 resolved_type: expr.resolved_type.clone(),
2898 span: expr.span,
2899 })
2900 }
2901 TypedExprKind::Like {
2902 expr: inner,
2903 pattern,
2904 escape,
2905 negated,
2906 kind,
2907 } => {
2908 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2909 let pattern =
2910 rewrite_expr_with_maps(pattern, group_key_map, aggregate_map, output_names)?;
2911 let escape = if let Some(esc) = escape {
2912 Some(Box::new(rewrite_expr_with_maps(
2913 esc,
2914 group_key_map,
2915 aggregate_map,
2916 output_names,
2917 )?))
2918 } else {
2919 None
2920 };
2921 Ok(TypedExpr {
2922 kind: TypedExprKind::Like {
2923 expr: Box::new(inner),
2924 pattern: Box::new(pattern),
2925 escape,
2926 negated: *negated,
2927 kind: *kind,
2928 },
2929 resolved_type: expr.resolved_type.clone(),
2930 span: expr.span,
2931 })
2932 }
2933 TypedExprKind::InList {
2934 expr: inner,
2935 list,
2936 negated,
2937 } => {
2938 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2939 let mut rewritten_list = Vec::with_capacity(list.len());
2940 for item in list {
2941 rewritten_list.push(rewrite_expr_with_maps(
2942 item,
2943 group_key_map,
2944 aggregate_map,
2945 output_names,
2946 )?);
2947 }
2948 Ok(TypedExpr {
2949 kind: TypedExprKind::InList {
2950 expr: Box::new(inner),
2951 list: rewritten_list,
2952 negated: *negated,
2953 },
2954 resolved_type: expr.resolved_type.clone(),
2955 span: expr.span,
2956 })
2957 }
2958 TypedExprKind::IsNull {
2959 expr: inner,
2960 negated,
2961 } => {
2962 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2963 Ok(TypedExpr {
2964 kind: TypedExprKind::IsNull {
2965 expr: Box::new(inner),
2966 negated: *negated,
2967 },
2968 resolved_type: expr.resolved_type.clone(),
2969 span: expr.span,
2970 })
2971 }
2972 TypedExprKind::Literal(_) | TypedExprKind::VectorLiteral(_) => Ok(expr.clone()),
2973 TypedExprKind::ColumnRef { .. } => Err(PlannerError::invalid_expression(
2974 "column reference must appear in GROUP BY or be aggregated".to_string(),
2975 )),
2976 TypedExprKind::Cast {
2977 expr: inner,
2978 target_type,
2979 } => {
2980 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2981 Ok(TypedExpr {
2982 kind: TypedExprKind::Cast {
2983 expr: Box::new(inner),
2984 target_type: target_type.clone(),
2985 },
2986 resolved_type: expr.resolved_type.clone(),
2987 span: expr.span,
2988 })
2989 }
2990 TypedExprKind::ScalarSubquery(_)
2991 | TypedExprKind::InSubquery { .. }
2992 | TypedExprKind::Exists { .. }
2993 | TypedExprKind::Quantified { .. } => Ok(expr.clone()),
2994 }
2995}
2996
2997fn make_output_column_ref(
2998 index: usize,
2999 output_names: &[String],
3000 resolved_type: ResolvedType,
3001 span: crate::ast::Span,
3002) -> TypedExpr {
3003 let name = output_names
3004 .get(index)
3005 .cloned()
3006 .unwrap_or_else(|| format!("col_{index}"));
3007 TypedExpr::column_ref("__agg__".to_string(), name, index, resolved_type, span)
3008}