1717
1818//! Analyzed rule to replace TableScan references
1919//! such as DataFrames and Views and inlines the LogicalPlan.
20- use std:: sync:: Arc ;
2120
2221use crate :: analyzer:: AnalyzerRule ;
2322
2423use datafusion_common:: config:: ConfigOptions ;
2524use datafusion_common:: tree_node:: { Transformed , TransformedResult , TreeNode } ;
2625use 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 } ;
3127
3228/// Analyzed rule that inlines TableScan that provide a [`LogicalPlan`]
3329/// (DataFrame / ViewTable)
@@ -51,65 +47,38 @@ impl AnalyzerRule for InlineTableScan {
5147}
5248
5349fn 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) ) ,
8078 }
81- _ => Ok ( Transformed :: no ( plan) ) ,
82- }
83- }
79+ } ) ?;
8480
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)
11382}
11483
11584fn generate_projection_expr (
0 commit comments