Merge pull request #72 from raisadz/feat/top_k

feat: add `top_k`

Author: ym-pett

narwhals_daft/dataframe.py

@@ -11,6 +11,7 @@
     ValidateBackendVersion,
     Version,
     check_column_names_are_unique,
+    extend_bool,
     generate_temporary_column_name,
     not_implemented,
     parse_columns_to_drop,
@@ -26,7 +27,7 @@
 from narwhals_daft.utils import evaluate_exprs, lit, native_to_narwhals_dtype
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Mapping, Sequence
+    from collections.abc import Iterable, Iterator, Mapping, Sequence
     from types import ModuleType
 
     from narwhals._compliant.typing import CompliantDataFrameAny
@@ -229,6 +230,16 @@ def sort(
             )
         )
 
+    def top_k(
+        self, k: int, *, by: Iterable[str], reverse: bool | Sequence[bool]
+    ) -> Self:
+        descending = [not x for x in extend_bool(reverse, len(list(by)))]
+        return self._with_native(
+            self._native_frame.sort(list(by), desc=descending, nulls_first=False).limit(
+                k
+            )
+        )
+
     def drop_nulls(self, subset: Sequence[str] | None) -> Self:
         if subset:
             return self._with_native(self._native_frame.drop_null(*subset))

run_tests.py

@@ -125,8 +125,6 @@
     "test_to_datetime_infer_fmt",
     "test_to_datetime_infer_fmt_from_date",
     "test_to_datetime_tz_aware",
-    "test_top_k",
-    "test_top_k_by_multiple",
     "test_truncate",
     "test_truncate_invalid_interval",
     "test_truncate_invalid_multiple",