diff --git a/datafusion/sql/src/unparser/rewrite.rs b/datafusion/sql/src/unparser/rewrite.rs index a95289d89f2e0..8f9135dcb7a91 100644 --- a/datafusion/sql/src/unparser/rewrite.rs +++ b/datafusion/sql/src/unparser/rewrite.rs @@ -184,6 +184,20 @@ pub(super) fn rewrite_qualify(plan: LogicalPlan) -> Result { /// 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 { @@ -221,6 +235,15 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields( }) .collect::>(); + // 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::>(); + 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::>(); - let inner_collects = inner_exprs - .iter() - .map(Expr::to_string) - .collect::>(); 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); } } diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index a9fc0716c4be2..540d07d5cbc64 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -1342,7 +1342,10 @@ fn table_scan_with_empty_projection_and_none_projection_helper( // yielding `Projection: -> 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(()) +}