Skip to content

Commit

Permalink
Improve SQLite subquery tables aliasing support
Browse files Browse the repository at this point in the history
  • Loading branch information
sgrebnov committed Sep 16, 2024
1 parent 7bd7747 commit 1146677
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 22 deletions.
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

0 comments on commit 1146677

Please sign in to comment.