polars_plan/plans/conversion/
ir_to_dsl.rs

1use super::*;
2
3/// converts a node from the AExpr arena to Expr
4#[recursive]
5pub fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
6    let expr = expr_arena.get(node).clone();
7
8    match expr {
9        AExpr::Explode(node) => Expr::Explode(Arc::new(node_to_expr(node, expr_arena))),
10        AExpr::Alias(expr, name) => {
11            let exp = node_to_expr(expr, expr_arena);
12            Expr::Alias(Arc::new(exp), name)
13        },
14        AExpr::Column(a) => Expr::Column(a),
15        AExpr::Literal(s) => Expr::Literal(s),
16        AExpr::BinaryExpr { left, op, right } => {
17            let l = node_to_expr(left, expr_arena);
18            let r = node_to_expr(right, expr_arena);
19            Expr::BinaryExpr {
20                left: Arc::new(l),
21                op,
22                right: Arc::new(r),
23            }
24        },
25        AExpr::Cast {
26            expr,
27            dtype,
28            options: strict,
29        } => {
30            let exp = node_to_expr(expr, expr_arena);
31            Expr::Cast {
32                expr: Arc::new(exp),
33                dtype,
34                options: strict,
35            }
36        },
37        AExpr::Sort { expr, options } => {
38            let exp = node_to_expr(expr, expr_arena);
39            Expr::Sort {
40                expr: Arc::new(exp),
41                options,
42            }
43        },
44        AExpr::Gather {
45            expr,
46            idx,
47            returns_scalar,
48        } => {
49            let expr = node_to_expr(expr, expr_arena);
50            let idx = node_to_expr(idx, expr_arena);
51            Expr::Gather {
52                expr: Arc::new(expr),
53                idx: Arc::new(idx),
54                returns_scalar,
55            }
56        },
57        AExpr::SortBy {
58            expr,
59            by,
60            sort_options,
61        } => {
62            let expr = node_to_expr(expr, expr_arena);
63            let by = by
64                .iter()
65                .map(|node| node_to_expr(*node, expr_arena))
66                .collect();
67            Expr::SortBy {
68                expr: Arc::new(expr),
69                by,
70                sort_options,
71            }
72        },
73        AExpr::Filter { input, by } => {
74            let input = node_to_expr(input, expr_arena);
75            let by = node_to_expr(by, expr_arena);
76            Expr::Filter {
77                input: Arc::new(input),
78                by: Arc::new(by),
79            }
80        },
81        AExpr::Agg(agg) => match agg {
82            IRAggExpr::Min {
83                input,
84                propagate_nans,
85            } => {
86                let exp = node_to_expr(input, expr_arena);
87                AggExpr::Min {
88                    input: Arc::new(exp),
89                    propagate_nans,
90                }
91                .into()
92            },
93            IRAggExpr::Max {
94                input,
95                propagate_nans,
96            } => {
97                let exp = node_to_expr(input, expr_arena);
98                AggExpr::Max {
99                    input: Arc::new(exp),
100                    propagate_nans,
101                }
102                .into()
103            },
104
105            IRAggExpr::Median(expr) => {
106                let exp = node_to_expr(expr, expr_arena);
107                AggExpr::Median(Arc::new(exp)).into()
108            },
109            IRAggExpr::NUnique(expr) => {
110                let exp = node_to_expr(expr, expr_arena);
111                AggExpr::NUnique(Arc::new(exp)).into()
112            },
113            IRAggExpr::First(expr) => {
114                let exp = node_to_expr(expr, expr_arena);
115                AggExpr::First(Arc::new(exp)).into()
116            },
117            IRAggExpr::Last(expr) => {
118                let exp = node_to_expr(expr, expr_arena);
119                AggExpr::Last(Arc::new(exp)).into()
120            },
121            IRAggExpr::Mean(expr) => {
122                let exp = node_to_expr(expr, expr_arena);
123                AggExpr::Mean(Arc::new(exp)).into()
124            },
125            IRAggExpr::Implode(expr) => {
126                let exp = node_to_expr(expr, expr_arena);
127                AggExpr::Implode(Arc::new(exp)).into()
128            },
129            IRAggExpr::Quantile {
130                expr,
131                quantile,
132                method,
133            } => {
134                let expr = node_to_expr(expr, expr_arena);
135                let quantile = node_to_expr(quantile, expr_arena);
136                AggExpr::Quantile {
137                    expr: Arc::new(expr),
138                    quantile: Arc::new(quantile),
139                    method,
140                }
141                .into()
142            },
143            IRAggExpr::Sum(expr) => {
144                let exp = node_to_expr(expr, expr_arena);
145                AggExpr::Sum(Arc::new(exp)).into()
146            },
147            IRAggExpr::Std(expr, ddof) => {
148                let exp = node_to_expr(expr, expr_arena);
149                AggExpr::Std(Arc::new(exp), ddof).into()
150            },
151            IRAggExpr::Var(expr, ddof) => {
152                let exp = node_to_expr(expr, expr_arena);
153                AggExpr::Var(Arc::new(exp), ddof).into()
154            },
155            IRAggExpr::AggGroups(expr) => {
156                let exp = node_to_expr(expr, expr_arena);
157                AggExpr::AggGroups(Arc::new(exp)).into()
158            },
159            IRAggExpr::Count(expr, include_nulls) => {
160                let expr = node_to_expr(expr, expr_arena);
161                AggExpr::Count(Arc::new(expr), include_nulls).into()
162            },
163        },
164        AExpr::Ternary {
165            predicate,
166            truthy,
167            falsy,
168        } => {
169            let p = node_to_expr(predicate, expr_arena);
170            let t = node_to_expr(truthy, expr_arena);
171            let f = node_to_expr(falsy, expr_arena);
172
173            Expr::Ternary {
174                predicate: Arc::new(p),
175                truthy: Arc::new(t),
176                falsy: Arc::new(f),
177            }
178        },
179        AExpr::AnonymousFunction {
180            input,
181            function,
182            output_type,
183            options,
184        } => Expr::AnonymousFunction {
185            input: expr_irs_to_exprs(input, expr_arena),
186            function,
187            output_type,
188            options,
189        },
190        AExpr::Function {
191            input,
192            function,
193            options,
194        } => Expr::Function {
195            input: expr_irs_to_exprs(input, expr_arena),
196            function,
197            options,
198        },
199        AExpr::Window {
200            function,
201            partition_by,
202            order_by,
203            options,
204        } => {
205            let function = Arc::new(node_to_expr(function, expr_arena));
206            let partition_by = nodes_to_exprs(&partition_by, expr_arena);
207            let order_by =
208                order_by.map(|(n, options)| (Arc::new(node_to_expr(n, expr_arena)), options));
209            Expr::Window {
210                function,
211                partition_by,
212                order_by,
213                options,
214            }
215        },
216        AExpr::Slice {
217            input,
218            offset,
219            length,
220        } => Expr::Slice {
221            input: Arc::new(node_to_expr(input, expr_arena)),
222            offset: Arc::new(node_to_expr(offset, expr_arena)),
223            length: Arc::new(node_to_expr(length, expr_arena)),
224        },
225        AExpr::Len => Expr::Len,
226    }
227}
228
229fn nodes_to_exprs(nodes: &[Node], expr_arena: &Arena<AExpr>) -> Vec<Expr> {
230    nodes.iter().map(|n| node_to_expr(*n, expr_arena)).collect()
231}
232
233pub fn node_to_lp_cloned(
234    node: Node,
235    expr_arena: &Arena<AExpr>,
236    mut lp_arena: &Arena<IR>,
237) -> DslPlan {
238    // we borrow again mutably only to make the types happy
239    // we want to initialize `to_lp` from a mutable and a immutable lp_arena
240    // by borrowing an immutable mutably, we still are immutable down the line.
241    let alp = lp_arena.get(node).clone();
242    alp.into_lp(
243        &|node, lp_arena: &mut &Arena<IR>| lp_arena.get(node).clone(),
244        &mut lp_arena,
245        expr_arena,
246    )
247}
248
249/// converts a node from the IR arena to a LogicalPlan
250pub fn node_to_lp(node: Node, expr_arena: &Arena<AExpr>, lp_arena: &mut Arena<IR>) -> DslPlan {
251    let alp = lp_arena.get_mut(node);
252    let alp = std::mem::take(alp);
253    alp.into_lp(
254        &|node, lp_arena: &mut Arena<IR>| {
255            let lp = lp_arena.get_mut(node);
256            std::mem::take(lp)
257        },
258        lp_arena,
259        expr_arena,
260    )
261}
polars_plan/plans/conversion/ir_to_dsl.rs

polars_plan/plans/conversion/
ir_to_dsl.rs
