Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid copies in InlineTableScan via TreeNode API #10038

Merged
merged 3 commits into from
Apr 12, 2024
Merged

Avoid copies in InlineTableScan via TreeNode API #10038

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ab128a6
Avoid copies in `InlineTableScan` via TreeNode API
alamb Apr 10, 2024
0e0ac8e
Improve variable name
alamb Apr 11, 2024
627ae7f
Merge remote-tracking branch 'apache/main' into alamb/less_analyzer_copy
alamb Apr 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 31 additions & 62 deletions datafusion/optimizer/src/analyzer/inline_table_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,13 @@

//! Analyzed rule to replace TableScan references
//! such as DataFrames and Views and inlines the LogicalPlan.
use std::sync::Arc;

use crate::analyzer::AnalyzerRule;

use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::{Column, Result};
use datafusion_expr::expr::{Exists, InSubquery};
use datafusion_expr::{
logical_plan::LogicalPlan, Expr, Filter, LogicalPlanBuilder, TableScan,
};
use datafusion_expr::{logical_plan::LogicalPlan, Expr, LogicalPlanBuilder, TableScan};

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

fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
match plan {
// Match only on scans without filter / projection / fetch
// Views and DataFrames won't have those added
// during the early stage of planning
LogicalPlan::TableScan(TableScan {
table_name,
source,
projection,
filters,
..
}) if filters.is_empty() && source.get_logical_plan().is_some() => {
let sub_plan = source.get_logical_plan().unwrap();
let projection_exprs = generate_projection_expr(&projection, sub_plan)?;
LogicalPlanBuilder::from(sub_plan.clone())
.project(projection_exprs)?
// Ensures that the reference to the inlined table remains the
// same, meaning we don't have to change any of the parent nodes
// that reference this table.
.alias(table_name)?
.build()
.map(Transformed::yes)
}
LogicalPlan::Filter(filter) => {
let new_expr = filter.predicate.transform(&rewrite_subquery).data()?;
Filter::try_new(new_expr, filter.input)
.map(|e| Transformed::yes(LogicalPlan::Filter(e)))
// rewrite any subqueries in the plan first
let transformed_plan =
plan.map_subqueries(|plan| plan.transform_up(&analyze_internal))?;

let transformed_plan = transformed_plan.transform_data(|plan| {
match plan {
// Match only on scans without filter / projection / fetch
// Views and DataFrames won't have those added
// during the early stage of planning
LogicalPlan::TableScan(TableScan {
table_name,
source,
projection,
filters,
..
}) if filters.is_empty() && source.get_logical_plan().is_some() => {
let sub_plan = source.get_logical_plan().unwrap();
let projection_exprs = generate_projection_expr(&projection, sub_plan)?;
LogicalPlanBuilder::from(sub_plan.clone())
.project(projection_exprs)?
// Ensures that the reference to the inlined table remains the
// same, meaning we don't have to change any of the parent nodes
// that reference this table.
.alias(table_name)?
.build()
.map(Transformed::yes)
}
_ => Ok(Transformed::no(plan)),
}
_ => Ok(Transformed::no(plan)),
}
}
})?;

fn rewrite_subquery(expr: Expr) -> Result<Transformed<Expr>> {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is now handled by apply_subqueries

match expr {
Expr::Exists(Exists { subquery, negated }) => {
let plan = subquery.subquery.as_ref().clone();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these clones are avoided in the current formulation

let new_plan = plan.transform_up(&analyze_internal).data()?;
let subquery = subquery.with_plan(Arc::new(new_plan));
Ok(Transformed::yes(Expr::Exists(Exists { subquery, negated })))
}
Expr::InSubquery(InSubquery {
expr,
subquery,
negated,
}) => {
let plan = subquery.subquery.as_ref().clone();
let new_plan = plan.transform_up(&analyze_internal).data()?;
let subquery = subquery.with_plan(Arc::new(new_plan));
Ok(Transformed::yes(Expr::InSubquery(InSubquery::new(
expr, subquery, negated,
))))
}
Expr::ScalarSubquery(subquery) => {
let plan = subquery.subquery.as_ref().clone();
let new_plan = plan.transform_up(&analyze_internal).data()?;
let subquery = subquery.with_plan(Arc::new(new_plan));
Ok(Transformed::yes(Expr::ScalarSubquery(subquery)))
}
_ => Ok(Transformed::no(expr)),
}
Ok(transformed_plan)
}

fn generate_projection_expr(
Expand Down
Loading