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