This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8555e41a90 Improve SQLite subquery tables aliasing support (#12482)
8555e41a90 is described below
commit 8555e41a90be6691a30e237f9168c5948810476f
Author: Sergei Grebnov <[email protected]>
AuthorDate: Tue Sep 17 10:44:16 2024 -0700
Improve SQLite subquery tables aliasing support (#12482)
---
datafusion/sql/src/unparser/plan.rs | 22 +++++++------
datafusion/sql/src/unparser/rewrite.rs | 53 ++++++++++++++++++++++++-------
datafusion/sql/tests/cases/plan_to_sql.rs | 12 +++++++
3 files changed, 65 insertions(+), 22 deletions(-)
diff --git a/datafusion/sql/src/unparser/plan.rs
b/datafusion/sql/src/unparser/plan.rs
index dc746b472a..c376a83ce7 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 e43c2eae23..9b4eaca834 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<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<_>>();
diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs
b/datafusion/sql/tests/cases/plan_to_sql.rs
index caec1e9c9d..48a3a4f360 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 {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
