SNOW-2268220: support ORDER BY ALL (#3912)

Author: sfc-gh-mayliu

CHANGELOG.md

@@ -11,6 +11,7 @@
 
 #### Improvements
 
+- Enhanced `DataFrame.sort()` to support `ORDER BY ALL` when no columns are specified.
 - Catalog API now uses SQL commands instead of SnowAPI calls. This new implementation is more reliable now.
 
 #### Dependency Updates
@@ -103,6 +104,9 @@
     - `st_geometryfromwkt`
     - `try_to_geography`
     - `try_to_geometry`
+
+#### Improvements
+
 - Added a parameter to enable and disable automatic column name aliasing for `interval_day_time_from_parts` and `interval_year_month_from_parts` functions.
 
 #### Bug Fixes

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -114,7 +114,10 @@
     SnowflakeTable,
     SnowflakeValues,
 )
-from snowflake.snowpark._internal.analyzer.sort_expression import SortOrder
+from snowflake.snowpark._internal.analyzer.sort_expression import (
+    SortOrder,
+    SortByAllOrder,
+)
 from snowflake.snowpark._internal.analyzer.table_function import (
     FlattenFunction,
     GeneratorTableFunction,
@@ -558,6 +561,13 @@ def analyze(
                 expr.null_ordering.sql,
             )
 
+        if isinstance(expr, SortByAllOrder):
+            return order_expression(
+                "ALL",
+                expr.direction.sql,
+                expr.null_ordering.sql,
+            )
+
         if isinstance(expr, ScalarSubquery):
             self.subquery_plans.append(expr.plan)
             return subquery_expression(expr.plan.queries[-1].sql)

src/snowflake/snowpark/_internal/analyzer/sort_expression.py

@@ -59,3 +59,23 @@ def dependent_column_names(self) -> Optional[AbstractSet[str]]:
 
     def dependent_column_names_with_duplication(self) -> List[str]:
         return derive_dependent_columns_with_duplication(self.child)
+
+
+class SortByAllOrder(Expression):
+    def __init__(
+        self,
+        direction: SortDirection,
+        null_ordering: Optional[NullOrdering] = None,
+    ) -> None:
+        super().__init__()
+        self.child: Expression
+        self.direction = direction
+        self.null_ordering = (
+            null_ordering if null_ordering else direction.default_null_ordering
+        )
+
+    def dependent_column_names(self) -> Optional[AbstractSet[str]]:
+        return derive_dependent_columns(self.child)
+
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.child)

src/snowflake/snowpark/_internal/analyzer/unary_plan_node.py

@@ -17,7 +17,10 @@
 from snowflake.snowpark._internal.analyzer.snowflake_plan_node import (
     DynamicTableCreateMode,
 )
-from snowflake.snowpark._internal.analyzer.sort_expression import SortOrder
+from snowflake.snowpark._internal.analyzer.sort_expression import (
+    SortOrder,
+    SortByAllOrder,
+)
 
 
 class UnaryNode(LogicalPlan):
@@ -90,7 +93,7 @@ def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
 class Sort(UnaryNode):
     def __init__(
         self,
-        order: List[SortOrder],
+        order: Union[List[SortOrder], List[SortByAllOrder]],
         child: LogicalPlan,
         is_order_by_append: bool = False,
     ) -> None:

src/snowflake/snowpark/dataframe.py

@@ -84,6 +84,7 @@
     Ascending,
     Descending,
     SortOrder,
+    SortByAllOrder,
 )
 from snowflake.snowpark._internal.analyzer.table_function import (
     FlattenFunction,
@@ -2103,6 +2104,8 @@ def sort(
     ) -> "DataFrame":
         """Sorts a DataFrame by the specified expressions (similar to ORDER BY in SQL).
 
+        When called with no column arguments, sorts by all columns (ORDER BY ALL).
+
         Examples::
 
             >>> from snowflake.snowpark.functions import col
@@ -2139,22 +2142,58 @@ def sort(
             -------------
             <BLANKLINE>
 
+            >>> # Sort by all columns (ORDER BY ALL) - no columns specified
+            >>> df.sort().show()
+            -------------
+            |"A"  |"B"  |
+            -------------
+            |1    |2    |
+            |1    |4    |
+            |3    |4    |
+            -------------
+            <BLANKLINE>
+
+            >>> df.sort(ascending=False).show()
+            -------------
+            |"A"  |"B"  |
+            -------------
+            |3    |4    |
+            |1    |4    |
+            |1    |2    |
+            -------------
+            <BLANKLINE>
+
         Args:
-            *cols: A column name as :class:`str` or :class:`Column`, or a list of
-             columns to sort by.
-            ascending: A :class:`bool` or a list of :class:`bool` for sorting the
-             DataFrame, where ``True`` sorts a column in ascending order and ``False``
-             sorts a column in descending order . If you specify a list of multiple
-             sort orders, the length of the list must equal the number of columns.
+            *cols: Column names as :class:`str`, :class:`Column` objects, or a list of
+             columns to sort by. If no columns are provided, the DataFrame is sorted
+             by all columns in the order they appear (equivalent to ``ORDER BY ALL`` in SQL).
+            ascending: Sort order specification.
+
+             - When sorting **specific columns**: A :class:`bool`, :class:`int`, or list of
+               :class:`bool`/:class:`int` values. ``True`` (or 1) for ascending, ``False``
+               (or 0) for descending. If a list is provided, its length must match the number
+               of columns.
+             - When sorting **all columns** (no columns specified): Must be a single
+               :class:`bool` or :class:`int`, not a list. Applies the same sort order to
+               all columns.
+             - Defaults to ``True`` (ascending) when not specified.
+
+        Note:
+            The aliases ``order_by()`` and ``orderBy()`` have the same behavior.
         """
-        if not cols:
-            raise ValueError("sort() needs at least one sort expression.")
-        # This code performs additional type checks, run first.
-        exprs = self._convert_cols_to_exprs("sort()", *cols)
-        if not exprs:
-            raise ValueError("sort() needs at least one sort expression.")
 
-        # AST.
+        is_order_by_all = not cols
+        if (
+            is_order_by_all
+            and ascending is not None
+            and not isinstance(ascending, (bool, int))
+        ):
+            raise TypeError(
+                "When no columns are specified (ORDER BY ALL), "
+                "ascending must be bool or int, not a list. "
+                "To sort specific columns with different orders, specify the columns."
+            )
+
         stmt = None
         if _emit_ast:
             stmt = self._session._ast_batch.bind()
@@ -2167,59 +2206,68 @@ def sort(
             ast.cols.variadic = is_variadic
             self._set_ast_ref(ast.df)
 
-        orders = []
-        # `ascending` is represented by Expr in the AST.
-        # Therefore, construct the required bool, int, or list and copy from that.
-        asc_expr_ast = None
-        if _emit_ast:
+            # Populate ascending as Expr
             asc_expr_ast = proto.Expr()
-        if ascending is not None:
-            if isinstance(ascending, (list, tuple)):
-                orders = [Ascending() if asc else Descending() for asc in ascending]
-                if _emit_ast:
-                    # Here asc_expr_ast is a list of bools and ints.
-                    for asc in ascending:
-                        asc_ast = proto.Expr()
-                        if isinstance(asc, bool):
-                            asc_ast.bool_val.v = asc
-                        else:
-                            asc_ast.int64_val.v = asc
-                        asc_expr_ast.list_val.vs.append(asc_ast)
-            elif isinstance(ascending, (bool, int)):
-                orders = [Ascending() if ascending else Descending()]
-                if _emit_ast:
-                    # Here asc_expr_ast is either a bool or an int.
-                    if isinstance(ascending, bool):
-                        asc_expr_ast.bool_val.v = ascending
+            asc_value = True if ascending is None else ascending
+            if isinstance(asc_value, (list, tuple)):
+                for asc in asc_value:
+                    asc_ast = proto.Expr()
+                    if isinstance(asc, bool):
+                        asc_ast.bool_val.v = asc
                     else:
-                        asc_expr_ast.int64_val.v = ascending
-            else:
-                raise TypeError(
-                    "ascending can only be boolean or list,"
-                    " but got {}".format(str(type(ascending)))
-                )
-            if _emit_ast:
-                ast.ascending.CopyFrom(asc_expr_ast)
-            if len(exprs) != len(orders):
-                raise ValueError(
-                    "The length of col ({}) should be same with"
-                    " the length of ascending ({}).".format(len(exprs), len(orders))
-                )
+                        asc_ast.int64_val.v = asc
+                    asc_expr_ast.list_val.vs.append(asc_ast)
+            elif isinstance(asc_value, (bool, int)):
+                if isinstance(asc_value, bool):
+                    asc_expr_ast.bool_val.v = asc_value
+                else:
+                    asc_expr_ast.int64_val.v = asc_value
+            ast.ascending.CopyFrom(asc_expr_ast)
+
+        # Build sort expressions
+        if is_order_by_all:
+            asc_value = True if ascending is None else ascending
+            order = Ascending() if bool(asc_value) else Descending()
+            sort_exprs = [SortByAllOrder(order)]
+        else:
+            exprs = self._convert_cols_to_exprs("sort()", *cols)
+            if not exprs:
+                raise ValueError("sort() needs at least one sort expression.")
+
+            orders = []
+            if ascending is not None:
+                if isinstance(ascending, (list, tuple)):
+                    orders = [Ascending() if asc else Descending() for asc in ascending]
+                elif isinstance(ascending, (bool, int)):
+                    orders = [Ascending() if ascending else Descending()]
+                else:
+                    raise TypeError(
+                        "ascending can only be boolean or list,"
+                        " but got {}".format(str(type(ascending)))
+                    )
 
-        sort_exprs = []
-        for idx in range(len(exprs)):
-            # orders will overwrite current orders in expression (but will not overwrite null ordering)
-            # if no order is provided, use ascending order
-            if isinstance(exprs[idx], SortOrder):
-                sort_exprs.append(
-                    SortOrder(exprs[idx].child, orders[idx], exprs[idx].null_ordering)
-                    if orders
-                    else exprs[idx]
-                )
-            else:
-                sort_exprs.append(
-                    SortOrder(exprs[idx], orders[idx] if orders else Ascending())
-                )
+                if len(exprs) != len(orders):
+                    raise ValueError(
+                        "The length of col ({}) should be same with"
+                        " the length of ascending ({}).".format(len(exprs), len(orders))
+                    )
+
+            sort_exprs = []
+            for idx in range(len(exprs)):
+                # orders will overwrite current orders in expression (but will not overwrite null ordering)
+                # if no order is provided, use ascending order
+                if isinstance(exprs[idx], SortOrder):
+                    sort_exprs.append(
+                        SortOrder(
+                            exprs[idx].child, orders[idx], exprs[idx].null_ordering
+                        )
+                        if orders
+                        else exprs[idx]
+                    )
+                else:
+                    sort_exprs.append(
+                        SortOrder(exprs[idx], orders[idx] if orders else Ascending())
+                    )
 
         # In snowpark_connect_compatible mode, we need to handle
         # the sorting for dataframe after aggregation without nesting