From 626ac3450991dfc76e4c3edab6b911fda69c2803 Mon Sep 17 00:00:00 2001 From: Chongchen Chen Date: Sat, 8 Mar 2025 01:15:14 +0800 Subject: [PATCH 1/3] fix: unparse for subqueryalias --- datafusion/core/tests/sql/select.rs | 46 +++++++++++++++++++++++++++++ datafusion/sql/src/unparser/plan.rs | 11 +++++-- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 6e81bf6410c11..f8cab6ab10d20 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -17,6 +17,7 @@ use super::*; use datafusion_common::ScalarValue; +use datafusion_sql::unparser::plan_to_sql; #[tokio::test] async fn test_list_query_parameters() -> Result<()> { @@ -350,3 +351,48 @@ async fn test_version_function() { assert_eq!(version.value(0), expected_version); } + +#[tokio::test] +async fn test_unparse_subqueryalias() -> Result<()> { + let ctx = SessionContext::new(); + let sql = r#" +SELECT + customer_view.c_custkey, + customer_view.c_name, + customer_view.custkey_plus +FROM + ( + SELECT + customer.c_custkey, + customer.c_name, + customer.custkey_plus + FROM + ( + SELECT + customer.c_custkey, + CAST(customer.c_custkey AS BIGINT) + 1 AS custkey_plus, + customer.c_name + FROM + ( + SELECT + customer.c_custkey AS c_custkey, + customer.c_name AS c_name + FROM + customer + ) AS customer + ) AS customer + ) AS customer_view + "#; + /// Return a RecordBatch with made up data about customer + fn customer() -> RecordBatch { + let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let name: ArrayRef = Arc::new(StringArray::from_iter_values([""])); + RecordBatch::try_from_iter(vec![("c_custkey", custkey), ("c_name", name)]) + .unwrap() + } + ctx.register_batch("customer", customer())?; + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?; + assert_eq!(sql.to_string(), "SELECT customer_view.c_custkey, customer_view.c_name, customer_view.custkey_plus FROM (SELECT customer.c_custkey, (CAST(customer.c_custkey AS BIGINT) + 1) AS custkey_plus, customer.c_name FROM (SELECT customer.c_custkey, customer.c_name FROM customer AS customer) AS customer) AS customer_view"); + Ok(()) +} diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 0fa203c60b7ba..a4f7a831fd5f7 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -984,11 +984,18 @@ impl Unparser<'_> { Ok(Some(builder.build()?)) } LogicalPlan::SubqueryAlias(subquery_alias) => { - Self::unparse_table_scan_pushdown( + let ret = Self::unparse_table_scan_pushdown( &subquery_alias.input, Some(subquery_alias.alias.clone()), already_projected, - ) + ); + if let Some(alias) = alias { + if let Ok(Some(plan)) = ret { + let plan = LogicalPlanBuilder::new(plan).alias(alias)?.build()?; + return Ok(Some(plan)); + } + } + ret } // SubqueryAlias could be rewritten to a plan with a projection as the top node by [rewrite::subquery_alias_inner_query_and_columns]. // The inner table scan could be a scan with pushdown operations. From fe48e0087163b83ee3fd7a69645e5415239c8a8d Mon Sep 17 00:00:00 2001 From: Chongchen Chen Date: Sun, 9 Mar 2025 21:36:44 +0800 Subject: [PATCH 2/3] update --- datafusion/sql/src/unparser/plan.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index a4f7a831fd5f7..b14fbdff236f5 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -988,14 +988,14 @@ impl Unparser<'_> { &subquery_alias.input, Some(subquery_alias.alias.clone()), already_projected, - ); + )?; if let Some(alias) = alias { - if let Ok(Some(plan)) = ret { + if let Some(plan) = ret { let plan = LogicalPlanBuilder::new(plan).alias(alias)?.build()?; return Ok(Some(plan)); } } - ret + Ok(ret) } // SubqueryAlias could be rewritten to a plan with a projection as the top node by [rewrite::subquery_alias_inner_query_and_columns]. // The inner table scan could be a scan with pushdown operations. From 2144399f2b104449f4631d5442bad8c91c2b501b Mon Sep 17 00:00:00 2001 From: Chongchen Chen Date: Tue, 11 Mar 2025 07:41:31 +0800 Subject: [PATCH 3/3] move test --- datafusion/core/tests/sql/select.rs | 46 ---------------- datafusion/sql/tests/cases/plan_to_sql.rs | 66 ++++++++++++++++++++++- 2 files changed, 65 insertions(+), 47 deletions(-) diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index f8cab6ab10d20..6e81bf6410c11 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -17,7 +17,6 @@ use super::*; use datafusion_common::ScalarValue; -use datafusion_sql::unparser::plan_to_sql; #[tokio::test] async fn test_list_query_parameters() -> Result<()> { @@ -351,48 +350,3 @@ async fn test_version_function() { assert_eq!(version.value(0), expected_version); } - -#[tokio::test] -async fn test_unparse_subqueryalias() -> Result<()> { - let ctx = SessionContext::new(); - let sql = r#" -SELECT - customer_view.c_custkey, - customer_view.c_name, - customer_view.custkey_plus -FROM - ( - SELECT - customer.c_custkey, - customer.c_name, - customer.custkey_plus - FROM - ( - SELECT - customer.c_custkey, - CAST(customer.c_custkey AS BIGINT) + 1 AS custkey_plus, - customer.c_name - FROM - ( - SELECT - customer.c_custkey AS c_custkey, - customer.c_name AS c_name - FROM - customer - ) AS customer - ) AS customer - ) AS customer_view - "#; - /// Return a RecordBatch with made up data about customer - fn customer() -> RecordBatch { - let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1])); - let name: ArrayRef = Arc::new(StringArray::from_iter_values([""])); - RecordBatch::try_from_iter(vec![("c_custkey", custkey), ("c_name", name)]) - .unwrap() - } - ctx.register_batch("customer", customer())?; - let plan = ctx.sql(sql).await?.into_optimized_plan()?; - let sql = plan_to_sql(&plan)?; - assert_eq!(sql.to_string(), "SELECT customer_view.c_custkey, customer_view.c_name, customer_view.custkey_plus FROM (SELECT customer.c_custkey, (CAST(customer.c_custkey AS BIGINT) + 1) AS custkey_plus, customer.c_name FROM (SELECT customer.c_custkey, customer.c_name FROM customer AS customer) AS customer) AS customer_view"); - Ok(()) -} diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 5af93a01e6c99..fc0b7a26baafa 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -21,7 +21,7 @@ use datafusion_expr::test::function_stub::{ count_udaf, max_udaf, min_udaf, sum, sum_udaf, }; use datafusion_expr::{ - col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, + cast, col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; use datafusion_functions::unicode; @@ -37,6 +37,7 @@ use datafusion_sql::unparser::dialect::{ use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser}; use sqlparser::ast::Statement; use std::hash::Hash; +use std::ops::Add; use std::sync::Arc; use std::{fmt, vec}; @@ -1687,3 +1688,66 @@ fn test_unparse_optimized_multi_union() -> Result<()> { Ok(()) } + +/// Test unparse the optimized plan from the following SQL: +/// ``` +/// SELECT +/// customer_view.c_custkey, +/// customer_view.c_name, +/// customer_view.custkey_plus +/// FROM +/// ( +/// SELECT +/// customer.c_custkey, +/// customer.c_name, +/// customer.custkey_plus +/// FROM +/// ( +/// SELECT +/// customer.c_custkey, +/// CAST(customer.c_custkey AS BIGINT) + 1 AS custkey_plus, +/// customer.c_name +/// FROM +/// ( +/// SELECT +/// customer.c_custkey AS c_custkey, +/// customer.c_name AS c_name +/// FROM +/// customer +/// ) AS customer +/// ) AS customer +/// ) AS customer_view +/// ``` +#[test] +fn test_unparse_subquery_alias_with_table_pushdown() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c_custkey", DataType::Int32, false), + Field::new("c_name", DataType::Utf8, false), + ]); + + let table_scan = table_scan(Some("customer"), &schema, Some(vec![0, 1]))?.build()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .alias("customer")? + .project(vec![ + col("customer.c_custkey"), + cast(col("customer.c_custkey"), DataType::Int64) + .add(lit(1)) + .alias("custkey_plus"), + col("customer.c_name"), + ])? + .alias("customer")? + .project(vec![ + col("customer.c_custkey"), + col("customer.c_name"), + col("customer.custkey_plus"), + ])? + .alias("customer_view")? + .build()?; + + let unparser = Unparser::default(); + let sql = unparser.plan_to_sql(&plan)?; + let expected = "SELECT customer_view.c_custkey, customer_view.c_name, customer_view.custkey_plus FROM (SELECT customer.c_custkey, (CAST(customer.c_custkey AS BIGINT) + 1) AS custkey_plus, customer.c_name FROM (SELECT customer.c_custkey, customer.c_name FROM customer AS customer) AS customer) AS customer_view"; + assert_eq!(sql.to_string(), expected); + Ok(()) +}