refactor(unparser): centralize aggregate-scope rendering in the SQL unparser

datafusion/core/tests/sql/unparser.rs

@@ -468,6 +468,59 @@ QUALIFY
     rn = 1 AND count(DISTINCT cs.customer_id) > 0
 "#;
 
+// https://github.com/apache/datafusion/issues/23668
+//
+// Extends the #23317 aggregate-scope fix to the window and ORDER BY clauses.
+// Reuses issue_23317_context() (same derived-projection shape).
+
+// Window sorting by an aggregate, over a derived-projection input. Already
+// correct today; this locks the OVER clause against keeping the out-of-scope
+// `cs` qualifier across the refactor.
+const ISSUE_23668_WINDOW_QUERY: &str = r#"
+SELECT
+    date_part('year', c.signup_date) AS signup_year,
+    count(DISTINCT cs.customer_id) AS customers,
+    row_number() OVER (ORDER BY count(DISTINCT cs.customer_id) DESC) AS rn
+FROM
+    "warehouse"."main"."sales" cs
+    JOIN "warehouse"."main"."customers" c USING (customer_id)
+GROUP BY
+    1
+"#;
+
+// ORDER BY an aggregate that is NOT selected, so it can't use a select alias
+// and is unprojected through the Aggregate. It must be normalized like the
+// SELECT list, not keep the out-of-scope `cs` qualifier.
+const ISSUE_23668_ORDER_BY_QUERY: &str = r#"
+SELECT
+    date_part('year', c.signup_date) AS signup_year,
+    count(DISTINCT cs.customer_id) AS customers
+FROM
+    "warehouse"."main"."sales" cs
+    JOIN "warehouse"."main"."customers" c USING (customer_id)
+GROUP BY
+    1
+ORDER BY
+    round(sum(cs.total_revenue), 2) DESC
+"#;
+
+// ORDER BY a selected aggregate keeps a top-level Sort (the direct `Sort` arm,
+// vs the projection-absorbed one above). It resolves to the select alias, so
+// this covers routing only -- the normalization in that arm isn't reachable
+// from SQL (an unselected aggregate takes the absorbed path above instead).
+const ISSUE_23668_TOP_LEVEL_SORT_QUERY: &str = r#"
+SELECT
+    date_part('year', c.signup_date) AS signup_year,
+    count(DISTINCT cs.customer_id) AS customers
+FROM
+    "warehouse"."main"."sales" cs
+    JOIN "warehouse"."main"."customers" c USING (customer_id)
+GROUP BY
+    1
+ORDER BY
+    customers DESC
+"#;
+
 fn issue_23317_context() -> Result<SessionContext> {
     let ctx = SessionContext::new();
 
@@ -612,6 +665,88 @@ async fn optimized_duckdb_unparse_qualify_unqualifies_agg_input() -> Result<()>
     Ok(())
 }
 
+#[tokio::test]
+async fn optimized_duckdb_unparse_window_over_agg_unqualifies_input() -> Result<()> {
+    let ctx = issue_23317_context()?;
+    assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name()));
+
+    let plan = ctx
+        .sql(ISSUE_23668_WINDOW_QUERY)
+        .await?
+        .into_optimized_plan()?;
+    let dialect = DuckDBDialect::new();
+    let unparser = Unparser::new(&dialect);
+    let sql = unparser.plan_to_sql(&plan)?.to_string();
+
+    assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?;
+
+    assert!(
+        sql.contains(r#"OVER (ORDER BY count(DISTINCT "customer_id")"#),
+        "window ORDER BY aggregate should resolve against the derived projection output: {sql}",
+    );
+    assert!(
+        !sql.contains(r#"count(DISTINCT "cs"."customer_id")"#),
+        "window OVER clause must not reference out-of-scope alias cs: {sql}",
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn optimized_duckdb_unparse_order_by_unqualifies_agg_input() -> Result<()> {
+    let ctx = issue_23317_context()?;
+    assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name()));
+
+    let plan = ctx
+        .sql(ISSUE_23668_ORDER_BY_QUERY)
+        .await?
+        .into_optimized_plan()?;
+    let dialect = DuckDBDialect::new();
+    let unparser = Unparser::new(&dialect);
+    let sql = unparser.plan_to_sql(&plan)?.to_string();
+
+    assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?;
+
+    assert!(
+        sql.contains(r#"ORDER BY round(sum("total_revenue"), 2)"#),
+        "ORDER BY aggregate should resolve against the derived projection output: {sql}",
+    );
+    assert!(
+        !sql.contains(r#"sum("cs"."total_revenue")"#),
+        "ORDER BY must not reference out-of-scope alias cs: {sql}",
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn optimized_duckdb_unparse_top_level_sort_over_agg_uses_select_alias() -> Result<()>
+{
+    let ctx = issue_23317_context()?;
+    assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name()));
+
+    let plan = ctx
+        .sql(ISSUE_23668_TOP_LEVEL_SORT_QUERY)
+        .await?
+        .into_optimized_plan()?;
+    let dialect = DuckDBDialect::new();
+    let unparser = Unparser::new(&dialect);
+    let sql = unparser.plan_to_sql(&plan)?.to_string();
+
+    assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?;
+
+    assert!(
+        sql.contains(r#"ORDER BY "customers""#),
+        "top-level ORDER BY should resolve to the select alias: {sql}",
+    );
+    assert!(
+        !sql.contains(r#""cs"."customer_id") AS "customers""#),
+        "aggregate output must not reference out-of-scope alias cs: {sql}",
+    );
+
+    Ok(())
+}
+
 /// The outcome of running a single roundtrip test.
 ///
 /// A successful test produces [`TestCaseResult::Success`].

datafusion/sql/src/unparser/plan.rs

@@ -101,6 +101,69 @@ pub fn plan_to_sql(plan: &LogicalPlan) -> Result<ast::Statement> {
     unparser.plan_to_sql(plan)
 }
 
+/// Aggregate-expression scope for one rendered SELECT block.
+///
+/// When an aggregate's input is itself emitted as a derived subquery (a
+/// projection sits between the aggregate and its relation), the input columns
+/// are only reachable by that derived table's output names. Base-table
+/// qualifiers like `t.col` name a relation that is out of scope above the
+/// boundary, so emitting them produces SQL a strict engine rejects.
+///
+/// Every clause that renders an aggregate expression (SELECT / GROUP BY /
+/// HAVING / QUALIFY / ORDER BY) has to apply the same rule. Detect the
+/// boundary once here and reuse it, so the clauses can't drift apart (which is
+/// how earlier fixes left some clauses correct and others not).
+struct UnparserAggScope<'a> {
+    agg: &'a Aggregate,
+    /// `agg.input` renders as a derived projection, so out-of-scope qualifiers
+    /// must be stripped from expressions in this scope.
+    input_is_derived_projection: bool,
+}
+
+impl<'a> UnparserAggScope<'a> {
+    fn new(agg: &'a Aggregate) -> Self {
+        Self {
+            agg,
+            input_is_derived_projection: Unparser::contains_projection_before_relation(
+                agg.input.as_ref(),
+            ),
+        }
+    }
+
+    /// Prepare a projected column or predicate that still references the
+    /// aggregate by its output columns: unproject it back onto the aggregate
+    /// (and `windows`) expressions, then normalize it for this scope.
+    fn prepare(&self, expr: Expr, windows: Option<&[&Window]>) -> Result<Expr> {
+        self.normalize(unproject_agg_exprs(expr, self.agg, windows)?)
+    }
+
+    /// Normalize an expression that is already in aggregate form (group / aggr
+    /// exprs, or an unprojected sort expr): strip the qualifiers that fall out
+    /// of scope once the input is a derived projection. No-op otherwise.
+    fn normalize(&self, expr: Expr) -> Result<Expr> {
+        if self.input_is_derived_projection {
+            Unparser::strip_column_qualifiers_for_schema(
+                expr,
+                self.agg.input.schema().as_ref(),
+            )
+        } else {
+            Ok(expr)
+        }
+    }
+
+    /// Unproject a sort expression onto this aggregate, then normalize it so
+    /// ORDER BY uses the same scope as the other clauses.
+    fn prepare_sort_expr(
+        &self,
+        sort_expr: SortExpr,
+        input: &LogicalPlan,
+    ) -> Result<SortExpr> {
+        let mut sort_expr = unproject_sort_expr(sort_expr, Some(self.agg), input)?;
+        sort_expr.expr = self.normalize(sort_expr.expr)?;
+        Ok(sort_expr)
+    }
+}
+
 impl Unparser<'_> {
     pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result<ast::Statement> {
         let mut plan = normalize_union_schema(plan)?;
@@ -312,17 +375,12 @@ impl Unparser<'_> {
         match (agg, window) {
             (Some(agg), window) => {
                 let window_option = window.as_deref();
-                let agg_input_has_derived_projection =
-                    Self::contains_projection_before_relation(agg.input.as_ref());
+                let unparser_agg_scope = UnparserAggScope::new(agg);
                 let items = exprs
                     .into_iter()
                     .map(|proj_expr| {
-                        let unproj = unproject_agg_exprs(proj_expr, agg, window_option)?;
-                        let unproj = Self::normalize_agg_input_columns(
-                            unproj,
-                            agg,
-                            agg_input_has_derived_projection,
-                        )?;
+                        let unproj =
+                            unparser_agg_scope.prepare(proj_expr, window_option)?;
                         self.select_item_to_sql(&unproj)
                     })
                     .collect::<Result<Vec<_>>>()?;
@@ -333,12 +391,7 @@ impl Unparser<'_> {
                         .iter()
                         .cloned()
                         .map(|expr| {
-                            let expr = Self::normalize_agg_input_columns(
-                                expr,
-                                agg,
-                                agg_input_has_derived_projection,
-                            )?;
-                            self.expr_to_sql(&expr)
+                            self.expr_to_sql(&unparser_agg_scope.normalize(expr)?)
                         })
                         .collect::<Result<Vec<_>>>()?,
                     vec![],
@@ -379,18 +432,6 @@ impl Unparser<'_> {
         }
     }
 
-    fn normalize_agg_input_columns(
-        expr: Expr,
-        agg: &Aggregate,
-        input_has_derived_projection: bool,
-    ) -> Result<Expr> {
-        if input_has_derived_projection {
-            Self::strip_column_qualifiers_for_schema(expr, agg.input.schema().as_ref())
-        } else {
-            Ok(expr)
-        }
-    }
-
     fn contains_projection_before_relation(plan: &LogicalPlan) -> bool {
         match plan {
             LogicalPlan::Projection(_) => true,
@@ -429,6 +470,19 @@ impl Unparser<'_> {
         }
     }
 
+    /// Unproject a sort expression; normalize it when the sort is above an
+    /// aggregate, otherwise just unproject (no scope to normalize against).
+    fn unproject_sort_expr_in_scope(
+        sort_expr: SortExpr,
+        agg: Option<&Aggregate>,
+        input: &LogicalPlan,
+    ) -> Result<SortExpr> {
+        match agg {
+            Some(agg) => UnparserAggScope::new(agg).prepare_sort_expr(sort_expr, input),
+            None => unproject_sort_expr(sort_expr, None, input),
+        }
+    }
+
     fn derive(
         &self,
         plan: &LogicalPlan,
@@ -592,6 +646,9 @@ impl Unparser<'_> {
             window_expr
                 .iter()
                 .map(|expr| {
+                    // No normalization: this agg branch is only reachable from a
+                    // hand-built plan. SQL wraps windows in a projection, which
+                    // reconstruct_select_statement handles (and normalizes).
                     let expr = if let Some(agg) = agg {
                         unproject_agg_exprs(expr.clone(), agg, None)?
                     } else {
@@ -977,7 +1034,7 @@ impl Unparser<'_> {
                         sort.expr
                             .iter()
                             .map(|sort_expr| {
-                                unproject_sort_expr(
+                                Self::unproject_sort_expr_in_scope(
                                     sort_expr.clone(),
                                     agg,
                                     sort.input.as_ref(),
@@ -1028,23 +1085,14 @@ impl Unparser<'_> {
                     let mut unprojected =
                         unproject_window_exprs(filter.predicate.clone(), window)?;
                     if let Some(agg) = agg {
-                        unprojected = unproject_agg_exprs(unprojected, agg, None)?;
-                        unprojected = Self::normalize_agg_input_columns(
-                            unprojected,
-                            agg,
-                            Self::contains_projection_before_relation(agg.input.as_ref()),
-                        )?;
+                        unprojected =
+                            UnparserAggScope::new(agg).prepare(unprojected, None)?;
                     }
                     let filter_expr = self.expr_to_sql(&unprojected)?;
                     select.qualify(Some(filter_expr));
                 } else if let Some(agg) = agg {
-                    let unprojected =
-                        unproject_agg_exprs(filter.predicate.clone(), agg, None)?;
-                    let unprojected = Self::normalize_agg_input_columns(
-                        unprojected,
-                        agg,
-                        Self::contains_projection_before_relation(agg.input.as_ref()),
-                    )?;
+                    let unprojected = UnparserAggScope::new(agg)
+                        .prepare(filter.predicate.clone(), None)?;
                     let filter_expr = self.expr_to_sql(&unprojected)?;
                     select.having(Some(filter_expr));
                 } else {
@@ -1130,7 +1178,11 @@ impl Unparser<'_> {
                     .expr
                     .iter()
                     .map(|sort_expr| {
-                        unproject_sort_expr(sort_expr.clone(), agg, sort.input.as_ref())
+                        Self::unproject_sort_expr_in_scope(
+                            sort_expr.clone(),
+                            agg,
+                            sort.input.as_ref(),
+                        )
                     })
                     .collect::<Result<Vec<_>>>()?;
 
@@ -1146,8 +1198,7 @@ impl Unparser<'_> {
             LogicalPlan::Aggregate(agg) => {
                 // Aggregation can be already handled in the projection case
                 if !select.already_projected() {
-                    let agg_input_has_derived_projection =
-                        Self::contains_projection_before_relation(agg.input.as_ref());
+                    let unparser_agg_scope = UnparserAggScope::new(agg);
                     // The query returns aggregate and group expressions. If that weren't the case,
                     // the aggregate would have been placed inside a projection, making the check above^ false
                     let exprs: Vec<_> = agg
@@ -1156,12 +1207,7 @@ impl Unparser<'_> {
                         .chain(agg.group_expr.iter())
                         .cloned()
                         .map(|expr| {
-                            let expr = Self::normalize_agg_input_columns(
-                                expr,
-                                agg,
-                                agg_input_has_derived_projection,
-                            )?;
-                            self.select_item_to_sql(&expr)
+                            self.select_item_to_sql(&unparser_agg_scope.normalize(expr)?)
                         })
                         .collect::<Result<Vec<_>>>()?;
                     select.projection(exprs);
@@ -1171,12 +1217,7 @@ impl Unparser<'_> {
                             .iter()
                             .cloned()
                             .map(|expr| {
-                                let expr = Self::normalize_agg_input_columns(
-                                    expr,
-                                    agg,
-                                    agg_input_has_derived_projection,
-                                )?;
-                                self.expr_to_sql(&expr)
+                                self.expr_to_sql(&unparser_agg_scope.normalize(expr)?)
                             })
                             .collect::<Result<Vec<_>>>()?,
                         vec![],