fix(unparser): keep a DISTINCT ON's sort expressions when naming its outputs

Rebuilding a node from the expressions a plan reports drops what those
expressions do not carry. For `Distinct::On` that is the `ORDER BY`, and
`with_new_exprs` asserts it was never asked to carry it — so naming the outputs
of a projection under a `DISTINCT ON ... ORDER BY` panicked while unparsing a
plan that is otherwise valid.

Replace only the input on each node walked through, leaving every other field as
it stands, which cannot drop a field the walk does not know about.

Author: claudespice

datafusion/sql/src/unparser/utils.rs

@@ -27,8 +27,8 @@ use datafusion_common::{
     tree_node::{Transformed, TransformedResult, TreeNode},
 };
 use datafusion_expr::{
-    Aggregate, Expr, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr, Unnest,
-    Window, expr, utils::grouping_set_to_exprlist,
+    Aggregate, Distinct, Expr, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr,
+    Unnest, Window, expr, utils::grouping_set_to_exprlist,
 };
 
 use indexmap::IndexSet;
@@ -337,29 +337,60 @@ pub(crate) fn name_derived_scope_outputs(
 /// the same way for a different purpose and stops at a narrower set, because a
 /// predicate can only reach a projection that stays in its own `SELECT`.
 fn name_scope_projection_outputs(plan: &LogicalPlan) -> Result<Option<LogicalPlan>> {
+    // Each node is rebuilt by replacing its input and leaving every other field as
+    // it stands. `LogicalPlan::with_new_exprs` is the shorter spelling and the
+    // wrong one: it reconstructs a node from the expressions a plan reports, and
+    // that round trip does not carry a `DISTINCT ON`'s sort expressions — it
+    // panics on a plan that has them.
     match plan {
         LogicalPlan::Projection(projection) => name_projection_outputs(projection),
-        // These carry the projection's output names through to the derived table
-        // unchanged, so the projection below is still the one exposing its columns.
-        LogicalPlan::Filter(_)
-        | LogicalPlan::Sort(_)
-        | LogicalPlan::Limit(_)
-        | LogicalPlan::Distinct(_)
-        | LogicalPlan::SubqueryAlias(_) => {
-            let inputs = plan.inputs();
-            let [input] = inputs.as_slice() else {
-                return Ok(None);
-            };
-            let Some(named_input) = name_scope_projection_outputs(input)? else {
-                return Ok(None);
-            };
-            plan.with_new_exprs(plan.expressions(), vec![named_input])
-                .map(Some)
+        // The nodes below carry the projection's output names out to the derived
+        // table unchanged, so the projection under them is still the one exposing
+        // its columns.
+        LogicalPlan::Filter(filter) => with_named_input(&filter.input, |input| {
+            let mut filter = filter.clone();
+            filter.input = input;
+            LogicalPlan::Filter(filter)
+        }),
+        LogicalPlan::Sort(sort) => with_named_input(&sort.input, |input| {
+            let mut sort = sort.clone();
+            sort.input = input;
+            LogicalPlan::Sort(sort)
+        }),
+        LogicalPlan::Limit(limit) => with_named_input(&limit.input, |input| {
+            let mut limit = limit.clone();
+            limit.input = input;
+            LogicalPlan::Limit(limit)
+        }),
+        LogicalPlan::Distinct(Distinct::All(input)) => {
+            with_named_input(input, |input| LogicalPlan::Distinct(Distinct::All(input)))
+        }
+        LogicalPlan::Distinct(Distinct::On(distinct_on)) => {
+            with_named_input(&distinct_on.input, |input| {
+                let mut distinct_on = distinct_on.clone();
+                distinct_on.input = input;
+                LogicalPlan::Distinct(Distinct::On(distinct_on))
+            })
+        }
+        LogicalPlan::SubqueryAlias(subquery_alias) => {
+            with_named_input(&subquery_alias.input, |input| {
+                let mut subquery_alias = subquery_alias.clone();
+                subquery_alias.input = input;
+                LogicalPlan::SubqueryAlias(subquery_alias)
+            })
         }
         _ => Ok(None),
     }
 }
 
+/// Rebuilds a node around a named input, or `None` where the input needs no naming.
+fn with_named_input(
+    input: &Arc<LogicalPlan>,
+    rebuild: impl FnOnce(Arc<LogicalPlan>) -> LogicalPlan,
+) -> Result<Option<LogicalPlan>> {
+    Ok(name_scope_projection_outputs(input)?.map(|named| rebuild(Arc::new(named))))
+}
+
 /// The projection with each unnamed output aliased to the name its schema reports.
 fn name_projection_outputs(projection: &Projection) -> Result<Option<LogicalPlan>> {
     if !projection.expr.iter().any(output_is_unnamed) {

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -6281,6 +6281,34 @@ fn test_derived_output_under_subquery_alias_is_named() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_derived_output_under_distinct_on_with_sort_is_named() -> Result<()> {
+    // `DISTINCT ON` carries its own `ORDER BY`, and rebuilding the node from the
+    // expressions a plan reports drops it — `LogicalPlan::with_new_exprs` asserts
+    // it was not asked to. Replacing only the input keeps every other field, so a
+    // plan with sort expressions is named rather than panicked on.
+    let schema = Schema::new(vec![
+        Field::new("a", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+    ]);
+    let plan = table_scan(Some("t"), &schema, Some(vec![0, 1]))?
+        .project(vec![col("t.a").add(col("t.b")), col("t.a")])?
+        .distinct_on(
+            vec![col("t.a")],
+            vec![col("t.a + t.b"), col("t.a")],
+            Some(vec![col("t.a").sort(true, false)]),
+        )?
+        .limit(0, Some(5))?
+        .project(vec![col("t.a + t.b")])?
+        .build()?;
+
+    assert_snapshot!(
+        plan_to_sql(&plan)?,
+        @r#"SELECT "t.a + t.b" FROM (SELECT DISTINCT ON (t.a) "t.a + t.b", a FROM (SELECT (t.a + t.b) AS "t.a + t.b", t.a FROM t) ORDER BY a ASC NULLS LAST LIMIT 5)"#
+    );
+    Ok(())
+}
+
 #[test]
 fn test_named_derived_projection_outputs_are_unchanged() -> Result<()> {
     // An alias and a bare column already carry the name the enclosing scope