Skip to content

Commit c5cbff5

Browse files
committed
Improve unparsing after optimize_projections optimization
1 parent eca0e07 commit c5cbff5

File tree

2 files changed

+50
-25
lines changed

2 files changed

+50
-25
lines changed

datafusion/sql/src/unparser/plan.rs

+36-24
Original file line numberDiff line numberDiff line change
@@ -276,9 +276,11 @@ impl Unparser<'_> {
276276
) -> Result<()> {
277277
match plan {
278278
LogicalPlan::TableScan(scan) => {
279-
if let Some(unparsed_table_scan) =
280-
Self::unparse_table_scan_pushdown(plan, None)?
281-
{
279+
if let Some(unparsed_table_scan) = Self::unparse_table_scan_pushdown(
280+
plan,
281+
None,
282+
select.already_projected(),
283+
)? {
282284
return self.select_to_sql_recursively(
283285
&unparsed_table_scan,
284286
query,
@@ -585,6 +587,7 @@ impl Unparser<'_> {
585587
let unparsed_table_scan = Self::unparse_table_scan_pushdown(
586588
plan,
587589
Some(plan_alias.alias.clone()),
590+
select.already_projected(),
588591
)?;
589592
// if the child plan is a TableScan with pushdown operations, we don't need to
590593
// create an additional subquery for it
@@ -714,6 +717,7 @@ impl Unparser<'_> {
714717
fn unparse_table_scan_pushdown(
715718
plan: &LogicalPlan,
716719
alias: Option<TableReference>,
720+
already_projected: bool,
717721
) -> Result<Option<LogicalPlan>> {
718722
match plan {
719723
LogicalPlan::TableScan(table_scan) => {
@@ -743,24 +747,29 @@ impl Unparser<'_> {
743747
}
744748
}
745749

746-
if let Some(project_vec) = &table_scan.projection {
747-
let project_columns = project_vec
748-
.iter()
749-
.cloned()
750-
.map(|i| {
751-
let schema = table_scan.source.schema();
752-
let field = schema.field(i);
753-
if alias.is_some() {
754-
Column::new(alias.clone(), field.name().clone())
755-
} else {
756-
Column::new(
757-
Some(table_scan.table_name.clone()),
758-
field.name().clone(),
759-
)
760-
}
761-
})
762-
.collect::<Vec<_>>();
763-
builder = builder.project(project_columns)?;
750+
// Avoid creating a duplicate Projection node, which would result in an additional subquery if a projection already exists.
751+
// For example, if the `optimize_projection` rule is applied, there will be a Projection node, and duplicate projection
752+
// information included in the TableScan node.
753+
if !already_projected {
754+
if let Some(project_vec) = &table_scan.projection {
755+
let project_columns = project_vec
756+
.iter()
757+
.cloned()
758+
.map(|i| {
759+
let schema = table_scan.source.schema();
760+
let field = schema.field(i);
761+
if alias.is_some() {
762+
Column::new(alias.clone(), field.name().clone())
763+
} else {
764+
Column::new(
765+
Some(table_scan.table_name.clone()),
766+
field.name().clone(),
767+
)
768+
}
769+
})
770+
.collect::<Vec<_>>();
771+
builder = builder.project(project_columns)?;
772+
}
764773
}
765774

766775
let filter_expr: Result<Option<Expr>> = table_scan
@@ -805,14 +814,17 @@ impl Unparser<'_> {
805814
Self::unparse_table_scan_pushdown(
806815
&subquery_alias.input,
807816
Some(subquery_alias.alias.clone()),
817+
already_projected,
808818
)
809819
}
810820
// SubqueryAlias could be rewritten to a plan with a projection as the top node by [rewrite::subquery_alias_inner_query_and_columns].
811821
// The inner table scan could be a scan with pushdown operations.
812822
LogicalPlan::Projection(projection) => {
813-
if let Some(plan) =
814-
Self::unparse_table_scan_pushdown(&projection.input, alias.clone())?
815-
{
823+
if let Some(plan) = Self::unparse_table_scan_pushdown(
824+
&projection.input,
825+
alias.clone(),
826+
already_projected,
827+
)? {
816828
let exprs = if alias.is_some() {
817829
let mut alias_rewriter =
818830
alias.as_ref().map(|alias_name| TableAliasRewriter {

datafusion/sql/tests/cases/plan_to_sql.rs

+14-1
Original file line numberDiff line numberDiff line change
@@ -926,12 +926,25 @@ fn test_table_scan_pushdown() -> Result<()> {
926926
let query_from_table_scan_with_projection = LogicalPlanBuilder::from(
927927
table_scan(Some("t1"), &schema, Some(vec![0, 1]))?.build()?,
928928
)
929-
.project(vec![wildcard()])?
929+
.project(vec![col("id"), col("age")])?
930930
.build()?;
931931
let query_from_table_scan_with_projection =
932932
plan_to_sql(&query_from_table_scan_with_projection)?;
933933
assert_eq!(
934934
query_from_table_scan_with_projection.to_string(),
935+
"SELECT t1.id, t1.age FROM t1"
936+
);
937+
938+
let query_from_table_scan_with_two_projections = LogicalPlanBuilder::from(
939+
table_scan(Some("t1"), &schema, Some(vec![0, 1]))?.build()?,
940+
)
941+
.project(vec![col("id"), col("age")])?
942+
.project(vec![wildcard()])?
943+
.build()?;
944+
let query_from_table_scan_with_two_projections =
945+
plan_to_sql(&query_from_table_scan_with_two_projections)?;
946+
assert_eq!(
947+
query_from_table_scan_with_two_projections.to_string(),
935948
"SELECT * FROM (SELECT t1.id, t1.age FROM t1)"
936949
);
937950

0 commit comments

Comments
 (0)