polyglot_sql/optimizer/
optimizer.rs

1//! Optimizer Orchestration Module
2//!
3//! This module provides the main entry point for SQL optimization,
4//! coordinating multiple optimization passes in the correct order.
5//!
6//! Ported from sqlglot's optimizer/optimizer.py
7
8use crate::dialects::DialectType;
9use crate::expressions::Expression;
10use crate::schema::Schema;
11
12use super::annotate_types::annotate_types;
13use super::canonicalize::canonicalize;
14use super::eliminate_ctes::eliminate_ctes;
15use super::eliminate_joins::eliminate_joins;
16use super::normalize::normalize;
17use super::optimize_joins::optimize_joins;
18use super::pushdown_predicates::pushdown_predicates;
19use super::pushdown_projections::pushdown_projections;
20use super::qualify_columns::{qualify_columns, quote_identifiers};
21use super::simplify::simplify;
22use super::subquery::{merge_subqueries, unnest_subqueries};
23
24/// Optimizer configuration
25pub struct OptimizerConfig<'a> {
26    /// Database schema for type inference and column resolution
27    pub schema: Option<&'a dyn Schema>,
28    /// Default database name
29    pub db: Option<String>,
30    /// Default catalog name
31    pub catalog: Option<String>,
32    /// Dialect for dialect-specific optimizations
33    pub dialect: Option<DialectType>,
34    /// Whether to keep tables isolated (don't merge from multiple tables)
35    pub isolate_tables: bool,
36    /// Whether to quote identifiers
37    pub quote_identifiers: bool,
38}
39
40impl<'a> Default for OptimizerConfig<'a> {
41    fn default() -> Self {
42        Self {
43            schema: None,
44            db: None,
45            catalog: None,
46            dialect: None,
47            isolate_tables: true,
48            quote_identifiers: false,
49        }
50    }
51}
52
53/// Optimization rule type
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum OptimizationRule {
56    /// Qualify columns and tables with their full names
57    Qualify,
58    /// Push projections down to eliminate unused columns early
59    PushdownProjections,
60    /// Normalize boolean expressions
61    Normalize,
62    /// Unnest correlated subqueries into joins
63    UnnestSubqueries,
64    /// Push predicates down to filter data early
65    PushdownPredicates,
66    /// Optimize join order and remove cross joins
67    OptimizeJoins,
68    /// Eliminate derived tables by converting to CTEs
69    EliminateSubqueries,
70    /// Merge subqueries into outer queries
71    MergeSubqueries,
72    /// Eliminate unused joins after join optimization and subquery merges
73    EliminateJoins,
74    /// Remove unused CTEs
75    EliminateCtes,
76    /// Quote identifiers that require quoting for the target dialect
77    QuoteIdentifiers,
78    /// Annotate expressions with type information
79    AnnotateTypes,
80    /// Convert expressions to canonical form
81    Canonicalize,
82    /// Simplify expressions
83    Simplify,
84}
85
86/// Default optimization rules in order of execution
87pub const DEFAULT_RULES: &[OptimizationRule] = &[
88    OptimizationRule::Qualify,
89    OptimizationRule::PushdownProjections,
90    OptimizationRule::Normalize,
91    OptimizationRule::UnnestSubqueries,
92    OptimizationRule::PushdownPredicates,
93    OptimizationRule::OptimizeJoins,
94    OptimizationRule::EliminateSubqueries,
95    OptimizationRule::MergeSubqueries,
96    OptimizationRule::EliminateJoins,
97    OptimizationRule::EliminateCtes,
98    OptimizationRule::QuoteIdentifiers,
99    OptimizationRule::AnnotateTypes,
100    OptimizationRule::Canonicalize,
101    OptimizationRule::Simplify,
102];
103
104const QUICK_RULES: &[OptimizationRule] =
105    &[OptimizationRule::Simplify, OptimizationRule::Canonicalize];
106const FAST_PATH_MAX_DEPTH: usize = 768;
107const FAST_PATH_MAX_CONNECTORS: usize = 10_000;
108const FAST_PATH_MAX_CONNECTOR_DEPTH: usize = 1024;
109const FAST_PATH_MAX_NODES: usize = 50_000;
110const CLONE_HEAVY_RULE_SKIP_NODES: usize = 20_000;
111
112#[derive(Debug, Clone, Copy)]
113struct ExpressionComplexity {
114    node_count: usize,
115    max_depth: usize,
116    connector_count: usize,
117    max_connector_depth: usize,
118}
119
120/// Optimize a SQL expression using the default set of rules.
121///
122/// This function coordinates multiple optimization passes in the correct order
123/// to produce an optimized query plan.
124///
125/// # Arguments
126/// * `expression` - The expression to optimize
127/// * `config` - Optimizer configuration
128///
129/// # Returns
130/// The optimized expression
131pub fn optimize(expression: Expression, config: &OptimizerConfig<'_>) -> Expression {
132    optimize_with_rules(expression, config, DEFAULT_RULES)
133}
134
135/// Optimize a SQL expression using a custom set of rules.
136///
137/// # Arguments
138/// * `expression` - The expression to optimize
139/// * `config` - Optimizer configuration
140/// * `rules` - The optimization rules to apply
141///
142/// # Returns
143/// The optimized expression
144pub fn optimize_with_rules(
145    mut expression: Expression,
146    config: &OptimizerConfig<'_>,
147    rules: &[OptimizationRule],
148) -> Expression {
149    let complexity = analyze_expression_complexity(&expression);
150    if rules == DEFAULT_RULES && should_skip_all_optimization(&complexity) {
151        return expression;
152    }
153
154    let active_rules = if rules == DEFAULT_RULES && should_use_quick_path(&complexity) {
155        QUICK_RULES
156    } else {
157        rules
158    };
159
160    for rule in active_rules {
161        if complexity.node_count >= CLONE_HEAVY_RULE_SKIP_NODES
162            && matches!(
163                rule,
164                OptimizationRule::Qualify | OptimizationRule::Normalize
165            )
166        {
167            continue;
168        }
169        expression = apply_rule(expression, *rule, config);
170    }
171    expression
172}
173
174fn should_skip_all_optimization(complexity: &ExpressionComplexity) -> bool {
175    complexity.max_depth >= FAST_PATH_MAX_DEPTH
176        || complexity.max_connector_depth >= FAST_PATH_MAX_CONNECTOR_DEPTH
177}
178
179fn should_use_quick_path(complexity: &ExpressionComplexity) -> bool {
180    complexity.connector_count >= FAST_PATH_MAX_CONNECTORS
181        || complexity.max_connector_depth >= FAST_PATH_MAX_CONNECTOR_DEPTH
182        || complexity.node_count >= FAST_PATH_MAX_NODES
183}
184
185fn analyze_expression_complexity(expression: &Expression) -> ExpressionComplexity {
186    let mut node_count = 0usize;
187    let mut max_depth = 0usize;
188    let mut connector_count = 0usize;
189    let mut max_connector_depth = 0usize;
190    let mut stack: Vec<(&Expression, usize, usize)> = vec![(expression, 0, 0)];
191
192    while let Some((node, depth, connector_depth)) = stack.pop() {
193        node_count += 1;
194        max_depth = max_depth.max(depth);
195
196        match node {
197            Expression::And(op) | Expression::Or(op) => {
198                connector_count += 1;
199                let next_connector_depth = connector_depth + 1;
200                max_connector_depth = max_connector_depth.max(next_connector_depth);
201                stack.push((&op.right, depth + 1, next_connector_depth));
202                stack.push((&op.left, depth + 1, next_connector_depth));
203            }
204            Expression::Paren(paren) => {
205                stack.push((&paren.this, depth + 1, connector_depth));
206            }
207            _ => {
208                let child_start = stack.len();
209                crate::ast_children::for_each_child(node, |_, child| {
210                    stack.push((child, depth + 1, 0));
211                });
212                stack[child_start..].reverse();
213            }
214        }
215    }
216
217    ExpressionComplexity {
218        node_count,
219        max_depth,
220        connector_count,
221        max_connector_depth,
222    }
223}
224
225/// Apply a single optimization rule
226fn apply_rule(
227    expression: Expression,
228    rule: OptimizationRule,
229    config: &OptimizerConfig<'_>,
230) -> Expression {
231    match rule {
232        OptimizationRule::Qualify => {
233            // Qualify columns with table references
234            if let Some(schema) = config.schema {
235                let options = super::qualify_columns::QualifyColumnsOptions {
236                    dialect: config.dialect,
237                    ..Default::default()
238                };
239                let original = expression.clone();
240                qualify_columns(expression, schema, &options).unwrap_or(original)
241            } else {
242                // Without schema, skip qualification
243                expression
244            }
245        }
246        OptimizationRule::PushdownProjections => {
247            pushdown_projections(expression, config.dialect, true)
248        }
249        OptimizationRule::Normalize => {
250            // Use CNF (dnf=false) with default max distance
251            let original = expression.clone();
252            normalize(expression, false, super::normalize::DEFAULT_MAX_DISTANCE).unwrap_or(original)
253        }
254        OptimizationRule::UnnestSubqueries => unnest_subqueries(expression),
255        OptimizationRule::PushdownPredicates => pushdown_predicates(expression, config.dialect),
256        OptimizationRule::OptimizeJoins => optimize_joins(expression),
257        OptimizationRule::EliminateSubqueries => eliminate_subqueries_opt(expression),
258        OptimizationRule::MergeSubqueries => merge_subqueries(expression, config.isolate_tables),
259        OptimizationRule::EliminateJoins => eliminate_joins(expression),
260        OptimizationRule::EliminateCtes => eliminate_ctes(expression),
261        OptimizationRule::QuoteIdentifiers => {
262            if config.quote_identifiers {
263                quote_identifiers(expression, config.dialect)
264            } else {
265                expression
266            }
267        }
268        OptimizationRule::AnnotateTypes => {
269            let mut expr = expression;
270            annotate_types(&mut expr, config.schema, config.dialect);
271            expr
272        }
273        OptimizationRule::Canonicalize => canonicalize(expression, config.dialect),
274        OptimizationRule::Simplify => simplify(expression, config.dialect),
275    }
276}
277
278// Re-import from subquery module with different name to avoid conflict
279use super::subquery::eliminate_subqueries as eliminate_subqueries_opt;
280
281/// Quick optimization that only applies essential passes.
282///
283/// This is faster than full optimization but may miss some opportunities.
284pub fn quick_optimize(expression: Expression, dialect: Option<DialectType>) -> Expression {
285    let config = OptimizerConfig {
286        dialect,
287        ..Default::default()
288    };
289
290    optimize_with_rules(expression, &config, QUICK_RULES)
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use crate::generator::Generator;
297    use crate::parser::Parser;
298
299    fn gen(expr: &Expression) -> String {
300        Generator::new().generate(expr).unwrap()
301    }
302
303    fn parse(sql: &str) -> Expression {
304        Parser::parse_sql(sql).expect("Failed to parse")[0].clone()
305    }
306
307    #[test]
308    fn test_optimize_simple() {
309        let expr = parse("SELECT a FROM t");
310        let config = OptimizerConfig::default();
311        let result = optimize(expr, &config);
312        let sql = gen(&result);
313        assert!(sql.contains("SELECT"));
314    }
315
316    #[test]
317    fn test_optimize_with_where() {
318        let expr = parse("SELECT a FROM t WHERE b = 1");
319        let config = OptimizerConfig::default();
320        let result = optimize(expr, &config);
321        let sql = gen(&result);
322        assert!(sql.contains("WHERE"));
323    }
324
325    #[test]
326    fn test_optimize_with_join() {
327        let expr = parse("SELECT t.a FROM t JOIN s ON t.id = s.id");
328        let config = OptimizerConfig::default();
329        let result = optimize(expr, &config);
330        let sql = gen(&result);
331        assert!(sql.contains("JOIN"));
332    }
333
334    #[test]
335    fn test_quick_optimize() {
336        let expr = parse("SELECT 1 + 0 FROM t");
337        let result = quick_optimize(expr, None);
338        let sql = gen(&result);
339        assert!(sql.contains("SELECT"));
340    }
341
342    #[test]
343    fn test_optimize_with_custom_rules() {
344        let expr = parse("SELECT a FROM t WHERE NOT NOT b = 1");
345        let config = OptimizerConfig::default();
346        let rules = &[OptimizationRule::Simplify];
347        let result = optimize_with_rules(expr, &config, rules);
348        let sql = gen(&result);
349        assert!(sql.contains("SELECT"));
350    }
351
352    #[test]
353    fn test_optimizer_config_default() {
354        let config = OptimizerConfig::default();
355        assert!(config.schema.is_none());
356        assert!(config.dialect.is_none());
357        assert!(config.isolate_tables);
358        assert!(!config.quote_identifiers);
359    }
360
361    #[test]
362    fn test_default_rules() {
363        assert_eq!(
364            DEFAULT_RULES,
365            &[
366                OptimizationRule::Qualify,
367                OptimizationRule::PushdownProjections,
368                OptimizationRule::Normalize,
369                OptimizationRule::UnnestSubqueries,
370                OptimizationRule::PushdownPredicates,
371                OptimizationRule::OptimizeJoins,
372                OptimizationRule::EliminateSubqueries,
373                OptimizationRule::MergeSubqueries,
374                OptimizationRule::EliminateJoins,
375                OptimizationRule::EliminateCtes,
376                OptimizationRule::QuoteIdentifiers,
377                OptimizationRule::AnnotateTypes,
378                OptimizationRule::Canonicalize,
379                OptimizationRule::Simplify,
380            ]
381        );
382    }
383
384    #[test]
385    fn test_quote_identifiers_rule_respects_config_flag() {
386        let mut expr = parse("SELECT a FROM t");
387        if let Expression::Select(ref mut select) = expr {
388            if let Expression::Column(ref mut col) = select.expressions[0] {
389                col.name.name = "select".to_string();
390            } else {
391                panic!("expected column projection");
392            }
393            if let Some(ref mut from) = select.from {
394                if let Expression::Table(ref mut table) = from.expressions[0] {
395                    table.name.name = "from".to_string();
396                } else {
397                    panic!("expected table reference");
398                }
399            } else {
400                panic!("expected FROM clause");
401            }
402        } else {
403            panic!("expected select expression");
404        }
405        let config = OptimizerConfig {
406            quote_identifiers: true,
407            dialect: Some(DialectType::PostgreSQL),
408            ..Default::default()
409        };
410        let result = optimize_with_rules(expr, &config, &[OptimizationRule::QuoteIdentifiers]);
411        let sql = gen(&result);
412        assert!(sql.contains("\"select\""), "{sql}");
413        assert!(sql.contains("\"from\""), "{sql}");
414    }
415
416    #[test]
417    fn test_quote_identifiers_rule_noop_by_default() {
418        let expr = parse("SELECT a FROM t");
419        let config = OptimizerConfig::default();
420        let result =
421            optimize_with_rules(expr.clone(), &config, &[OptimizationRule::QuoteIdentifiers]);
422        assert_eq!(gen(&result), gen(&expr));
423    }
424
425    #[test]
426    fn test_optimize_subquery() {
427        let expr = parse("SELECT * FROM (SELECT a FROM t) AS sub");
428        let config = OptimizerConfig::default();
429        let result = optimize(expr, &config);
430        let sql = gen(&result);
431        assert!(sql.contains("SELECT"));
432    }
433
434    #[test]
435    fn test_optimize_cte() {
436        let expr = parse("WITH cte AS (SELECT a FROM t) SELECT * FROM cte");
437        let config = OptimizerConfig::default();
438        let result = optimize(expr, &config);
439        let sql = gen(&result);
440        assert!(sql.contains("WITH"));
441    }
442
443    #[test]
444    fn test_optimize_preserves_semantics() {
445        let expr = parse("SELECT a, b FROM t WHERE c > 1 ORDER BY a");
446        let config = OptimizerConfig::default();
447        let result = optimize(expr, &config);
448        let sql = gen(&result);
449        assert!(sql.contains("ORDER BY"));
450    }
451
452    #[test]
453    fn test_analyze_expression_complexity_deep_connector_chain() {
454        let mut expr = Expression::Eq(Box::new(crate::expressions::BinaryOp::new(
455            Expression::column("c0"),
456            Expression::number(0),
457        )));
458
459        for i in 1..1500 {
460            let predicate = Expression::Eq(Box::new(crate::expressions::BinaryOp::new(
461                Expression::column(format!("c{i}")),
462                Expression::number(i as i64),
463            )));
464            expr = Expression::And(Box::new(crate::expressions::BinaryOp::new(expr, predicate)));
465        }
466
467        let complexity = analyze_expression_complexity(&expr);
468        assert!(complexity.max_connector_depth >= 1499);
469        assert!(complexity.connector_count >= 1499);
470    }
471
472    #[test]
473    fn test_optimize_handles_deep_connector_chain() {
474        let mut expr = Expression::Eq(Box::new(crate::expressions::BinaryOp::new(
475            Expression::column("c0"),
476            Expression::number(0),
477        )));
478
479        for i in 1..2200 {
480            let predicate = Expression::Eq(Box::new(crate::expressions::BinaryOp::new(
481                Expression::column(format!("c{i}")),
482                Expression::number(i as i64),
483            )));
484            expr = Expression::And(Box::new(crate::expressions::BinaryOp::new(expr, predicate)));
485        }
486
487        let config = OptimizerConfig::default();
488        let optimized = optimize(expr, &config);
489        let sql = gen(&optimized);
490        assert!(sql.contains("c2199 = 2199"), "{sql}");
491    }
492}
polyglot_sql/optimizer/optimizer.rs

polyglot_sql/optimizer/
optimizer.rs
