Skip to content

Commit a5cf0b8

Browse files
authored
Avoid copies in InlineTableScan via TreeNode API (#10038)
* Avoid copies in `InlineTableScan` via TreeNode API * Improve variable name
1 parent e161cd6 commit a5cf0b8

File tree

1 file changed

+31
-62
lines changed

1 file changed

+31
-62
lines changed

datafusion/optimizer/src/analyzer/inline_table_scan.rs

Lines changed: 31 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,13 @@
1717

1818
//! Analyzed rule to replace TableScan references
1919
//! such as DataFrames and Views and inlines the LogicalPlan.
20-
use std::sync::Arc;
2120
2221
use crate::analyzer::AnalyzerRule;
2322

2423
use datafusion_common::config::ConfigOptions;
2524
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
2625
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};
3127

3228
/// Analyzed rule that inlines TableScan that provide a [`LogicalPlan`]
3329
/// (DataFrame / ViewTable)
@@ -51,65 +47,38 @@ impl AnalyzerRule for InlineTableScan {
5147
}
5248

5349
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)),
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

11584
fn generate_projection_expr(

0 commit comments

Comments
 (0)