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, HashSet};
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, &[])?;
1286 let alias = alias.clone().ok_or_else(|| {
1287 PlannerError::invalid_expression("derived table requires an alias".to_string())
1288 })?;
1289 relation.plan = LogicalPlan::Project {
1290 input: Box::new(relation.plan),
1291 projection: Projection::All(
1292 relation.schema.iter().map(|col| col.name.clone()).collect(),
1293 ),
1294 };
1295 relation.scope = vec![ScopedTable::new(
1296 TableMetadata::new(alias, relation.schema.clone()),
1297 start_index,
1298 )];
1299 Ok(relation)
1300 }
1301 }
1302 }
1303
1304 fn combine_join_relation(
1305 &self,
1306 left: PlannedRelation,
1307 right: PlannedRelation,
1308 join_type: JoinType,
1309 condition: Option<TypedExpr>,
1310 using: Option<Vec<String>>,
1311 _span: crate::ast::Span,
1312 ) -> Result<PlannedRelation, PlannerError> {
1313 let mut schema = left.schema.clone();
1314 schema.extend(right.schema.clone());
1315 let mut scope = left.scope.clone();
1316 let mut right_scope = right.scope.clone();
1317 if let Some(columns) = &using {
1318 for column in columns {
1322 let right_index = right_scope.iter().find_map(|table| {
1323 table
1324 .table
1325 .get_column_index(column)
1326 .map(|index| table.start_index + index)
1327 });
1328 let Some(right_index) = right_index else {
1329 continue;
1330 };
1331 for table in &mut scope {
1332 if table.table.get_column_index(column).is_some() {
1333 table.merge_column_with(column, right_index);
1334 }
1335 }
1336 }
1337 for table in &mut right_scope {
1338 table.hide_unqualified_columns(columns);
1339 }
1340 }
1341 scope.extend(right_scope);
1342 Ok(PlannedRelation {
1343 plan: LogicalPlan::Join {
1344 left: Box::new(left.plan),
1345 right: Box::new(right.plan),
1346 join_type,
1347 condition,
1348 using,
1349 },
1350 schema,
1351 scope,
1352 })
1353 }
1354
1355 fn build_using_condition(
1356 &self,
1357 using: Option<&[String]>,
1358 left: &PlannedRelation,
1359 right: &PlannedRelation,
1360 span: crate::ast::Span,
1361 ) -> Result<Option<TypedExpr>, PlannerError> {
1362 let Some(columns) = using else {
1363 return Ok(None);
1364 };
1365 let mut condition = None;
1366 for column in columns {
1367 let left_col = find_scoped_column(&left.scope, column, span)?;
1368 let right_col = find_scoped_column(&right.scope, column, span)?;
1369 let left_expr = merged_scoped_column_expr(&left_col, column, span);
1370 let right_expr = merged_scoped_column_expr(&right_col, column, span);
1371 self.type_checker
1372 .check_comparison_op(&left_col.ty, &right_col.ty, span)?;
1373 let eq = TypedExpr::binary_op(
1374 left_expr,
1375 crate::ast::expr::BinaryOp::Eq,
1376 right_expr,
1377 ResolvedType::Boolean,
1378 span,
1379 );
1380 condition = Some(match condition {
1381 Some(prev) => TypedExpr::binary_op(
1382 prev,
1383 crate::ast::expr::BinaryOp::And,
1384 eq,
1385 ResolvedType::Boolean,
1386 span,
1387 ),
1388 None => eq,
1389 });
1390 }
1391 Ok(condition)
1392 }
1393
1394 fn infer_expr_with_scope(
1395 &self,
1396 expr: &crate::ast::expr::Expr,
1397 scope: &[ScopedTable],
1398 ) -> Result<TypedExpr, PlannerError> {
1399 self.type_checker
1400 .infer_type_with_scope(expr, scope, &|stmt, outer_scope| {
1401 let crate::ast::StatementKind::Select(select) = &stmt.kind else {
1402 return Err(PlannerError::unsupported_feature(
1403 "non-SELECT subquery",
1404 "v0.6.0-subquery Phase 6",
1405 stmt.span(),
1406 ));
1407 };
1408 let relation = self.plan_select_relation(select, outer_scope)?;
1409 Ok((relation.plan, relation.schema))
1410 })
1411 }
1412
1413 #[allow(dead_code)]
1414 fn build_projection(
1415 &self,
1416 items: &[SelectItem],
1417 table: &TableMetadata,
1418 ) -> Result<Projection, PlannerError> {
1419 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1421 let columns = self.name_resolver.expand_wildcard(table);
1422 return Ok(Projection::All(columns));
1423 }
1424
1425 let mut projected_columns = Vec::new();
1427 for item in items {
1428 match item {
1429 SelectItem::Wildcard { span } => {
1430 for col in &table.columns {
1432 let column_index = table.get_column_index(&col.name).unwrap();
1433 let typed_expr = TypedExpr::column_ref(
1434 table.name.clone(),
1435 col.name.clone(),
1436 column_index,
1437 col.data_type.clone(),
1438 *span,
1439 );
1440 projected_columns.push(ProjectedColumn::new(typed_expr));
1441 }
1442 }
1443 SelectItem::QualifiedWildcard {
1444 table: qualifier,
1445 span,
1446 } => {
1447 if qualifier != &table.name {
1448 return Err(PlannerError::invalid_expression(format!(
1449 "table '{qualifier}' is not available for wildcard projection"
1450 )));
1451 }
1452 for col in &table.columns {
1453 let column_index = table.get_column_index(&col.name).unwrap();
1454 projected_columns.push(ProjectedColumn::new(TypedExpr::column_ref(
1455 table.name.clone(),
1456 col.name.clone(),
1457 column_index,
1458 col.data_type.clone(),
1459 *span,
1460 )));
1461 }
1462 }
1463 SelectItem::Expr { expr, alias, .. } => {
1464 let typed_expr = self.type_checker.infer_type(expr, table)?;
1465 let projected = if let Some(alias) = alias {
1466 ProjectedColumn::with_alias(typed_expr, alias.clone())
1467 } else {
1468 ProjectedColumn::new(typed_expr)
1469 };
1470 projected_columns.push(projected);
1471 }
1472 }
1473 }
1474
1475 Ok(Projection::Columns(projected_columns))
1476 }
1477
1478 fn build_projection_with_scope(
1479 &self,
1480 items: &[SelectItem],
1481 schema: &[ColumnMetadata],
1482 scope: &[ScopedTable],
1483 ) -> Result<Projection, PlannerError> {
1484 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1485 return Ok(Projection::All(visible_wildcard_columns(schema, scope)));
1486 }
1487
1488 let mut projected_columns = Vec::new();
1489 for item in items {
1490 match item {
1491 SelectItem::Wildcard { span } => {
1492 for scoped in scope {
1493 for (local_idx, col) in scoped.table.columns.iter().enumerate() {
1494 projected_columns.push(ProjectedColumn::new(TypedExpr::column_ref(
1495 scoped.table.name.clone(),
1496 col.name.clone(),
1497 scoped.start_index + local_idx,
1498 col.data_type.clone(),
1499 *span,
1500 )));
1501 }
1502 }
1503 }
1504 SelectItem::QualifiedWildcard { table, span } => {
1505 let scoped = scope
1506 .iter()
1507 .filter(|scoped| scoped.table.name == *table)
1508 .collect::<Vec<_>>();
1509 match scoped.as_slice() {
1510 [] => {
1511 return Err(PlannerError::invalid_expression(format!(
1512 "table '{table}' is not available for wildcard projection"
1513 )));
1514 }
1515 [scoped] => {
1516 for (local_idx, col) in scoped.table.columns.iter().enumerate() {
1517 projected_columns.push(ProjectedColumn::new(
1518 TypedExpr::column_ref(
1519 scoped.table.name.clone(),
1520 col.name.clone(),
1521 scoped.start_index + local_idx,
1522 col.data_type.clone(),
1523 *span,
1524 ),
1525 ));
1526 }
1527 }
1528 _ => {
1529 return Err(PlannerError::ambiguous_column(
1530 table,
1531 scoped
1532 .iter()
1533 .map(|scoped| scoped.table.name.clone())
1534 .collect(),
1535 *span,
1536 ));
1537 }
1538 }
1539 }
1540 SelectItem::Expr { expr, alias, .. } => {
1541 let typed_expr = self.infer_expr_with_scope(expr, scope)?;
1542 let projected = if let Some(alias) = alias {
1543 ProjectedColumn::with_alias(typed_expr, alias.clone())
1544 } else {
1545 ProjectedColumn::new(typed_expr)
1546 };
1547 projected_columns.push(projected);
1548 }
1549 }
1550 }
1551
1552 Ok(Projection::Columns(projected_columns))
1553 }
1554
1555 #[allow(dead_code)]
1557 fn build_sort_exprs(
1558 &self,
1559 order_by: &[OrderByExpr],
1560 table: &TableMetadata,
1561 ) -> Result<Vec<SortExpr>, PlannerError> {
1562 let mut sort_exprs = Vec::new();
1563
1564 for order_expr in order_by {
1565 let typed_expr = self.type_checker.infer_type(&order_expr.expr, table)?;
1566
1567 let asc = order_expr.asc.unwrap_or(true);
1569
1570 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1572
1573 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1574 }
1575
1576 Ok(sort_exprs)
1577 }
1578
1579 fn build_sort_exprs_with_scope(
1580 &self,
1581 order_by: &[OrderByExpr],
1582 scope: &[ScopedTable],
1583 ) -> Result<Vec<SortExpr>, PlannerError> {
1584 let mut sort_exprs = Vec::new();
1585 for order_expr in order_by {
1586 let typed_expr = self.infer_expr_with_scope(&order_expr.expr, scope)?;
1587 let asc = order_expr.asc.unwrap_or(true);
1588 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1589 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1590 }
1591 Ok(sort_exprs)
1592 }
1593
1594 fn select_contains_aggregate(&self, stmt: &Select) -> bool {
1595 stmt.projection.iter().any(|item| match item {
1596 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => false,
1597 SelectItem::Expr { expr, .. } => expr_contains_aggregate(expr),
1598 }) || stmt
1599 .group_by
1600 .as_ref()
1601 .map(|items| items.iter().any(expr_contains_aggregate))
1602 .unwrap_or(false)
1603 || stmt
1604 .having
1605 .as_ref()
1606 .map(expr_contains_aggregate)
1607 .unwrap_or(false)
1608 || stmt
1609 .order_by
1610 .iter()
1611 .any(|order| expr_contains_aggregate(&order.expr))
1612 }
1613
1614 #[allow(dead_code)]
1615 fn build_group_keys(
1616 &self,
1617 stmt: &Select,
1618 table: &TableMetadata,
1619 ) -> Result<Vec<TypedExpr>, PlannerError> {
1620 let mut keys = Vec::new();
1621 if let Some(items) = &stmt.group_by {
1622 for expr in items {
1623 let typed = self.type_checker.infer_type(expr, table)?;
1624 if typed_expr_contains_aggregate(&typed) {
1625 return Err(PlannerError::invalid_expression(
1626 "GROUP BY cannot contain aggregate functions".to_string(),
1627 ));
1628 }
1629 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1630 return Err(PlannerError::invalid_expression(
1631 "GROUP BY expressions must be column references".to_string(),
1632 ));
1633 }
1634 keys.push(typed);
1635 }
1636 }
1637 Ok(keys)
1638 }
1639
1640 fn build_group_keys_with_scope(
1641 &self,
1642 stmt: &Select,
1643 scope: &[ScopedTable],
1644 ) -> Result<Vec<TypedExpr>, PlannerError> {
1645 let mut keys = Vec::new();
1646 if let Some(items) = &stmt.group_by {
1647 for expr in items {
1648 let typed = self.infer_expr_with_scope(expr, scope)?;
1649 if typed_expr_contains_aggregate(&typed) {
1650 return Err(PlannerError::invalid_expression(
1651 "GROUP BY cannot contain aggregate functions".to_string(),
1652 ));
1653 }
1654 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1655 return Err(PlannerError::invalid_expression(
1656 "GROUP BY expressions must be column references".to_string(),
1657 ));
1658 }
1659 keys.push(typed);
1660 }
1661 }
1662 Ok(keys)
1663 }
1664
1665 #[allow(dead_code)]
1666 fn build_projected_columns_for_aggregate(
1667 &self,
1668 items: &[SelectItem],
1669 table: &TableMetadata,
1670 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1671 let mut projected = Vec::new();
1672 for item in items {
1673 match item {
1674 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => {
1675 return Err(PlannerError::invalid_expression(
1676 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1677 ));
1678 }
1679 SelectItem::Expr { expr, alias, .. } => {
1680 let typed = self.type_checker.infer_type(expr, table)?;
1681 projected.push(ProjectedColumn {
1682 expr: typed,
1683 alias: alias.clone(),
1684 });
1685 }
1686 }
1687 }
1688 Ok(projected)
1689 }
1690
1691 fn build_projected_columns_for_aggregate_with_scope(
1692 &self,
1693 items: &[SelectItem],
1694 scope: &[ScopedTable],
1695 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1696 let mut projected = Vec::new();
1697 for item in items {
1698 match item {
1699 SelectItem::Wildcard { .. } | SelectItem::QualifiedWildcard { .. } => {
1700 return Err(PlannerError::invalid_expression(
1701 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1702 ));
1703 }
1704 SelectItem::Expr { expr, alias, .. } => {
1705 let typed = self.infer_expr_with_scope(expr, scope)?;
1706 projected.push(ProjectedColumn {
1707 expr: typed,
1708 alias: alias.clone(),
1709 });
1710 }
1711 }
1712 }
1713 Ok(projected)
1714 }
1715
1716 #[allow(dead_code)]
1717 fn build_projected_columns_for_distinct(
1718 &self,
1719 items: &[SelectItem],
1720 table: &TableMetadata,
1721 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1722 let projection = self.build_projection(items, table)?;
1723 match projection {
1724 Projection::All(columns) => {
1725 let mut projected = Vec::with_capacity(columns.len());
1726 for column in columns {
1727 let column_index = table.get_column_index(&column).ok_or_else(|| {
1728 PlannerError::invalid_expression(format!(
1729 "column '{column}' not found for DISTINCT projection"
1730 ))
1731 })?;
1732 let column_meta = table.get_column(&column).ok_or_else(|| {
1733 PlannerError::invalid_expression(format!(
1734 "column '{column}' not found for DISTINCT projection"
1735 ))
1736 })?;
1737 let typed_expr = TypedExpr::column_ref(
1738 table.name.clone(),
1739 column.clone(),
1740 column_index,
1741 column_meta.data_type.clone(),
1742 crate::ast::Span::default(),
1743 );
1744 projected.push(ProjectedColumn::new(typed_expr));
1745 }
1746 Ok(projected)
1747 }
1748 Projection::Columns(columns) => Ok(columns),
1749 }
1750 }
1751
1752 fn build_projected_columns_for_distinct_with_scope(
1753 &self,
1754 items: &[SelectItem],
1755 schema: &[ColumnMetadata],
1756 scope: &[ScopedTable],
1757 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1758 let projection = self.build_projection_with_scope(items, schema, scope)?;
1759 match projection {
1760 Projection::All(columns) => {
1761 let mut projected = Vec::with_capacity(columns.len());
1762 for (idx, column) in columns.into_iter().enumerate() {
1763 let column_meta = schema.get(idx).ok_or_else(|| {
1764 PlannerError::invalid_expression(format!(
1765 "column '{column}' not found for DISTINCT projection"
1766 ))
1767 })?;
1768 projected.push(ProjectedColumn::new(TypedExpr::column_ref(
1769 LITERAL_TABLE.to_string(),
1770 column,
1771 idx,
1772 column_meta.data_type.clone(),
1773 crate::ast::Span::default(),
1774 )));
1775 }
1776 Ok(projected)
1777 }
1778 Projection::Columns(columns) => Ok(columns),
1779 }
1780 }
1781
1782 fn collect_aggregates_from_typed_expr(
1783 &self,
1784 expr: &TypedExpr,
1785 aggregates: &mut Vec<AggregateExpr>,
1786 aggregate_map: &mut HashMap<AggregateSignature, usize>,
1787 ) -> Result<(), PlannerError> {
1788 match &expr.kind {
1789 TypedExprKind::FunctionCall {
1790 name,
1791 args,
1792 distinct,
1793 star,
1794 } if is_aggregate_function(name) => {
1795 for arg in args {
1796 if typed_expr_contains_aggregate(arg) {
1797 return Err(PlannerError::invalid_expression(
1798 "nested aggregate functions are not supported".to_string(),
1799 ));
1800 }
1801 }
1802 let (agg, signature) =
1803 self.build_aggregate_expr_from_typed(expr, name, args, *distinct, *star)?;
1804 aggregate_map.entry(signature).or_insert_with(|| {
1805 aggregates.push(agg);
1806 aggregates.len() - 1
1807 });
1808 Ok(())
1809 }
1810 TypedExprKind::BinaryOp { left, right, .. } => {
1811 self.collect_aggregates_from_typed_expr(left, aggregates, aggregate_map)?;
1812 self.collect_aggregates_from_typed_expr(right, aggregates, aggregate_map)?;
1813 Ok(())
1814 }
1815 TypedExprKind::UnaryOp { operand, .. } => {
1816 self.collect_aggregates_from_typed_expr(operand, aggregates, aggregate_map)
1817 }
1818 TypedExprKind::FunctionCall { args, .. } => {
1819 for arg in args {
1820 self.collect_aggregates_from_typed_expr(arg, aggregates, aggregate_map)?;
1821 }
1822 Ok(())
1823 }
1824 TypedExprKind::Between {
1825 expr, low, high, ..
1826 } => {
1827 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1828 self.collect_aggregates_from_typed_expr(low, aggregates, aggregate_map)?;
1829 self.collect_aggregates_from_typed_expr(high, aggregates, aggregate_map)?;
1830 Ok(())
1831 }
1832 TypedExprKind::Like {
1833 expr,
1834 pattern,
1835 escape,
1836 ..
1837 } => {
1838 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1839 self.collect_aggregates_from_typed_expr(pattern, aggregates, aggregate_map)?;
1840 if let Some(esc) = escape {
1841 self.collect_aggregates_from_typed_expr(esc, aggregates, aggregate_map)?;
1842 }
1843 Ok(())
1844 }
1845 TypedExprKind::InList { expr, list, .. } => {
1846 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1847 for item in list {
1848 self.collect_aggregates_from_typed_expr(item, aggregates, aggregate_map)?;
1849 }
1850 Ok(())
1851 }
1852 TypedExprKind::IsNull { expr, .. } => {
1853 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)
1854 }
1855 _ => Ok(()),
1856 }
1857 }
1858
1859 fn build_aggregate_expr_from_typed(
1860 &self,
1861 expr: &TypedExpr,
1862 name: &str,
1863 args: &[TypedExpr],
1864 distinct: bool,
1865 star: bool,
1866 ) -> Result<(AggregateExpr, AggregateSignature), PlannerError> {
1867 let lower = name.to_lowercase();
1868 match lower.as_str() {
1869 "count" => {
1870 if star {
1871 let agg = AggregateExpr::count_star();
1872 let signature = aggregate_signature(name, distinct, star, None, None, expr);
1873 return Ok((agg, signature));
1874 }
1875 if args.len() != 1 {
1876 return Err(PlannerError::type_mismatch(
1877 "1 argument",
1878 format!("{} arguments", args.len()),
1879 expr.span,
1880 ));
1881 }
1882 let agg = AggregateExpr {
1883 function: AggregateFunction::Count,
1884 arg: Some(args[0].clone()),
1885 distinct,
1886 result_type: ResolvedType::BigInt,
1887 };
1888 let signature =
1889 aggregate_signature(name, distinct, star, Some(&args[0]), None, expr);
1890 Ok((agg, signature))
1891 }
1892 "sum" => {
1893 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1894 let agg = AggregateExpr {
1895 function: AggregateFunction::Sum,
1896 arg: Some(arg.clone()),
1897 distinct,
1898 result_type: crate::planner::aggregate_expr::sum_result_type(
1899 &arg.resolved_type,
1900 ),
1901 };
1902 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1903 Ok((agg, signature))
1904 }
1905 "total" => {
1906 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1907 let agg = AggregateExpr {
1908 function: AggregateFunction::Total,
1909 arg: Some(arg.clone()),
1910 distinct: false,
1911 result_type: ResolvedType::Double,
1912 };
1913 let signature = aggregate_signature(name, false, star, Some(arg), None, expr);
1914 Ok((agg, signature))
1915 }
1916 "avg" => {
1917 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1918 let agg = AggregateExpr {
1919 function: AggregateFunction::Avg,
1920 arg: Some(arg.clone()),
1921 distinct,
1922 result_type: ResolvedType::Double,
1923 };
1924 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1925 Ok((agg, signature))
1926 }
1927 "min" => {
1928 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1929 let agg = AggregateExpr {
1930 function: AggregateFunction::Min,
1931 arg: Some(arg.clone()),
1932 distinct,
1933 result_type: arg.resolved_type.clone(),
1934 };
1935 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1936 Ok((agg, signature))
1937 }
1938 "max" => {
1939 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1940 let agg = AggregateExpr {
1941 function: AggregateFunction::Max,
1942 arg: Some(arg.clone()),
1943 distinct,
1944 result_type: arg.resolved_type.clone(),
1945 };
1946 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1947 Ok((agg, signature))
1948 }
1949 "group_concat" => {
1950 if args.is_empty() || args.len() > 2 {
1951 return Err(PlannerError::type_mismatch(
1952 "1 or 2 arguments",
1953 format!("{} arguments", args.len()),
1954 expr.span,
1955 ));
1956 }
1957 let arg = &args[0];
1958 let mut separator = None;
1959 if args.len() == 2 {
1960 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
1961 separator = Some(value.clone());
1962 } else {
1963 return Err(PlannerError::invalid_expression(
1964 "GROUP_CONCAT separator must be a string literal".to_string(),
1965 ));
1966 }
1967 }
1968 let agg = AggregateExpr {
1969 function: AggregateFunction::GroupConcat { separator },
1970 arg: Some(arg.clone()),
1971 distinct,
1972 result_type: ResolvedType::Text,
1973 };
1974 let signature = aggregate_signature(
1975 name,
1976 distinct,
1977 star,
1978 Some(arg),
1979 match &agg.function {
1980 AggregateFunction::GroupConcat { separator } => separator.as_ref(),
1981 _ => None,
1982 },
1983 expr,
1984 );
1985 Ok((agg, signature))
1986 }
1987 "string_agg" => {
1988 if args.len() != 2 {
1989 return Err(PlannerError::type_mismatch(
1990 "2 arguments",
1991 format!("{} arguments", args.len()),
1992 expr.span,
1993 ));
1994 }
1995 let arg = &args[0];
1996 let separator =
1997 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
1998 Some(value.clone())
1999 } else {
2000 return Err(PlannerError::invalid_expression(
2001 "STRING_AGG separator must be a string literal".to_string(),
2002 ));
2003 };
2004 let agg = AggregateExpr {
2005 function: AggregateFunction::StringAgg { separator },
2006 arg: Some(arg.clone()),
2007 distinct,
2008 result_type: ResolvedType::Text,
2009 };
2010 let signature = aggregate_signature(
2011 name,
2012 distinct,
2013 star,
2014 Some(arg),
2015 match &agg.function {
2016 AggregateFunction::StringAgg { separator } => separator.as_ref(),
2017 _ => None,
2018 },
2019 expr,
2020 );
2021 Ok((agg, signature))
2022 }
2023 _ => Err(PlannerError::unsupported_feature(
2024 format!("function '{}'", name),
2025 "future",
2026 expr.span,
2027 )),
2028 }
2029 }
2030
2031 fn require_single_aggregate_arg<'b>(
2032 &self,
2033 args: &'b [TypedExpr],
2034 span: crate::ast::Span,
2035 ) -> Result<&'b TypedExpr, PlannerError> {
2036 if args.len() != 1 {
2037 return Err(PlannerError::type_mismatch(
2038 "1 argument",
2039 format!("{} arguments", args.len()),
2040 span,
2041 ));
2042 }
2043 Ok(&args[0])
2044 }
2045
2046 fn build_aggregate_projection(
2047 &self,
2048 projected: Vec<ProjectedColumn>,
2049 group_keys: &[TypedExpr],
2050 aggregates: &[AggregateExpr],
2051 output_names: &[String],
2052 ) -> Result<Projection, PlannerError> {
2053 let mut columns = Vec::new();
2054 for col in projected {
2055 let rewritten =
2056 self.rewrite_expr_for_aggregate(&col.expr, group_keys, aggregates, output_names)?;
2057 columns.push(ProjectedColumn {
2058 expr: rewritten,
2059 alias: col.alias,
2060 });
2061 }
2062 Ok(Projection::Columns(columns))
2063 }
2064
2065 fn rewrite_expr_for_aggregate(
2066 &self,
2067 expr: &TypedExpr,
2068 group_keys: &[TypedExpr],
2069 aggregates: &[AggregateExpr],
2070 output_names: &[String],
2071 ) -> Result<TypedExpr, PlannerError> {
2072 let group_key_map = build_group_key_map(group_keys);
2073 let aggregate_map = build_aggregate_map(aggregates);
2074
2075 rewrite_expr_with_maps(expr, &group_key_map, &aggregate_map, output_names)
2076 }
2077
2078 fn extract_limit_value(
2082 &self,
2083 expr: &Option<crate::ast::expr::Expr>,
2084 stmt_span: crate::ast::Span,
2085 ) -> Result<Option<u64>, PlannerError> {
2086 match expr {
2087 None => Ok(None),
2088 Some(e) => {
2089 if let crate::ast::expr::ExprKind::Literal {
2091 literal: Literal::Number(s),
2092 } = &e.kind
2093 {
2094 s.parse::<u64>().map(Some).map_err(|_| {
2095 PlannerError::type_mismatch("unsigned integer", s.clone(), e.span)
2096 })
2097 } else {
2098 Err(PlannerError::unsupported_feature(
2099 "non-literal LIMIT/OFFSET",
2100 "v0.3.0+",
2101 stmt_span,
2102 ))
2103 }
2104 }
2105 }
2106 }
2107
2108 fn plan_insert(&self, stmt: &Insert) -> Result<LogicalPlan, PlannerError> {
2113 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2115
2116 let columns: Vec<String> = if let Some(ref cols) = stmt.columns {
2118 for col in cols {
2120 self.name_resolver.resolve_column(table, col, stmt.span)?;
2121 }
2122 cols.clone()
2123 } else {
2124 table.column_names().into_iter().map(String::from).collect()
2126 };
2127
2128 match &stmt.source {
2129 InsertSource::Values { values } => {
2130 let mut typed_values: Vec<Vec<TypedExpr>> = Vec::new();
2131
2132 for row in values {
2133 if row.len() != columns.len() {
2134 return Err(PlannerError::column_value_count_mismatch(
2135 columns.len(),
2136 row.len(),
2137 stmt.span,
2138 ));
2139 }
2140
2141 typed_values.push(self.type_check_insert_values(row, &columns, table)?);
2142 }
2143
2144 Ok(LogicalPlan::Insert {
2145 table: table.name.clone(),
2146 columns,
2147 values: typed_values,
2148 })
2149 }
2150 InsertSource::Select { select } => {
2151 let source = self.plan_select_relation(select, &[])?;
2152 if source.schema.len() != columns.len() {
2153 return Err(PlannerError::column_value_count_mismatch(
2154 columns.len(),
2155 source.schema.len(),
2156 stmt.span,
2157 ));
2158 }
2159
2160 for (source_column, target_column) in source.schema.iter().zip(&columns) {
2161 let target = table
2162 .get_column(target_column)
2163 .expect("validated target column");
2164 if target.not_null && source_column.data_type == ResolvedType::Null {
2165 return Err(PlannerError::null_constraint_violation(
2166 target_column,
2167 stmt.span,
2168 ));
2169 }
2170 self.validate_resolved_type_assignment(
2171 &source_column.data_type,
2172 &target.data_type,
2173 stmt.span,
2174 )?;
2175 }
2176
2177 Ok(LogicalPlan::InsertSelect {
2178 table: table.name.clone(),
2179 columns,
2180 source: Box::new(source.plan),
2181 })
2182 }
2183 }
2184 }
2185
2186 fn type_check_insert_values(
2188 &self,
2189 values: &[crate::ast::expr::Expr],
2190 columns: &[String],
2191 table: &TableMetadata,
2192 ) -> Result<Vec<TypedExpr>, PlannerError> {
2193 let mut typed_values = Vec::new();
2194
2195 for (i, value) in values.iter().enumerate() {
2196 let column_name = &columns[i];
2197 let column_meta = table.get_column(column_name).ok_or_else(|| {
2198 PlannerError::column_not_found(column_name, &table.name, value.span)
2199 })?;
2200
2201 let typed_value = self.type_checker.infer_type(value, table)?;
2203
2204 if column_meta.not_null
2206 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2207 {
2208 return Err(PlannerError::null_constraint_violation(
2209 column_name,
2210 value.span,
2211 ));
2212 }
2213
2214 self.validate_type_assignment(&typed_value, &column_meta.data_type, value.span)?;
2216
2217 let typed_value =
2218 self.coerce_assignment_value(typed_value, &column_meta.data_type, value.span);
2219
2220 typed_values.push(typed_value);
2221 }
2222
2223 Ok(typed_values)
2224 }
2225
2226 fn validate_type_assignment(
2228 &self,
2229 value: &TypedExpr,
2230 target_type: &ResolvedType,
2231 span: crate::ast::Span,
2232 ) -> Result<(), PlannerError> {
2233 self.validate_resolved_type_assignment(&value.resolved_type, target_type, span)
2234 }
2235
2236 fn validate_resolved_type_assignment(
2237 &self,
2238 source_type: &ResolvedType,
2239 target_type: &ResolvedType,
2240 span: crate::ast::Span,
2241 ) -> Result<(), PlannerError> {
2242 if *source_type == ResolvedType::Null {
2244 return Ok(());
2245 }
2246
2247 if self.types_compatible(source_type, target_type) {
2249 return Ok(());
2250 }
2251
2252 Err(PlannerError::type_mismatch(
2253 target_type.to_string(),
2254 source_type.to_string(),
2255 span,
2256 ))
2257 }
2258
2259 fn types_compatible(&self, source: &ResolvedType, target: &ResolvedType) -> bool {
2261 use ResolvedType::*;
2262
2263 if source == target {
2265 return true;
2266 }
2267
2268 match (source, target) {
2270 (Integer, BigInt) | (Integer, Float) | (Integer, Double) => true,
2272 (BigInt, Float) | (BigInt, Double) => true,
2274 (Float, Double) => true,
2276 (Double, Float) => true,
2279 (Text | Integer | BigInt | Float | Double, Timestamp) => true,
2282 (Vector { dimension: d1, .. }, Vector { dimension: d2, .. }) => d1 == d2,
2284 _ => false,
2285 }
2286 }
2287
2288 fn coerce_assignment_value(
2289 &self,
2290 value: TypedExpr,
2291 target_type: &ResolvedType,
2292 span: crate::ast::Span,
2293 ) -> TypedExpr {
2294 if matches!(target_type, ResolvedType::Timestamp)
2295 && !matches!(
2296 value.resolved_type,
2297 ResolvedType::Timestamp | ResolvedType::Null
2298 )
2299 {
2300 TypedExpr::cast(value, ResolvedType::Timestamp, span)
2301 } else {
2302 value
2303 }
2304 }
2305
2306 fn plan_update(&self, stmt: &Update) -> Result<LogicalPlan, PlannerError> {
2310 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2312
2313 let mut typed_assignments = Vec::new();
2315
2316 for assignment in &stmt.assignments {
2317 let column_meta =
2319 self.name_resolver
2320 .resolve_column(table, &assignment.column, assignment.span)?;
2321 let column_index = table.get_column_index(&assignment.column).unwrap();
2322
2323 let typed_value = self.type_checker.infer_type(&assignment.value, table)?;
2325
2326 if column_meta.not_null
2328 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2329 {
2330 return Err(PlannerError::null_constraint_violation(
2331 &assignment.column,
2332 assignment.value.span,
2333 ));
2334 }
2335
2336 self.validate_type_assignment(
2338 &typed_value,
2339 &column_meta.data_type,
2340 assignment.value.span,
2341 )?;
2342
2343 let typed_value = self.coerce_assignment_value(
2344 typed_value,
2345 &column_meta.data_type,
2346 assignment.value.span,
2347 );
2348
2349 typed_assignments.push(TypedAssignment::new(
2350 assignment.column.clone(),
2351 column_index,
2352 typed_value,
2353 ));
2354 }
2355
2356 let filter = if let Some(ref selection) = stmt.selection {
2358 let predicate = self.type_checker.infer_type(selection, table)?;
2359
2360 if predicate.resolved_type != ResolvedType::Boolean {
2362 return Err(PlannerError::type_mismatch(
2363 "Boolean",
2364 predicate.resolved_type.to_string(),
2365 selection.span,
2366 ));
2367 }
2368
2369 Some(predicate)
2370 } else {
2371 None
2372 };
2373
2374 Ok(LogicalPlan::Update {
2375 table: table.name.clone(),
2376 assignments: typed_assignments,
2377 filter,
2378 })
2379 }
2380
2381 fn plan_delete(&self, stmt: &Delete) -> Result<LogicalPlan, PlannerError> {
2385 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2387
2388 let filter = if let Some(ref selection) = stmt.selection {
2390 let predicate = self.type_checker.infer_type(selection, table)?;
2391
2392 if predicate.resolved_type != ResolvedType::Boolean {
2394 return Err(PlannerError::type_mismatch(
2395 "Boolean",
2396 predicate.resolved_type.to_string(),
2397 selection.span,
2398 ));
2399 }
2400
2401 Some(predicate)
2402 } else {
2403 None
2404 };
2405
2406 Ok(LogicalPlan::Delete {
2407 table: table.name.clone(),
2408 filter,
2409 })
2410 }
2411}
2412
2413#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2414struct AggregateSignature {
2415 name: String,
2416 distinct: bool,
2417 star: bool,
2418 arg_key: Option<String>,
2419 separator: Option<String>,
2420}
2421
2422fn expr_contains_aggregate(expr: &crate::ast::expr::Expr) -> bool {
2423 use crate::ast::expr::ExprKind;
2424
2425 match &expr.kind {
2426 ExprKind::FunctionCall { name, args, .. } => {
2427 if is_aggregate_function(name) {
2428 return true;
2429 }
2430 args.iter().any(expr_contains_aggregate)
2431 }
2432 ExprKind::BinaryOp { left, right, .. } => {
2433 expr_contains_aggregate(left) || expr_contains_aggregate(right)
2434 }
2435 ExprKind::UnaryOp { operand, .. } => expr_contains_aggregate(operand),
2436 ExprKind::Cast { expr, .. } => expr_contains_aggregate(expr),
2437 ExprKind::Between {
2438 expr, low, high, ..
2439 } => {
2440 expr_contains_aggregate(expr)
2441 || expr_contains_aggregate(low)
2442 || expr_contains_aggregate(high)
2443 }
2444 ExprKind::Like {
2445 expr,
2446 pattern,
2447 escape,
2448 ..
2449 } => {
2450 expr_contains_aggregate(expr)
2451 || expr_contains_aggregate(pattern)
2452 || escape.as_deref().is_some_and(expr_contains_aggregate)
2453 }
2454 ExprKind::InList { expr, list, .. } => {
2455 expr_contains_aggregate(expr) || list.iter().any(expr_contains_aggregate)
2456 }
2457 ExprKind::IsNull { expr, .. } => expr_contains_aggregate(expr),
2458 ExprKind::ScalarSubquery { .. }
2459 | ExprKind::InSubquery { .. }
2460 | ExprKind::Exists { .. }
2461 | ExprKind::Quantified { .. }
2462 | ExprKind::Literal { .. }
2463 | ExprKind::VectorLiteral { .. }
2464 | ExprKind::ColumnRef { .. } => false,
2465 }
2466}
2467
2468fn typed_expr_contains_aggregate(expr: &TypedExpr) -> bool {
2469 match &expr.kind {
2470 TypedExprKind::FunctionCall { name, args, .. } => {
2471 if is_aggregate_function(name) {
2472 return true;
2473 }
2474 args.iter().any(typed_expr_contains_aggregate)
2475 }
2476 TypedExprKind::BinaryOp { left, right, .. } => {
2477 typed_expr_contains_aggregate(left) || typed_expr_contains_aggregate(right)
2478 }
2479 TypedExprKind::UnaryOp { operand, .. } => typed_expr_contains_aggregate(operand),
2480 TypedExprKind::Between {
2481 expr, low, high, ..
2482 } => {
2483 typed_expr_contains_aggregate(expr)
2484 || typed_expr_contains_aggregate(low)
2485 || typed_expr_contains_aggregate(high)
2486 }
2487 TypedExprKind::Like {
2488 expr,
2489 pattern,
2490 escape,
2491 ..
2492 } => {
2493 typed_expr_contains_aggregate(expr)
2494 || typed_expr_contains_aggregate(pattern)
2495 || escape
2496 .as_ref()
2497 .is_some_and(|inner| typed_expr_contains_aggregate(inner))
2498 }
2499 TypedExprKind::InList { expr, list, .. } => {
2500 typed_expr_contains_aggregate(expr) || list.iter().any(typed_expr_contains_aggregate)
2501 }
2502 TypedExprKind::IsNull { expr, .. } => typed_expr_contains_aggregate(expr),
2503 TypedExprKind::InSubquery { expr, .. } => typed_expr_contains_aggregate(expr),
2504 TypedExprKind::Quantified { expr, .. } => typed_expr_contains_aggregate(expr),
2505 TypedExprKind::ScalarSubquery(_) | TypedExprKind::Exists { .. } => false,
2506 _ => false,
2507 }
2508}
2509
2510fn map_join_type(join_type: crate::ast::dml::JoinType) -> JoinType {
2511 match join_type {
2512 crate::ast::dml::JoinType::Inner => JoinType::Inner,
2513 crate::ast::dml::JoinType::Left => JoinType::Left,
2514 crate::ast::dml::JoinType::Right => JoinType::Right,
2515 crate::ast::dml::JoinType::Full => JoinType::Full,
2516 crate::ast::dml::JoinType::Cross => JoinType::Cross,
2517 }
2518}
2519
2520struct FoundScopedColumn {
2521 table: String,
2522 index: usize,
2523 ty: ResolvedType,
2524 partner_indices: Vec<usize>,
2525}
2526
2527fn find_scoped_column(
2528 scope: &[ScopedTable],
2529 column: &str,
2530 span: crate::ast::Span,
2531) -> Result<FoundScopedColumn, PlannerError> {
2532 let mut matches = Vec::new();
2533 for table in scope {
2534 if table.hidden_unqualified_columns.contains(column) {
2535 continue;
2536 }
2537 if let Some(local_idx) = table.table.get_column_index(column) {
2538 let meta = &table.table.columns[local_idx];
2539 matches.push(FoundScopedColumn {
2540 table: table.table.name.clone(),
2541 index: table.start_index + local_idx,
2542 ty: meta.data_type.clone(),
2543 partner_indices: table
2544 .merged_column_partners
2545 .get(column)
2546 .cloned()
2547 .unwrap_or_default(),
2548 });
2549 }
2550 }
2551 match matches.len() {
2552 0 => Err(PlannerError::column_not_found(column, "JOIN input", span)),
2553 1 => Ok(matches.remove(0)),
2554 _ => Err(PlannerError::ambiguous_column(
2555 column,
2556 scope.iter().map(|s| s.table.name.clone()).collect(),
2557 span,
2558 )),
2559 }
2560}
2561
2562fn merged_scoped_column_expr(
2563 found: &FoundScopedColumn,
2564 column: &str,
2565 span: crate::ast::Span,
2566) -> TypedExpr {
2567 let own = TypedExpr::column_ref(
2568 found.table.clone(),
2569 column.to_string(),
2570 found.index,
2571 found.ty.clone(),
2572 span,
2573 );
2574 if found.partner_indices.is_empty() {
2575 return own;
2576 }
2577
2578 let mut args = Vec::with_capacity(found.partner_indices.len() + 1);
2579 args.push(own);
2580 args.extend(found.partner_indices.iter().map(|&index| {
2581 TypedExpr::column_ref(
2582 found.table.clone(),
2583 column.to_string(),
2584 index,
2585 found.ty.clone(),
2586 span,
2587 )
2588 }));
2589 TypedExpr {
2590 kind: TypedExprKind::FunctionCall {
2591 name: "coalesce".to_string(),
2592 args,
2593 distinct: false,
2594 star: false,
2595 },
2596 resolved_type: found.ty.clone(),
2597 span,
2598 }
2599}
2600
2601fn projection_schema(
2602 projection: &Projection,
2603 input_schema: &[ColumnMetadata],
2604) -> Vec<ColumnMetadata> {
2605 match projection {
2606 Projection::All(names) => names
2607 .iter()
2608 .enumerate()
2609 .map(|(idx, name)| {
2610 let ty = (names.len() == input_schema.len())
2611 .then(|| input_schema.get(idx))
2612 .flatten()
2613 .or_else(|| input_schema.iter().find(|col| &col.name == name))
2614 .map(|col| col.data_type.clone())
2615 .unwrap_or(ResolvedType::Null);
2616 ColumnMetadata::new(name.clone(), ty)
2617 })
2618 .collect(),
2619 Projection::Columns(columns) => columns
2620 .iter()
2621 .enumerate()
2622 .map(|(idx, col)| {
2623 let name = col
2624 .alias
2625 .clone()
2626 .or_else(|| match &col.expr.kind {
2627 TypedExprKind::ColumnRef { column, .. } => Some(column.clone()),
2628 TypedExprKind::FunctionCall { name, args, .. }
2632 if name == "coalesce" && !args.is_empty() =>
2633 {
2634 let first_column = match &args[0].kind {
2635 TypedExprKind::ColumnRef { column, .. } => Some(column),
2636 _ => None,
2637 };
2638 first_column
2639 .filter(|column| {
2640 args.iter().all(|arg| {
2641 matches!(
2642 &arg.kind,
2643 TypedExprKind::ColumnRef { column: other, .. }
2644 if other == *column
2645 )
2646 })
2647 })
2648 .cloned()
2649 }
2650 _ => None,
2651 })
2652 .unwrap_or_else(|| format!("col_{idx}"));
2653 ColumnMetadata::new(name, col.expr.resolved_type.clone())
2654 })
2655 .collect(),
2656 }
2657}
2658
2659fn visible_wildcard_columns(schema: &[ColumnMetadata], scope: &[ScopedTable]) -> Vec<String> {
2660 schema
2661 .iter()
2662 .enumerate()
2663 .filter(|(index, column)| {
2664 !scope.iter().any(|table| {
2665 *index >= table.start_index
2666 && *index < table.start_index + table.table.columns.len()
2667 && table.hidden_unqualified_columns.contains(&column.name)
2668 })
2669 })
2670 .map(|(_, column)| column.name.clone())
2671 .collect()
2672}
2673
2674fn offset_scope(scope: &[ScopedTable], offset: usize) -> Vec<ScopedTable> {
2675 scope
2676 .iter()
2677 .cloned()
2678 .map(|mut table| {
2679 table.start_index += offset;
2680 table.scope_level += 1;
2681 table
2682 })
2683 .collect()
2684}
2685
2686fn natural_join_columns(
2687 left_schema: &[ColumnMetadata],
2688 right_schema: &[ColumnMetadata],
2689) -> Vec<String> {
2690 let right_names = right_schema
2694 .iter()
2695 .map(|column| column.name.as_str())
2696 .collect::<HashSet<_>>();
2697 left_schema
2698 .iter()
2699 .filter(|left| right_names.contains(left.name.as_str()))
2700 .map(|column| column.name.clone())
2701 .collect()
2702}
2703
2704fn install_base_projection(plan: &mut LogicalPlan, projection: &Projection) {
2705 match plan {
2706 LogicalPlan::Scan {
2707 projection: scan_projection,
2708 ..
2709 } => *scan_projection = projection.clone(),
2710 LogicalPlan::Filter { input, .. } => install_base_projection(input, projection),
2711 _ => {}
2712 }
2713}
2714
2715fn is_aggregate_function(name: &str) -> bool {
2716 matches!(
2717 name.to_ascii_lowercase().as_str(),
2718 "count" | "sum" | "total" | "avg" | "min" | "max" | "group_concat" | "string_agg"
2719 )
2720}
2721
2722fn expr_key(expr: &TypedExpr) -> String {
2723 format!("{:?}", expr.kind)
2724}
2725
2726fn aggregate_signature(
2727 name: &str,
2728 distinct: bool,
2729 star: bool,
2730 arg: Option<&TypedExpr>,
2731 separator: Option<&String>,
2732 _expr: &TypedExpr,
2733) -> AggregateSignature {
2734 AggregateSignature {
2735 name: name.to_ascii_lowercase(),
2736 distinct,
2737 star,
2738 arg_key: arg.map(expr_key),
2739 separator: separator.cloned(),
2740 }
2741}
2742
2743fn build_group_key_map(group_keys: &[TypedExpr]) -> HashMap<String, usize> {
2744 let mut map = HashMap::new();
2745 for (idx, key) in group_keys.iter().enumerate() {
2746 map.insert(expr_key(key), idx);
2747 }
2748 map
2749}
2750
2751fn build_aggregate_map(aggregates: &[AggregateExpr]) -> HashMap<AggregateSignature, usize> {
2752 let mut map = HashMap::new();
2753 for (idx, agg) in aggregates.iter().enumerate() {
2754 let (name, separator, star, arg) = match &agg.function {
2755 AggregateFunction::Count => (
2756 "count".to_string(),
2757 None,
2758 agg.arg.is_none(),
2759 agg.arg.as_ref(),
2760 ),
2761 AggregateFunction::Sum => ("sum".to_string(), None, false, agg.arg.as_ref()),
2762 AggregateFunction::Total => ("total".to_string(), None, false, agg.arg.as_ref()),
2763 AggregateFunction::Avg => ("avg".to_string(), None, false, agg.arg.as_ref()),
2764 AggregateFunction::Min => ("min".to_string(), None, false, agg.arg.as_ref()),
2765 AggregateFunction::Max => ("max".to_string(), None, false, agg.arg.as_ref()),
2766 AggregateFunction::GroupConcat { separator } => (
2767 "group_concat".to_string(),
2768 separator.clone(),
2769 false,
2770 agg.arg.as_ref(),
2771 ),
2772 AggregateFunction::StringAgg { separator } => (
2773 "string_agg".to_string(),
2774 separator.clone(),
2775 false,
2776 agg.arg.as_ref(),
2777 ),
2778 };
2779 let signature = AggregateSignature {
2780 name,
2781 distinct: agg.distinct,
2782 star,
2783 arg_key: arg.map(expr_key),
2784 separator,
2785 };
2786 map.insert(signature, idx);
2787 }
2788 map
2789}
2790
2791fn build_aggregate_schema(
2792 group_keys: &[TypedExpr],
2793 aggregates: &[AggregateExpr],
2794) -> Vec<ColumnMetadata> {
2795 let mut schema = Vec::new();
2796 for (idx, key) in group_keys.iter().enumerate() {
2797 let name = match &key.kind {
2798 TypedExprKind::ColumnRef { column, .. } => column.clone(),
2799 _ => format!("group_{idx}"),
2800 };
2801 schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
2802 }
2803 for (idx, agg) in aggregates.iter().enumerate() {
2804 let name = match &agg.function {
2805 AggregateFunction::Count => format!("count_{idx}"),
2806 AggregateFunction::Sum => format!("sum_{idx}"),
2807 AggregateFunction::Total => format!("total_{idx}"),
2808 AggregateFunction::Avg => format!("avg_{idx}"),
2809 AggregateFunction::Min => format!("min_{idx}"),
2810 AggregateFunction::Max => format!("max_{idx}"),
2811 AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
2812 AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
2813 };
2814 schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
2815 }
2816 schema
2817}
2818
2819fn rewrite_expr_with_maps(
2820 expr: &TypedExpr,
2821 group_key_map: &HashMap<String, usize>,
2822 aggregate_map: &HashMap<AggregateSignature, usize>,
2823 output_names: &[String],
2824) -> Result<TypedExpr, PlannerError> {
2825 let group_key_count = output_names.len().saturating_sub(aggregate_map.len());
2826 let key = expr_key(expr);
2827 if let Some(idx) = group_key_map.get(&key) {
2828 return Ok(make_output_column_ref(
2829 *idx,
2830 output_names,
2831 expr.resolved_type.clone(),
2832 expr.span,
2833 ));
2834 }
2835
2836 match &expr.kind {
2837 TypedExprKind::FunctionCall {
2838 name,
2839 args,
2840 distinct,
2841 star,
2842 } if is_aggregate_function(name) => {
2843 let separator = if name.eq_ignore_ascii_case("group_concat") && args.len() == 2 {
2844 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2845 Some(value.clone())
2846 } else {
2847 return Err(PlannerError::invalid_expression(
2848 "GROUP_CONCAT separator must be a string literal".to_string(),
2849 ));
2850 }
2851 } else if name.eq_ignore_ascii_case("string_agg") && args.len() == 2 {
2852 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2853 Some(value.clone())
2854 } else {
2855 return Err(PlannerError::invalid_expression(
2856 "STRING_AGG separator must be a string literal".to_string(),
2857 ));
2858 }
2859 } else {
2860 None
2861 };
2862 let signature = AggregateSignature {
2863 name: name.to_ascii_lowercase(),
2864 distinct: *distinct,
2865 star: *star,
2866 arg_key: args.first().map(expr_key),
2867 separator,
2868 };
2869 let idx = aggregate_map.get(&signature).ok_or_else(|| {
2870 PlannerError::invalid_expression(
2871 "aggregate in expression is not part of plan".to_string(),
2872 )
2873 })?;
2874 let output_index = group_key_count + idx;
2875 Ok(make_output_column_ref(
2876 output_index,
2877 output_names,
2878 expr.resolved_type.clone(),
2879 expr.span,
2880 ))
2881 }
2882 TypedExprKind::FunctionCall {
2883 name,
2884 args,
2885 distinct,
2886 star,
2887 } => {
2888 if *distinct || *star {
2889 return Err(PlannerError::invalid_expression(
2890 "DISTINCT/STAR modifiers are only supported for aggregates".to_string(),
2891 ));
2892 }
2893 let mut rewritten_args = Vec::with_capacity(args.len());
2894 for arg in args {
2895 rewritten_args.push(rewrite_expr_with_maps(
2896 arg,
2897 group_key_map,
2898 aggregate_map,
2899 output_names,
2900 )?);
2901 }
2902 Ok(TypedExpr {
2903 kind: TypedExprKind::FunctionCall {
2904 name: name.clone(),
2905 args: rewritten_args,
2906 distinct: false,
2907 star: false,
2908 },
2909 resolved_type: expr.resolved_type.clone(),
2910 span: expr.span,
2911 })
2912 }
2913 TypedExprKind::BinaryOp { left, op, right } => {
2914 let left = rewrite_expr_with_maps(left, group_key_map, aggregate_map, output_names)?;
2915 let right = rewrite_expr_with_maps(right, group_key_map, aggregate_map, output_names)?;
2916 Ok(TypedExpr {
2917 kind: TypedExprKind::BinaryOp {
2918 left: Box::new(left),
2919 op: *op,
2920 right: Box::new(right),
2921 },
2922 resolved_type: expr.resolved_type.clone(),
2923 span: expr.span,
2924 })
2925 }
2926 TypedExprKind::UnaryOp { op, operand } => {
2927 let operand =
2928 rewrite_expr_with_maps(operand, group_key_map, aggregate_map, output_names)?;
2929 Ok(TypedExpr {
2930 kind: TypedExprKind::UnaryOp {
2931 op: *op,
2932 operand: Box::new(operand),
2933 },
2934 resolved_type: expr.resolved_type.clone(),
2935 span: expr.span,
2936 })
2937 }
2938 TypedExprKind::Between {
2939 expr: inner,
2940 low,
2941 high,
2942 negated,
2943 } => {
2944 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2945 let low = rewrite_expr_with_maps(low, group_key_map, aggregate_map, output_names)?;
2946 let high = rewrite_expr_with_maps(high, group_key_map, aggregate_map, output_names)?;
2947 Ok(TypedExpr {
2948 kind: TypedExprKind::Between {
2949 expr: Box::new(inner),
2950 low: Box::new(low),
2951 high: Box::new(high),
2952 negated: *negated,
2953 },
2954 resolved_type: expr.resolved_type.clone(),
2955 span: expr.span,
2956 })
2957 }
2958 TypedExprKind::Like {
2959 expr: inner,
2960 pattern,
2961 escape,
2962 negated,
2963 kind,
2964 } => {
2965 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2966 let pattern =
2967 rewrite_expr_with_maps(pattern, group_key_map, aggregate_map, output_names)?;
2968 let escape = if let Some(esc) = escape {
2969 Some(Box::new(rewrite_expr_with_maps(
2970 esc,
2971 group_key_map,
2972 aggregate_map,
2973 output_names,
2974 )?))
2975 } else {
2976 None
2977 };
2978 Ok(TypedExpr {
2979 kind: TypedExprKind::Like {
2980 expr: Box::new(inner),
2981 pattern: Box::new(pattern),
2982 escape,
2983 negated: *negated,
2984 kind: *kind,
2985 },
2986 resolved_type: expr.resolved_type.clone(),
2987 span: expr.span,
2988 })
2989 }
2990 TypedExprKind::InList {
2991 expr: inner,
2992 list,
2993 negated,
2994 } => {
2995 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2996 let mut rewritten_list = Vec::with_capacity(list.len());
2997 for item in list {
2998 rewritten_list.push(rewrite_expr_with_maps(
2999 item,
3000 group_key_map,
3001 aggregate_map,
3002 output_names,
3003 )?);
3004 }
3005 Ok(TypedExpr {
3006 kind: TypedExprKind::InList {
3007 expr: Box::new(inner),
3008 list: rewritten_list,
3009 negated: *negated,
3010 },
3011 resolved_type: expr.resolved_type.clone(),
3012 span: expr.span,
3013 })
3014 }
3015 TypedExprKind::IsNull {
3016 expr: inner,
3017 negated,
3018 } => {
3019 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
3020 Ok(TypedExpr {
3021 kind: TypedExprKind::IsNull {
3022 expr: Box::new(inner),
3023 negated: *negated,
3024 },
3025 resolved_type: expr.resolved_type.clone(),
3026 span: expr.span,
3027 })
3028 }
3029 TypedExprKind::Literal(_) | TypedExprKind::VectorLiteral(_) => Ok(expr.clone()),
3030 TypedExprKind::ColumnRef { .. } => Err(PlannerError::invalid_expression(
3031 "column reference must appear in GROUP BY or be aggregated".to_string(),
3032 )),
3033 TypedExprKind::Cast {
3034 expr: inner,
3035 target_type,
3036 } => {
3037 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
3038 Ok(TypedExpr {
3039 kind: TypedExprKind::Cast {
3040 expr: Box::new(inner),
3041 target_type: target_type.clone(),
3042 },
3043 resolved_type: expr.resolved_type.clone(),
3044 span: expr.span,
3045 })
3046 }
3047 TypedExprKind::ScalarSubquery(_)
3048 | TypedExprKind::InSubquery { .. }
3049 | TypedExprKind::Exists { .. }
3050 | TypedExprKind::Quantified { .. } => Ok(expr.clone()),
3051 }
3052}
3053
3054fn make_output_column_ref(
3055 index: usize,
3056 output_names: &[String],
3057 resolved_type: ResolvedType,
3058 span: crate::ast::Span,
3059) -> TypedExpr {
3060 let name = output_names
3061 .get(index)
3062 .cloned()
3063 .unwrap_or_else(|| format!("col_{index}"));
3064 TypedExpr::column_ref("__agg__".to_string(), name, index, resolved_type, span)
3065}