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,
 };
 use sqlparser::ast::{
-    Expr as SQLExpr, ExprWithAliasAndOrderBy, Ident, LimitClause, Offset, OffsetRows,
-    OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, SelectInto, SetExpr,
-    SetOperator, SetQuantifier, TableAlias,
+    Expr as SQLExpr, ExprWithAliasAndOrderBy, 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
@@ -83,7 +82,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(),
@@ -379,47 +379,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,86 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# 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
+
+# Swapped aliases must sort by output positions, not resolve back to input names.
+query II
+SELECT b AS a, a AS b
+FROM (VALUES (2, 1), (1, 2), (1, 1)) AS t(a, b)
+ORDER BY ALL;
+----
+1 1
+1 2
+2 1
+
+# Duplicate output names from a self-join remain unambiguous through ordinals.
+query II
+WITH t(id) AS (VALUES (2), (1))
+SELECT l.id, r.id
+FROM t AS l
+CROSS JOIN t AS r
+ORDER BY ALL;
+----
+1 1
+1 2
+2 1
+2 2
+
+statement ok
+set datafusion.sql_parser.dialect = 'Generic';