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    if crate::traversal::is_aggregate(expr) {
355        return true;
356    }
357
358    match expr {
359        Expression::Alias(alias) => contains_aggregate(&alias.this),
360        Expression::Add(op) | Expression::Sub(op) | Expression::Mul(op) | Expression::Div(op) => {
361            contains_aggregate(&op.left) || contains_aggregate(&op.right)
362        }
363        Expression::Function(func) => func.args.iter().any(contains_aggregate),
364        _ => false,
365    }
366}
367
368/// Extract aggregate expressions from a list
369fn extract_aggregations(expressions: &[Expression]) -> Vec<Expression> {
370    let mut aggs = Vec::new();
371    for expr in expressions {
372        collect_aggregations(expr, &mut aggs);
373    }
374    aggs
375}
376
377fn collect_aggregations(expr: &Expression, aggs: &mut Vec<Expression>) {
378    if crate::traversal::is_aggregate(expr) {
379        aggs.push(expr.clone());
380        return;
381    }
382
383    match expr {
384        Expression::Alias(alias) => {
385            collect_aggregations(&alias.this, aggs);
386        }
387        Expression::Add(op) | Expression::Sub(op) | Expression::Mul(op) | Expression::Div(op) => {
388            collect_aggregations(&op.left, aggs);
389            collect_aggregations(&op.right, aggs);
390        }
391        Expression::Function(func) => {
392            for arg in &func.args {
393                collect_aggregations(arg, aggs);
394            }
395        }
396        _ => {}
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403    use crate::dialects::{Dialect, DialectType};
404
405    fn parse(sql: &str) -> Expression {
406        let dialect = Dialect::get(DialectType::Generic);
407        let ast = dialect.parse(sql).unwrap();
408        ast.into_iter().next().unwrap()
409    }
410
411    #[test]
412    fn test_simple_scan() {
413        let sql = "SELECT a, b FROM t";
414        let expr = parse(sql);
415        let plan = Plan::from_expression(&expr);
416
417        assert!(plan.is_some());
418        let plan = plan.unwrap();
419        assert_eq!(plan.root.kind, StepKind::Scan);
420        assert_eq!(plan.root.name, "t");
421    }
422
423    #[test]
424    fn test_join() {
425        let sql = "SELECT t1.a, t2.b FROM t1 JOIN t2 ON t1.id = t2.id";
426        let expr = parse(sql);
427        let plan = Plan::from_expression(&expr);
428
429        assert!(plan.is_some());
430        let plan = plan.unwrap();
431        assert!(matches!(plan.root.kind, StepKind::Join(_)));
432        assert_eq!(plan.root.dependencies.len(), 2);
433    }
434
435    #[test]
436    fn test_nested_join_dag_has_unique_preorder_ids() {
437        let expr = parse(
438            "SELECT t1.a FROM t1 \
439             JOIN t2 ON t1.id = t2.id \
440             JOIN t3 ON t2.id = t3.id",
441        );
442        let mut plan = Plan::from_expression(&expr).unwrap();
443        let dag = plan.dag();
444
445        assert_eq!(dag.len(), 5);
446        assert_eq!(dag.get(&0), Some(&HashSet::from([1, 4])));
447        assert_eq!(dag.get(&1), Some(&HashSet::from([2, 3])));
448        assert_eq!(dag.get(&2), Some(&HashSet::new()));
449        assert_eq!(dag.get(&3), Some(&HashSet::new()));
450        assert_eq!(dag.get(&4), Some(&HashSet::new()));
451
452        assert!(dag
453            .values()
454            .flat_map(|dependencies| dependencies.iter())
455            .all(|dependency| dag.contains_key(dependency)));
456    }
457
458    #[test]
459    fn test_aggregate() {
460        let sql = "SELECT x, SUM(y) FROM t GROUP BY x";
461        let expr = parse(sql);
462        let plan = Plan::from_expression(&expr);
463
464        assert!(plan.is_some());
465        let plan = plan.unwrap();
466        assert_eq!(plan.root.kind, StepKind::Aggregate);
467    }
468
469    #[test]
470    fn test_union() {
471        let sql = "SELECT a FROM t1 UNION SELECT b FROM t2";
472        let expr = parse(sql);
473        let plan = Plan::from_expression(&expr);
474
475        assert!(plan.is_some());
476        let plan = plan.unwrap();
477        assert!(matches!(
478            plan.root.kind,
479            StepKind::SetOperation(SetOperationType::Union)
480        ));
481    }
482
483    #[test]
484    fn test_contains_aggregate() {
485        // Parse a SELECT with an aggregate function and check the expression
486        let select_with_agg = parse("SELECT SUM(x) FROM t");
487        if let Expression::Select(ref sel) = select_with_agg {
488            assert!(!sel.expressions.is_empty());
489            assert!(
490                contains_aggregate(&sel.expressions[0]),
491                "Expected SUM to be detected as aggregate function"
492            );
493        } else {
494            panic!("Expected SELECT expression");
495        }
496
497        // Parse a SELECT with a non-aggregate expression
498        let select_without_agg = parse("SELECT x + 1 FROM t");
499        if let Expression::Select(ref sel) = select_without_agg {
500            assert!(!sel.expressions.is_empty());
501            assert!(
502                !contains_aggregate(&sel.expressions[0]),
503                "Expected x + 1 to not be an aggregate function"
504            );
505        } else {
506            panic!("Expected SELECT expression");
507        }
508    }
509}
polyglot_sql/planner.rs

polyglot_sql/
planner.rs
