/Users/andrewlamb/Software/datafusion/datafusion/expr/src/tree_node.rs
Line | Count | Source (jump to first uncovered line) |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | //! Tree node implementation for logical expr |
19 | | |
20 | | use crate::expr::{ |
21 | | AggregateFunction, Alias, Between, BinaryExpr, Case, Cast, GroupingSet, InList, |
22 | | InSubquery, Like, Placeholder, ScalarFunction, Sort, TryCast, Unnest, WindowFunction, |
23 | | }; |
24 | | use crate::{Expr, ExprFunctionExt}; |
25 | | |
26 | | use datafusion_common::tree_node::{ |
27 | | Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, |
28 | | }; |
29 | | use datafusion_common::{map_until_stop_and_collect, Result}; |
30 | | |
31 | | impl TreeNode for Expr { |
32 | 0 | fn apply_children<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>( |
33 | 0 | &'n self, |
34 | 0 | f: F, |
35 | 0 | ) -> Result<TreeNodeRecursion> { |
36 | 0 | let children = match self { |
37 | 0 | Expr::Alias(Alias{expr,..}) |
38 | 0 | | Expr::Unnest(Unnest{expr}) |
39 | 0 | | Expr::Not(expr) |
40 | 0 | | Expr::IsNotNull(expr) |
41 | 0 | | Expr::IsTrue(expr) |
42 | 0 | | Expr::IsFalse(expr) |
43 | 0 | | Expr::IsUnknown(expr) |
44 | 0 | | Expr::IsNotTrue(expr) |
45 | 0 | | Expr::IsNotFalse(expr) |
46 | 0 | | Expr::IsNotUnknown(expr) |
47 | 0 | | Expr::IsNull(expr) |
48 | 0 | | Expr::Negative(expr) |
49 | 0 | | Expr::Cast(Cast { expr, .. }) |
50 | 0 | | Expr::TryCast(TryCast { expr, .. }) |
51 | 0 | | Expr::InSubquery(InSubquery{ expr, .. }) => vec![expr.as_ref()], |
52 | 0 | Expr::GroupingSet(GroupingSet::Rollup(exprs)) |
53 | 0 | | Expr::GroupingSet(GroupingSet::Cube(exprs)) => exprs.iter().collect(), |
54 | 0 | Expr::ScalarFunction (ScalarFunction{ args, .. } ) => { |
55 | 0 | args.iter().collect() |
56 | | } |
57 | 0 | Expr::GroupingSet(GroupingSet::GroupingSets(lists_of_exprs)) => { |
58 | 0 | lists_of_exprs.iter().flatten().collect() |
59 | | } |
60 | | Expr::Column(_) |
61 | | // Treat OuterReferenceColumn as a leaf expression |
62 | | | Expr::OuterReferenceColumn(_, _) |
63 | | | Expr::ScalarVariable(_, _) |
64 | | | Expr::Literal(_) |
65 | | | Expr::Exists {..} |
66 | | | Expr::ScalarSubquery(_) |
67 | | | Expr::Wildcard {..} |
68 | 0 | | Expr::Placeholder (_) => vec![], |
69 | 0 | Expr::BinaryExpr(BinaryExpr { left, right, .. }) => { |
70 | 0 | vec![left.as_ref(), right.as_ref()] |
71 | | } |
72 | 0 | Expr::Like(Like { expr, pattern, .. }) |
73 | 0 | | Expr::SimilarTo(Like { expr, pattern, .. }) => { |
74 | 0 | vec![expr.as_ref(), pattern.as_ref()] |
75 | | } |
76 | | Expr::Between(Between { |
77 | 0 | expr, low, high, .. |
78 | 0 | }) => vec![expr.as_ref(), low.as_ref(), high.as_ref()], |
79 | 0 | Expr::Case(case) => { |
80 | 0 | let mut expr_vec = vec![]; |
81 | 0 | if let Some(expr) = case.expr.as_ref() { |
82 | 0 | expr_vec.push(expr.as_ref()); |
83 | 0 | }; |
84 | 0 | for (when, then) in case.when_then_expr.iter() { |
85 | 0 | expr_vec.push(when.as_ref()); |
86 | 0 | expr_vec.push(then.as_ref()); |
87 | 0 | } |
88 | 0 | if let Some(else_expr) = case.else_expr.as_ref() { |
89 | 0 | expr_vec.push(else_expr.as_ref()); |
90 | 0 | } |
91 | 0 | expr_vec |
92 | | } |
93 | 0 | Expr::AggregateFunction(AggregateFunction { args, filter, order_by, .. }) |
94 | 0 | => { |
95 | 0 | let mut expr_vec = args.iter().collect::<Vec<_>>(); |
96 | 0 | if let Some(f) = filter { |
97 | 0 | expr_vec.push(f.as_ref()); |
98 | 0 | } |
99 | 0 | if let Some(order_by) = order_by { |
100 | 0 | expr_vec.extend(order_by.iter().map(|sort| &sort.expr)); |
101 | 0 | } |
102 | 0 | expr_vec |
103 | | } |
104 | | Expr::WindowFunction(WindowFunction { |
105 | 0 | args, |
106 | 0 | partition_by, |
107 | 0 | order_by, |
108 | 0 | .. |
109 | 0 | }) => { |
110 | 0 | let mut expr_vec = args.iter().collect::<Vec<_>>(); |
111 | 0 | expr_vec.extend(partition_by); |
112 | 0 | expr_vec.extend(order_by.iter().map(|sort| &sort.expr)); |
113 | 0 | expr_vec |
114 | | } |
115 | 0 | Expr::InList(InList { expr, list, .. }) => { |
116 | 0 | let mut expr_vec = vec![expr.as_ref()]; |
117 | 0 | expr_vec.extend(list); |
118 | 0 | expr_vec |
119 | | } |
120 | | }; |
121 | | |
122 | 0 | children.into_iter().apply_until_stop(f) |
123 | 0 | } |
124 | | |
125 | 0 | fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>( |
126 | 0 | self, |
127 | 0 | mut f: F, |
128 | 0 | ) -> Result<Transformed<Self>> { |
129 | 0 | Ok(match self { |
130 | | Expr::Column(_) |
131 | | | Expr::Wildcard { .. } |
132 | | | Expr::Placeholder(Placeholder { .. }) |
133 | | | Expr::OuterReferenceColumn(_, _) |
134 | | | Expr::Exists { .. } |
135 | | | Expr::ScalarSubquery(_) |
136 | | | Expr::ScalarVariable(_, _) |
137 | 0 | | Expr::Literal(_) => Transformed::no(self), |
138 | 0 | Expr::Unnest(Unnest { expr, .. }) => transform_box(expr, &mut f)? |
139 | 0 | .update_data(|be| Expr::Unnest(Unnest::new_boxed(be))), |
140 | | Expr::Alias(Alias { |
141 | 0 | expr, |
142 | 0 | relation, |
143 | 0 | name, |
144 | 0 | }) => f(*expr)?.update_data(|e| Expr::Alias(Alias::new(e, relation, name))), |
145 | | Expr::InSubquery(InSubquery { |
146 | 0 | expr, |
147 | 0 | subquery, |
148 | 0 | negated, |
149 | 0 | }) => transform_box(expr, &mut f)?.update_data(|be| { |
150 | 0 | Expr::InSubquery(InSubquery::new(be, subquery, negated)) |
151 | 0 | }), |
152 | 0 | Expr::BinaryExpr(BinaryExpr { left, op, right }) => { |
153 | 0 | map_until_stop_and_collect!( |
154 | 0 | transform_box(left, &mut f), |
155 | 0 | right, |
156 | 0 | transform_box(right, &mut f) |
157 | 0 | )? |
158 | 0 | .update_data(|(new_left, new_right)| { |
159 | 0 | Expr::BinaryExpr(BinaryExpr::new(new_left, op, new_right)) |
160 | 0 | }) |
161 | | } |
162 | | Expr::Like(Like { |
163 | 0 | negated, |
164 | 0 | expr, |
165 | 0 | pattern, |
166 | 0 | escape_char, |
167 | 0 | case_insensitive, |
168 | 0 | }) => map_until_stop_and_collect!( |
169 | 0 | transform_box(expr, &mut f), |
170 | 0 | pattern, |
171 | 0 | transform_box(pattern, &mut f) |
172 | 0 | )? |
173 | 0 | .update_data(|(new_expr, new_pattern)| { |
174 | 0 | Expr::Like(Like::new( |
175 | 0 | negated, |
176 | 0 | new_expr, |
177 | 0 | new_pattern, |
178 | 0 | escape_char, |
179 | 0 | case_insensitive, |
180 | 0 | )) |
181 | 0 | }), |
182 | | Expr::SimilarTo(Like { |
183 | 0 | negated, |
184 | 0 | expr, |
185 | 0 | pattern, |
186 | 0 | escape_char, |
187 | 0 | case_insensitive, |
188 | 0 | }) => map_until_stop_and_collect!( |
189 | 0 | transform_box(expr, &mut f), |
190 | 0 | pattern, |
191 | 0 | transform_box(pattern, &mut f) |
192 | 0 | )? |
193 | 0 | .update_data(|(new_expr, new_pattern)| { |
194 | 0 | Expr::SimilarTo(Like::new( |
195 | 0 | negated, |
196 | 0 | new_expr, |
197 | 0 | new_pattern, |
198 | 0 | escape_char, |
199 | 0 | case_insensitive, |
200 | 0 | )) |
201 | 0 | }), |
202 | 0 | Expr::Not(expr) => transform_box(expr, &mut f)?.update_data(Expr::Not), |
203 | 0 | Expr::IsNotNull(expr) => { |
204 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsNotNull) |
205 | | } |
206 | 0 | Expr::IsNull(expr) => transform_box(expr, &mut f)?.update_data(Expr::IsNull), |
207 | 0 | Expr::IsTrue(expr) => transform_box(expr, &mut f)?.update_data(Expr::IsTrue), |
208 | 0 | Expr::IsFalse(expr) => { |
209 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsFalse) |
210 | | } |
211 | 0 | Expr::IsUnknown(expr) => { |
212 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsUnknown) |
213 | | } |
214 | 0 | Expr::IsNotTrue(expr) => { |
215 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsNotTrue) |
216 | | } |
217 | 0 | Expr::IsNotFalse(expr) => { |
218 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsNotFalse) |
219 | | } |
220 | 0 | Expr::IsNotUnknown(expr) => { |
221 | 0 | transform_box(expr, &mut f)?.update_data(Expr::IsNotUnknown) |
222 | | } |
223 | 0 | Expr::Negative(expr) => { |
224 | 0 | transform_box(expr, &mut f)?.update_data(Expr::Negative) |
225 | | } |
226 | | Expr::Between(Between { |
227 | 0 | expr, |
228 | 0 | negated, |
229 | 0 | low, |
230 | 0 | high, |
231 | 0 | }) => map_until_stop_and_collect!( |
232 | 0 | transform_box(expr, &mut f), |
233 | 0 | low, |
234 | 0 | transform_box(low, &mut f), |
235 | 0 | high, |
236 | 0 | transform_box(high, &mut f) |
237 | 0 | )? |
238 | 0 | .update_data(|(new_expr, new_low, new_high)| { |
239 | 0 | Expr::Between(Between::new(new_expr, negated, new_low, new_high)) |
240 | 0 | }), |
241 | | Expr::Case(Case { |
242 | 0 | expr, |
243 | 0 | when_then_expr, |
244 | 0 | else_expr, |
245 | 0 | }) => map_until_stop_and_collect!( |
246 | 0 | transform_option_box(expr, &mut f), |
247 | 0 | when_then_expr, |
248 | 0 | when_then_expr |
249 | 0 | .into_iter() |
250 | 0 | .map_until_stop_and_collect(|(when, then)| { |
251 | 0 | map_until_stop_and_collect!( |
252 | 0 | transform_box(when, &mut f), |
253 | 0 | then, |
254 | 0 | transform_box(then, &mut f) |
255 | 0 | ) |
256 | 0 | }), |
257 | 0 | else_expr, |
258 | 0 | transform_option_box(else_expr, &mut f) |
259 | 0 | )? |
260 | 0 | .update_data(|(new_expr, new_when_then_expr, new_else_expr)| { |
261 | 0 | Expr::Case(Case::new(new_expr, new_when_then_expr, new_else_expr)) |
262 | 0 | }), |
263 | 0 | Expr::Cast(Cast { expr, data_type }) => transform_box(expr, &mut f)? |
264 | 0 | .update_data(|be| Expr::Cast(Cast::new(be, data_type))), |
265 | 0 | Expr::TryCast(TryCast { expr, data_type }) => transform_box(expr, &mut f)? |
266 | 0 | .update_data(|be| Expr::TryCast(TryCast::new(be, data_type))), |
267 | 0 | Expr::ScalarFunction(ScalarFunction { func, args }) => { |
268 | 0 | transform_vec(args, &mut f)?.map_data(|new_args| { |
269 | 0 | Ok(Expr::ScalarFunction(ScalarFunction::new_udf( |
270 | 0 | func, new_args, |
271 | 0 | ))) |
272 | 0 | })? |
273 | | } |
274 | | Expr::WindowFunction(WindowFunction { |
275 | 0 | args, |
276 | 0 | fun, |
277 | 0 | partition_by, |
278 | 0 | order_by, |
279 | 0 | window_frame, |
280 | 0 | null_treatment, |
281 | 0 | }) => map_until_stop_and_collect!( |
282 | 0 | transform_vec(args, &mut f), |
283 | 0 | partition_by, |
284 | 0 | transform_vec(partition_by, &mut f), |
285 | 0 | order_by, |
286 | 0 | transform_sort_vec(order_by, &mut f) |
287 | 0 | )? |
288 | 0 | .update_data(|(new_args, new_partition_by, new_order_by)| { |
289 | 0 | Expr::WindowFunction(WindowFunction::new(fun, new_args)) |
290 | 0 | .partition_by(new_partition_by) |
291 | 0 | .order_by(new_order_by) |
292 | 0 | .window_frame(window_frame) |
293 | 0 | .null_treatment(null_treatment) |
294 | 0 | .build() |
295 | 0 | .unwrap() |
296 | 0 | }), |
297 | | Expr::AggregateFunction(AggregateFunction { |
298 | 0 | args, |
299 | 0 | func, |
300 | 0 | distinct, |
301 | 0 | filter, |
302 | 0 | order_by, |
303 | 0 | null_treatment, |
304 | 0 | }) => map_until_stop_and_collect!( |
305 | 0 | transform_vec(args, &mut f), |
306 | 0 | filter, |
307 | 0 | transform_option_box(filter, &mut f), |
308 | 0 | order_by, |
309 | 0 | transform_sort_option_vec(order_by, &mut f) |
310 | 0 | )? |
311 | 0 | .map_data(|(new_args, new_filter, new_order_by)| { |
312 | 0 | Ok(Expr::AggregateFunction(AggregateFunction::new_udf( |
313 | 0 | func, |
314 | 0 | new_args, |
315 | 0 | distinct, |
316 | 0 | new_filter, |
317 | 0 | new_order_by, |
318 | 0 | null_treatment, |
319 | 0 | ))) |
320 | 0 | })?, |
321 | 0 | Expr::GroupingSet(grouping_set) => match grouping_set { |
322 | 0 | GroupingSet::Rollup(exprs) => transform_vec(exprs, &mut f)? |
323 | 0 | .update_data(|ve| Expr::GroupingSet(GroupingSet::Rollup(ve))), |
324 | 0 | GroupingSet::Cube(exprs) => transform_vec(exprs, &mut f)? |
325 | 0 | .update_data(|ve| Expr::GroupingSet(GroupingSet::Cube(ve))), |
326 | 0 | GroupingSet::GroupingSets(lists_of_exprs) => lists_of_exprs |
327 | 0 | .into_iter() |
328 | 0 | .map_until_stop_and_collect(|exprs| transform_vec(exprs, &mut f))? |
329 | 0 | .update_data(|new_lists_of_exprs| { |
330 | 0 | Expr::GroupingSet(GroupingSet::GroupingSets(new_lists_of_exprs)) |
331 | 0 | }), |
332 | | }, |
333 | | Expr::InList(InList { |
334 | 0 | expr, |
335 | 0 | list, |
336 | 0 | negated, |
337 | 0 | }) => map_until_stop_and_collect!( |
338 | 0 | transform_box(expr, &mut f), |
339 | 0 | list, |
340 | 0 | transform_vec(list, &mut f) |
341 | 0 | )? |
342 | 0 | .update_data(|(new_expr, new_list)| { |
343 | 0 | Expr::InList(InList::new(new_expr, new_list, negated)) |
344 | 0 | }), |
345 | | }) |
346 | 0 | } |
347 | | } |
348 | | |
349 | 0 | fn transform_box<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
350 | 0 | be: Box<Expr>, |
351 | 0 | f: &mut F, |
352 | 0 | ) -> Result<Transformed<Box<Expr>>> { |
353 | 0 | Ok(f(*be)?.update_data(Box::new)) |
354 | 0 | } |
355 | | |
356 | 0 | fn transform_option_box<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
357 | 0 | obe: Option<Box<Expr>>, |
358 | 0 | f: &mut F, |
359 | 0 | ) -> Result<Transformed<Option<Box<Expr>>>> { |
360 | 0 | obe.map_or(Ok(Transformed::no(None)), |be| { |
361 | 0 | Ok(transform_box(be, f)?.update_data(Some)) |
362 | 0 | }) |
363 | 0 | } |
364 | | |
365 | | /// &mut transform a Option<`Vec` of `Expr`s> |
366 | 0 | pub fn transform_option_vec<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
367 | 0 | ove: Option<Vec<Expr>>, |
368 | 0 | f: &mut F, |
369 | 0 | ) -> Result<Transformed<Option<Vec<Expr>>>> { |
370 | 0 | ove.map_or(Ok(Transformed::no(None)), |ve| { |
371 | 0 | Ok(transform_vec(ve, f)?.update_data(Some)) |
372 | 0 | }) |
373 | 0 | } |
374 | | |
375 | | /// &mut transform a `Vec` of `Expr`s |
376 | 0 | fn transform_vec<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
377 | 0 | ve: Vec<Expr>, |
378 | 0 | f: &mut F, |
379 | 0 | ) -> Result<Transformed<Vec<Expr>>> { |
380 | 0 | ve.into_iter().map_until_stop_and_collect(f) |
381 | 0 | } |
382 | | |
383 | 0 | pub fn transform_sort_option_vec<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
384 | 0 | sorts_option: Option<Vec<Sort>>, |
385 | 0 | f: &mut F, |
386 | 0 | ) -> Result<Transformed<Option<Vec<Sort>>>> { |
387 | 0 | sorts_option.map_or(Ok(Transformed::no(None)), |sorts| { |
388 | 0 | Ok(transform_sort_vec(sorts, f)?.update_data(Some)) |
389 | 0 | }) |
390 | 0 | } |
391 | | |
392 | 0 | pub fn transform_sort_vec<F: FnMut(Expr) -> Result<Transformed<Expr>>>( |
393 | 0 | sorts: Vec<Sort>, |
394 | 0 | mut f: &mut F, |
395 | 0 | ) -> Result<Transformed<Vec<Sort>>> { |
396 | 0 | Ok(sorts |
397 | 0 | .iter() |
398 | 0 | .map(|sort| sort.expr.clone()) |
399 | 0 | .map_until_stop_and_collect(&mut f)? |
400 | 0 | .update_data(|transformed_exprs| { |
401 | 0 | replace_sort_expressions(sorts, transformed_exprs) |
402 | 0 | })) |
403 | 0 | } |
404 | | |
405 | 0 | pub fn replace_sort_expressions(sorts: Vec<Sort>, new_expr: Vec<Expr>) -> Vec<Sort> { |
406 | 0 | assert_eq!(sorts.len(), new_expr.len()); |
407 | 0 | sorts |
408 | 0 | .into_iter() |
409 | 0 | .zip(new_expr) |
410 | 0 | .map(|(sort, expr)| replace_sort_expression(sort, expr)) |
411 | 0 | .collect() |
412 | 0 | } |
413 | | |
414 | 0 | pub fn replace_sort_expression(sort: Sort, new_expr: Expr) -> Sort { |
415 | 0 | Sort { |
416 | 0 | expr: new_expr, |
417 | 0 | ..sort |
418 | 0 | } |
419 | 0 | } |