polars_plan/plans/optimizer/
stack_opt.rs

1use polars_core::prelude::PolarsResult;
2
3use crate::plans::aexpr::AExpr;
4use crate::plans::ir::IR;
5use crate::prelude::{Arena, Node};
6
7/// Optimizer that uses a stack and memory arenas in favor of recursion
8pub struct StackOptimizer {}
9
10impl StackOptimizer {
11    pub fn optimize_loop(
12        &self,
13        rules: &mut [Box<dyn OptimizationRule>],
14        expr_arena: &mut Arena<AExpr>,
15        lp_arena: &mut Arena<IR>,
16        lp_top: Node,
17    ) -> PolarsResult<Node> {
18        let mut changed = true;
19
20        // Nodes of expressions and lp node from which the expressions are a member of.
21        let mut plans = vec![];
22        let mut exprs = vec![];
23        let mut scratch = vec![];
24
25        // Run loop until reaching fixed point.
26        while changed {
27            // Recurse into sub plans and expressions and apply rules.
28            changed = false;
29            plans.push(lp_top);
30            while let Some(current_node) = plans.pop() {
31                // Apply rules
32                for rule in rules.iter_mut() {
33                    // keep iterating over same rule
34                    while let Some(x) = rule.optimize_plan(lp_arena, expr_arena, current_node)? {
35                        lp_arena.replace(current_node, x);
36                        changed = true;
37                    }
38                }
39
40                let plan = lp_arena.get(current_node);
41
42                // traverse subplans and expressions and add to the stack
43                plan.copy_exprs(&mut scratch);
44                plan.copy_inputs(&mut plans);
45
46                if scratch.is_empty() {
47                    continue;
48                }
49
50                while let Some(expr_ir) = scratch.pop() {
51                    exprs.push(expr_ir.node());
52                }
53
54                // process the expressions on the stack and apply optimizations.
55                while let Some(current_expr_node) = exprs.pop() {
56                    {
57                        let expr = unsafe { expr_arena.get_unchecked(current_expr_node) };
58                        if expr.is_leaf() {
59                            continue;
60                        }
61                    }
62                    for rule in rules.iter_mut() {
63                        // keep iterating over same rule
64                        while let Some(x) = rule.optimize_expr(
65                            expr_arena,
66                            current_expr_node,
67                            lp_arena,
68                            current_node,
69                        )? {
70                            expr_arena.replace(current_expr_node, x);
71                            changed = true;
72                        }
73                    }
74
75                    let expr = unsafe { expr_arena.get_unchecked(current_expr_node) };
76                    // traverse subexpressions and add to the stack
77                    expr.inputs_rev(&mut exprs)
78                }
79            }
80        }
81        Ok(lp_top)
82    }
83}
84
85pub trait OptimizationRule {
86    ///  Optimize (subplan) in LogicalPlan
87    ///
88    /// * `lp_arena` - LogicalPlan memory arena
89    /// * `expr_arena` - Expression memory arena
90    /// * `node` - node of the current LogicalPlan node
91    fn optimize_plan(
92        &mut self,
93        _lp_arena: &mut Arena<IR>,
94        _expr_arena: &mut Arena<AExpr>,
95        _node: Node,
96    ) -> PolarsResult<Option<IR>> {
97        Ok(None)
98    }
99    fn optimize_expr(
100        &mut self,
101        _expr_arena: &mut Arena<AExpr>,
102        _expr_node: Node,
103        _lp_arena: &Arena<IR>,
104        _lp_node: Node,
105    ) -> PolarsResult<Option<AExpr>> {
106        Ok(None)
107    }
108}
polars_plan/plans/optimizer/stack_opt.rs

polars_plan/plans/optimizer/
stack_opt.rs
