17
17
18
18
//! Analyzed rule to replace TableScan references
19
19
//! such as DataFrames and Views and inlines the LogicalPlan.
20
- use std:: sync:: Arc ;
21
20
22
21
use crate :: analyzer:: AnalyzerRule ;
23
22
24
23
use datafusion_common:: config:: ConfigOptions ;
25
24
use datafusion_common:: tree_node:: { Transformed , TransformedResult , TreeNode } ;
26
25
use datafusion_common:: { Column , Result } ;
27
- use datafusion_expr:: expr:: { Exists , InSubquery } ;
28
- use datafusion_expr:: {
29
- logical_plan:: LogicalPlan , Expr , Filter , LogicalPlanBuilder , TableScan ,
30
- } ;
26
+ use datafusion_expr:: { logical_plan:: LogicalPlan , Expr , LogicalPlanBuilder , TableScan } ;
31
27
32
28
/// Analyzed rule that inlines TableScan that provide a [`LogicalPlan`]
33
29
/// (DataFrame / ViewTable)
@@ -51,65 +47,38 @@ impl AnalyzerRule for InlineTableScan {
51
47
}
52
48
53
49
fn analyze_internal ( plan : LogicalPlan ) -> Result < Transformed < LogicalPlan > > {
54
- match plan {
55
- // Match only on scans without filter / projection / fetch
56
- // Views and DataFrames won't have those added
57
- // during the early stage of planning
58
- LogicalPlan :: TableScan ( TableScan {
59
- table_name,
60
- source,
61
- projection,
62
- filters,
63
- ..
64
- } ) if filters. is_empty ( ) && source. get_logical_plan ( ) . is_some ( ) => {
65
- let sub_plan = source. get_logical_plan ( ) . unwrap ( ) ;
66
- let projection_exprs = generate_projection_expr ( & projection, sub_plan) ?;
67
- LogicalPlanBuilder :: from ( sub_plan. clone ( ) )
68
- . project ( projection_exprs) ?
69
- // Ensures that the reference to the inlined table remains the
70
- // same, meaning we don't have to change any of the parent nodes
71
- // that reference this table.
72
- . alias ( table_name) ?
73
- . build ( )
74
- . map ( Transformed :: yes)
75
- }
76
- LogicalPlan :: Filter ( filter) => {
77
- let new_expr = filter. predicate . transform ( & rewrite_subquery) . data ( ) ?;
78
- Filter :: try_new ( new_expr, filter. input )
79
- . map ( |e| Transformed :: yes ( LogicalPlan :: Filter ( e) ) )
50
+ // rewrite any subqueries in the plan first
51
+ let transformed_plan =
52
+ plan. map_subqueries ( |plan| plan. transform_up ( & analyze_internal) ) ?;
53
+
54
+ let transformed_plan = transformed_plan. transform_data ( |plan| {
55
+ match plan {
56
+ // Match only on scans without filter / projection / fetch
57
+ // Views and DataFrames won't have those added
58
+ // during the early stage of planning
59
+ LogicalPlan :: TableScan ( TableScan {
60
+ table_name,
61
+ source,
62
+ projection,
63
+ filters,
64
+ ..
65
+ } ) if filters. is_empty ( ) && source. get_logical_plan ( ) . is_some ( ) => {
66
+ let sub_plan = source. get_logical_plan ( ) . unwrap ( ) ;
67
+ let projection_exprs = generate_projection_expr ( & projection, sub_plan) ?;
68
+ LogicalPlanBuilder :: from ( sub_plan. clone ( ) )
69
+ . project ( projection_exprs) ?
70
+ // Ensures that the reference to the inlined table remains the
71
+ // same, meaning we don't have to change any of the parent nodes
72
+ // that reference this table.
73
+ . alias ( table_name) ?
74
+ . build ( )
75
+ . map ( Transformed :: yes)
76
+ }
77
+ _ => Ok ( Transformed :: no ( plan) ) ,
80
78
}
81
- _ => Ok ( Transformed :: no ( plan) ) ,
82
- }
83
- }
79
+ } ) ?;
84
80
85
- fn rewrite_subquery ( expr : Expr ) -> Result < Transformed < Expr > > {
86
- match expr {
87
- Expr :: Exists ( Exists { subquery, negated } ) => {
88
- let plan = subquery. subquery . as_ref ( ) . clone ( ) ;
89
- let new_plan = plan. transform_up ( & analyze_internal) . data ( ) ?;
90
- let subquery = subquery. with_plan ( Arc :: new ( new_plan) ) ;
91
- Ok ( Transformed :: yes ( Expr :: Exists ( Exists { subquery, negated } ) ) )
92
- }
93
- Expr :: InSubquery ( InSubquery {
94
- expr,
95
- subquery,
96
- negated,
97
- } ) => {
98
- let plan = subquery. subquery . as_ref ( ) . clone ( ) ;
99
- let new_plan = plan. transform_up ( & analyze_internal) . data ( ) ?;
100
- let subquery = subquery. with_plan ( Arc :: new ( new_plan) ) ;
101
- Ok ( Transformed :: yes ( Expr :: InSubquery ( InSubquery :: new (
102
- expr, subquery, negated,
103
- ) ) ) )
104
- }
105
- Expr :: ScalarSubquery ( subquery) => {
106
- let plan = subquery. subquery . as_ref ( ) . clone ( ) ;
107
- let new_plan = plan. transform_up ( & analyze_internal) . data ( ) ?;
108
- let subquery = subquery. with_plan ( Arc :: new ( new_plan) ) ;
109
- Ok ( Transformed :: yes ( Expr :: ScalarSubquery ( subquery) ) )
110
- }
111
- _ => Ok ( Transformed :: no ( expr) ) ,
112
- }
81
+ Ok ( transformed_plan)
113
82
}
114
83
115
84
fn generate_projection_expr (
0 commit comments