Support ORDER BY ALL for projected expressions

datafusion/sql/src/query.rs

@@ -27,11 +27,10 @@ use datafusion_expr::{
     CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder, lit,
 };
 use sqlparser::ast::{
-    Expr as SQLExpr, ExprWithAliasAndOrderBy, Fetch, Ident, LimitClause, Offset,
-    OffsetRows, OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, SelectInto,
-    SetExpr, SetOperator, SetQuantifier, TableAlias,
+    Expr as SQLExpr, ExprWithAliasAndOrderBy, Fetch, LimitClause, Offset, OffsetRows,
+    OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, SelectInto, SetExpr,
+    SetOperator, SetQuantifier, TableAlias, Value,
 };
-use sqlparser::tokenizer::Span;
 
 impl<S: ContextProvider> SqlToRel<'_, S> {
     /// Generate a logical plan from an SQL query/subquery
@@ -84,7 +83,8 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 let plan = crate::stack::maybe_grow(|| {
                     self.set_expr_to_plan(other, planner_context)
                 })?;
-                let oby_exprs = to_order_by_exprs(order_by)?;
+                let oby_exprs =
+                    to_order_by_exprs(order_by, plan.schema().fields().len())?;
                 let order_by_rex = self.order_by_to_sort_expr(
                     oby_exprs,
                     plan.schema(),
@@ -406,47 +406,35 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
 }
 
 /// Returns the order by expressions from the query.
-fn to_order_by_exprs(order_by: Option<OrderBy>) -> Result<Vec<OrderByExpr>> {
-    to_order_by_exprs_with_select(order_by, None)
-}
-
-/// Returns the order by expressions from the query with the select expressions.
-pub(crate) fn to_order_by_exprs_with_select(
+fn to_order_by_exprs(
     order_by: Option<OrderBy>,
-    select_exprs: Option<&Vec<Expr>>,
+    order_by_all_column_count: usize,
 ) -> Result<Vec<OrderByExpr>> {
     let Some(OrderBy { kind, interpolate }) = order_by else {
-        // If no order by, return an empty array.
         return Ok(vec![]);
     };
     if let Some(_interpolate) = interpolate {
         return not_impl_err!("ORDER BY INTERPOLATE is not supported");
     }
     match kind {
-        OrderByKind::All(order_by_options) => {
-            let Some(exprs) = select_exprs else {
-                return Ok(vec![]);
-            };
-            let order_by_exprs = exprs
-                .iter()
-                .map(|select_expr| match select_expr {
-                    Expr::Column(column) => Ok(OrderByExpr {
-                        expr: SQLExpr::Identifier(Ident {
-                            value: column.name.clone(),
-                            quote_style: None,
-                            span: Span::empty(),
-                        }),
-                        options: order_by_options,
-                        with_fill: None,
-                    }),
-                    // TODO: Support other types of expressions
-                    _ => not_impl_err!(
-                        "ORDER BY ALL is not supported for non-column expressions"
-                    ),
-                })
-                .collect::<Result<Vec<_>>>()?;
-            Ok(order_by_exprs)
-        }
+        OrderByKind::All(options) => Ok((1..=order_by_all_column_count)
+            .map(|position| OrderByExpr {
+                expr: SQLExpr::value(Value::Number(position.to_string(), false)),
+                options,
+                with_fill: None,
+            })
+            .collect()),
         OrderByKind::Expressions(order_by_exprs) => Ok(order_by_exprs),
     }
 }
+
+/// Returns the order by expressions from the query with the select expressions.
+pub(crate) fn to_order_by_exprs_with_select(
+    order_by: Option<OrderBy>,
+    select_exprs: Option<&Vec<Expr>>,
+) -> Result<Vec<OrderByExpr>> {
+    let Some(select_exprs) = select_exprs else {
+        return to_order_by_exprs(order_by, 0);
+    };
+    to_order_by_exprs(order_by, select_exprs.len())
+}

datafusion/sqllogictest/test_files/order_by_all.slt

@@ -0,0 +1,46 @@
+# ORDER BY ALL expands computed expressions and aliases in select-list order.
+statement ok
+set datafusion.sql_parser.dialect = 'DuckDB';
+
+query II
+SELECT column1 + 1 AS computed, length(column2) AS text_length
+FROM (VALUES (2, 'x'), (1, 'zzz'), (1, 'a')) AS t(column1, column2)
+ORDER BY ALL;
+----
+2 1
+2 3
+3 1
+
+# Direction and null ordering apply to every expanded sort key.
+query IT
+SELECT column1 + 1 AS computed, column2 AS label
+FROM (VALUES (2, 'b'), (NULL, 'n'), (1, 'a')) AS t(column1, column2)
+ORDER BY ALL DESC NULLS FIRST;
+----
+NULL n
+3 b
+2 a
+
+# ORDER BY ALL also applies to the output columns of a set operation.
+query IT
+SELECT 2 AS x, 'b' AS y
+UNION ALL SELECT 1, 'z'
+UNION ALL SELECT 1, 'a'
+ORDER BY ALL;
+----
+1 a
+1 z
+2 b
+
+# Aggregate outputs are sorted by their projected positions too.
+query II
+SELECT column1 AS value, count(*) AS occurrences
+FROM (VALUES (2), (1), (1)) AS t(column1)
+GROUP BY column1
+ORDER BY ALL DESC;
+----
+2 1
+1 2
+
+statement ok
+set datafusion.sql_parser.dialect = 'Generic';