Skip to content

Commit e843bb9

Browse files
committed
fix(optimizer): optimize_projections rule bug
1 parent 945902d commit e843bb9

File tree

1 file changed

+70
-1
lines changed
  • datafusion/optimizer/src/optimize_projections

1 file changed

+70
-1
lines changed

datafusion/optimizer/src/optimize_projections/mod.rs

+70-1
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,26 @@ fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Project
469469
return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
470470
};
471471

472+
// In order to pass the detection of `assert_schema_is_the_same` function, avoid changing
473+
// the plan schema here. e.g.
474+
// Projection: public.test.a -- proj
475+
// Projection: public.test.a AS public.test.a -- prev_projection
476+
// TableScan: public.test
477+
if prev_projection.expr.iter().any(|expr| match expr {
478+
Expr::Alias(Alias {
479+
expr,
480+
name: alias_name,
481+
..
482+
}) => match expr.as_ref() {
483+
Expr::Column(column) => &column.flat_name() == alias_name,
484+
_ => false,
485+
},
486+
_ => false,
487+
}) {
488+
// no change
489+
return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
490+
}
491+
472492
// Count usages (referrals) of each projection expression in its input fields:
473493
let mut column_referral_map = HashMap::<&Column, usize>::new();
474494
expr.iter()
@@ -796,7 +816,7 @@ mod tests {
796816
use std::vec;
797817

798818
use crate::optimize_projections::OptimizeProjections;
799-
use crate::optimizer::Optimizer;
819+
use crate::optimizer::{assert_schema_is_the_same, Optimizer};
800820
use crate::test::{
801821
assert_fields_eq, assert_optimized_plan_eq, scan_empty, test_table_scan,
802822
test_table_scan_fields, test_table_scan_with_name,
@@ -1954,4 +1974,53 @@ mod tests {
19541974
optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
19551975
Ok(optimized_plan)
19561976
}
1977+
1978+
#[test]
1979+
fn test_subquery() -> Result<()> {
1980+
fn assert_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
1981+
// check schema
1982+
let prev_schema = plan.schema().clone();
1983+
let new_plan = optimize(plan)?;
1984+
assert_schema_is_the_same("", &prev_schema, &new_plan)?;
1985+
1986+
let actual = format!("{new_plan}");
1987+
assert_eq!(&actual, expected);
1988+
1989+
Ok(())
1990+
}
1991+
1992+
// SELECT "test.a" FROM (SELECT a AS "test.a" FROM test)
1993+
let table_scan = test_table_scan_with_name("test")?;
1994+
let plan = LogicalPlanBuilder::from(table_scan)
1995+
.project(vec![col("a").alias("test.a")])?
1996+
.project(vec![Expr::Column(Column::new_unqualified("test.a"))])?
1997+
.build()?;
1998+
let expected = "Projection: test.a AS test.a\
1999+
\n TableScan: test projection=[a]";
2000+
assert_eq(plan, expected)?;
2001+
2002+
// SELECT "public.test.a" FROM (SELECT a AS "public.test.a" FROM public.test)
2003+
let table_scan = test_table_scan_with_name("public.test")?;
2004+
let plan = LogicalPlanBuilder::from(table_scan)
2005+
.project(vec![col("a").alias("public.test.a")])?
2006+
.project(vec![Expr::Column(Column::new_unqualified("public.test.a"))])?
2007+
.build()?;
2008+
let expected = "Projection: public.test.a AS public.test.a\
2009+
\n TableScan: public.test projection=[a]";
2010+
assert_eq(plan, expected)?;
2011+
2012+
// SELECT "datafusion.public.test.a" FROM (SELECT a AS "datafusion.public.test.a" FROM datafusion.public.test)
2013+
let table_scan = test_table_scan_with_name("datafusion.public.test")?;
2014+
let plan = LogicalPlanBuilder::from(table_scan)
2015+
.project(vec![col("a").alias("datafusion.public.test.a")])?
2016+
.project(vec![Expr::Column(Column::new_unqualified(
2017+
"datafusion.public.test.a",
2018+
))])?
2019+
.build()?;
2020+
let expected = "Projection: datafusion.public.test.a AS datafusion.public.test.a\
2021+
\n TableScan: datafusion.public.test projection=[a]";
2022+
assert_eq(plan, expected)?;
2023+
2024+
Ok(())
2025+
}
19572026
}

0 commit comments

Comments
 (0)