From e843bb9ceeec8c7bcf545081b5d007ca9b696f59 Mon Sep 17 00:00:00 2001
From: JasonLi-cn <lijingxuan92@126.com>
Date: Tue, 27 Aug 2024 09:58:15 +0800
Subject: [PATCH] fix(optimizer): optimize_projections rule bug

---
 .../optimizer/src/optimize_projections/mod.rs | 71 ++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs
index ac4ed87a4a1a..af563eae2800 100644
--- a/datafusion/optimizer/src/optimize_projections/mod.rs
+++ b/datafusion/optimizer/src/optimize_projections/mod.rs
@@ -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()
@@ -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,
@@ -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(())
+    }
 }