polyglot_sql/
planner.rs

1//! Query Execution Planner
2//!
3//! This module provides functionality to convert SQL AST into an execution plan
4//! represented as a DAG (Directed Acyclic Graph) of steps.
5//!
6
7use crate::expressions::{Expression, JoinKind};
8use serde::{Deserialize, Serialize};
9use std::collections::{HashMap, HashSet};
10
11/// A query execution plan
12#[derive(Debug)]
13pub struct Plan {
14    /// The root step of the plan DAG
15    pub root: Step,
16    /// Cached DAG representation
17    dag: Option<HashMap<usize, HashSet<usize>>>,
18}
19
20impl Plan {
21    /// Create a new plan from an expression
22    pub fn from_expression(expression: &Expression) -> Option<Self> {
23        let root = Step::from_expression(expression, &HashMap::new())?;
24        Some(Self { root, dag: None })
25    }
26
27    /// Get the DAG representation of the plan
28    pub fn dag(&mut self) -> &HashMap<usize, HashSet<usize>> {
29        if self.dag.is_none() {
30            let mut dag = HashMap::new();
31            let mut next_id = 0;
32            Self::build_dag(&self.root, &mut dag, &mut next_id);
33            self.dag = Some(dag);
34        }
35        self.dag.as_ref().unwrap()
36    }
37
38    fn build_dag(
39        step: &Step,
40        dag: &mut HashMap<usize, HashSet<usize>>,
41        next_id: &mut usize,
42    ) -> usize {
43        let id = *next_id;
44        *next_id += 1;
45
46        let dependencies = step
47            .dependencies
48            .iter()
49            .map(|dependency| Self::build_dag(dependency, dag, next_id))
50            .collect();
51        dag.insert(id, dependencies);
52        id
53    }
54
55    /// Get all leaf steps (steps with no dependencies)
56    pub fn leaves(&self) -> Vec<&Step> {
57        let mut leaves = Vec::new();
58        self.collect_leaves(&self.root, &mut leaves);
59        leaves
60    }
61
62    fn collect_leaves<'a>(&'a self, step: &'a Step, leaves: &mut Vec<&'a Step>) {
63        if step.dependencies.is_empty() {
64            leaves.push(step);
65        } else {
66            for dep in &step.dependencies {
67                self.collect_leaves(dep, leaves);
68            }
69        }
70    }
71}
72
73/// A step in the execution plan
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct Step {
76    /// Name of this step
77    pub name: String,
78    /// Type of step
79    pub kind: StepKind,
80    /// Projections to output
81    pub projections: Vec<Expression>,
82    /// Dependencies (other steps that must complete first)
83    pub dependencies: Vec<Step>,
84    /// Aggregation expressions (for Aggregate steps)
85    pub aggregations: Vec<Expression>,
86    /// Group by expressions (for Aggregate steps)
87    pub group_by: Vec<Expression>,
88    /// Join condition (for Join steps)
89    pub condition: Option<Expression>,
90    /// Sort expressions (for Sort steps)
91    pub order_by: Vec<Expression>,
92    /// Limit value (for Scan/other steps)
93    pub limit: Option<Expression>,
94}
95
96/// Types of execution steps
97#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
98#[serde(rename_all = "snake_case")]
99pub enum StepKind {
100    /// Scan a table
101    Scan,
102    /// Join multiple inputs
103    Join(JoinType),
104    /// Aggregate rows
105    Aggregate,
106    /// Sort rows
107    Sort,
108    /// Set operation (UNION, INTERSECT, EXCEPT)
109    SetOperation(SetOperationType),
110}
111
112/// Types of joins in execution plans
113#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
114#[serde(rename_all = "snake_case")]
115pub enum JoinType {
116    Inner,
117    Left,
118    Right,
119    Full,
120    Cross,
121}
122
123/// Types of set operations
124#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
125#[serde(rename_all = "snake_case")]
126pub enum SetOperationType {
127    Union,
128    UnionAll,
129    Intersect,
130    Except,
131}
132
133impl Step {
134    /// Create a new step
135    pub fn new(name: impl Into<String>, kind: StepKind) -> Self {
136        Self {
137            name: name.into(),
138            kind,
139            projections: Vec::new(),
140            dependencies: Vec::new(),
141            aggregations: Vec::new(),
142            group_by: Vec::new(),
143            condition: None,
144            order_by: Vec::new(),
145            limit: None,
146        }
147    }
148
149    /// Build a step from an expression
150    pub fn from_expression(expression: &Expression, ctes: &HashMap<String, Step>) -> Option<Self> {
151        match expression {
152            Expression::Select(select) => {
153                let mut step = Self::from_select(select, ctes)?;
154
155                // Handle ORDER BY
156                if let Some(ref order_by) = select.order_by {
157                    let sort_step = Step {
158                        name: step.name.clone(),
159                        kind: StepKind::Sort,
160                        projections: Vec::new(),
161                        dependencies: vec![step],
162                        aggregations: Vec::new(),
163                        group_by: Vec::new(),
164                        condition: None,
165                        order_by: order_by
166                            .expressions
167                            .iter()
168                            .map(|o| o.this.clone())
169                            .collect(),
170                        limit: None,
171                    };
172                    step = sort_step;
173                }
174
175                // Handle LIMIT
176                if let Some(ref limit) = select.limit {
177                    step.limit = Some(limit.this.clone());
178                }
179
180                Some(step)
181            }
182            Expression::Union(union) => {
183                let left = Self::from_expression(&union.left, ctes)?;
184                let right = Self::from_expression(&union.right, ctes)?;
185
186                let op_type = if union.all {
187                    SetOperationType::UnionAll
188                } else {
189                    SetOperationType::Union
190                };
191
192                Some(Step {
193                    name: "UNION".to_string(),
194                    kind: StepKind::SetOperation(op_type),
195                    projections: Vec::new(),
196                    dependencies: vec![left, right],
197                    aggregations: Vec::new(),
198                    group_by: Vec::new(),
199                    condition: None,
200                    order_by: Vec::new(),
201                    limit: None,
202                })
203            }
204            Expression::Intersect(intersect) => {
205                let left = Self::from_expression(&intersect.left, ctes)?;
206                let right = Self::from_expression(&intersect.right, ctes)?;
207
208                Some(Step {
209                    name: "INTERSECT".to_string(),
210                    kind: StepKind::SetOperation(SetOperationType::Intersect),
211                    projections: Vec::new(),
212                    dependencies: vec![left, right],
213                    aggregations: Vec::new(),
214                    group_by: Vec::new(),
215                    condition: None,
216                    order_by: Vec::new(),
217                    limit: None,
218                })
219            }
220            Expression::Except(except) => {
221                let left = Self::from_expression(&except.left, ctes)?;
222                let right = Self::from_expression(&except.right, ctes)?;
223
224                Some(Step {
225                    name: "EXCEPT".to_string(),
226                    kind: StepKind::SetOperation(SetOperationType::Except),
227                    projections: Vec::new(),
228                    dependencies: vec![left, right],
229                    aggregations: Vec::new(),
230                    group_by: Vec::new(),
231                    condition: None,
232                    order_by: Vec::new(),
233                    limit: None,
234                })
235            }
236            _ => None,
237        }
238    }
239
240    fn from_select(
241        select: &crate::expressions::Select,
242        ctes: &HashMap<String, Step>,
243    ) -> Option<Self> {
244        // Process CTEs first
245        let mut ctes = ctes.clone();
246        if let Some(ref with) = select.with {
247            for cte in &with.ctes {
248                if let Some(step) = Self::from_expression(&cte.this, &ctes) {
249                    ctes.insert(cte.alias.name.clone(), step);
250                }
251            }
252        }
253
254        // Start with the FROM clause
255        let mut step = if let Some(ref from) = select.from {
256            if let Some(table_expr) = from.expressions.first() {
257                Self::from_table_expression(table_expr, &ctes)?
258            } else {
259                return None;
260            }
261        } else {
262            // SELECT without FROM (e.g., SELECT 1)
263            Step::new("", StepKind::Scan)
264        };
265
266        // Process JOINs
267        for join in &select.joins {
268            let right = Self::from_table_expression(&join.this, &ctes)?;
269
270            let join_type = match join.kind {
271                JoinKind::Inner => JoinType::Inner,
272                JoinKind::Left | JoinKind::NaturalLeft => JoinType::Left,
273                JoinKind::Right | JoinKind::NaturalRight => JoinType::Right,
274                JoinKind::Full | JoinKind::NaturalFull => JoinType::Full,
275                JoinKind::Cross | JoinKind::Natural => JoinType::Cross,
276                _ => JoinType::Inner,
277            };
278
279            let join_step = Step {
280                name: step.name.clone(),
281                kind: StepKind::Join(join_type),
282                projections: Vec::new(),
283                dependencies: vec![step, right],
284                aggregations: Vec::new(),
285                group_by: Vec::new(),
286                condition: join.on.clone(),
287                order_by: Vec::new(),
288                limit: None,
289            };
290            step = join_step;
291        }
292
293        // Check for aggregations
294        let has_aggregations = select.expressions.iter().any(|e| contains_aggregate(e));
295        let has_group_by = select.group_by.is_some();
296
297        if has_aggregations || has_group_by {
298            // Create aggregate step
299            let agg_step = Step {
300                name: step.name.clone(),
301                kind: StepKind::Aggregate,
302                projections: select.expressions.clone(),
303                dependencies: vec![step],
304                aggregations: extract_aggregations(&select.expressions),
305                group_by: select
306                    .group_by
307                    .as_ref()
308                    .map(|g| g.expressions.clone())
309                    .unwrap_or_default(),
310                condition: None,
311                order_by: Vec::new(),
312                limit: None,
313            };
314            step = agg_step;
315        } else {
316            step.projections = select.expressions.clone();
317        }
318
319        Some(step)
320    }
321
322    fn from_table_expression(expr: &Expression, ctes: &HashMap<String, Step>) -> Option<Self> {
323        match expr {
324            Expression::Table(table) => {
325                // Check if this references a CTE
326                if let Some(cte_step) = ctes.get(&table.name.name) {
327                    return Some(cte_step.clone());
328                }
329
330                // Regular table scan
331                Some(Step::new(&table.name.name, StepKind::Scan))
332            }
333            Expression::Alias(alias) => {
334                let mut step = Self::from_table_expression(&alias.this, ctes)?;
335                step.name = alias.alias.name.clone();
336                Some(step)
337            }
338            Expression::Subquery(sq) => {
339                let step = Self::from_expression(&sq.this, ctes)?;
340                Some(step)
341            }
342            _ => None,
343        }
344    }
345
346    /// Add a dependency to this step
347    pub fn add_dependency(&mut self, dep: Step) {
348        self.dependencies.push(dep);
349    }
350}
351
352/// Check if an expression contains an aggregate function
353fn contains_aggregate(expr: &Expression) -> bool {
354    match expr {
355        // Specific aggregate function variants
356        Expression::Sum(_)
357        | Expression::Count(_)
358        | Expression::Avg(_)
359        | Expression::Min(_)
360        | Expression::Max(_)
361        | Expression::ArrayAgg(_)
362        | Expression::StringAgg(_)
363        | Expression::ListAgg(_)
364        | Expression::Stddev(_)
365        | Expression::StddevPop(_)
366        | Expression::StddevSamp(_)
367        | Expression::Variance(_)
368        | Expression::VarPop(_)
369        | Expression::VarSamp(_)
370        | Expression::Median(_)
371        | Expression::Mode(_)
372        | Expression::First(_)
373        | Expression::Last(_)
374        | Expression::AnyValue(_)
375        | Expression::ApproxDistinct(_)
376        | Expression::ApproxCountDistinct(_)
377        | Expression::LogicalAnd(_)
378        | Expression::LogicalOr(_)
379        | Expression::AggregateFunction(_) => true,
380
381        Expression::Alias(alias) => contains_aggregate(&alias.this),
382        Expression::Add(op) | Expression::Sub(op) | Expression::Mul(op) | Expression::Div(op) => {
383            contains_aggregate(&op.left) || contains_aggregate(&op.right)
384        }
385        Expression::Function(func) => {
386            // Check for aggregate function names (fallback)
387            let name = func.name.to_uppercase();
388            matches!(
389                name.as_str(),
390                "SUM"
391                    | "COUNT"
392                    | "AVG"
393                    | "MIN"
394                    | "MAX"
395                    | "ARRAY_AGG"
396                    | "STRING_AGG"
397                    | "GROUP_CONCAT"
398            )
399        }
400        _ => false,
401    }
402}
403
404/// Extract aggregate expressions from a list
405fn extract_aggregations(expressions: &[Expression]) -> Vec<Expression> {
406    let mut aggs = Vec::new();
407    for expr in expressions {
408        collect_aggregations(expr, &mut aggs);
409    }
410    aggs
411}
412
413fn collect_aggregations(expr: &Expression, aggs: &mut Vec<Expression>) {
414    match expr {
415        // Specific aggregate function variants
416        Expression::Sum(_)
417        | Expression::Count(_)
418        | Expression::Avg(_)
419        | Expression::Min(_)
420        | Expression::Max(_)
421        | Expression::ArrayAgg(_)
422        | Expression::StringAgg(_)
423        | Expression::ListAgg(_)
424        | Expression::Stddev(_)
425        | Expression::StddevPop(_)
426        | Expression::StddevSamp(_)
427        | Expression::Variance(_)
428        | Expression::VarPop(_)
429        | Expression::VarSamp(_)
430        | Expression::Median(_)
431        | Expression::Mode(_)
432        | Expression::First(_)
433        | Expression::Last(_)
434        | Expression::AnyValue(_)
435        | Expression::ApproxDistinct(_)
436        | Expression::ApproxCountDistinct(_)
437        | Expression::LogicalAnd(_)
438        | Expression::LogicalOr(_)
439        | Expression::AggregateFunction(_) => {
440            aggs.push(expr.clone());
441        }
442        Expression::Alias(alias) => {
443            collect_aggregations(&alias.this, aggs);
444        }
445        Expression::Add(op) | Expression::Sub(op) | Expression::Mul(op) | Expression::Div(op) => {
446            collect_aggregations(&op.left, aggs);
447            collect_aggregations(&op.right, aggs);
448        }
449        Expression::Function(func) => {
450            let name = func.name.to_uppercase();
451            if matches!(
452                name.as_str(),
453                "SUM"
454                    | "COUNT"
455                    | "AVG"
456                    | "MIN"
457                    | "MAX"
458                    | "ARRAY_AGG"
459                    | "STRING_AGG"
460                    | "GROUP_CONCAT"
461            ) {
462                aggs.push(expr.clone());
463            } else {
464                for arg in &func.args {
465                    collect_aggregations(arg, aggs);
466                }
467            }
468        }
469        _ => {}
470    }
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476    use crate::dialects::{Dialect, DialectType};
477
478    fn parse(sql: &str) -> Expression {
479        let dialect = Dialect::get(DialectType::Generic);
480        let ast = dialect.parse(sql).unwrap();
481        ast.into_iter().next().unwrap()
482    }
483
484    #[test]
485    fn test_simple_scan() {
486        let sql = "SELECT a, b FROM t";
487        let expr = parse(sql);
488        let plan = Plan::from_expression(&expr);
489
490        assert!(plan.is_some());
491        let plan = plan.unwrap();
492        assert_eq!(plan.root.kind, StepKind::Scan);
493        assert_eq!(plan.root.name, "t");
494    }
495
496    #[test]
497    fn test_join() {
498        let sql = "SELECT t1.a, t2.b FROM t1 JOIN t2 ON t1.id = t2.id";
499        let expr = parse(sql);
500        let plan = Plan::from_expression(&expr);
501
502        assert!(plan.is_some());
503        let plan = plan.unwrap();
504        assert!(matches!(plan.root.kind, StepKind::Join(_)));
505        assert_eq!(plan.root.dependencies.len(), 2);
506    }
507
508    #[test]
509    fn test_nested_join_dag_has_unique_preorder_ids() {
510        let expr = parse(
511            "SELECT t1.a FROM t1 \
512             JOIN t2 ON t1.id = t2.id \
513             JOIN t3 ON t2.id = t3.id",
514        );
515        let mut plan = Plan::from_expression(&expr).unwrap();
516        let dag = plan.dag();
517
518        assert_eq!(dag.len(), 5);
519        assert_eq!(dag.get(&0), Some(&HashSet::from([1, 4])));
520        assert_eq!(dag.get(&1), Some(&HashSet::from([2, 3])));
521        assert_eq!(dag.get(&2), Some(&HashSet::new()));
522        assert_eq!(dag.get(&3), Some(&HashSet::new()));
523        assert_eq!(dag.get(&4), Some(&HashSet::new()));
524
525        assert!(dag
526            .values()
527            .flat_map(|dependencies| dependencies.iter())
528            .all(|dependency| dag.contains_key(dependency)));
529    }
530
531    #[test]
532    fn test_aggregate() {
533        let sql = "SELECT x, SUM(y) FROM t GROUP BY x";
534        let expr = parse(sql);
535        let plan = Plan::from_expression(&expr);
536
537        assert!(plan.is_some());
538        let plan = plan.unwrap();
539        assert_eq!(plan.root.kind, StepKind::Aggregate);
540    }
541
542    #[test]
543    fn test_union() {
544        let sql = "SELECT a FROM t1 UNION SELECT b FROM t2";
545        let expr = parse(sql);
546        let plan = Plan::from_expression(&expr);
547
548        assert!(plan.is_some());
549        let plan = plan.unwrap();
550        assert!(matches!(
551            plan.root.kind,
552            StepKind::SetOperation(SetOperationType::Union)
553        ));
554    }
555
556    #[test]
557    fn test_contains_aggregate() {
558        // Parse a SELECT with an aggregate function and check the expression
559        let select_with_agg = parse("SELECT SUM(x) FROM t");
560        if let Expression::Select(ref sel) = select_with_agg {
561            assert!(!sel.expressions.is_empty());
562            assert!(
563                contains_aggregate(&sel.expressions[0]),
564                "Expected SUM to be detected as aggregate function"
565            );
566        } else {
567            panic!("Expected SELECT expression");
568        }
569
570        // Parse a SELECT with a non-aggregate expression
571        let select_without_agg = parse("SELECT x + 1 FROM t");
572        if let Expression::Select(ref sel) = select_without_agg {
573            assert!(!sel.expressions.is_empty());
574            assert!(
575                !contains_aggregate(&sel.expressions[0]),
576                "Expected x + 1 to not be an aggregate function"
577            );
578        } else {
579            panic!("Expected SELECT expression");
580        }
581    }
582}
polyglot_sql/planner.rs

polyglot_sql/
planner.rs
