Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 12 additions & 10 deletions datafusion/sql/src/unparser/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand All @@ -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,
},
Expand Down Expand Up @@ -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(
Expand Down
53 changes: 41 additions & 12 deletions datafusion/sql/src/unparser/rewrite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<_>>();

Expand Down
12 changes: 12 additions & 0 deletions datafusion/sql/tests/cases/plan_to_sql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down