From 4abdeef1e03e7d37be7a0dd1683985e9e3c8a7d1 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Sat, 19 Sep 2026 17:27:22 -0700 Subject: [PATCH 1/2] fix: preserve computed names in derived SQL projections Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/sql/src/unparser/plan.rs | 48 +++++++++++++++++------ datafusion/sql/tests/cases/plan_to_sql.rs | 22 +++++++++++ 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 18af08fc18361..eeed711a290e0 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -593,10 +593,37 @@ impl Unparser<'_> { alias: Option, lateral: bool, ) -> Result<()> { + let preserve_names = matches!(plan, LogicalPlan::Projection(_)) + && alias.as_ref().is_some_and(|alias| alias.columns.is_empty()); let mut derived_builder = DerivedRelationBuilder::default(); derived_builder.lateral(lateral).alias(alias).subquery({ let inner_statement = self.plan_to_sql(plan)?; - if let ast::Statement::Query(inner_query) = inner_statement { + if let ast::Statement::Query(mut inner_query) = inner_statement { + if preserve_names + && let SetExpr::Select(select) = inner_query.body.as_mut() + && select.projection.len() == plan.schema().fields().len() + { + for (item, field) in + select.projection.iter_mut().zip(plan.schema().fields()) + { + if let ast::SelectItem::UnnamedExpr(expr) = item { + let alias = self.column_alias_to_sql(field.name())?; + let preserves_name = match expr { + ast::Expr::Identifier(name) => name.value == alias.value, + ast::Expr::CompoundIdentifier(names) => names + .last() + .is_some_and(|name| name.value == alias.value), + _ => false, + }; + if !preserves_name { + *item = ast::SelectItem::ExprWithAlias { + expr: expr.clone(), + alias, + }; + } + } + } + } inner_query } else { return internal_err!( @@ -2689,18 +2716,9 @@ impl Unparser<'_> { Expr::Alias(Alias { expr, name, .. }) => { let inner = self.expr_to_sql(expr)?; - // Determine the alias name to use - let col_name = if let Some(rewritten_name) = - self.dialect.col_alias_overrides(name)? - { - rewritten_name.to_string() - } else { - name.to_string() - }; - Ok(ast::SelectItem::ExprWithAlias { expr: inner, - alias: self.new_ident_quoted_if_needs(col_name), + alias: self.column_alias_to_sql(name)?, }) } _ => { @@ -2711,6 +2729,14 @@ impl Unparser<'_> { } } + fn column_alias_to_sql(&self, name: &str) -> Result { + let name = self + .dialect + .col_alias_overrides(name)? + .unwrap_or_else(|| name.to_string()); + Ok(self.new_ident_quoted_if_needs(name)) + } + fn sorts_to_sql(&self, sort_exprs: &[SortExpr]) -> Result { Ok(OrderByKind::Expressions( sort_exprs diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 191425416c119..d6da1fc17eecb 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -392,6 +392,28 @@ fn roundtrip_statement_with_dialect_4() -> Result<(), DataFusionError> { Ok(()) } +#[test] +fn unparse_preserves_derived_aggregate_output_name() -> Result<()> { + let schema = Schema::new(vec![Field::new("j1_id", DataType::Int32, false)]); + let aggregate = sum(col("j1.j1_id")); + let output = Expr::Column(Column::from_name(aggregate.schema_name().to_string())); + let plan = table_scan(Some("j1"), &schema, None)? + .aggregate(Vec::::new(), vec![aggregate])? + .project(vec![output.clone().alias("visible"), output.clone()])? + .project(vec![output])? + .build()?; + + let sql = Unparser::new(&UnparserPostgreSqlDialect {}) + .plan_to_sql(&plan)? + .to_string(); + println!("UNPARSED_SQL={sql}"); + assert_snapshot!( + sql, + @r#"SELECT "sum(j1.j1_id)" FROM (SELECT sum("j1"."j1_id") AS "visible", sum("j1"."j1_id") AS "sum(j1.j1_id)" FROM "j1") AS "derived_projection""# + ); + Ok(()) +} + #[test] fn roundtrip_rebases_derived_projection_references() -> Result<(), DataFusionError> { roundtrip_statement_with_dialect_helper!( From 91ca248ad866dc1a19e1a99f8de0108927e5d229 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Sun, 20 Sep 2026 15:00:20 -0700 Subject: [PATCH 2/2] fix: preserve output names without derived table aliases Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/sql/src/unparser/ast.rs | 45 ++++++++++++++++++++--- datafusion/sql/src/unparser/plan.rs | 38 ++++++------------- datafusion/sql/tests/cases/plan_to_sql.rs | 35 +++++++++++++++++- 3 files changed, 84 insertions(+), 34 deletions(-) diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index c335d3ee4c16d..38b51afe3e262 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -677,6 +677,7 @@ pub struct DerivedRelationBuilder { lateral: Option, subquery: Option>, alias: Option, + projection_names: Vec, } impl DerivedRelationBuilder { @@ -692,18 +693,49 @@ impl DerivedRelationBuilder { self.alias = value; self } + pub(super) fn projection_names(&mut self, value: Vec) -> &mut Self { + self.projection_names = value; + self + } fn build(&self) -> Result { + let mut subquery = match self.subquery { + Some(ref value) => value.clone(), + None => { + return Err(Into::into(UninitializedFieldError::from("subquery"))); + } + }; + if self + .alias + .as_ref() + .is_none_or(|alias| alias.columns.is_empty()) + && let ast::SetExpr::Select(select) = subquery.body.as_mut() + && select.projection.len() == self.projection_names.len() + { + for (item, alias) in select.projection.iter_mut().zip(&self.projection_names) + { + if let ast::SelectItem::UnnamedExpr(expr) = item { + let preserves_name = match expr { + ast::Expr::Identifier(name) => name.value == alias.value, + ast::Expr::CompoundIdentifier(names) => { + names.last().is_some_and(|name| name.value == alias.value) + } + _ => false, + }; + if !preserves_name { + *item = ast::SelectItem::ExprWithAlias { + expr: expr.clone(), + alias: alias.clone(), + }; + } + } + } + } Ok(ast::TableFactor::Derived { lateral: match self.lateral { Some(ref value) => *value, None => return Err(Into::into(UninitializedFieldError::from("lateral"))), }, - subquery: match self.subquery { - Some(ref value) => value.clone(), - None => { - return Err(Into::into(UninitializedFieldError::from("subquery"))); - } - }, + subquery, alias: self.alias.clone(), sample: None, }) @@ -713,6 +745,7 @@ impl DerivedRelationBuilder { lateral: Default::default(), subquery: Default::default(), alias: Default::default(), + projection_names: Default::default(), } } } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index eeed711a290e0..f51d76f93a444 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -594,36 +594,20 @@ impl Unparser<'_> { lateral: bool, ) -> Result<()> { let preserve_names = matches!(plan, LogicalPlan::Projection(_)) - && alias.as_ref().is_some_and(|alias| alias.columns.is_empty()); + && alias.as_ref().is_none_or(|alias| alias.columns.is_empty()); let mut derived_builder = DerivedRelationBuilder::default(); + if preserve_names { + derived_builder.projection_names( + plan.schema() + .fields() + .iter() + .map(|field| self.column_alias_to_sql(field.name())) + .collect::>>()?, + ); + } derived_builder.lateral(lateral).alias(alias).subquery({ let inner_statement = self.plan_to_sql(plan)?; - if let ast::Statement::Query(mut inner_query) = inner_statement { - if preserve_names - && let SetExpr::Select(select) = inner_query.body.as_mut() - && select.projection.len() == plan.schema().fields().len() - { - for (item, field) in - select.projection.iter_mut().zip(plan.schema().fields()) - { - if let ast::SelectItem::UnnamedExpr(expr) = item { - let alias = self.column_alias_to_sql(field.name())?; - let preserves_name = match expr { - ast::Expr::Identifier(name) => name.value == alias.value, - ast::Expr::CompoundIdentifier(names) => names - .last() - .is_some_and(|name| name.value == alias.value), - _ => false, - }; - if !preserves_name { - *item = ast::SelectItem::ExprWithAlias { - expr: expr.clone(), - alias, - }; - } - } - } - } + if let ast::Statement::Query(inner_query) = inner_statement { inner_query } else { return internal_err!( diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index d6da1fc17eecb..de22321327d75 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -17,6 +17,7 @@ use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::tree_node::{Transformed, TransformedResult}; use datafusion_common::{ Column, DFSchema, DFSchemaRef, DataFusionError, Result, TableReference, assert_contains, @@ -111,6 +112,25 @@ fn roundtrip_expr(table: TableReference, sql: &str) -> Result { Ok(ast.to_string()) } +fn remove_column_self_aliases(plan: LogicalPlan) -> Result { + plan.transform_up_with_subqueries(|plan| { + plan.map_expressions(|expr| { + if let Expr::Alias(alias) = &expr + && alias.relation.is_none() + && alias.metadata.is_none() + && let Expr::Column(column) = alias.expr.as_ref() + && column.relation.is_none() + && column.name == alias.name + { + Ok(Transformed::yes(*alias.expr.clone())) + } else { + Ok(Transformed::no(expr)) + } + }) + }) + .data() +} + #[test] fn roundtrip_statement() -> Result<()> { let tests: Vec<&str> = vec![ @@ -254,7 +274,11 @@ fn roundtrip_statement() -> Result<()> { .sql_statement_to_plan(roundtrip_statement.clone()) .unwrap(); - assert_eq!(plan, plan_roundtrip); + // Explicit output names can add unqualified self-aliases without changing the plan's meaning. + assert_eq!( + remove_column_self_aliases(plan)?, + remove_column_self_aliases(plan_roundtrip)?, + ); } Ok(()) @@ -411,6 +435,15 @@ fn unparse_preserves_derived_aggregate_output_name() -> Result<()> { sql, @r#"SELECT "sum(j1.j1_id)" FROM (SELECT sum("j1"."j1_id") AS "visible", sum("j1"."j1_id") AS "sum(j1.j1_id)" FROM "j1") AS "derived_projection""# ); + + let sql = Unparser::new(&BigQueryDialect {}) + .plan_to_sql(&plan)? + .to_string(); + println!("BIGQUERY_SQL={sql}"); + assert_snapshot!( + sql, + @r#"SELECT `sum_40j1_46j1_id_41` FROM (SELECT sum(`j1`.`j1_id`) AS `visible`, sum(`j1`.`j1_id`) AS `sum_40j1_46j1_id_41` FROM `j1`)"# + ); Ok(()) }