Improve SQLite subquery tables aliasing support (#12482)

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(

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<Ident>,
+) -> Result<LogicalPlan> {
+    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::<Vec<_>>();
 

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 {