Merge pull request #74 from raisadz/feat/rank

feat: add `rank`

Author: MarcoGorelli

narwhals_daft/expr.py

@@ -23,6 +23,7 @@
     from narwhals._compliant.typing import AliasNames, EvalNames, EvalSeries
     from narwhals._utils import Version, _LimitedContext
     from narwhals.dtypes import DType
+    from narwhals.typing import RankMethod
     from typing_extensions import TypeIs
 
     from narwhals_daft.dataframe import DaftLazyFrame
@@ -783,6 +784,67 @@ def func(df: DaftLazyFrame, inputs: WindowInputs) -> Sequence[Expression]:
 
         return self._with_window_function(func)
 
+    def rank(self, method: RankMethod, *, descending: bool) -> DaftExpr:
+        if method in {"min", "max", "average"}:
+            func = F.rank()
+        elif method == "dense":
+            func = F.dense_rank()
+        else:  # method == "ordinal"
+            func = F.row_number()
+
+        def _rank(
+            expr: Expression,
+            partition_by: Sequence[str | Expression] = (),
+            *,
+            descending: bool,
+        ) -> Expression:
+            count_expr = expr.count(mode="all")
+            window_kwargs: dict[str, Any] = {
+                "partition_by": partition_by,
+                "order_by": (expr,),
+                "descending": [descending],
+                "nulls_first": [False],
+            }
+            count_window_kwargs: dict[str, Any] = {
+                "partition_by": (*partition_by, expr)
+            }
+            window = self._window_expression
+            if method == "max":
+                rank_expr = op.sub(
+                    op.add(
+                        window(func, **window_kwargs),
+                        window(count_expr, **count_window_kwargs),
+                    ),
+                    lit(1),
+                )
+            elif method == "average":
+                rank_expr = op.add(
+                    window(func, **window_kwargs),
+                    op.truediv(
+                        op.sub(window(count_expr, **count_window_kwargs), lit(1)),
+                        lit(2.0),
+                    ),
+                )
+            else:
+                rank_expr = window(func, **window_kwargs)
+            return F.when(F.is_null(expr), expr).otherwise(rank_expr)
+
+        def _unpartitioned_rank(expr: Expression) -> Expression:
+            return _rank(expr, descending=descending)
+
+        def _partitioned_rank(
+            df: DaftLazyFrame, inputs: WindowInputs
+        ) -> Sequence[Expression]:
+            if inputs.order_by:
+                msg = "`rank` followed by `over` with `order_by` specified is not supported for Daft."
+                raise NotImplementedError(msg)
+            return [
+                _rank(expr, inputs.partition_by, descending=descending)
+                for expr in self(df)
+            ]
+
+        return self._with_callable(_unpartitioned_rank, _partitioned_rank)
+
     @property
     def name(self) -> ExprNameNamespace:
         return ExprNameNamespace(self)
@@ -796,7 +858,6 @@ def str(self) -> ExprStringNamespace:
     filter = not_implemented()
     ewm_mean = not_implemented()
     kurtosis = not_implemented()
-    rank = not_implemented()
     map_batches = not_implemented()
     median = not_implemented()
     mode = not_implemented()

run_tests.py

@@ -57,8 +57,6 @@
     "test_joinasof_suffix",
     "test_joinasof_time",
     "test_kurtosis_expr",
-    "test_lazy_rank_expr",
-    "test_lazy_rank_expr_desc",
     "test_left_to_right_broadcasting",
     "test_len_expr",
     "test_mapping_key_not_in_expr",
@@ -89,8 +87,6 @@
     "test_parse_weight",
     "test_pipe_expr",
     "test_quantile_expr",
-    "test_rank_expr_in_over_context",
-    "test_rank_expr_in_over_desc",
     "test_rank_with_order_by",
     "test_rank_with_order_by_and_partition_by",
     "test_replace_strict_expr_basic",