diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index dc746b472a7ec..c376a83ce70b7 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -17,7 +17,7 @@ use crate::unparser::utils::unproject_agg_exprs; use datafusion_common::{ - internal_err, not_impl_err, plan_err, + internal_err, not_impl_err, tree_node::{TransformedResult, TreeNode}, Column, DataFusionError, Result, TableReference, }; @@ -34,7 +34,7 @@ use super::{ SelectBuilder, TableRelationBuilder, TableWithJoinsBuilder, }, rewrite::{ - inject_column_aliases, normalize_union_schema, + inject_column_aliases_into_subquery, normalize_union_schema, rewrite_plan_for_sort_on_non_projected_fields, subquery_alias_inner_query_and_columns, TableAliasRewriter, }, @@ -477,15 +477,17 @@ impl Unparser<'_> { if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() { - // if columns are returned then the plan corresponds to a projection - let LogicalPlan::Projection(inner_p) = plan else { - return plan_err!( - "Inner projection for subquery alias is expected" - ); - }; - // Instead of specifying column aliases as part of the outer table, inject them directly into the inner projection - let rewritten_plan = inject_column_aliases(&inner_p, columns); + let rewritten_plan = + match inject_column_aliases_into_subquery(plan, columns) { + Ok(p) => p, + Err(e) => { + return internal_err!( + "Failed to transform SubqueryAlias plan: {e}" + ) + } + }; + columns = vec![]; self.select_to_sql_recursively( diff --git a/datafusion/sql/src/unparser/rewrite.rs b/datafusion/sql/src/unparser/rewrite.rs index e43c2eae238e4..9b4eaca834f16 100644 --- a/datafusion/sql/src/unparser/rewrite.rs +++ b/datafusion/sql/src/unparser/rewrite.rs @@ -258,8 +258,36 @@ pub(super) fn subquery_alias_inner_query_and_columns( (outer_projections.input.as_ref(), columns) } -/// Injects column aliases into the projection of a logical plan by wrapping `Expr::Column` expressions -/// with `Expr::Alias` using the provided list of aliases. Non-column expressions are left unchanged. +/// Injects column aliases into a subquery's logical plan. The function searches for a `Projection` +/// within the given plan, which may be wrapped by other operators (e.g., LIMIT, SORT). +/// If the top-level plan is a `Projection`, it directly injects the column aliases. +/// Otherwise, it iterates through the plan's children to locate and transform the `Projection`. +/// +/// Example: +/// - `SELECT col1, col2 FROM table LIMIT 10` plan with aliases `["alias_1", "some_alias_2"]` will be transformed to +/// - `SELECT col1 AS alias_1, col2 AS some_alias_2 FROM table LIMIT 10` +pub(super) fn inject_column_aliases_into_subquery( + plan: LogicalPlan, + aliases: Vec, +) -> Result { + match &plan { + LogicalPlan::Projection(inner_p) => Ok(inject_column_aliases(inner_p, aliases)), + _ => { + // projection is wrapped by other operator (LIMIT, SORT, etc), iterate through the plan to find it + plan.map_children(|child| { + if let LogicalPlan::Projection(p) = &child { + Ok(Transformed::yes(inject_column_aliases(p, aliases.clone()))) + } else { + Ok(Transformed::no(child)) + } + }) + .map(|plan| plan.data) + } + } +} + +/// Injects column aliases into the projection of a logical plan by wrapping expressions +/// with `Expr::Alias` using the provided list of aliases. /// /// Example: /// - `SELECT col1, col2 FROM table` with aliases `["alias_1", "some_alias_2"]` will be transformed to @@ -274,16 +302,17 @@ pub(super) fn inject_column_aliases( .expr .into_iter() .zip(aliases) - .map(|(expr, col_alias)| match expr { - Expr::Column(col) => { - let relation = col.relation.clone(); - Expr::Alias(Alias { - expr: Box::new(Expr::Column(col)), - relation, - name: col_alias.value, - }) - } - _ => expr, + .map(|(expr, col_alias)| { + let relation = match &expr { + Expr::Column(col) => col.relation.clone(), + _ => None, + }; + + Expr::Alias(Alias { + expr: Box::new(expr.clone()), + relation, + name: col_alias.value, + }) }) .collect::>(); diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index caec1e9c9d70e..48a3a4f36024f 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -421,6 +421,18 @@ fn roundtrip_statement_with_dialect() -> Result<()> { parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(SqliteDialect {}), }, + TestStatementWithDialect { + sql: "SELECT * FROM (SELECT j1_id + 1 FROM j1) AS temp_j(id2)", + expected: r#"SELECT * FROM (SELECT (`j1`.`j1_id` + 1) AS `id2` FROM `j1`) AS `temp_j`"#, + parser_dialect: Box::new(GenericDialect {}), + unparser_dialect: Box::new(SqliteDialect {}), + }, + TestStatementWithDialect { + sql: "SELECT * FROM (SELECT j1_id FROM j1 LIMIT 1) AS temp_j(id2)", + expected: r#"SELECT * FROM (SELECT `j1`.`j1_id` AS `id2` FROM `j1` LIMIT 1) AS `temp_j`"#, + parser_dialect: Box::new(GenericDialect {}), + unparser_dialect: Box::new(SqliteDialect {}), + }, ]; for query in tests {