enh: allow for over(order_by=...) for rank (#2746)

Author: MarcoGorelli

narwhals/_duckdb/dataframe.py

@@ -477,6 +477,9 @@ def unpivot(
 
     @requires.backend_version((1, 3))
     def with_row_index(self, name: str, order_by: Sequence[str]) -> Self:
+        if order_by is None:
+            msg = "Cannot pass `order_by` to `with_row_index` for DuckDB"
+            raise TypeError(msg)
         expr = (window_expression(F("row_number"), order_by=order_by) - lit(1)).alias(
             name
         )

narwhals/_duckdb/expr.py

@@ -112,8 +112,8 @@ def func(df: DuckDBLazyFrame, inputs: DuckDBWindowInputs) -> list[Expression]:
                     F(func_name, expr),
                     inputs.partition_by,
                     inputs.order_by,
-                    descending=reverse,
-                    nulls_last=reverse,
+                    descending=[reverse] * len(inputs.order_by),
+                    nulls_last=[reverse] * len(inputs.order_by),
                     rows_start="unbounded preceding",
                     rows_end="current row",
                 )
@@ -564,8 +564,8 @@ def func(df: DuckDBLazyFrame, inputs: DuckDBWindowInputs) -> Sequence[Expression
                     F("row_number"),
                     (*inputs.partition_by, expr),
                     inputs.order_by,
-                    descending=True,
-                    nulls_last=True,
+                    descending=[True] * len(inputs.order_by),
+                    nulls_last=[True] * len(inputs.order_by),
                 )
                 == lit(1)
                 for expr in self(df)
@@ -731,16 +731,18 @@ def rank(self, method: RankMethod, *, descending: bool) -> Self:
 
         def _rank(
             expr: Expression,
+            partition_by: Sequence[str | Expression] = (),
+            order_by: Sequence[str | Expression] = (),
             *,
-            descending: bool,
-            partition_by: Sequence[str | Expression],
+            descending: Sequence[bool],
+            nulls_last: Sequence[bool],
         ) -> Expression:
             count_expr = F("count", StarExpression())
             window_kwargs: WindowExpressionKwargs = {
                 "partition_by": partition_by,
-                "order_by": (expr,),
+                "order_by": (expr, *order_by),
                 "descending": descending,
-                "nulls_last": True,
+                "nulls_last": nulls_last,
             }
             count_window_kwargs: WindowExpressionKwargs = {
                 "partition_by": (*partition_by, expr)
@@ -760,14 +762,21 @@ def _rank(
             return when(expr.isnotnull(), rank_expr)
 
         def _unpartitioned_rank(expr: Expression) -> Expression:
-            return _rank(expr, partition_by=(), descending=descending)
+            return _rank(expr, descending=[descending], nulls_last=[True])
 
         def _partitioned_rank(
             df: DuckDBLazyFrame, inputs: DuckDBWindowInputs
         ) -> Sequence[Expression]:
-            assert not inputs.order_by  # noqa: S101
+            # node: when `descending` / `nulls_last` are supported in `.over`, they should be respected here
+            # https://github.com/narwhals-dev/narwhals/issues/2790
             return [
-                _rank(expr, descending=descending, partition_by=inputs.partition_by)
+                _rank(
+                    expr,
+                    inputs.partition_by,
+                    inputs.order_by,
+                    descending=[descending] + [False] * len(inputs.order_by),
+                    nulls_last=[True] + [False] * len(inputs.order_by),
+                )
                 for expr in self(df)
             ]
 

narwhals/_duckdb/typing.py

@@ -13,6 +13,6 @@ class WindowExpressionKwargs(TypedDict, total=False):
     order_by: Sequence[str | Expression]
     rows_start: str
     rows_end: str
-    descending: bool
-    nulls_last: bool
+    descending: Sequence[bool]
+    nulls_last: Sequence[bool]
     ignore_nulls: bool

narwhals/_duckdb/utils.py

@@ -36,6 +36,8 @@
     "us": "TIMESTAMP",
     "ns": "TIMESTAMP_NS",
 }
+DESCENDING_TO_ORDER = {True: "desc", False: "asc"}
+NULLS_LAST_TO_NULLS_POS = {True: "nulls last", False: "nulls first"}
 
 col = duckdb.ColumnExpression
 """Alias for `duckdb.ColumnExpression`."""
@@ -304,15 +306,14 @@ def generate_partition_by_sql(*partition_by: str | Expression) -> str:
 
 
 def generate_order_by_sql(
-    *order_by: str | Expression, descending: bool, nulls_last: bool
+    *order_by: str | Expression, descending: Sequence[bool], nulls_last: Sequence[bool]
 ) -> str:
     if not order_by:
         return ""
-    nulls = "nulls last" if nulls_last else "nulls first"
-    if descending:
-        by_sql = ", ".join([f"{parse_into_expression(x)} desc {nulls}" for x in order_by])
-    else:
-        by_sql = ", ".join([f"{parse_into_expression(x)} asc {nulls}" for x in order_by])
+    by_sql = ",".join(
+        f"{parse_into_expression(x)} {DESCENDING_TO_ORDER[_descending]} {NULLS_LAST_TO_NULLS_POS[_nulls_last]}"
+        for x, _descending, _nulls_last in zip(order_by, descending, nulls_last)
+    )
     return f"order by {by_sql}"
 
 
@@ -323,8 +324,8 @@ def window_expression(
     rows_start: str = "",
     rows_end: str = "",
     *,
-    descending: bool = False,
-    nulls_last: bool = False,
+    descending: Sequence[bool] | None = None,
+    nulls_last: Sequence[bool] | None = None,
     ignore_nulls: bool = False,
 ) -> Expression:
     # TODO(unassigned): Replace with `duckdb.WindowExpression` when they release it.
@@ -335,6 +336,8 @@ def window_expression(
         msg = f"DuckDB>=1.3.0 is required for this operation. Found: DuckDB {duckdb.__version__}"
         raise NotImplementedError(msg) from exc
     pb = generate_partition_by_sql(*partition_by)
+    descending = descending or [False] * len(order_by)
+    nulls_last = nulls_last or [False] * len(order_by)
     ob = generate_order_by_sql(*order_by, descending=descending, nulls_last=nulls_last)
 
     if rows_start and rows_end:

narwhals/_expression_parsing.py

@@ -131,8 +131,8 @@ class ExprKind(Enum):
     ORDERABLE_WINDOW = auto()
     """Depends on the rows around it and on their order, e.g. `diff`."""
 
-    UNORDERABLE_WINDOW = auto()
-    """Depends on the rows around it but not on their order, e.g. `rank`."""
+    WINDOW = auto()
+    """Depends on the rows around it and possibly their order, e.g. `rank`."""
 
     FILTRATION = auto()
     """Changes length, not affected by row order, e.g. `drop_nulls`."""
@@ -217,6 +217,22 @@ def __and__(self, other: ExpansionKind) -> Literal[ExpansionKind.MULTI_UNNAMED]:
 
 
 class ExprMetadata:
+    """Expression metadata.
+
+    Parameters:
+        expansion_kind: What kind of expansion the expression performs.
+        has_windows: Whether it already contains window functions.
+        is_elementwise: Whether it can operate row-by-row without context
+            of the other rows around it.
+        is_literal: Whether it is just a literal wrapped in an expression.
+        is_scalar_like: Whether it is a literal or an aggregation.
+        last_node: The ExprKind of the last node.
+        n_orderable_ops: The number of order-dependent operations. In the
+            lazy case, this number must be `0` by the time the expression
+            is evaluated.
+        preserves_length: Whether the expression preserves the input length.
+    """
+
     __slots__ = (
         "expansion_kind",
         "has_windows",
@@ -317,14 +333,17 @@ def with_elementwise_op(self) -> ExprMetadata:
             is_literal=self.is_literal,
         )
 
-    def with_unorderable_window(self) -> ExprMetadata:
+    def with_window(self) -> ExprMetadata:
+        # Window function which may (but doesn't have to) be used with `over(order_by=...)`.
         if self.is_scalar_like:
-            msg = "Can't apply unorderable window (`rank`, `is_unique`) to scalar-like expression."
+            msg = "Can't apply window (e.g. `rank`) to scalar-like expression."
             raise InvalidOperationError(msg)
         return ExprMetadata(
             self.expansion_kind,
-            ExprKind.UNORDERABLE_WINDOW,
+            ExprKind.WINDOW,
             has_windows=self.has_windows,
+            # The function isn't order-dependent (but, users can still use `order_by` if they wish!),
+            # so we don't increment `n_orderable_ops`.
             n_orderable_ops=self.n_orderable_ops,
             preserves_length=self.preserves_length,
             is_elementwise=False,
@@ -333,6 +352,7 @@ def with_unorderable_window(self) -> ExprMetadata:
         )
 
     def with_orderable_window(self) -> ExprMetadata:
+        # Window function which must be used with `over(order_by=...)`.
         if self.is_scalar_like:
             msg = "Can't apply orderable window (e.g. `diff`, `shift`) to scalar-like expression."
             raise InvalidOperationError(msg)
@@ -358,8 +378,16 @@ def with_ordered_over(self) -> ExprMetadata:
             )
             raise InvalidOperationError(msg)
         n_orderable_ops = self.n_orderable_ops
-        if not n_orderable_ops:
-            msg = "Cannot use `order_by` in `over` on expression which isn't orderable."
+        if not n_orderable_ops and self.last_node is not ExprKind.WINDOW:
+            msg = (
+                "Cannot use `order_by` in `over` on expression which isn't orderable.\n"
+                "If your expression is orderable, then make sure that `over(order_by=...)`\n"
+                "comes immediately after the order-dependent expression.\n\n"
+                "Hint: instead of\n"
+                "  - `(nw.col('price').diff() + 1).over(order_by='date')`\n"
+                "write:\n"
+                "  + `nw.col('price').diff().over(order_by='date') + 1`\n"
+            )
             raise InvalidOperationError(msg)
         if self.last_node.is_orderable_window:
             n_orderable_ops -= 1

narwhals/_ibis/expr.py

@@ -65,9 +65,13 @@ def window_function(self) -> IbisWindowFunction:
         def default_window_func(
             df: IbisLazyFrame, window_inputs: IbisWindowInputs
         ) -> list[ir.Value]:
-            assert not window_inputs.order_by  # noqa: S101
             return [
-                expr.over(ibis.window(group_by=window_inputs.partition_by))
+                expr.over(
+                    ibis.window(
+                        group_by=window_inputs.partition_by,
+                        order_by=self._sort(*window_inputs.order_by),
+                    )
+                )
                 for expr in self(df)
             ]
 

narwhals/_spark_like/dataframe.py

@@ -509,6 +509,9 @@ def unpivot(
         return self._with_native(unpivoted_native_frame)
 
     def with_row_index(self, name: str, order_by: Sequence[str]) -> Self:
+        if order_by is None:
+            msg = "Cannot pass `order_by` to `with_row_index` for PySpark-like"
+            raise TypeError(msg)
         row_index_expr = (
             self._F.row_number().over(
                 self._Window.partitionBy(self._F.lit(1)).orderBy(*order_by)

narwhals/_spark_like/expr.py

@@ -129,17 +129,24 @@ def _Window(self) -> type[Window]:  # noqa: N802
             return import_window(self._implementation)
 
     def _sort(
-        self, *cols: Column | str, descending: bool = False, nulls_last: bool = False
+        self,
+        *cols: Column | str,
+        descending: Sequence[bool] | None = None,
+        nulls_last: Sequence[bool] | None = None,
     ) -> Iterator[Column]:
         F = self._F  # noqa: N806
+        descending = descending or [False] * len(cols)
+        nulls_last = nulls_last or [False] * len(cols)
         mapping = {
             (False, False): F.asc_nulls_first,
             (False, True): F.asc_nulls_last,
             (True, False): F.desc_nulls_first,
             (True, True): F.desc_nulls_last,
         }
-        sort = mapping[(descending, nulls_last)]
-        yield from (sort(col) for col in cols)
+        yield from (
+            mapping[(_desc, _nulls_last)](col)
+            for col, _desc, _nulls_last in zip(cols, descending, nulls_last)
+        )
 
     def partition_by(self, *cols: Column | str) -> WindowSpec:
         """Wraps `Window().paritionBy`, with default and `WindowInputs` handling."""
@@ -178,7 +185,9 @@ def func(df: SparkLikeLazyFrame, inputs: SparkWindowInputs) -> Sequence[Column]:
             window = (
                 self.partition_by(*inputs.partition_by)
                 .orderBy(
-                    *self._sort(*inputs.order_by, descending=reverse, nulls_last=reverse)
+                    *self._sort(
+                        *inputs.order_by, descending=[reverse], nulls_last=[reverse]
+                    )
                 )
                 .rowsBetween(self._Window.unboundedPreceding, 0)
             )
@@ -695,7 +704,9 @@ def func(df: SparkLikeLazyFrame, inputs: SparkWindowInputs) -> Sequence[Column]:
             return [
                 self._F.row_number().over(
                     self.partition_by(*inputs.partition_by, expr).orderBy(
-                        *self._sort(*inputs.order_by, descending=True, nulls_last=True)
+                        *self._sort(
+                            *inputs.order_by, descending=[True], nulls_last=[True]
+                        )
                     )
                 )
                 == 1
@@ -823,17 +834,17 @@ def rank(self, method: RankMethod, *, descending: bool) -> Self:
 
         def _rank(
             expr: Column,
+            partition_by: Sequence[str | Column] = (),
+            order_by: Sequence[str | Column] = (),
             *,
-            descending: bool,
-            partition_by: Sequence[str | Column] | None = None,
+            descending: Sequence[bool],
+            nulls_last: Sequence[bool],
         ) -> Column:
-            order_by = self._sort(expr, descending=descending, nulls_last=True)
-            if partition_by is not None:
-                window = self.partition_by(*partition_by).orderBy(*order_by)
-                count_window = self.partition_by(*partition_by, expr)
-            else:
-                window = self.partition_by().orderBy(*order_by)
-                count_window = self.partition_by(expr)
+            _order_by = self._sort(
+                expr, *order_by, descending=descending, nulls_last=nulls_last
+            )
+            window = self.partition_by(*partition_by).orderBy(*_order_by)
+            count_window = self.partition_by(*partition_by, expr)
             if method == "max":
                 rank_expr = (
                     getattr(self._F, func_name)().over(window)
@@ -852,14 +863,21 @@ def _rank(
             return self._F.when(expr.isNotNull(), rank_expr)
 
         def _unpartitioned_rank(expr: Column) -> Column:
-            return _rank(expr, descending=descending)
+            return _rank(expr, descending=[descending], nulls_last=[True])
 
         def _partitioned_rank(
             df: SparkLikeLazyFrame, inputs: SparkWindowInputs
         ) -> Sequence[Column]:
-            assert not inputs.order_by  # noqa: S101
+            # node: when `descending` / `nulls_last` are supported in `.over`, they should be respected here
+            # https://github.com/narwhals-dev/narwhals/issues/2790
             return [
-                _rank(expr, descending=descending, partition_by=inputs.partition_by)
+                _rank(
+                    expr,
+                    inputs.partition_by,
+                    inputs.order_by,
+                    descending=[descending] + [False] * len(inputs.order_by),
+                    nulls_last=[True] + [False] * len(inputs.order_by),
+                )
                 for expr in self(df)
             ]
 

narwhals/expr.py

@@ -79,8 +79,8 @@ def _with_orderable_aggregation(
     def _with_orderable_window(self, to_compliant_expr: Callable[[Any], Any]) -> Self:
         return self.__class__(to_compliant_expr, self._metadata.with_orderable_window())
 
-    def _with_unorderable_window(self, to_compliant_expr: Callable[[Any], Any]) -> Self:
-        return self.__class__(to_compliant_expr, self._metadata.with_unorderable_window())
+    def _with_window(self, to_compliant_expr: Callable[[Any], Any]) -> Self:
+        return self.__class__(to_compliant_expr, self._metadata.with_window())
 
     def _with_filtration(self, to_compliant_expr: Callable[[Any], Any]) -> Self:
         return self.__class__(to_compliant_expr, self._metadata.with_filtration())
@@ -1646,9 +1646,7 @@ def is_unique(self) -> Self:
             |3  1  c        False         True|
             └─────────────────────────────────┘
         """
-        return self._with_unorderable_window(
-            lambda plx: self._to_compliant_expr(plx).is_unique()
-        )
+        return self._with_window(lambda plx: self._to_compliant_expr(plx).is_unique())
 
     def null_count(self) -> Self:
         r"""Count null values.
@@ -2472,7 +2470,7 @@ def rank(self, method: RankMethod = "average", *, descending: bool = False) -> S
             )
             raise ValueError(msg)
 
-        return self._with_unorderable_window(
+        return self._with_window(
             lambda plx: self._to_compliant_expr(plx).rank(
                 method=method, descending=descending
             )

narwhals/stable/v1/__init__.py

@@ -405,7 +405,7 @@ def sort(self, *, descending: bool = False, nulls_last: bool = False) -> Self:
         Returns:
             A new expression.
         """
-        return self._with_unorderable_window(
+        return self._with_window(
             lambda plx: self._to_compliant_expr(plx).sort(
                 descending=descending, nulls_last=nulls_last
             )