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, 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::Update {
378 table,
379 assignments,
380 filter,
381 } => {
382 push_table_reference(
383 references,
384 table,
385 root_access,
386 TableReferenceSource::LogicalPlanMutationTarget,
387 );
388 for assignment in assignments {
389 self.extract_typed_expr(&assignment.value, diagnostics, references);
390 }
391 if let Some(filter) = filter {
392 self.extract_typed_expr(filter, diagnostics, references);
393 }
394 }
395 LogicalPlan::Delete { table, filter } => {
396 push_table_reference(
397 references,
398 table,
399 root_access,
400 TableReferenceSource::LogicalPlanMutationTarget,
401 );
402 if let Some(filter) = filter {
403 self.extract_typed_expr(filter, diagnostics, references);
404 }
405 }
406 LogicalPlan::CreateTable { table, .. } => push_table_reference(
407 references,
408 &table.name,
409 root_access,
410 TableReferenceSource::LogicalPlanDdlTarget,
411 ),
412 LogicalPlan::DropTable { name, .. } => push_table_reference(
413 references,
414 name,
415 root_access,
416 TableReferenceSource::LogicalPlanDdlTarget,
417 ),
418 LogicalPlan::CreateIndex { index, .. } => push_table_reference(
419 references,
420 &index.table,
421 root_access,
422 TableReferenceSource::LogicalPlanIndexTarget,
423 ),
424 LogicalPlan::DropIndex { name, .. } => diagnostics.push(PlanningDiagnostic::warning(
425 "ALOPEX-PLAN-ROUTE-003",
426 format!(
427 "DROP INDEX {name} does not expose a target table in the current logical plan"
428 ),
429 )),
430 LogicalPlan::Pragma { .. } => {}
431 }
432 }
433
434 fn extract_projection(
435 &self,
436 projection: &Projection,
437 diagnostics: &mut Vec<PlanningDiagnostic>,
438 references: &mut Vec<TableReference>,
439 ) {
440 if let Projection::Columns(columns) = projection {
441 for column in columns {
442 self.extract_typed_expr(&column.expr, diagnostics, references);
443 }
444 }
445 }
446
447 fn extract_typed_expr(
448 &self,
449 expr: &TypedExpr,
450 diagnostics: &mut Vec<PlanningDiagnostic>,
451 references: &mut Vec<TableReference>,
452 ) {
453 match &expr.kind {
454 TypedExprKind::Literal(_)
455 | TypedExprKind::ColumnRef { .. }
456 | TypedExprKind::VectorLiteral(_) => {}
457 TypedExprKind::BinaryOp { left, right, .. } => {
458 self.extract_typed_expr(left, diagnostics, references);
459 self.extract_typed_expr(right, diagnostics, references);
460 }
461 TypedExprKind::UnaryOp { operand, .. }
462 | TypedExprKind::Cast { expr: operand, .. }
463 | TypedExprKind::IsNull { expr: operand, .. } => {
464 self.extract_typed_expr(operand, diagnostics, references);
465 }
466 TypedExprKind::FunctionCall { args, .. } => {
467 for arg in args {
468 self.extract_typed_expr(arg, diagnostics, references);
469 }
470 }
471 TypedExprKind::Between {
472 expr, low, high, ..
473 } => {
474 self.extract_typed_expr(expr, diagnostics, references);
475 self.extract_typed_expr(low, diagnostics, references);
476 self.extract_typed_expr(high, diagnostics, references);
477 }
478 TypedExprKind::Like {
479 expr,
480 pattern,
481 escape,
482 ..
483 } => {
484 self.extract_typed_expr(expr, diagnostics, references);
485 self.extract_typed_expr(pattern, diagnostics, references);
486 if let Some(escape) = escape {
487 self.extract_typed_expr(escape, diagnostics, references);
488 }
489 }
490 TypedExprKind::InList { expr, list, .. } => {
491 self.extract_typed_expr(expr, diagnostics, references);
492 for item in list {
493 self.extract_typed_expr(item, diagnostics, references);
494 }
495 }
496 TypedExprKind::ScalarSubquery(subquery) => self.extract_plan(
497 subquery,
498 TableReferenceAccess::Read,
499 TableReferenceSource::TypedExprSubquery,
500 diagnostics,
501 references,
502 ),
503 TypedExprKind::InSubquery { expr, subquery, .. } => {
504 self.extract_typed_expr(expr, diagnostics, references);
505 self.extract_plan(
506 subquery,
507 TableReferenceAccess::Read,
508 TableReferenceSource::TypedExprSubquery,
509 diagnostics,
510 references,
511 );
512 }
513 TypedExprKind::Exists { subquery, .. } => self.extract_plan(
514 subquery,
515 TableReferenceAccess::Read,
516 TableReferenceSource::TypedExprSubquery,
517 diagnostics,
518 references,
519 ),
520 TypedExprKind::Quantified { expr, subquery, .. } => {
521 self.extract_typed_expr(expr, diagnostics, references);
522 self.extract_plan(
523 subquery,
524 TableReferenceAccess::Read,
525 TableReferenceSource::TypedExprSubquery,
526 diagnostics,
527 references,
528 );
529 }
530 }
531 }
532}
533
534fn push_table_reference(
535 references: &mut Vec<TableReference>,
536 table_name: &str,
537 access: TableReferenceAccess,
538 source: TableReferenceSource,
539) {
540 if !references.iter().any(|reference| {
541 reference.table_name == table_name
542 && reference.access == access
543 && reference.source == source
544 }) {
545 references.push(TableReference::new(table_name, access, source));
546 }
547}
548
549fn table_reference_access(statement_kind: &StatementKind) -> TableReferenceAccess {
550 match statement_kind {
551 StatementKind::Select(_) => TableReferenceAccess::Read,
552 StatementKind::Insert(_) | StatementKind::Update(_) | StatementKind::Delete(_) => {
553 TableReferenceAccess::Write
554 }
555 StatementKind::CreateTable(_) => TableReferenceAccess::Create,
556 StatementKind::DropTable(_) => TableReferenceAccess::Drop,
557 StatementKind::CreateIndex(_) | StatementKind::DropIndex(_) => {
558 TableReferenceAccess::Metadata
559 }
560 StatementKind::Pragma { .. } => TableReferenceAccess::Metadata,
561 }
562}
563
564pub struct Planner<'a, C: Catalog + ?Sized> {
591 catalog: &'a C,
592 name_resolver: NameResolver<'a, C>,
593 type_checker: TypeChecker<'a, C>,
594}
595
596impl<'a, C: Catalog + ?Sized> Planner<'a, C> {
597 pub fn new(catalog: &'a C) -> Self {
599 Self {
600 catalog,
601 name_resolver: NameResolver::new(catalog),
602 type_checker: TypeChecker::new(catalog),
603 }
604 }
605
606 pub fn plan(&self, stmt: &Statement) -> Result<LogicalPlan, PlannerError> {
617 match &stmt.kind {
618 StatementKind::CreateTable(ct) => self.plan_create_table(ct),
620 StatementKind::DropTable(dt) => self.plan_drop_table(dt),
621 StatementKind::CreateIndex(ci) => self.plan_create_index(ci),
622 StatementKind::DropIndex(di) => self.plan_drop_index(di),
623 StatementKind::Pragma { name, value } => self.plan_pragma(name, value),
624
625 StatementKind::Select(sel) => self.plan_select(sel),
627 StatementKind::Insert(ins) => self.plan_insert(ins),
628 StatementKind::Update(upd) => self.plan_update(upd),
629 StatementKind::Delete(del) => self.plan_delete(del),
630 }
631 }
632
633 fn plan_pragma(
634 &self,
635 raw_name: &str,
636 value: &Option<PragmaValue>,
637 ) -> Result<LogicalPlan, PlannerError> {
638 let name = raw_name.to_ascii_lowercase();
639 if !matches!(name.as_str(), "cache_size" | "memory_limit" | "io_stats") {
640 return Err(PlannerError::InvalidPragma {
641 name,
642 reason: "supported names are cache_size, memory_limit, and io_stats".to_string(),
643 });
644 }
645 match name.as_str() {
646 "cache_size" => match value {
647 Some(PragmaValue::Int(v)) if *v > 0 => {}
648 Some(PragmaValue::Int(_)) => {
649 return Err(PlannerError::InvalidPragma {
650 name,
651 reason: "cache_size must be a positive page count".to_string(),
652 });
653 }
654 Some(PragmaValue::Text(_)) => {
655 return Err(PlannerError::InvalidPragma {
656 name,
657 reason: "cache_size requires an integer page count".to_string(),
658 });
659 }
660 None => {}
661 },
662 "memory_limit" => {
663 if let Some(PragmaValue::Int(v)) = value
664 && *v < 0
665 {
666 return Err(PlannerError::InvalidPragma {
667 name,
668 reason: "memory_limit cannot be negative".to_string(),
669 });
670 }
671 }
672 "io_stats" => {
673 if value.is_some() {
674 return Err(PlannerError::InvalidPragma {
675 name,
676 reason: "io_stats does not accept a value".to_string(),
677 });
678 }
679 }
680 _ => unreachable!(),
681 }
682 Ok(LogicalPlan::Pragma {
683 name,
684 value: value.clone(),
685 })
686 }
687
688 fn plan_create_table(&self, stmt: &CreateTable) -> Result<LogicalPlan, PlannerError> {
697 if !stmt.if_not_exists && self.catalog.table_exists(&stmt.name) {
699 return Err(PlannerError::table_already_exists(&stmt.name));
700 }
701
702 let columns: Vec<ColumnMetadata> = stmt
704 .columns
705 .iter()
706 .map(|col| self.convert_column_def(col))
707 .collect();
708
709 let primary_key = Self::extract_primary_key(stmt);
711
712 let mut table = TableMetadata::new(stmt.name.clone(), columns);
715 if let Some(pk) = primary_key {
716 table = table.with_primary_key(pk);
717 }
718 table.catalog_name = "default".to_string();
719 table.namespace_name = "default".to_string();
720 table.table_type = TableType::Managed;
721 table.data_source_format = DataSourceFormat::Alopex;
722 table.properties = HashMap::new();
723
724 Ok(LogicalPlan::CreateTable {
725 table,
726 if_not_exists: stmt.if_not_exists,
727 with_options: stmt
728 .with_options
729 .iter()
730 .map(|opt| (opt.key.clone(), opt.value.clone()))
731 .collect(),
732 })
733 }
734
735 fn convert_column_def(&self, col: &ColumnDef) -> ColumnMetadata {
737 let data_type = ResolvedType::from_ast(&col.data_type);
738 let mut meta = ColumnMetadata::new(col.name.clone(), data_type);
739
740 for constraint in &col.constraints {
742 meta = Self::apply_column_constraint(meta, constraint);
743 }
744
745 meta
746 }
747
748 fn apply_column_constraint(
750 mut meta: ColumnMetadata,
751 constraint: &ColumnConstraint,
752 ) -> ColumnMetadata {
753 match constraint {
754 ColumnConstraint::NotNull { .. } => {
755 meta.not_null = true;
756 }
757 ColumnConstraint::PrimaryKey { .. } => {
758 meta.primary_key = true;
759 meta.not_null = true; }
761 ColumnConstraint::Unique { .. } => {
762 meta.unique = true;
763 }
764 ColumnConstraint::Default { value: expr, .. } => {
765 meta.default = Some(expr.clone());
766 }
767 }
768 meta
769 }
770
771 fn extract_primary_key(stmt: &CreateTable) -> Option<Vec<String>> {
773 use crate::ast::ddl::TableConstraint;
774
775 if let Some(TableConstraint::PrimaryKey { columns, .. }) = stmt.constraints.first() {
779 return Some(columns.clone());
780 }
781
782 let pk_columns: Vec<String> = stmt
784 .columns
785 .iter()
786 .filter(|col| col.constraints.iter().any(Self::is_primary_key_constraint))
787 .map(|col| col.name.clone())
788 .collect();
789
790 if pk_columns.is_empty() {
791 None
792 } else {
793 Some(pk_columns)
794 }
795 }
796
797 fn is_primary_key_constraint(constraint: &ColumnConstraint) -> bool {
799 matches!(constraint, ColumnConstraint::PrimaryKey { .. })
800 }
801
802 fn plan_drop_table(&self, stmt: &DropTable) -> Result<LogicalPlan, PlannerError> {
806 if !stmt.if_exists && !self.table_exists_in_default(&stmt.name) {
808 return Err(PlannerError::TableNotFound {
809 name: stmt.name.clone(),
810 line: stmt.span.start.line,
811 column: stmt.span.start.column,
812 });
813 }
814
815 Ok(LogicalPlan::DropTable {
816 name: stmt.name.clone(),
817 if_exists: stmt.if_exists,
818 })
819 }
820
821 fn table_exists_in_default(&self, name: &str) -> bool {
822 match self.catalog.get_table(name) {
823 Some(table) => table.catalog_name == "default" && table.namespace_name == "default",
824 None => false,
825 }
826 }
827
828 fn plan_create_index(&self, stmt: &CreateIndex) -> Result<LogicalPlan, PlannerError> {
835 if !stmt.if_not_exists && self.catalog.index_exists(&stmt.name) {
837 return Err(PlannerError::index_already_exists(&stmt.name));
838 }
839
840 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
842
843 self.name_resolver
845 .resolve_column(table, &stmt.column, stmt.span)?;
846
847 let mut index = IndexMetadata::new(
851 0,
852 stmt.name.clone(),
853 stmt.table.clone(),
854 vec![stmt.column.clone()],
855 );
856
857 if let Some(method) = stmt.method {
858 index = index.with_method(method);
859 }
860
861 let options: Vec<(String, String)> = stmt
862 .options
863 .iter()
864 .map(|opt| (opt.key.clone(), opt.value.clone()))
865 .collect();
866 if !options.is_empty() {
867 index = index.with_options(options);
868 }
869
870 Ok(LogicalPlan::CreateIndex {
871 index,
872 if_not_exists: stmt.if_not_exists,
873 })
874 }
875
876 fn plan_drop_index(&self, stmt: &DropIndex) -> Result<LogicalPlan, PlannerError> {
880 if !stmt.if_exists && !self.index_exists_in_default(&stmt.name) {
882 return Err(PlannerError::index_not_found(&stmt.name));
883 }
884
885 Ok(LogicalPlan::DropIndex {
886 name: stmt.name.clone(),
887 if_exists: stmt.if_exists,
888 })
889 }
890
891 fn index_exists_in_default(&self, name: &str) -> bool {
892 match self.catalog.get_index(name) {
893 Some(index) => index.catalog_name == "default" && index.namespace_name == "default",
894 None => false,
895 }
896 }
897
898 fn plan_select(&self, stmt: &Select) -> Result<LogicalPlan, PlannerError> {
907 self.plan_select_relation(stmt, &[])
908 .map(|relation| relation.plan)
909 }
910
911 fn plan_select_relation(
912 &self,
913 stmt: &Select,
914 outer_scope: &[ScopedTable],
915 ) -> Result<PlannedRelation, PlannerError> {
916 let mut relation = self.plan_from_items(&stmt.from, stmt.span, outer_scope)?;
917 let expr_scope = relation
918 .scope
919 .iter()
920 .cloned()
921 .chain(offset_scope(outer_scope, relation.schema.len()))
922 .collect::<Vec<_>>();
923
924 let has_group_by = stmt
925 .group_by
926 .as_ref()
927 .is_some_and(|items| !items.is_empty());
928 let has_aggregate = self.select_contains_aggregate(stmt);
929 let distinct_only =
930 stmt.distinct && !has_group_by && !has_aggregate && stmt.having.is_none();
931
932 let final_projection =
933 self.build_projection_with_scope(&stmt.projection, &relation.schema, &expr_scope)?;
934 install_base_projection(&mut relation.plan, &final_projection);
935 let needs_project_boundary = !matches!(relation.plan, LogicalPlan::Scan { .. });
936 let mut plan = relation.plan;
937
938 if let Some(ref selection) = stmt.selection {
940 let predicate = self.infer_expr_with_scope(selection, &expr_scope)?;
941
942 if predicate.resolved_type != ResolvedType::Boolean {
944 return Err(PlannerError::type_mismatch(
945 "Boolean",
946 predicate.resolved_type.to_string(),
947 selection.span,
948 ));
949 }
950
951 plan = LogicalPlan::Filter {
952 input: Box::new(plan),
953 predicate,
954 };
955 }
956
957 if has_group_by || has_aggregate || stmt.having.is_some() || stmt.distinct {
958 if !has_group_by && !has_aggregate && stmt.having.is_some() {
959 return Err(PlannerError::invalid_expression(
960 "HAVING requires GROUP BY or aggregate functions".to_string(),
961 ));
962 }
963
964 let (group_keys, projected) = if distinct_only {
965 let projected = self.build_projected_columns_for_distinct_with_scope(
966 &stmt.projection,
967 &relation.schema,
968 &expr_scope,
969 )?;
970 let group_keys = projected.iter().map(|col| col.expr.clone()).collect();
971 (group_keys, projected)
972 } else {
973 let group_keys = self.build_group_keys_with_scope(stmt, &expr_scope)?;
974 let projected = self.build_projected_columns_for_aggregate_with_scope(
975 &stmt.projection,
976 &expr_scope,
977 )?;
978 (group_keys, projected)
979 };
980 let mut aggregates = Vec::new();
981 let mut agg_map = HashMap::new();
982
983 for col in &projected {
984 self.collect_aggregates_from_typed_expr(&col.expr, &mut aggregates, &mut agg_map)?;
985 }
986
987 let having_typed = if let Some(having) = &stmt.having {
988 let typed = self.infer_expr_with_scope(having, &expr_scope)?;
989 if typed.resolved_type != ResolvedType::Boolean {
990 return Err(PlannerError::type_mismatch(
991 "Boolean",
992 typed.resolved_type.type_name().to_string(),
993 typed.span,
994 ));
995 }
996 self.collect_aggregates_from_typed_expr(&typed, &mut aggregates, &mut agg_map)?;
997 Some(typed)
998 } else {
999 None
1000 };
1001
1002 let mut order_by = Vec::new();
1003 if !stmt.order_by.is_empty() {
1004 for order_expr in &stmt.order_by {
1005 let typed = self.infer_expr_with_scope(&order_expr.expr, &expr_scope)?;
1006 self.collect_aggregates_from_typed_expr(&typed, &mut aggregates, &mut agg_map)?;
1007 let asc = order_expr.asc.unwrap_or(true);
1008 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1009 order_by.push(SortExpr::new(typed, asc, nulls_first));
1010 }
1011 }
1012
1013 if let Some(ref having) = having_typed {
1014 self.type_checker
1015 .validate_having_expr(having, &group_keys, &aggregates)?;
1016 }
1017
1018 let output_schema = build_aggregate_schema(&group_keys, &aggregates);
1019 let output_names: Vec<String> = output_schema.iter().map(|c| c.name.clone()).collect();
1020
1021 let projection = self.build_aggregate_projection(
1022 projected,
1023 &group_keys,
1024 &aggregates,
1025 &output_names,
1026 )?;
1027
1028 let having = if let Some(having) = having_typed {
1029 Some(self.rewrite_expr_for_aggregate(
1030 &having,
1031 &group_keys,
1032 &aggregates,
1033 &output_names,
1034 )?)
1035 } else {
1036 None
1037 };
1038
1039 let order_by = order_by
1040 .into_iter()
1041 .map(|expr| {
1042 let rewritten = self.rewrite_expr_for_aggregate(
1043 &expr.expr,
1044 &group_keys,
1045 &aggregates,
1046 &output_names,
1047 )?;
1048 Ok(SortExpr::new(rewritten, expr.asc, expr.nulls_first))
1049 })
1050 .collect::<Result<Vec<_>, PlannerError>>()?;
1051
1052 let schema = projection_schema(&projection, &output_schema);
1053 plan = LogicalPlan::Aggregate {
1054 input: Box::new(plan),
1055 group_keys,
1056 aggregates,
1057 having,
1058 projection,
1059 };
1060
1061 if !order_by.is_empty() {
1062 plan = LogicalPlan::Sort {
1063 input: Box::new(plan),
1064 order_by,
1065 };
1066 }
1067
1068 if stmt.limit.is_some() || stmt.offset.is_some() {
1069 let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
1070 let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
1071 plan = LogicalPlan::Limit {
1072 input: Box::new(plan),
1073 limit,
1074 offset,
1075 };
1076 }
1077
1078 return Ok(PlannedRelation {
1079 plan,
1080 schema: schema.clone(),
1081 scope: vec![ScopedTable::new(
1082 TableMetadata::new(LITERAL_TABLE, schema),
1083 0,
1084 )],
1085 });
1086 }
1087
1088 if !stmt.order_by.is_empty() {
1090 let order_by = self.build_sort_exprs_with_scope(&stmt.order_by, &expr_scope)?;
1091 plan = LogicalPlan::Sort {
1092 input: Box::new(plan),
1093 order_by,
1094 };
1095 }
1096
1097 if stmt.limit.is_some() || stmt.offset.is_some() {
1098 let limit = self.extract_limit_value(&stmt.limit, stmt.span)?;
1099 let offset = self.extract_limit_value(&stmt.offset, stmt.span)?;
1100 plan = LogicalPlan::Limit {
1101 input: Box::new(plan),
1102 limit,
1103 offset,
1104 };
1105 }
1106
1107 let output_schema = projection_schema(&final_projection, &relation.schema);
1108 if needs_project_boundary {
1109 plan = LogicalPlan::Project {
1110 input: Box::new(plan),
1111 projection: final_projection,
1112 };
1113 }
1114 Ok(PlannedRelation {
1115 plan,
1116 schema: output_schema.clone(),
1117 scope: vec![ScopedTable::new(
1118 TableMetadata::new(LITERAL_TABLE, output_schema),
1119 0,
1120 )],
1121 })
1122 }
1123
1124 fn plan_from_items(
1128 &self,
1129 items: &[FromItem],
1130 select_span: crate::ast::Span,
1131 outer_scope: &[ScopedTable],
1132 ) -> Result<PlannedRelation, PlannerError> {
1133 match items {
1134 [] => {
1135 let schema = Vec::new();
1136 Ok(PlannedRelation {
1137 plan: LogicalPlan::Scan {
1138 table: LITERAL_TABLE.to_string(),
1139 projection: Projection::All(Vec::new()),
1140 },
1141 schema: schema.clone(),
1142 scope: vec![ScopedTable::new(
1143 TableMetadata::new(LITERAL_TABLE, schema),
1144 0,
1145 )],
1146 })
1147 }
1148 [single] => self.plan_from_item(single, 0, outer_scope),
1149 [first, rest @ ..] => {
1150 let mut relation = self.plan_from_item(first, 0, outer_scope)?;
1151 for item in rest {
1152 let right = self.plan_from_item(item, relation.schema.len(), outer_scope)?;
1153 relation = self.combine_join_relation(
1154 relation,
1155 right,
1156 JoinType::Cross,
1157 None,
1158 None,
1159 select_span,
1160 )?;
1161 }
1162 Ok(relation)
1163 }
1164 }
1165 }
1166
1167 fn plan_from_item(
1168 &self,
1169 item: &FromItem,
1170 start_index: usize,
1171 outer_scope: &[ScopedTable],
1172 ) -> Result<PlannedRelation, PlannerError> {
1173 match item {
1174 FromItem::Table { name, alias, span } => {
1175 let table = self.name_resolver.resolve_table(name, *span)?.clone();
1176 let mut scope_table = table.clone();
1177 if let Some(alias) = alias {
1178 scope_table.name = alias.clone();
1179 }
1180 let schema = table.columns.clone();
1181 Ok(PlannedRelation {
1182 plan: LogicalPlan::Scan {
1183 table: name.clone(),
1184 projection: Projection::All(
1185 schema.iter().map(|col| col.name.clone()).collect(),
1186 ),
1187 },
1188 schema,
1189 scope: vec![ScopedTable::new(scope_table, start_index)],
1190 })
1191 }
1192 FromItem::Join {
1193 left,
1194 right,
1195 join_type,
1196 condition,
1197 using,
1198 span,
1199 } => {
1200 let left_relation = self.plan_from_item(left, start_index, outer_scope)?;
1201 let right_relation = self.plan_from_item(
1202 right,
1203 start_index + left_relation.schema.len(),
1204 outer_scope,
1205 )?;
1206 let expr_scope = left_relation
1207 .scope
1208 .iter()
1209 .cloned()
1210 .chain(right_relation.scope.iter().cloned())
1211 .chain(offset_scope(
1212 outer_scope,
1213 left_relation.schema.len() + right_relation.schema.len(),
1214 ))
1215 .collect::<Vec<_>>();
1216 let typed_condition = if let Some(expr) = condition {
1217 let typed = self.infer_expr_with_scope(expr, &expr_scope)?;
1218 if typed.resolved_type != ResolvedType::Boolean {
1219 return Err(PlannerError::type_mismatch(
1220 "Boolean",
1221 typed.resolved_type.to_string(),
1222 expr.span,
1223 ));
1224 }
1225 Some(typed)
1226 } else {
1227 self.build_using_condition(
1228 using.as_deref(),
1229 &left_relation,
1230 &right_relation,
1231 *span,
1232 )?
1233 };
1234 self.combine_join_relation(
1235 left_relation,
1236 right_relation,
1237 map_join_type(*join_type),
1238 typed_condition,
1239 using.clone(),
1240 *span,
1241 )
1242 }
1243 FromItem::Derived {
1244 subquery,
1245 alias,
1246 span,
1247 } => {
1248 let crate::ast::StatementKind::Select(select) = &subquery.kind else {
1249 return Err(PlannerError::unsupported_feature(
1250 "non-SELECT derived table",
1251 "v0.6.0-subquery Phase 6",
1252 *span,
1253 ));
1254 };
1255 let mut relation = self.plan_select_relation(select, outer_scope)?;
1256 let alias = alias.clone().ok_or_else(|| {
1257 PlannerError::invalid_expression("derived table requires an alias".to_string())
1258 })?;
1259 relation.plan = LogicalPlan::Project {
1260 input: Box::new(relation.plan),
1261 projection: Projection::All(
1262 relation.schema.iter().map(|col| col.name.clone()).collect(),
1263 ),
1264 };
1265 relation.scope = vec![ScopedTable::new(
1266 TableMetadata::new(alias, relation.schema.clone()),
1267 start_index,
1268 )];
1269 Ok(relation)
1270 }
1271 }
1272 }
1273
1274 fn combine_join_relation(
1275 &self,
1276 left: PlannedRelation,
1277 right: PlannedRelation,
1278 join_type: JoinType,
1279 condition: Option<TypedExpr>,
1280 using: Option<Vec<String>>,
1281 _span: crate::ast::Span,
1282 ) -> Result<PlannedRelation, PlannerError> {
1283 let mut schema = left.schema.clone();
1284 schema.extend(right.schema.clone());
1285 let mut scope = left.scope.clone();
1286 scope.extend(right.scope.clone());
1287 Ok(PlannedRelation {
1288 plan: LogicalPlan::Join {
1289 left: Box::new(left.plan),
1290 right: Box::new(right.plan),
1291 join_type,
1292 condition,
1293 using,
1294 },
1295 schema,
1296 scope,
1297 })
1298 }
1299
1300 fn build_using_condition(
1301 &self,
1302 using: Option<&[String]>,
1303 left: &PlannedRelation,
1304 right: &PlannedRelation,
1305 span: crate::ast::Span,
1306 ) -> Result<Option<TypedExpr>, PlannerError> {
1307 let Some(columns) = using else {
1308 return Ok(None);
1309 };
1310 let mut condition = None;
1311 for column in columns {
1312 let left_col = find_scoped_column(&left.scope, column, span)?;
1313 let right_col = find_scoped_column(&right.scope, column, span)?;
1314 let left_expr = TypedExpr::column_ref(
1315 left_col.table,
1316 column.clone(),
1317 left_col.index,
1318 left_col.ty.clone(),
1319 span,
1320 );
1321 let right_expr = TypedExpr::column_ref(
1322 right_col.table,
1323 column.clone(),
1324 right_col.index,
1325 right_col.ty.clone(),
1326 span,
1327 );
1328 self.type_checker
1329 .check_comparison_op(&left_col.ty, &right_col.ty, span)?;
1330 let eq = TypedExpr::binary_op(
1331 left_expr,
1332 crate::ast::expr::BinaryOp::Eq,
1333 right_expr,
1334 ResolvedType::Boolean,
1335 span,
1336 );
1337 condition = Some(match condition {
1338 Some(prev) => TypedExpr::binary_op(
1339 prev,
1340 crate::ast::expr::BinaryOp::And,
1341 eq,
1342 ResolvedType::Boolean,
1343 span,
1344 ),
1345 None => eq,
1346 });
1347 }
1348 Ok(condition)
1349 }
1350
1351 fn infer_expr_with_scope(
1352 &self,
1353 expr: &crate::ast::expr::Expr,
1354 scope: &[ScopedTable],
1355 ) -> Result<TypedExpr, PlannerError> {
1356 self.type_checker
1357 .infer_type_with_scope(expr, scope, &|stmt, outer_scope| {
1358 let crate::ast::StatementKind::Select(select) = &stmt.kind else {
1359 return Err(PlannerError::unsupported_feature(
1360 "non-SELECT subquery",
1361 "v0.6.0-subquery Phase 6",
1362 stmt.span(),
1363 ));
1364 };
1365 let relation = self.plan_select_relation(select, outer_scope)?;
1366 Ok((relation.plan, relation.schema))
1367 })
1368 }
1369
1370 #[allow(dead_code)]
1371 fn build_projection(
1372 &self,
1373 items: &[SelectItem],
1374 table: &TableMetadata,
1375 ) -> Result<Projection, PlannerError> {
1376 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1378 let columns = self.name_resolver.expand_wildcard(table);
1379 return Ok(Projection::All(columns));
1380 }
1381
1382 let mut projected_columns = Vec::new();
1384 for item in items {
1385 match item {
1386 SelectItem::Wildcard { span } => {
1387 for col in &table.columns {
1389 let column_index = table.get_column_index(&col.name).unwrap();
1390 let typed_expr = TypedExpr::column_ref(
1391 table.name.clone(),
1392 col.name.clone(),
1393 column_index,
1394 col.data_type.clone(),
1395 *span,
1396 );
1397 projected_columns.push(ProjectedColumn::new(typed_expr));
1398 }
1399 }
1400 SelectItem::Expr { expr, alias, .. } => {
1401 let typed_expr = self.type_checker.infer_type(expr, table)?;
1402 let projected = if let Some(alias) = alias {
1403 ProjectedColumn::with_alias(typed_expr, alias.clone())
1404 } else {
1405 ProjectedColumn::new(typed_expr)
1406 };
1407 projected_columns.push(projected);
1408 }
1409 }
1410 }
1411
1412 Ok(Projection::Columns(projected_columns))
1413 }
1414
1415 fn build_projection_with_scope(
1416 &self,
1417 items: &[SelectItem],
1418 schema: &[ColumnMetadata],
1419 scope: &[ScopedTable],
1420 ) -> Result<Projection, PlannerError> {
1421 if items.len() == 1 && matches!(&items[0], SelectItem::Wildcard { .. }) {
1422 return Ok(Projection::All(
1423 schema.iter().map(|col| col.name.clone()).collect(),
1424 ));
1425 }
1426
1427 let mut projected_columns = Vec::new();
1428 for item in items {
1429 match item {
1430 SelectItem::Wildcard { span } => {
1431 for scoped in scope {
1432 for (local_idx, col) in scoped.table.columns.iter().enumerate() {
1433 projected_columns.push(ProjectedColumn::new(TypedExpr::column_ref(
1434 scoped.table.name.clone(),
1435 col.name.clone(),
1436 scoped.start_index + local_idx,
1437 col.data_type.clone(),
1438 *span,
1439 )));
1440 }
1441 }
1442 }
1443 SelectItem::Expr { expr, alias, .. } => {
1444 let typed_expr = self.infer_expr_with_scope(expr, scope)?;
1445 let projected = if let Some(alias) = alias {
1446 ProjectedColumn::with_alias(typed_expr, alias.clone())
1447 } else {
1448 ProjectedColumn::new(typed_expr)
1449 };
1450 projected_columns.push(projected);
1451 }
1452 }
1453 }
1454
1455 Ok(Projection::Columns(projected_columns))
1456 }
1457
1458 #[allow(dead_code)]
1460 fn build_sort_exprs(
1461 &self,
1462 order_by: &[OrderByExpr],
1463 table: &TableMetadata,
1464 ) -> Result<Vec<SortExpr>, PlannerError> {
1465 let mut sort_exprs = Vec::new();
1466
1467 for order_expr in order_by {
1468 let typed_expr = self.type_checker.infer_type(&order_expr.expr, table)?;
1469
1470 let asc = order_expr.asc.unwrap_or(true);
1472
1473 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1475
1476 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1477 }
1478
1479 Ok(sort_exprs)
1480 }
1481
1482 fn build_sort_exprs_with_scope(
1483 &self,
1484 order_by: &[OrderByExpr],
1485 scope: &[ScopedTable],
1486 ) -> Result<Vec<SortExpr>, PlannerError> {
1487 let mut sort_exprs = Vec::new();
1488 for order_expr in order_by {
1489 let typed_expr = self.infer_expr_with_scope(&order_expr.expr, scope)?;
1490 let asc = order_expr.asc.unwrap_or(true);
1491 let nulls_first = order_expr.nulls_first.unwrap_or(false);
1492 sort_exprs.push(SortExpr::new(typed_expr, asc, nulls_first));
1493 }
1494 Ok(sort_exprs)
1495 }
1496
1497 fn select_contains_aggregate(&self, stmt: &Select) -> bool {
1498 stmt.projection.iter().any(|item| match item {
1499 SelectItem::Wildcard { .. } => false,
1500 SelectItem::Expr { expr, .. } => expr_contains_aggregate(expr),
1501 }) || stmt
1502 .group_by
1503 .as_ref()
1504 .map(|items| items.iter().any(expr_contains_aggregate))
1505 .unwrap_or(false)
1506 || stmt
1507 .having
1508 .as_ref()
1509 .map(expr_contains_aggregate)
1510 .unwrap_or(false)
1511 || stmt
1512 .order_by
1513 .iter()
1514 .any(|order| expr_contains_aggregate(&order.expr))
1515 }
1516
1517 #[allow(dead_code)]
1518 fn build_group_keys(
1519 &self,
1520 stmt: &Select,
1521 table: &TableMetadata,
1522 ) -> Result<Vec<TypedExpr>, PlannerError> {
1523 let mut keys = Vec::new();
1524 if let Some(items) = &stmt.group_by {
1525 for expr in items {
1526 let typed = self.type_checker.infer_type(expr, table)?;
1527 if typed_expr_contains_aggregate(&typed) {
1528 return Err(PlannerError::invalid_expression(
1529 "GROUP BY cannot contain aggregate functions".to_string(),
1530 ));
1531 }
1532 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1533 return Err(PlannerError::invalid_expression(
1534 "GROUP BY expressions must be column references".to_string(),
1535 ));
1536 }
1537 keys.push(typed);
1538 }
1539 }
1540 Ok(keys)
1541 }
1542
1543 fn build_group_keys_with_scope(
1544 &self,
1545 stmt: &Select,
1546 scope: &[ScopedTable],
1547 ) -> Result<Vec<TypedExpr>, PlannerError> {
1548 let mut keys = Vec::new();
1549 if let Some(items) = &stmt.group_by {
1550 for expr in items {
1551 let typed = self.infer_expr_with_scope(expr, scope)?;
1552 if typed_expr_contains_aggregate(&typed) {
1553 return Err(PlannerError::invalid_expression(
1554 "GROUP BY cannot contain aggregate functions".to_string(),
1555 ));
1556 }
1557 if !matches!(typed.kind, TypedExprKind::ColumnRef { .. }) {
1558 return Err(PlannerError::invalid_expression(
1559 "GROUP BY expressions must be column references".to_string(),
1560 ));
1561 }
1562 keys.push(typed);
1563 }
1564 }
1565 Ok(keys)
1566 }
1567
1568 #[allow(dead_code)]
1569 fn build_projected_columns_for_aggregate(
1570 &self,
1571 items: &[SelectItem],
1572 table: &TableMetadata,
1573 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1574 let mut projected = Vec::new();
1575 for item in items {
1576 match item {
1577 SelectItem::Wildcard { .. } => {
1578 return Err(PlannerError::invalid_expression(
1579 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1580 ));
1581 }
1582 SelectItem::Expr { expr, alias, .. } => {
1583 let typed = self.type_checker.infer_type(expr, table)?;
1584 projected.push(ProjectedColumn {
1585 expr: typed,
1586 alias: alias.clone(),
1587 });
1588 }
1589 }
1590 }
1591 Ok(projected)
1592 }
1593
1594 fn build_projected_columns_for_aggregate_with_scope(
1595 &self,
1596 items: &[SelectItem],
1597 scope: &[ScopedTable],
1598 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1599 let mut projected = Vec::new();
1600 for item in items {
1601 match item {
1602 SelectItem::Wildcard { .. } => {
1603 return Err(PlannerError::invalid_expression(
1604 "wildcard projection not supported with GROUP BY/aggregate".to_string(),
1605 ));
1606 }
1607 SelectItem::Expr { expr, alias, .. } => {
1608 let typed = self.infer_expr_with_scope(expr, scope)?;
1609 projected.push(ProjectedColumn {
1610 expr: typed,
1611 alias: alias.clone(),
1612 });
1613 }
1614 }
1615 }
1616 Ok(projected)
1617 }
1618
1619 #[allow(dead_code)]
1620 fn build_projected_columns_for_distinct(
1621 &self,
1622 items: &[SelectItem],
1623 table: &TableMetadata,
1624 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1625 let projection = self.build_projection(items, table)?;
1626 match projection {
1627 Projection::All(columns) => {
1628 let mut projected = Vec::with_capacity(columns.len());
1629 for column in columns {
1630 let column_index = table.get_column_index(&column).ok_or_else(|| {
1631 PlannerError::invalid_expression(format!(
1632 "column '{column}' not found for DISTINCT projection"
1633 ))
1634 })?;
1635 let column_meta = table.get_column(&column).ok_or_else(|| {
1636 PlannerError::invalid_expression(format!(
1637 "column '{column}' not found for DISTINCT projection"
1638 ))
1639 })?;
1640 let typed_expr = TypedExpr::column_ref(
1641 table.name.clone(),
1642 column.clone(),
1643 column_index,
1644 column_meta.data_type.clone(),
1645 crate::ast::Span::default(),
1646 );
1647 projected.push(ProjectedColumn::new(typed_expr));
1648 }
1649 Ok(projected)
1650 }
1651 Projection::Columns(columns) => Ok(columns),
1652 }
1653 }
1654
1655 fn build_projected_columns_for_distinct_with_scope(
1656 &self,
1657 items: &[SelectItem],
1658 schema: &[ColumnMetadata],
1659 scope: &[ScopedTable],
1660 ) -> Result<Vec<ProjectedColumn>, PlannerError> {
1661 let projection = self.build_projection_with_scope(items, schema, scope)?;
1662 match projection {
1663 Projection::All(columns) => {
1664 let mut projected = Vec::with_capacity(columns.len());
1665 for (idx, column) in columns.into_iter().enumerate() {
1666 let column_meta = schema.get(idx).ok_or_else(|| {
1667 PlannerError::invalid_expression(format!(
1668 "column '{column}' not found for DISTINCT projection"
1669 ))
1670 })?;
1671 projected.push(ProjectedColumn::new(TypedExpr::column_ref(
1672 LITERAL_TABLE.to_string(),
1673 column,
1674 idx,
1675 column_meta.data_type.clone(),
1676 crate::ast::Span::default(),
1677 )));
1678 }
1679 Ok(projected)
1680 }
1681 Projection::Columns(columns) => Ok(columns),
1682 }
1683 }
1684
1685 fn collect_aggregates_from_typed_expr(
1686 &self,
1687 expr: &TypedExpr,
1688 aggregates: &mut Vec<AggregateExpr>,
1689 aggregate_map: &mut HashMap<AggregateSignature, usize>,
1690 ) -> Result<(), PlannerError> {
1691 match &expr.kind {
1692 TypedExprKind::FunctionCall {
1693 name,
1694 args,
1695 distinct,
1696 star,
1697 } if is_aggregate_function(name) => {
1698 for arg in args {
1699 if typed_expr_contains_aggregate(arg) {
1700 return Err(PlannerError::invalid_expression(
1701 "nested aggregate functions are not supported".to_string(),
1702 ));
1703 }
1704 }
1705 let (agg, signature) =
1706 self.build_aggregate_expr_from_typed(expr, name, args, *distinct, *star)?;
1707 aggregate_map.entry(signature).or_insert_with(|| {
1708 aggregates.push(agg);
1709 aggregates.len() - 1
1710 });
1711 Ok(())
1712 }
1713 TypedExprKind::BinaryOp { left, right, .. } => {
1714 self.collect_aggregates_from_typed_expr(left, aggregates, aggregate_map)?;
1715 self.collect_aggregates_from_typed_expr(right, aggregates, aggregate_map)?;
1716 Ok(())
1717 }
1718 TypedExprKind::UnaryOp { operand, .. } => {
1719 self.collect_aggregates_from_typed_expr(operand, aggregates, aggregate_map)
1720 }
1721 TypedExprKind::FunctionCall { args, .. } => {
1722 for arg in args {
1723 self.collect_aggregates_from_typed_expr(arg, aggregates, aggregate_map)?;
1724 }
1725 Ok(())
1726 }
1727 TypedExprKind::Between {
1728 expr, low, high, ..
1729 } => {
1730 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1731 self.collect_aggregates_from_typed_expr(low, aggregates, aggregate_map)?;
1732 self.collect_aggregates_from_typed_expr(high, aggregates, aggregate_map)?;
1733 Ok(())
1734 }
1735 TypedExprKind::Like {
1736 expr,
1737 pattern,
1738 escape,
1739 ..
1740 } => {
1741 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1742 self.collect_aggregates_from_typed_expr(pattern, aggregates, aggregate_map)?;
1743 if let Some(esc) = escape {
1744 self.collect_aggregates_from_typed_expr(esc, aggregates, aggregate_map)?;
1745 }
1746 Ok(())
1747 }
1748 TypedExprKind::InList { expr, list, .. } => {
1749 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)?;
1750 for item in list {
1751 self.collect_aggregates_from_typed_expr(item, aggregates, aggregate_map)?;
1752 }
1753 Ok(())
1754 }
1755 TypedExprKind::IsNull { expr, .. } => {
1756 self.collect_aggregates_from_typed_expr(expr, aggregates, aggregate_map)
1757 }
1758 _ => Ok(()),
1759 }
1760 }
1761
1762 fn build_aggregate_expr_from_typed(
1763 &self,
1764 expr: &TypedExpr,
1765 name: &str,
1766 args: &[TypedExpr],
1767 distinct: bool,
1768 star: bool,
1769 ) -> Result<(AggregateExpr, AggregateSignature), PlannerError> {
1770 let lower = name.to_lowercase();
1771 match lower.as_str() {
1772 "count" => {
1773 if star {
1774 let agg = AggregateExpr::count_star();
1775 let signature = aggregate_signature(name, distinct, star, None, None, expr);
1776 return Ok((agg, signature));
1777 }
1778 if args.len() != 1 {
1779 return Err(PlannerError::type_mismatch(
1780 "1 argument",
1781 format!("{} arguments", args.len()),
1782 expr.span,
1783 ));
1784 }
1785 let agg = AggregateExpr {
1786 function: AggregateFunction::Count,
1787 arg: Some(args[0].clone()),
1788 distinct,
1789 result_type: ResolvedType::BigInt,
1790 };
1791 let signature =
1792 aggregate_signature(name, distinct, star, Some(&args[0]), None, expr);
1793 Ok((agg, signature))
1794 }
1795 "sum" => {
1796 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1797 let agg = AggregateExpr {
1798 function: AggregateFunction::Sum,
1799 arg: Some(arg.clone()),
1800 distinct,
1801 result_type: ResolvedType::Double,
1802 };
1803 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1804 Ok((agg, signature))
1805 }
1806 "total" => {
1807 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1808 let agg = AggregateExpr {
1809 function: AggregateFunction::Total,
1810 arg: Some(arg.clone()),
1811 distinct: false,
1812 result_type: ResolvedType::Double,
1813 };
1814 let signature = aggregate_signature(name, false, star, Some(arg), None, expr);
1815 Ok((agg, signature))
1816 }
1817 "avg" => {
1818 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1819 let agg = AggregateExpr {
1820 function: AggregateFunction::Avg,
1821 arg: Some(arg.clone()),
1822 distinct,
1823 result_type: ResolvedType::Double,
1824 };
1825 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1826 Ok((agg, signature))
1827 }
1828 "min" => {
1829 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1830 let agg = AggregateExpr {
1831 function: AggregateFunction::Min,
1832 arg: Some(arg.clone()),
1833 distinct,
1834 result_type: arg.resolved_type.clone(),
1835 };
1836 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1837 Ok((agg, signature))
1838 }
1839 "max" => {
1840 let arg = self.require_single_aggregate_arg(args, expr.span)?;
1841 let agg = AggregateExpr {
1842 function: AggregateFunction::Max,
1843 arg: Some(arg.clone()),
1844 distinct,
1845 result_type: arg.resolved_type.clone(),
1846 };
1847 let signature = aggregate_signature(name, distinct, star, Some(arg), None, expr);
1848 Ok((agg, signature))
1849 }
1850 "group_concat" => {
1851 if args.is_empty() || args.len() > 2 {
1852 return Err(PlannerError::type_mismatch(
1853 "1 or 2 arguments",
1854 format!("{} arguments", args.len()),
1855 expr.span,
1856 ));
1857 }
1858 let arg = &args[0];
1859 let mut separator = None;
1860 if args.len() == 2 {
1861 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
1862 separator = Some(value.clone());
1863 } else {
1864 return Err(PlannerError::invalid_expression(
1865 "GROUP_CONCAT separator must be a string literal".to_string(),
1866 ));
1867 }
1868 }
1869 let agg = AggregateExpr {
1870 function: AggregateFunction::GroupConcat { separator },
1871 arg: Some(arg.clone()),
1872 distinct,
1873 result_type: ResolvedType::Text,
1874 };
1875 let signature = aggregate_signature(
1876 name,
1877 distinct,
1878 star,
1879 Some(arg),
1880 match &agg.function {
1881 AggregateFunction::GroupConcat { separator } => separator.as_ref(),
1882 _ => None,
1883 },
1884 expr,
1885 );
1886 Ok((agg, signature))
1887 }
1888 "string_agg" => {
1889 if args.len() != 2 {
1890 return Err(PlannerError::type_mismatch(
1891 "2 arguments",
1892 format!("{} arguments", args.len()),
1893 expr.span,
1894 ));
1895 }
1896 let arg = &args[0];
1897 let separator =
1898 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
1899 Some(value.clone())
1900 } else {
1901 return Err(PlannerError::invalid_expression(
1902 "STRING_AGG separator must be a string literal".to_string(),
1903 ));
1904 };
1905 let agg = AggregateExpr {
1906 function: AggregateFunction::StringAgg { separator },
1907 arg: Some(arg.clone()),
1908 distinct,
1909 result_type: ResolvedType::Text,
1910 };
1911 let signature = aggregate_signature(
1912 name,
1913 distinct,
1914 star,
1915 Some(arg),
1916 match &agg.function {
1917 AggregateFunction::StringAgg { separator } => separator.as_ref(),
1918 _ => None,
1919 },
1920 expr,
1921 );
1922 Ok((agg, signature))
1923 }
1924 _ => Err(PlannerError::unsupported_feature(
1925 format!("function '{}'", name),
1926 "future",
1927 expr.span,
1928 )),
1929 }
1930 }
1931
1932 fn require_single_aggregate_arg<'b>(
1933 &self,
1934 args: &'b [TypedExpr],
1935 span: crate::ast::Span,
1936 ) -> Result<&'b TypedExpr, PlannerError> {
1937 if args.len() != 1 {
1938 return Err(PlannerError::type_mismatch(
1939 "1 argument",
1940 format!("{} arguments", args.len()),
1941 span,
1942 ));
1943 }
1944 Ok(&args[0])
1945 }
1946
1947 fn build_aggregate_projection(
1948 &self,
1949 projected: Vec<ProjectedColumn>,
1950 group_keys: &[TypedExpr],
1951 aggregates: &[AggregateExpr],
1952 output_names: &[String],
1953 ) -> Result<Projection, PlannerError> {
1954 let mut columns = Vec::new();
1955 for col in projected {
1956 let rewritten =
1957 self.rewrite_expr_for_aggregate(&col.expr, group_keys, aggregates, output_names)?;
1958 columns.push(ProjectedColumn {
1959 expr: rewritten,
1960 alias: col.alias,
1961 });
1962 }
1963 Ok(Projection::Columns(columns))
1964 }
1965
1966 fn rewrite_expr_for_aggregate(
1967 &self,
1968 expr: &TypedExpr,
1969 group_keys: &[TypedExpr],
1970 aggregates: &[AggregateExpr],
1971 output_names: &[String],
1972 ) -> Result<TypedExpr, PlannerError> {
1973 let group_key_map = build_group_key_map(group_keys);
1974 let aggregate_map = build_aggregate_map(aggregates);
1975
1976 rewrite_expr_with_maps(expr, &group_key_map, &aggregate_map, output_names)
1977 }
1978
1979 fn extract_limit_value(
1983 &self,
1984 expr: &Option<crate::ast::expr::Expr>,
1985 stmt_span: crate::ast::Span,
1986 ) -> Result<Option<u64>, PlannerError> {
1987 match expr {
1988 None => Ok(None),
1989 Some(e) => {
1990 if let crate::ast::expr::ExprKind::Literal {
1992 literal: Literal::Number(s),
1993 } = &e.kind
1994 {
1995 s.parse::<u64>().map(Some).map_err(|_| {
1996 PlannerError::type_mismatch("unsigned integer", s.clone(), e.span)
1997 })
1998 } else {
1999 Err(PlannerError::unsupported_feature(
2000 "non-literal LIMIT/OFFSET",
2001 "v0.3.0+",
2002 stmt_span,
2003 ))
2004 }
2005 }
2006 }
2007 }
2008
2009 fn plan_insert(&self, stmt: &Insert) -> Result<LogicalPlan, PlannerError> {
2014 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2016
2017 let columns: Vec<String> = if let Some(ref cols) = stmt.columns {
2019 for col in cols {
2021 self.name_resolver.resolve_column(table, col, stmt.span)?;
2022 }
2023 cols.clone()
2024 } else {
2025 table.column_names().into_iter().map(String::from).collect()
2027 };
2028
2029 let mut typed_values: Vec<Vec<TypedExpr>> = Vec::new();
2031
2032 for row in &stmt.values {
2033 if row.len() != columns.len() {
2035 return Err(PlannerError::column_value_count_mismatch(
2036 columns.len(),
2037 row.len(),
2038 stmt.span,
2039 ));
2040 }
2041
2042 let typed_row = self.type_check_insert_values(row, &columns, table)?;
2044 typed_values.push(typed_row);
2045 }
2046
2047 Ok(LogicalPlan::Insert {
2048 table: table.name.clone(),
2049 columns,
2050 values: typed_values,
2051 })
2052 }
2053
2054 fn type_check_insert_values(
2056 &self,
2057 values: &[crate::ast::expr::Expr],
2058 columns: &[String],
2059 table: &TableMetadata,
2060 ) -> Result<Vec<TypedExpr>, PlannerError> {
2061 let mut typed_values = Vec::new();
2062
2063 for (i, value) in values.iter().enumerate() {
2064 let column_name = &columns[i];
2065 let column_meta = table.get_column(column_name).ok_or_else(|| {
2066 PlannerError::column_not_found(column_name, &table.name, value.span)
2067 })?;
2068
2069 let typed_value = self.type_checker.infer_type(value, table)?;
2071
2072 if column_meta.not_null
2074 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2075 {
2076 return Err(PlannerError::null_constraint_violation(
2077 column_name,
2078 value.span,
2079 ));
2080 }
2081
2082 self.validate_type_assignment(&typed_value, &column_meta.data_type, value.span)?;
2084
2085 typed_values.push(typed_value);
2086 }
2087
2088 Ok(typed_values)
2089 }
2090
2091 fn validate_type_assignment(
2093 &self,
2094 value: &TypedExpr,
2095 target_type: &ResolvedType,
2096 span: crate::ast::Span,
2097 ) -> Result<(), PlannerError> {
2098 if value.resolved_type == ResolvedType::Null {
2100 return Ok(());
2101 }
2102
2103 if self.types_compatible(&value.resolved_type, target_type) {
2105 return Ok(());
2106 }
2107
2108 Err(PlannerError::type_mismatch(
2109 target_type.to_string(),
2110 value.resolved_type.to_string(),
2111 span,
2112 ))
2113 }
2114
2115 fn types_compatible(&self, source: &ResolvedType, target: &ResolvedType) -> bool {
2117 use ResolvedType::*;
2118
2119 if source == target {
2121 return true;
2122 }
2123
2124 match (source, target) {
2126 (Integer, BigInt) | (Integer, Float) | (Integer, Double) => true,
2128 (BigInt, Float) | (BigInt, Double) => true,
2130 (Float, Double) => true,
2132 (Vector { dimension: d1, .. }, Vector { dimension: d2, .. }) => d1 == d2,
2134 _ => false,
2135 }
2136 }
2137
2138 fn plan_update(&self, stmt: &Update) -> Result<LogicalPlan, PlannerError> {
2142 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2144
2145 let mut typed_assignments = Vec::new();
2147
2148 for assignment in &stmt.assignments {
2149 let column_meta =
2151 self.name_resolver
2152 .resolve_column(table, &assignment.column, assignment.span)?;
2153 let column_index = table.get_column_index(&assignment.column).unwrap();
2154
2155 let typed_value = self.type_checker.infer_type(&assignment.value, table)?;
2157
2158 if column_meta.not_null
2160 && matches!(&typed_value.kind, TypedExprKind::Literal(Literal::Null))
2161 {
2162 return Err(PlannerError::null_constraint_violation(
2163 &assignment.column,
2164 assignment.value.span,
2165 ));
2166 }
2167
2168 self.validate_type_assignment(
2170 &typed_value,
2171 &column_meta.data_type,
2172 assignment.value.span,
2173 )?;
2174
2175 typed_assignments.push(TypedAssignment::new(
2176 assignment.column.clone(),
2177 column_index,
2178 typed_value,
2179 ));
2180 }
2181
2182 let filter = if let Some(ref selection) = stmt.selection {
2184 let predicate = self.type_checker.infer_type(selection, table)?;
2185
2186 if predicate.resolved_type != ResolvedType::Boolean {
2188 return Err(PlannerError::type_mismatch(
2189 "Boolean",
2190 predicate.resolved_type.to_string(),
2191 selection.span,
2192 ));
2193 }
2194
2195 Some(predicate)
2196 } else {
2197 None
2198 };
2199
2200 Ok(LogicalPlan::Update {
2201 table: table.name.clone(),
2202 assignments: typed_assignments,
2203 filter,
2204 })
2205 }
2206
2207 fn plan_delete(&self, stmt: &Delete) -> Result<LogicalPlan, PlannerError> {
2211 let table = self.name_resolver.resolve_table(&stmt.table, stmt.span)?;
2213
2214 let filter = if let Some(ref selection) = stmt.selection {
2216 let predicate = self.type_checker.infer_type(selection, table)?;
2217
2218 if predicate.resolved_type != ResolvedType::Boolean {
2220 return Err(PlannerError::type_mismatch(
2221 "Boolean",
2222 predicate.resolved_type.to_string(),
2223 selection.span,
2224 ));
2225 }
2226
2227 Some(predicate)
2228 } else {
2229 None
2230 };
2231
2232 Ok(LogicalPlan::Delete {
2233 table: table.name.clone(),
2234 filter,
2235 })
2236 }
2237}
2238
2239#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2240struct AggregateSignature {
2241 name: String,
2242 distinct: bool,
2243 star: bool,
2244 arg_key: Option<String>,
2245 separator: Option<String>,
2246}
2247
2248fn expr_contains_aggregate(expr: &crate::ast::expr::Expr) -> bool {
2249 use crate::ast::expr::ExprKind;
2250
2251 match &expr.kind {
2252 ExprKind::FunctionCall { name, args, .. } => {
2253 if is_aggregate_function(name) {
2254 return true;
2255 }
2256 args.iter().any(expr_contains_aggregate)
2257 }
2258 ExprKind::BinaryOp { left, right, .. } => {
2259 expr_contains_aggregate(left) || expr_contains_aggregate(right)
2260 }
2261 ExprKind::UnaryOp { operand, .. } => expr_contains_aggregate(operand),
2262 ExprKind::Between {
2263 expr, low, high, ..
2264 } => {
2265 expr_contains_aggregate(expr)
2266 || expr_contains_aggregate(low)
2267 || expr_contains_aggregate(high)
2268 }
2269 ExprKind::Like {
2270 expr,
2271 pattern,
2272 escape,
2273 ..
2274 } => {
2275 expr_contains_aggregate(expr)
2276 || expr_contains_aggregate(pattern)
2277 || escape.as_deref().is_some_and(expr_contains_aggregate)
2278 }
2279 ExprKind::InList { expr, list, .. } => {
2280 expr_contains_aggregate(expr) || list.iter().any(expr_contains_aggregate)
2281 }
2282 ExprKind::IsNull { expr, .. } => expr_contains_aggregate(expr),
2283 ExprKind::ScalarSubquery { .. }
2284 | ExprKind::InSubquery { .. }
2285 | ExprKind::Exists { .. }
2286 | ExprKind::Quantified { .. }
2287 | ExprKind::Literal { .. }
2288 | ExprKind::VectorLiteral { .. }
2289 | ExprKind::ColumnRef { .. } => false,
2290 }
2291}
2292
2293fn typed_expr_contains_aggregate(expr: &TypedExpr) -> bool {
2294 match &expr.kind {
2295 TypedExprKind::FunctionCall { name, args, .. } => {
2296 if is_aggregate_function(name) {
2297 return true;
2298 }
2299 args.iter().any(typed_expr_contains_aggregate)
2300 }
2301 TypedExprKind::BinaryOp { left, right, .. } => {
2302 typed_expr_contains_aggregate(left) || typed_expr_contains_aggregate(right)
2303 }
2304 TypedExprKind::UnaryOp { operand, .. } => typed_expr_contains_aggregate(operand),
2305 TypedExprKind::Between {
2306 expr, low, high, ..
2307 } => {
2308 typed_expr_contains_aggregate(expr)
2309 || typed_expr_contains_aggregate(low)
2310 || typed_expr_contains_aggregate(high)
2311 }
2312 TypedExprKind::Like {
2313 expr,
2314 pattern,
2315 escape,
2316 ..
2317 } => {
2318 typed_expr_contains_aggregate(expr)
2319 || typed_expr_contains_aggregate(pattern)
2320 || escape
2321 .as_ref()
2322 .is_some_and(|inner| typed_expr_contains_aggregate(inner))
2323 }
2324 TypedExprKind::InList { expr, list, .. } => {
2325 typed_expr_contains_aggregate(expr) || list.iter().any(typed_expr_contains_aggregate)
2326 }
2327 TypedExprKind::IsNull { expr, .. } => typed_expr_contains_aggregate(expr),
2328 TypedExprKind::InSubquery { expr, .. } => typed_expr_contains_aggregate(expr),
2329 TypedExprKind::Quantified { expr, .. } => typed_expr_contains_aggregate(expr),
2330 TypedExprKind::ScalarSubquery(_) | TypedExprKind::Exists { .. } => false,
2331 _ => false,
2332 }
2333}
2334
2335fn map_join_type(join_type: crate::ast::dml::JoinType) -> JoinType {
2336 match join_type {
2337 crate::ast::dml::JoinType::Inner => JoinType::Inner,
2338 crate::ast::dml::JoinType::Left => JoinType::Left,
2339 crate::ast::dml::JoinType::Right => JoinType::Right,
2340 crate::ast::dml::JoinType::Full => JoinType::Full,
2341 crate::ast::dml::JoinType::Cross => JoinType::Cross,
2342 }
2343}
2344
2345struct FoundScopedColumn {
2346 table: String,
2347 index: usize,
2348 ty: ResolvedType,
2349}
2350
2351fn find_scoped_column(
2352 scope: &[ScopedTable],
2353 column: &str,
2354 span: crate::ast::Span,
2355) -> Result<FoundScopedColumn, PlannerError> {
2356 let mut matches = Vec::new();
2357 for table in scope {
2358 if let Some(local_idx) = table.table.get_column_index(column) {
2359 let meta = &table.table.columns[local_idx];
2360 matches.push(FoundScopedColumn {
2361 table: table.table.name.clone(),
2362 index: table.start_index + local_idx,
2363 ty: meta.data_type.clone(),
2364 });
2365 }
2366 }
2367 match matches.len() {
2368 0 => Err(PlannerError::column_not_found(column, "JOIN input", span)),
2369 1 => Ok(matches.remove(0)),
2370 _ => Err(PlannerError::ambiguous_column(
2371 column,
2372 scope.iter().map(|s| s.table.name.clone()).collect(),
2373 span,
2374 )),
2375 }
2376}
2377
2378fn projection_schema(
2379 projection: &Projection,
2380 input_schema: &[ColumnMetadata],
2381) -> Vec<ColumnMetadata> {
2382 match projection {
2383 Projection::All(names) => names
2384 .iter()
2385 .enumerate()
2386 .map(|(idx, name)| {
2387 let ty = input_schema
2388 .get(idx)
2389 .or_else(|| input_schema.iter().find(|col| &col.name == name))
2390 .map(|col| col.data_type.clone())
2391 .unwrap_or(ResolvedType::Null);
2392 ColumnMetadata::new(name.clone(), ty)
2393 })
2394 .collect(),
2395 Projection::Columns(columns) => columns
2396 .iter()
2397 .enumerate()
2398 .map(|(idx, col)| {
2399 let name = col
2400 .alias
2401 .clone()
2402 .or_else(|| match &col.expr.kind {
2403 TypedExprKind::ColumnRef { column, .. } => Some(column.clone()),
2404 _ => None,
2405 })
2406 .unwrap_or_else(|| format!("col_{idx}"));
2407 ColumnMetadata::new(name, col.expr.resolved_type.clone())
2408 })
2409 .collect(),
2410 }
2411}
2412
2413fn offset_scope(scope: &[ScopedTable], offset: usize) -> Vec<ScopedTable> {
2414 scope
2415 .iter()
2416 .cloned()
2417 .map(|mut table| {
2418 table.start_index += offset;
2419 table
2420 })
2421 .collect()
2422}
2423
2424fn install_base_projection(plan: &mut LogicalPlan, projection: &Projection) {
2425 match plan {
2426 LogicalPlan::Scan {
2427 projection: scan_projection,
2428 ..
2429 } => *scan_projection = projection.clone(),
2430 LogicalPlan::Filter { input, .. } => install_base_projection(input, projection),
2431 _ => {}
2432 }
2433}
2434
2435fn is_aggregate_function(name: &str) -> bool {
2436 matches!(
2437 name.to_ascii_lowercase().as_str(),
2438 "count" | "sum" | "total" | "avg" | "min" | "max" | "group_concat" | "string_agg"
2439 )
2440}
2441
2442fn expr_key(expr: &TypedExpr) -> String {
2443 format!("{:?}", expr.kind)
2444}
2445
2446fn aggregate_signature(
2447 name: &str,
2448 distinct: bool,
2449 star: bool,
2450 arg: Option<&TypedExpr>,
2451 separator: Option<&String>,
2452 _expr: &TypedExpr,
2453) -> AggregateSignature {
2454 AggregateSignature {
2455 name: name.to_ascii_lowercase(),
2456 distinct,
2457 star,
2458 arg_key: arg.map(expr_key),
2459 separator: separator.cloned(),
2460 }
2461}
2462
2463fn build_group_key_map(group_keys: &[TypedExpr]) -> HashMap<String, usize> {
2464 let mut map = HashMap::new();
2465 for (idx, key) in group_keys.iter().enumerate() {
2466 map.insert(expr_key(key), idx);
2467 }
2468 map
2469}
2470
2471fn build_aggregate_map(aggregates: &[AggregateExpr]) -> HashMap<AggregateSignature, usize> {
2472 let mut map = HashMap::new();
2473 for (idx, agg) in aggregates.iter().enumerate() {
2474 let (name, separator, star, arg) = match &agg.function {
2475 AggregateFunction::Count => (
2476 "count".to_string(),
2477 None,
2478 agg.arg.is_none(),
2479 agg.arg.as_ref(),
2480 ),
2481 AggregateFunction::Sum => ("sum".to_string(), None, false, agg.arg.as_ref()),
2482 AggregateFunction::Total => ("total".to_string(), None, false, agg.arg.as_ref()),
2483 AggregateFunction::Avg => ("avg".to_string(), None, false, agg.arg.as_ref()),
2484 AggregateFunction::Min => ("min".to_string(), None, false, agg.arg.as_ref()),
2485 AggregateFunction::Max => ("max".to_string(), None, false, agg.arg.as_ref()),
2486 AggregateFunction::GroupConcat { separator } => (
2487 "group_concat".to_string(),
2488 separator.clone(),
2489 false,
2490 agg.arg.as_ref(),
2491 ),
2492 AggregateFunction::StringAgg { separator } => (
2493 "string_agg".to_string(),
2494 separator.clone(),
2495 false,
2496 agg.arg.as_ref(),
2497 ),
2498 };
2499 let signature = AggregateSignature {
2500 name,
2501 distinct: agg.distinct,
2502 star,
2503 arg_key: arg.map(expr_key),
2504 separator,
2505 };
2506 map.insert(signature, idx);
2507 }
2508 map
2509}
2510
2511fn build_aggregate_schema(
2512 group_keys: &[TypedExpr],
2513 aggregates: &[AggregateExpr],
2514) -> Vec<ColumnMetadata> {
2515 let mut schema = Vec::new();
2516 for (idx, key) in group_keys.iter().enumerate() {
2517 let name = match &key.kind {
2518 TypedExprKind::ColumnRef { column, .. } => column.clone(),
2519 _ => format!("group_{idx}"),
2520 };
2521 schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
2522 }
2523 for (idx, agg) in aggregates.iter().enumerate() {
2524 let name = match &agg.function {
2525 AggregateFunction::Count => format!("count_{idx}"),
2526 AggregateFunction::Sum => format!("sum_{idx}"),
2527 AggregateFunction::Total => format!("total_{idx}"),
2528 AggregateFunction::Avg => format!("avg_{idx}"),
2529 AggregateFunction::Min => format!("min_{idx}"),
2530 AggregateFunction::Max => format!("max_{idx}"),
2531 AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
2532 AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
2533 };
2534 schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
2535 }
2536 schema
2537}
2538
2539fn rewrite_expr_with_maps(
2540 expr: &TypedExpr,
2541 group_key_map: &HashMap<String, usize>,
2542 aggregate_map: &HashMap<AggregateSignature, usize>,
2543 output_names: &[String],
2544) -> Result<TypedExpr, PlannerError> {
2545 let group_key_count = output_names.len().saturating_sub(aggregate_map.len());
2546 let key = expr_key(expr);
2547 if let Some(idx) = group_key_map.get(&key) {
2548 return Ok(make_output_column_ref(
2549 *idx,
2550 output_names,
2551 expr.resolved_type.clone(),
2552 expr.span,
2553 ));
2554 }
2555
2556 match &expr.kind {
2557 TypedExprKind::FunctionCall {
2558 name,
2559 args,
2560 distinct,
2561 star,
2562 } if is_aggregate_function(name) => {
2563 let separator = if name.eq_ignore_ascii_case("group_concat") && args.len() == 2 {
2564 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2565 Some(value.clone())
2566 } else {
2567 return Err(PlannerError::invalid_expression(
2568 "GROUP_CONCAT separator must be a string literal".to_string(),
2569 ));
2570 }
2571 } else if name.eq_ignore_ascii_case("string_agg") && args.len() == 2 {
2572 if let TypedExprKind::Literal(Literal::String(value)) = &args[1].kind {
2573 Some(value.clone())
2574 } else {
2575 return Err(PlannerError::invalid_expression(
2576 "STRING_AGG separator must be a string literal".to_string(),
2577 ));
2578 }
2579 } else {
2580 None
2581 };
2582 let signature = AggregateSignature {
2583 name: name.to_ascii_lowercase(),
2584 distinct: *distinct,
2585 star: *star,
2586 arg_key: args.first().map(expr_key),
2587 separator,
2588 };
2589 let idx = aggregate_map.get(&signature).ok_or_else(|| {
2590 PlannerError::invalid_expression(
2591 "aggregate in expression is not part of plan".to_string(),
2592 )
2593 })?;
2594 let output_index = group_key_count + idx;
2595 Ok(make_output_column_ref(
2596 output_index,
2597 output_names,
2598 expr.resolved_type.clone(),
2599 expr.span,
2600 ))
2601 }
2602 TypedExprKind::FunctionCall {
2603 name,
2604 args,
2605 distinct,
2606 star,
2607 } => {
2608 if *distinct || *star {
2609 return Err(PlannerError::invalid_expression(
2610 "DISTINCT/STAR modifiers are only supported for aggregates".to_string(),
2611 ));
2612 }
2613 let mut rewritten_args = Vec::with_capacity(args.len());
2614 for arg in args {
2615 rewritten_args.push(rewrite_expr_with_maps(
2616 arg,
2617 group_key_map,
2618 aggregate_map,
2619 output_names,
2620 )?);
2621 }
2622 Ok(TypedExpr {
2623 kind: TypedExprKind::FunctionCall {
2624 name: name.clone(),
2625 args: rewritten_args,
2626 distinct: false,
2627 star: false,
2628 },
2629 resolved_type: expr.resolved_type.clone(),
2630 span: expr.span,
2631 })
2632 }
2633 TypedExprKind::BinaryOp { left, op, right } => {
2634 let left = rewrite_expr_with_maps(left, group_key_map, aggregate_map, output_names)?;
2635 let right = rewrite_expr_with_maps(right, group_key_map, aggregate_map, output_names)?;
2636 Ok(TypedExpr {
2637 kind: TypedExprKind::BinaryOp {
2638 left: Box::new(left),
2639 op: *op,
2640 right: Box::new(right),
2641 },
2642 resolved_type: expr.resolved_type.clone(),
2643 span: expr.span,
2644 })
2645 }
2646 TypedExprKind::UnaryOp { op, operand } => {
2647 let operand =
2648 rewrite_expr_with_maps(operand, group_key_map, aggregate_map, output_names)?;
2649 Ok(TypedExpr {
2650 kind: TypedExprKind::UnaryOp {
2651 op: *op,
2652 operand: Box::new(operand),
2653 },
2654 resolved_type: expr.resolved_type.clone(),
2655 span: expr.span,
2656 })
2657 }
2658 TypedExprKind::Between {
2659 expr: inner,
2660 low,
2661 high,
2662 negated,
2663 } => {
2664 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2665 let low = rewrite_expr_with_maps(low, group_key_map, aggregate_map, output_names)?;
2666 let high = rewrite_expr_with_maps(high, group_key_map, aggregate_map, output_names)?;
2667 Ok(TypedExpr {
2668 kind: TypedExprKind::Between {
2669 expr: Box::new(inner),
2670 low: Box::new(low),
2671 high: Box::new(high),
2672 negated: *negated,
2673 },
2674 resolved_type: expr.resolved_type.clone(),
2675 span: expr.span,
2676 })
2677 }
2678 TypedExprKind::Like {
2679 expr: inner,
2680 pattern,
2681 escape,
2682 negated,
2683 kind,
2684 } => {
2685 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2686 let pattern =
2687 rewrite_expr_with_maps(pattern, group_key_map, aggregate_map, output_names)?;
2688 let escape = if let Some(esc) = escape {
2689 Some(Box::new(rewrite_expr_with_maps(
2690 esc,
2691 group_key_map,
2692 aggregate_map,
2693 output_names,
2694 )?))
2695 } else {
2696 None
2697 };
2698 Ok(TypedExpr {
2699 kind: TypedExprKind::Like {
2700 expr: Box::new(inner),
2701 pattern: Box::new(pattern),
2702 escape,
2703 negated: *negated,
2704 kind: *kind,
2705 },
2706 resolved_type: expr.resolved_type.clone(),
2707 span: expr.span,
2708 })
2709 }
2710 TypedExprKind::InList {
2711 expr: inner,
2712 list,
2713 negated,
2714 } => {
2715 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2716 let mut rewritten_list = Vec::with_capacity(list.len());
2717 for item in list {
2718 rewritten_list.push(rewrite_expr_with_maps(
2719 item,
2720 group_key_map,
2721 aggregate_map,
2722 output_names,
2723 )?);
2724 }
2725 Ok(TypedExpr {
2726 kind: TypedExprKind::InList {
2727 expr: Box::new(inner),
2728 list: rewritten_list,
2729 negated: *negated,
2730 },
2731 resolved_type: expr.resolved_type.clone(),
2732 span: expr.span,
2733 })
2734 }
2735 TypedExprKind::IsNull {
2736 expr: inner,
2737 negated,
2738 } => {
2739 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2740 Ok(TypedExpr {
2741 kind: TypedExprKind::IsNull {
2742 expr: Box::new(inner),
2743 negated: *negated,
2744 },
2745 resolved_type: expr.resolved_type.clone(),
2746 span: expr.span,
2747 })
2748 }
2749 TypedExprKind::Literal(_) | TypedExprKind::VectorLiteral(_) => Ok(expr.clone()),
2750 TypedExprKind::ColumnRef { .. } => Err(PlannerError::invalid_expression(
2751 "column reference must appear in GROUP BY or be aggregated".to_string(),
2752 )),
2753 TypedExprKind::Cast {
2754 expr: inner,
2755 target_type,
2756 } => {
2757 let inner = rewrite_expr_with_maps(inner, group_key_map, aggregate_map, output_names)?;
2758 Ok(TypedExpr {
2759 kind: TypedExprKind::Cast {
2760 expr: Box::new(inner),
2761 target_type: target_type.clone(),
2762 },
2763 resolved_type: expr.resolved_type.clone(),
2764 span: expr.span,
2765 })
2766 }
2767 TypedExprKind::ScalarSubquery(_)
2768 | TypedExprKind::InSubquery { .. }
2769 | TypedExprKind::Exists { .. }
2770 | TypedExprKind::Quantified { .. } => Ok(expr.clone()),
2771 }
2772}
2773
2774fn make_output_column_ref(
2775 index: usize,
2776 output_names: &[String],
2777 resolved_type: ResolvedType,
2778 span: crate::ast::Span,
2779) -> TypedExpr {
2780 let name = output_names
2781 .get(index)
2782 .cloned()
2783 .unwrap_or_else(|| format!("col_{index}"));
2784 TypedExpr::column_ref("__agg__".to_string(), name, index, resolved_type, span)
2785}