Unparse SubqueryAlias without projections to SQL (apache#12896)

* change pub function comment to doc

* unparse subquery alias without projections

* fix tests

* rollback the empty line

* rollback the empty line

* exclude the table_scan with pushdown case

* fmt and clippy

* simplify the ast to string and remove unused debug code

Author: goldmedal

datafusion/sql/src/unparser/plan.rs

@@ -46,6 +46,18 @@ use super::{
     },
     Unparser,
 };
+use crate::unparser::utils::unproject_agg_exprs;
+use datafusion_common::{
+    internal_err, not_impl_err,
+    tree_node::{TransformedResult, TreeNode},
+    Column, DataFusionError, Result, TableReference,
+};
+use datafusion_expr::{
+    expr::Alias, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
+    LogicalPlanBuilder, Projection, SortExpr, TableScan,
+};
+use sqlparser::ast::{self, Ident, SetExpr};
+use std::sync::Arc;
 
 /// Convert a DataFusion [`LogicalPlan`] to [`ast::Statement`]
 ///
@@ -279,12 +291,9 @@ impl Unparser<'_> {
     ) -> Result<()> {
         match plan {
             LogicalPlan::TableScan(scan) => {
-                if scan.projection.is_some()
-                    || !scan.filters.is_empty()
-                    || scan.fetch.is_some()
+                if let Some(unparsed_table_scan) =
+                    Self::unparse_table_scan_pushdown(plan, None)?
                 {
-                    let unparsed_table_scan =
-                        Self::unparse_table_scan_pushdown(plan, None)?;
                     return self.select_to_sql_recursively(
                         &unparsed_table_scan,
                         query,
@@ -556,10 +565,18 @@ impl Unparser<'_> {
             LogicalPlan::SubqueryAlias(plan_alias) => {
                 let (plan, mut columns) =
                     subquery_alias_inner_query_and_columns(plan_alias);
-                let plan = Self::unparse_table_scan_pushdown(
+                let unparsed_table_scan = Self::unparse_table_scan_pushdown(
                     plan,
                     Some(plan_alias.alias.clone()),
                 )?;
+                // if the child plan is a TableScan with pushdown operations, we don't need to
+                // create an additional subquery for it
+                if !select.already_projected() && unparsed_table_scan.is_none() {
+                    select.projection(vec![ast::SelectItem::Wildcard(
+                        ast::WildcardAdditionalOptions::default(),
+                    )]);
+                }
+                let plan = unparsed_table_scan.unwrap_or_else(|| plan.clone());
                 if !columns.is_empty()
                     && !self.dialect.supports_column_alias_in_table_alias()
                 {
@@ -668,12 +685,21 @@ impl Unparser<'_> {
         }
     }
 
+    fn is_scan_with_pushdown(scan: &TableScan) -> bool {
+        scan.projection.is_some() || !scan.filters.is_empty() || scan.fetch.is_some()
+    }
+
+    /// Try to unparse a table scan with pushdown operations into a new subquery plan.
+    /// If the table scan is without any pushdown operations, return None.
     fn unparse_table_scan_pushdown(
         plan: &LogicalPlan,
         alias: Option<TableReference>,
-    ) -> Result<LogicalPlan> {
+    ) -> Result<Option<LogicalPlan>> {
         match plan {
             LogicalPlan::TableScan(table_scan) => {
+                if !Self::is_scan_with_pushdown(table_scan) {
+                    return Ok(None);
+                }
                 let mut filter_alias_rewriter =
                     alias.as_ref().map(|alias_name| TableAliasRewriter {
                         table_schema: table_scan.source.schema(),
@@ -734,18 +760,15 @@ impl Unparser<'_> {
                     builder = builder.limit(0, Some(fetch))?;
                 }
 
-                builder.build()
+                Ok(Some(builder.build()?))
             }
             LogicalPlan::SubqueryAlias(subquery_alias) => {
-                let new_plan = Self::unparse_table_scan_pushdown(
+                Self::unparse_table_scan_pushdown(
                     &subquery_alias.input,
                     Some(subquery_alias.alias.clone()),
-                )?;
-                LogicalPlanBuilder::from(new_plan)
-                    .alias(subquery_alias.alias.clone())?
-                    .build()
+                )
             }
-            _ => Ok(plan.clone()),
+            _ => Ok(None),
         }
     }
 

datafusion/sql/src/unparser/rewrite.rs

@@ -101,25 +101,25 @@ fn rewrite_sort_expr_for_union(exprs: Vec<SortExpr>) -> Result<Vec<SortExpr>> {
     Ok(sort_exprs)
 }
 
-// Rewrite logic plan for query that order by columns are not in projections
-// Plan before rewrite:
-//
-// Projection: j1.j1_string, j2.j2_string
-//   Sort: j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST
-//     Projection: j1.j1_string, j2.j2_string, j1.j1_id, j2.j2_id
-//       Inner Join:  Filter: j1.j1_id = j2.j2_id
-//         TableScan: j1
-//         TableScan: j2
-//
-// Plan after rewrite
-//
-// Sort: j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST
-//   Projection: j1.j1_string, j2.j2_string
-//     Inner Join:  Filter: j1.j1_id = j2.j2_id
-//       TableScan: j1
-//       TableScan: j2
-//
-// This prevents the original plan generate query with derived table but missing alias.
+/// Rewrite logic plan for query that order by columns are not in projections
+/// Plan before rewrite:
+///
+/// Projection: j1.j1_string, j2.j2_string
+///   Sort: j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST
+///     Projection: j1.j1_string, j2.j2_string, j1.j1_id, j2.j2_id
+///       Inner Join:  Filter: j1.j1_id = j2.j2_id
+///         TableScan: j1
+///         TableScan: j2
+///
+/// Plan after rewrite
+///
+/// Sort: j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST
+///   Projection: j1.j1_string, j2.j2_string
+///     Inner Join:  Filter: j1.j1_id = j2.j2_id
+///       TableScan: j1
+///       TableScan: j2
+///
+/// This prevents the original plan generate query with derived table but missing alias.
 pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
     p: &Projection,
 ) -> Option<LogicalPlan> {
@@ -191,33 +191,33 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
     }
 }
 
-// This logic is to work out the columns and inner query for SubqueryAlias plan for both types of
-// subquery
-// - `(SELECT column_a as a from table) AS A`
-// - `(SELECT column_a from table) AS A (a)`
-//
-// A roundtrip example for table alias with columns
-//
-// query: SELECT id FROM (SELECT j1_id from j1) AS c (id)
-//
-// LogicPlan:
-// Projection: c.id
-//   SubqueryAlias: c
-//     Projection: j1.j1_id AS id
-//       Projection: j1.j1_id
-//         TableScan: j1
-//
-// Before introducing this logic, the unparsed query would be `SELECT c.id FROM (SELECT j1.j1_id AS
-// id FROM (SELECT j1.j1_id FROM j1)) AS c`.
-// The query is invalid as `j1.j1_id` is not a valid identifier in the derived table
-// `(SELECT j1.j1_id FROM j1)`
-//
-// With this logic, the unparsed query will be:
-// `SELECT c.id FROM (SELECT j1.j1_id FROM j1) AS c (id)`
-//
-// Caveat: this won't handle the case like `select * from (select 1, 2) AS a (b, c)`
-// as the parser gives a wrong plan which has mismatch `Int(1)` types: Literal and
-// Column in the Projections. Once the parser side is fixed, this logic should work
+/// This logic is to work out the columns and inner query for SubqueryAlias plan for both types of
+/// subquery
+/// - `(SELECT column_a as a from table) AS A`
+/// - `(SELECT column_a from table) AS A (a)`
+///
+/// A roundtrip example for table alias with columns
+///
+/// query: SELECT id FROM (SELECT j1_id from j1) AS c (id)
+///
+/// LogicPlan:
+/// Projection: c.id
+///   SubqueryAlias: c
+///     Projection: j1.j1_id AS id
+///       Projection: j1.j1_id
+///         TableScan: j1
+///
+/// Before introducing this logic, the unparsed query would be `SELECT c.id FROM (SELECT j1.j1_id AS
+/// id FROM (SELECT j1.j1_id FROM j1)) AS c`.
+/// The query is invalid as `j1.j1_id` is not a valid identifier in the derived table
+/// `(SELECT j1.j1_id FROM j1)`
+///
+/// With this logic, the unparsed query will be:
+/// `SELECT c.id FROM (SELECT j1.j1_id FROM j1) AS c (id)`
+///
+/// Caveat: this won't handle the case like `select * from (select 1, 2) AS a (b, c)`
+/// as the parser gives a wrong plan which has mismatch `Int(1)` types: Literal and
+/// Column in the Projections. Once the parser side is fixed, this logic should work
 pub(super) fn subquery_alias_inner_query_and_columns(
     subquery_alias: &datafusion_expr::SubqueryAlias,
 ) -> (&LogicalPlan, Vec<Ident>) {
@@ -330,6 +330,7 @@ fn find_projection(logical_plan: &LogicalPlan) -> Option<&Projection> {
         _ => None,
     }
 }
+
 /// A `TreeNodeRewriter` implementation that rewrites `Expr::Column` expressions by
 /// replacing the column's name with an alias if the column exists in the provided schema.
 ///