SELECT * FROM
(
    SELECT order_id
    FROM "warehouse"."main"."order_items"
) oi
JOIN (
    SELECT
        order_id,
        coalesce(discount_pct, 0) AS discount_pct_2
    FROM "warehouse"."main"."orders"
) o USING (order_id)

use std::sync::Arc;

use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::catalog::{
    CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider,
};
use datafusion::datasource::empty::EmptyTable;
use datafusion::prelude::*;
use datafusion::sql::unparser::dialect::DuckDBDialect;
use datafusion::sql::unparser::Unparser;
use duckdb::Connection;

const QUERY: &str = r#"
SELECT * FROM
(
SELECT
        order_id
    FROM
        "warehouse"."main"."order_items"
) oi
JOIN (
    SELECT
        order_id,
        coalesce(discount_pct, 0) AS discount_pct_2
    FROM
        "warehouse"."main"."orders"
) o USING (order_id)
"#;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let ctx = build_datafusion_context()?;
    let conn = Connection::open_in_memory()?;
    register_duckdb_tables(&conn)?;

    let dialect = DuckDBDialect::new();
    // Unparser 是用来把 DataFusion 的内部计划再转回 SQL 字符串 的
    let unparser = Unparser::new(&dialect);

    let unoptimized_plan = ctx.sql(QUERY).await?.into_unoptimized_plan();
    let unoptimized_sql = unparser.plan_to_sql(&unoptimized_plan)?;
    println!("=== Unoptimized SQL ===\n{unoptimized_sql}\n");

    match conn.prepare(&unoptimized_sql.to_string()) {
        Ok(_) => println!("unoptimized query is valid in DuckDB\n"),
        Err(e) => println!("unoptimized query is not valid in DuckDB: {e}\n"),
    }

    let optimized_plan = ctx.sql(QUERY).await?.into_optimized_plan()?;
    let optimized_sql = match unparser.plan_to_sql(&optimized_plan) {
        Ok(sql) => sql,
        Err(e) => {
            println!("reproduced #23138: optimized plan cannot be unparsed");
            println!("Unparser error: {e}");
            println!(
                "=== Optimized Plan ===\n{}",
                optimized_plan.display_indent()
            );
            return Ok(());
        }
    };

    println!("=== Optimized SQL ===\n{optimized_sql}\n");
    println!(
        "=== Optimized Plan ===\n{}\n",
        optimized_plan.display_indent()
    );

    match conn.prepare(&optimized_sql.to_string()) {
        Ok(_) => println!("unexpected: optimized query is valid in DuckDB"),
        Err(e) => {
            println!("reproduced #23138: optimized query is invalid in DuckDB");
            println!("DuckDB error: {e}");
        }
    }

    Ok(())
}

fn register_duckdb_tables(conn: &Connection) -> duckdb::Result<()> {
    conn.execute_batch(
        r#"
        ATTACH ':memory:' AS warehouse;
        CREATE TABLE warehouse.main.order_items(order_id INTEGER);
        CREATE TABLE warehouse.main.orders(
            order_id INTEGER NOT NULL,
            customer_id INTEGER,
            discount_pct DECIMAL(5, 2)
        );
        "#,
    )
}

/// catalog  ≈ 数据库实例 / 数据目录
/// schema   ≈ 库 / 命名空间
/// table    ≈ 表
fn build_datafusion_context() -> datafusion::common::Result<SessionContext> {
    // 构建订单项 schema
    let order_items_schema = Arc::new(Schema::new(vec![Field::new(
        "order_id",
        DataType::Int32,
        true,
    )]));

    // 构建订单 schema
    let orders_schema = Arc::new(Schema::new(vec![
        Field::new("order_id", DataType::Int32, false),
        Field::new("customer_id", DataType::Int32, true),
        Field::new("discount_pct", DataType::Decimal128(5, 2), true),
    ]));


    let schema_provider = Arc::new(MemorySchemaProvider::new());
    schema_provider.register_table(
        "order_items".to_string(),
        Arc::new(EmptyTable::new(order_items_schema)),
    )?;
    schema_provider.register_table(
        "orders".to_string(),
        Arc::new(EmptyTable::new(orders_schema)),
    )?;

    let catalog = Arc::new(MemoryCatalogProvider::new());
    catalog.register_schema("main", schema_provider)?;

    let ctx = SessionContext::new();
    ctx.register_catalog("warehouse", catalog);

    Ok(ctx)
}

=== Unoptimized SQL ===
SELECT "o"."order_id", "o"."discount_pct_2" FROM (SELECT "order_items"."order_id" FROM "warehouse"."main"."order_items") AS "oi" INNER JOIN (SELECT "orders"."order_id", coalesce("orders"."discount_pct", 0) AS "discount_pct_2" FROM "warehouse"."main"."orders") AS "o" USING("order_id")

unoptimized query is valid in DuckDB

=== Optimized SQL ===
SELECT "o"."order_id", "o"."discount_pct_2" FROM "warehouse"."main"."order_items" AS "oi" INNER JOIN (SELECT "o"."order_id", CASE WHEN "__common_expr_1" IS NOT NULL THEN "__common_expr_1" ELSE 0.00 END AS "discount_pct_2" FROM (SELECT CAST("o"."discount_pct" AS DECIMAL(22,2)) AS "__common_expr_1", "o"."order_id" FROM "warehouse"."main"."orders" AS "o")) AS "o" ON "oi"."order_id" = "o"."order_id"

=== Optimized Plan ===
Projection: o.order_id, o.discount_pct_2
  Inner Join: oi.order_id = o.order_id
    SubqueryAlias: oi
      TableScan: warehouse.main.order_items projection=[order_id]
    SubqueryAlias: o
      Projection: warehouse.main.orders.order_id, CASE WHEN __common_expr_1 IS NOT NULL THEN __common_expr_1 ELSE Decimal128(0.00,22,2) END AS discount_pct_2
        Projection: CAST(warehouse.main.orders.discount_pct AS Decimal128(22, 2)) AS __common_expr_1, warehouse.main.orders.order_id
          TableScan: warehouse.main.orders projection=[order_id, discount_pct]

reproduced #23138: optimized query is invalid in DuckDB
DuckDB error: Binder Error: Referenced table "o" not found!
Candidate tables: "oi"

LINE 1: ... "warehouse"."main"."order_items" AS "oi" INNER JOIN (SELECT "o"."order_id", CASE WHEN "__common_expr_1" IS NOT NULL...

// datafusion/sql/src/unparser/plan.rs
//
// AI 本地试修思路：
// 当 Projection 下面还是 Projection 时，说明外层 Projection 面对的是 derived table，
// 不能继续把列重写成 "o"."order_id"，而应该重写成 derived table 输出列 "order_id"。

use crate::unparser::rewrite::{
    ProjectionInputRewriter, TableAliasRewriter, inject_column_aliases_into_subquery,
};

// ...

LogicalPlan::Projection(projection) => {
    let input_is_projection =
        matches!(projection.input.as_ref(), LogicalPlan::Projection(_));

    if let Some(plan) = self.unparse_table_scan_pushdown(
        &projection.input,
        alias.clone(),
        already_projected,
    )? {
        let exprs = if input_is_projection {
            let mut input_rewriter = ProjectionInputRewriter {
                input_schema: plan.schema().as_ref(),
            };
            projection
                .expr
                .iter()
                .cloned()
                .map(|expr| expr.rewrite(&mut input_rewriter).data())
                .collect::<Result<Vec<_>>>()?
        } else if alias.is_some() {
            let mut alias_rewriter =
                alias.as_ref().map(|alias_name| TableAliasRewriter {
                    table_schema: plan.schema().as_ref(),
                    alias_name: alias_name.clone(),
                    rewrite_unqualified: false,
                });
            projection
                .expr
                .iter()
                .cloned()
                .map(|expr| {
                    if let Some(ref mut rewriter) = alias_rewriter {
                        expr.rewrite(rewriter).data()
                    } else {
                        Ok(expr)
                    }
                })
                .collect::<Result<Vec<_>>>()?
        } else {
            projection.expr.clone()
        };

        Ok(Some(LogicalPlan::Projection(
            Projection::try_new_with_schema(
                exprs,
                Arc::new(plan),
                Arc::clone(&projection.schema),
            )?,
        )))
    } else {
        Ok(None)
    }
}

// datafusion/sql/src/unparser/rewrite.rs
//
// AI 本地试修新增的 rewriter：
// 用 input projection 的 schema，把原始 qualified column 改成 derived table 输出列。

/// A `TreeNodeRewriter` implementation that rewrites `Expr::Column` expressions
/// to the field names produced by an input projection.
///
/// This is used when a projection is rebuilt on top of another projection that
/// will be unparsed as a derived subquery. In that SQL scope, the original table
/// qualifier is no longer visible, and the outer projection must refer to the
/// derived subquery's output columns by name.
pub struct ProjectionInputRewriter<'a> {
    pub input_schema: &'a DFSchema,
}

impl TreeNodeRewriter for ProjectionInputRewriter<'_> {
    type Node = Expr;

    fn f_down(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
        match expr {
            Expr::Column(column) => {
                match self
                    .input_schema
                    .qualified_field_from_column(&column)
                    .or_else(|_| {
                        self.input_schema
                            .qualified_field_with_unqualified_name(&column.name)
                    }) {
                    Ok((_qualifier, field)) => {
                        let new_column = Column::new_unqualified(field.name().clone());
                        Ok(Transformed::yes(Expr::Column(new_column)))
                    }
                    Err(_) => Ok(Transformed::no(Expr::Column(column))),
                }
            }
            _ => Ok(Transformed::no(expr)),
        }
    }
}

// datafusion/core/tests/sql/unparser.rs
//
// AI 本地试修加的最小回归测试：
// 确认 optimized SQL 不再生成 out-of-scope alias: "o"."order_id"。

const ISSUE_23138_QUERY: &str = r#"
SELECT * FROM
(
SELECT
        order_id
    FROM
        "warehouse"."main"."order_items"
) oi
JOIN (
    SELECT
        order_id,
        coalesce(discount_pct, 0) AS discount_pct_2
    FROM
        "warehouse"."main"."orders"
) o USING (order_id)
"#;

#[tokio::test]
async fn optimized_duckdb_unparse_resolves_columns_in_outer_query() -> Result<()> {
    let ctx = issue_22961_context()?;
    let plan = ctx.sql(ISSUE_23138_QUERY).await?.into_optimized_plan()?;
    let dialect = DuckDBDialect::new();
    let unparser = Unparser::new(&dialect);
    let sql = unparser.plan_to_sql(&plan)?.to_string();

    assert!(
        !sql.contains(r#"INNER JOIN (SELECT "o"."order_id""#),
        "unparsed SQL references derived table alias before it is in scope: {sql}"
    );
    assert!(
        sql.contains(r#"INNER JOIN (SELECT "order_id", CASE"#),
        "unparsed SQL should resolve columns against the derived query output: {sql}"
    );

    Ok(())
}

// datafusion/sql/src/unparser/plan.rs
//
// 官方 PR #23176:
// fix: unparse columns of stacked pushdown projections unqualified
//
// 核心：直接把 stacked pushdown projection 里的 column qualifier 去掉。

/// Strip the table qualifier from every column in a pushdown pass-through
/// projection expression, so it resolves against the unnamed derived table
/// rendered for the inner pushdown projection rather than a deeper table
/// alias that is out of scope at this nesting level.
fn strip_pushdown_column_qualifiers(expr: Expr) -> Result<Expr> {
    expr.transform(|e| match e {
        Expr::Column(mut column) => {
            column.relation = None;
            Ok(Transformed::yes(Expr::Column(column)))
        }
        other => Ok(Transformed::no(other)),
    })
    .data()
}

// datafusion/sql/src/unparser/plan.rs
//
// 官方 PR #23176 在 unparse_table_scan_pushdown 的 Projection 分支里加的逻辑。

LogicalPlan::Projection(projection) => {
    if let Some(plan) = self.unparse_table_scan_pushdown(
        &projection.input,
        alias.clone(),
        already_projected,
    )? {
        // The pushed-down scan alias is only in scope for the
        // projection directly above the aliased table scan. `plan`
        // is the result of pushing the alias further down: if it is
        // itself a `Projection`, the input was another projection
        // (e.g. common subexpression elimination stacked one), so
        // this projection sits over a derived table rather than
        // directly over the aliased scan, and the alias is out of
        // scope here. Its qualified pass-through columns must then
        // reference the derived table's output unqualified instead
        // of being rebased to the alias. Build it directly so the
        // unqualified columns are not re-normalized back to the
        // alias. (Otherwise `plan` is the scan-derived plan and we
        // fall through to rebase to the alias, correct one level
        // above the scan.)
        if alias.is_some() && matches!(plan, LogicalPlan::Projection(_)) {
            let exprs = projection
                .expr
                .iter()
                .cloned()
                .map(Self::strip_pushdown_column_qualifiers)
                .collect::<Result<Vec<_>>>()?;
            return Ok(Some(LogicalPlan::Projection(Projection::try_new(
                exprs,
                Arc::new(plan),
            )?)));
        }

        let exprs = if alias.is_some() {
            let mut alias_rewriter =
                alias.as_ref().map(|alias_name| TableAliasRewriter {
                    table_schema: plan.schema().as_ref(),
                    alias_name: alias_name.clone(),
                    rewrite_unqualified: false,
                });

            projection
                .expr
                .iter()
                .cloned()
                .map(|expr| {
                    if let Some(ref mut rewriter) = alias_rewriter {
                        expr.rewrite(rewriter).data()
                    } else {
                        Ok(expr)
                    }
                })
                .collect::<Result<Vec<_>>>()?
        } else {
            projection.expr.clone()
        };

        Ok(Some(LogicalPlan::Projection(Projection::try_new(
            exprs,
            Arc::new(plan),
        )?)))
    } else {
        Ok(None)
    }
}

// datafusion/core/tests/sql/unparser.rs
//
// 官方 PR #23176 加的回归测试。

// https://github.com/apache/datafusion/issues/23138
//
// CSE on `coalesce(discount_pct, 0)` factors a shared CAST into an extra inner
// projection, so `SubqueryAlias: o` ends up over two stacked projections. When
// the unparser renders that as nested derived tables it must qualify the
// pass-through `order_id` with a name in scope at each level -- it must not
// rebase it to the outer subquery alias `o`, which is not visible inside the
// inner derived table.
const ISSUE_23138_QUERY: &str = r#"
SELECT * FROM
(
    SELECT order_id FROM "warehouse"."main"."order_items"
) oi
JOIN (
    SELECT order_id, coalesce(discount_pct, 0) AS discount_pct_2
    FROM "warehouse"."main"."orders"
) o USING (order_id)
"#;

#[tokio::test]
async fn optimized_duckdb_unparse_qualifies_nested_passthrough_column() -> Result<()> {
    let ctx = issue_22961_context()?;
    let plan = ctx.sql(ISSUE_23138_QUERY).await?.into_optimized_plan()?;
    let dialect = DuckDBDialect::new();
    let unparser = Unparser::new(&dialect);
    let sql = unparser.plan_to_sql(&plan)?.to_string();

    // The intermediate derived table has no `o` in scope, so the pass-through
    // `order_id` must be unqualified there, not rebased to the subquery alias
    // `o` (which is only the base-table alias one level deeper). The bug emitted
    // `"o"."order_id"` inside that derived table; the fix emits a bare column.
    let expected = concat!(
        r#"SELECT "o"."order_id", "o"."discount_pct_2" "#,
        r#"FROM "warehouse"."main"."order_items" AS "oi" "#,
        r#"INNER JOIN (SELECT "order_id", "#,
        r#"CASE WHEN "__common_expr_1" IS NOT NULL "#,
        r#"THEN "__common_expr_1" ELSE 0.00 END AS "discount_pct_2" "#,
        r#"FROM (SELECT CAST("o"."discount_pct" AS DECIMAL(22,2)) "#,
        r#"AS "__common_expr_1", "o"."order_id" "#,
        r#"FROM "warehouse"."main"."orders" AS "o")) AS "o" "#,
        r#"ON "oi"."order_id" = "o"."order_id""#,
    );
    assert_eq!(sql, expected);

    assert!(
        sql.contains(r#"(SELECT "order_id", CASE WHEN"#),
        "pass-through order_id should be unqualified in derived table: {sql}"
    );
    assert!(
        !sql.contains(r#"(SELECT "o"."order_id", CASE WHEN"#),
        "derived table must not reference out-of-scope alias o: {sql}"
    );

    Ok(())
}

// #23138 的本质可以用这两个 SQL 片段总结。

// 修复前：错误。
// 这一层 SELECT 里面，"o" 还不在作用域里。
let bad_sql = r#"
INNER JOIN (
  SELECT "o"."order_id",
         CASE WHEN "__common_expr_1" IS NOT NULL
              THEN "__common_expr_1"
              ELSE 0.00
         END AS "discount_pct_2"
  FROM (
    SELECT CAST("o"."discount_pct" AS DECIMAL(22,2)) AS "__common_expr_1",
           "o"."order_id"
    FROM "warehouse"."main"."orders" AS "o"
  )
) AS "o" ON ...
"#;

// 修复后：正确。
// 外层 derived table 只能引用自己 input 的输出列，所以这里应该是 "order_id"。
let good_sql = r#"
INNER JOIN (
  SELECT "order_id",
         CASE WHEN "__common_expr_1" IS NOT NULL
              THEN "__common_expr_1"
              ELSE 0.00
         END AS "discount_pct_2"
  FROM (
    SELECT CAST("o"."discount_pct" AS DECIMAL(22,2)) AS "__common_expr_1",
           "o"."order_id"
    FROM "warehouse"."main"."orders" AS "o"
  )
) AS "o" ON ...
"#;

// 最终结论：
//
// LogicalPlan 没坏。
// Optimizer 引入 __common_expr_1 也没坏。
// 坏在 SQL Unparser 把 stacked projection 转回 SQL 时，
// 把内层 table alias 错误带到了外层 derived table。

fn issue_23138_summary() {
    let problem = "SubqueryAlias(o) + Projection + Projection + TableScan";
    let wrong = r#""o"."order_id""#;
    let right = r#""order_id""#;

    println!("problem shape: {problem}");
    println!("wrong generated column: {wrong}");
    println!("right generated column: {right}");
}