Skip to content

Commit

Permalink
fix(optimizer): optimize_projections rule bug
Browse files Browse the repository at this point in the history
  • Loading branch information
JasonLi-cn committed Aug 27, 2024
1 parent 945902d commit e843bb9
Showing 1 changed file with 70 additions and 1 deletion.
71 changes: 70 additions & 1 deletion datafusion/optimizer/src/optimize_projections/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -469,6 +469,26 @@ fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Project
return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
};

// In order to pass the detection of `assert_schema_is_the_same` function, avoid changing
// the plan schema here. e.g.
// Projection: public.test.a -- proj
// Projection: public.test.a AS public.test.a -- prev_projection
// TableScan: public.test
if prev_projection.expr.iter().any(|expr| match expr {
Expr::Alias(Alias {
expr,
name: alias_name,
..
}) => match expr.as_ref() {
Expr::Column(column) => &column.flat_name() == alias_name,
_ => false,
},
_ => false,
}) {
// no change
return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
}

// Count usages (referrals) of each projection expression in its input fields:
let mut column_referral_map = HashMap::<&Column, usize>::new();
expr.iter()
Expand Down Expand Up @@ -796,7 +816,7 @@ mod tests {
use std::vec;

use crate::optimize_projections::OptimizeProjections;
use crate::optimizer::Optimizer;
use crate::optimizer::{assert_schema_is_the_same, Optimizer};
use crate::test::{
assert_fields_eq, assert_optimized_plan_eq, scan_empty, test_table_scan,
test_table_scan_fields, test_table_scan_with_name,
Expand Down Expand Up @@ -1954,4 +1974,53 @@ mod tests {
optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
Ok(optimized_plan)
}

#[test]
fn test_subquery() -> Result<()> {
fn assert_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
// check schema
let prev_schema = plan.schema().clone();
let new_plan = optimize(plan)?;
assert_schema_is_the_same("", &prev_schema, &new_plan)?;

let actual = format!("{new_plan}");
assert_eq!(&actual, expected);

Ok(())
}

// SELECT "test.a" FROM (SELECT a AS "test.a" FROM test)
let table_scan = test_table_scan_with_name("test")?;
let plan = LogicalPlanBuilder::from(table_scan)
.project(vec![col("a").alias("test.a")])?
.project(vec![Expr::Column(Column::new_unqualified("test.a"))])?
.build()?;
let expected = "Projection: test.a AS test.a\
\n TableScan: test projection=[a]";
assert_eq(plan, expected)?;

// SELECT "public.test.a" FROM (SELECT a AS "public.test.a" FROM public.test)
let table_scan = test_table_scan_with_name("public.test")?;
let plan = LogicalPlanBuilder::from(table_scan)
.project(vec![col("a").alias("public.test.a")])?
.project(vec![Expr::Column(Column::new_unqualified("public.test.a"))])?
.build()?;
let expected = "Projection: public.test.a AS public.test.a\
\n TableScan: public.test projection=[a]";
assert_eq(plan, expected)?;

// SELECT "datafusion.public.test.a" FROM (SELECT a AS "datafusion.public.test.a" FROM datafusion.public.test)
let table_scan = test_table_scan_with_name("datafusion.public.test")?;
let plan = LogicalPlanBuilder::from(table_scan)
.project(vec![col("a").alias("datafusion.public.test.a")])?
.project(vec![Expr::Column(Column::new_unqualified(
"datafusion.public.test.a",
))])?
.build()?;
let expected = "Projection: datafusion.public.test.a AS datafusion.public.test.a\
\n TableScan: datafusion.public.test projection=[a]";
assert_eq(plan, expected)?;

Ok(())
}
}

0 comments on commit e843bb9

Please sign in to comment.