datafusion sql unparser 错误sql
datafusion sql unparser error 生成错误sql 问题描述 从订单明细表取订单id，从订单表取订单id和折扣百分比，使用USING连接两个子查询，生成的优化后的SQL在DuckDB中无法执行。 两个子查询join后，DataFusion 把优化后的 LogicalPlan 转回 SQL 时，生成的 SQL 里面 alias 作用域错了，所以 DuckDB 执行这个生成 SQL 会报错 复现 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... ai 如何解决的 // 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}"); }