fix(unparser): keep ORDER BY out of a derived table when the sort key is computed

datafusion/sql/src/unparser/rewrite.rs

@@ -184,6 +184,20 @@ pub(super) fn rewrite_qualify(plan: LogicalPlan) -> Result<LogicalPlan> {
 ///       TableScan: j2
 ///
 /// This prevents the original plan generate query with derived table but missing alias.
+///
+/// It also keeps the ORDER BY at the top level of the emitted statement. Left as
+/// `Projection -> Sort`, the `Sort` is unparsed as a derived table:
+///
+/// ```sql
+/// SELECT id FROM (SELECT person.id, person.age FROM person ORDER BY person.age)
+/// ```
+///
+/// SQL does not require the enclosing query to preserve the ordering of a derived
+/// table, so the rows come back in an arbitrary order.
+///
+/// A sort key does not have to *be* one of the inner Projection's outputs for the
+/// hoist to be valid -- it only has to be computable from them, as `age + 1` is from
+/// `age`.
 pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
     p: &Projection,
 ) -> Option<LogicalPlan> {
@@ -221,6 +235,15 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
         })
         .collect::<Vec<_>>();
 
+    // Compare outer collects Expr::to_string with inner collected transformed values
+    // alias -> alias column
+    // column -> remain
+    // others, extract schema field name
+    let inner_collects = inner_exprs
+        .iter()
+        .map(Expr::to_string)
+        .collect::<HashSet<_>>();
+
     let mut collects = p.expr.clone();
     for sort in &sort.expr {
         // Strip aliases from sort expressions so the comparison matches
@@ -231,18 +254,21 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
         while let Expr::Alias(alias) = expr {
             expr = *alias.expr;
         }
-        collects.push(expr);
+        if inner_collects.contains(&expr.to_string()) {
+            collects.push(expr);
+            continue;
+        }
+        // The sort key is not itself one of the inner Projection's outputs, but
+        // it may be an expression *over* them (`ORDER BY age + 1` while the
+        // inner Projection exposes `age`). Account for the columns it reads so
+        // an inner Projection that exists only to expose them still matches.
+        // Without this the rewrite bails out and the Sort is emitted as a
+        // derived table, where SQL does not guarantee the ORDER BY is honoured
+        // by the enclosing query -- the rows come back in an arbitrary order.
+        collects.extend(expr.column_refs().into_iter().cloned().map(Expr::Column));
     }
 
-    // Compare outer collects Expr::to_string with inner collected transformed values
-    // alias -> alias column
-    // column -> remain
-    // others, extract schema field name
     let outer_collects = collects.iter().map(Expr::to_string).collect::<HashSet<_>>();
-    let inner_collects = inner_exprs
-        .iter()
-        .map(Expr::to_string)
-        .collect::<HashSet<_>>();
 
     if outer_collects == inner_collects {
         let mut sort = sort.clone();
@@ -288,11 +314,21 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
                 while let Expr::Alias(alias) = expr {
                     expr = *alias.expr;
                 }
-                if let Expr::Column(ref col) = expr
-                    && let Some(underlying) = dropped_aliases.get(col.name())
-                {
-                    sort_expr.expr = underlying.clone();
-                }
+                // Substitute nested references too, not just a whole-expression
+                // one: `ORDER BY x + 1` over a dropped `expr AS x` has to become
+                // `ORDER BY expr + 1`.
+                sort_expr.expr = expr
+                    .clone()
+                    .transform_down(|e| {
+                        Ok(match &e {
+                            Expr::Column(col) => match dropped_aliases.get(col.name()) {
+                                Some(underlying) => Transformed::yes(underlying.clone()),
+                                None => Transformed::no(e),
+                            },
+                            _ => Transformed::no(e),
+                        })
+                    })
+                    .map_or(expr, |transformed| transformed.data);
             }
         }
 

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -1342,7 +1342,10 @@ fn table_scan_with_empty_projection_and_none_projection_helper(
 // yielding `Projection: <empty> -> TableScan`. It must not be unparsed as an
 // empty `SELECT` list for dialects that reject `SELECT FROM t` (e.g. DuckDB):
 // fall back to `SELECT 1` just like the bare empty-projection `TableScan`.
-fn empty_projection_over_table_helper(table_name: &str, table_schema: Schema) -> LogicalPlan {
+fn empty_projection_over_table_helper(
+    table_name: &str,
+    table_schema: Schema,
+) -> LogicalPlan {
     project(
         table_scan(Some(table_name), &table_schema, None)
             .unwrap()
@@ -4274,3 +4277,59 @@ fn test_unparse_chained_intersect_build_side_is_self_contained() -> Result<()> {
     );
     Ok(())
 }
+
+/// A `Sort` that has to be emitted below the SELECT list must not end up inside a
+/// derived table: SQL does not require an enclosing query to honour the ORDER BY of
+/// a derived table, so the rows come back in an arbitrary order.
+#[test]
+fn order_by_over_non_projected_field_stays_top_level() -> Result<()> {
+    // Sort key is an expression over a column the SELECT list does not project.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT id FROM person ORDER BY age + 1 DESC",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: UnparserDefaultDialect {},
+        expected: @"SELECT person.id FROM person ORDER BY (person.age + 1) DESC NULLS FIRST",
+    );
+
+    // Same, with the sort key an expression over an aggregate that is not selected.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT id, first_name FROM person GROUP BY id, first_name ORDER BY max(age) + 1 DESC",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: UnparserDefaultDialect {},
+        expected: @"SELECT person.id, person.first_name FROM person GROUP BY person.id, person.first_name ORDER BY (max(person.age) + 1) DESC NULLS FIRST",
+    );
+
+    // A plain non-projected column already worked; keep it covered so the
+    // generalisation above cannot regress it.
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT id FROM person ORDER BY age DESC",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: UnparserDefaultDialect {},
+        expected: @"SELECT person.id FROM person ORDER BY person.age DESC NULLS FIRST",
+    );
+
+    Ok(())
+}
+
+/// The inner `Projection` may define an alias that only the `Sort` uses. Hoisting the
+/// `Sort` drops that alias, so every reference to it -- including one nested inside a
+/// larger sort expression -- has to be replaced by the expression it named.
+#[test]
+fn order_by_nested_reference_to_dropped_alias() -> Result<()> {
+    let schema = Schema::new(vec![
+        Field::new("id", DataType::Int64, false),
+        Field::new("age", DataType::Int64, false),
+    ]);
+    let plan = table_scan(Some("person"), &schema, None)?
+        .project(vec![col("id"), (col("age") * lit(2)).alias("doubled")])?
+        .sort(vec![(col("doubled") + lit(1)).sort(false, true)])?
+        .project(vec![col("id")])?
+        .build()?;
+
+    assert_snapshot!(
+        plan_to_sql(&plan)?,
+        @"SELECT person.id FROM person ORDER BY ((person.age * 2) + 1) DESC NULLS FIRST"
+    );
+
+    Ok(())
+}