fix(python): Add {top, bottom}_k_by to Series (#22902)

Co-authored-by: TareqKomboz <[email protected]>

Author: coastalwhite

py-polars/polars/series/series.py

@@ -3381,7 +3381,9 @@ def top_k(self, k: int = 5) -> Series:
 
         See Also
         --------
+        top_k_by
         bottom_k
+        bottom_k_by
 
         Examples
         --------
@@ -3396,6 +3398,56 @@ def top_k(self, k: int = 5) -> Series:
         ]
         """
 
+    def top_k_by(
+        self,
+        by: IntoExpr | Iterable[IntoExpr],
+        k: int = 5,
+        *,
+        reverse: bool | Sequence[bool] = False,
+    ) -> Series:
+        r"""
+        Return the `k` largest elements of the `by` column.
+
+        Non-null elements are always preferred over null elements, regardless of
+        the value of `reverse`. The output is not guaranteed to be in any
+        particular order, call :func:`sort` after this function if you wish the
+        output to be sorted.
+
+        This has time complexity:
+
+        .. math:: O(n \log{n})
+
+        Parameters
+        ----------
+        by
+            Column used to determine the largest elements.
+            Accepts expression input. Strings are parsed as column names.
+        k
+            Number of elements to return.
+        reverse
+            Consider the `k` smallest elements of the `by` column (instead of the `k`
+            largest). This can be specified per column by passing a sequence of
+            booleans.
+
+        See Also
+        --------
+        top_k
+        bottom_k
+        bottom_k_by
+
+        Examples
+        --------
+        >>> s = pl.Series("a", [2, 5, 1, 4, 3])
+        >>> s.top_k_by("a", 3)
+        shape: (3,)
+        Series: 'a' [i64]
+        [
+            5
+            4
+            3
+        ]
+        """
+
     def bottom_k(self, k: int = 5) -> Series:
         r"""
         Return the `k` smallest elements.
@@ -3416,6 +3468,8 @@ def bottom_k(self, k: int = 5) -> Series:
         See Also
         --------
         top_k
+        top_k_by
+        bottom_k_by
 
         Examples
         --------
@@ -3430,6 +3484,56 @@ def bottom_k(self, k: int = 5) -> Series:
         ]
         """
 
+    def bottom_k_by(
+        self,
+        by: IntoExpr | Iterable[IntoExpr],
+        k: int = 5,
+        *,
+        reverse: bool | Sequence[bool] = False,
+    ) -> Series:
+        r"""
+        Return the `k` smallest elements of the `by` column.
+
+        Non-null elements are always preferred over null elements, regardless of
+        the value of `reverse`. The output is not guaranteed to be in any
+        particular order, call :func:`sort` after this function if you wish the
+        output to be sorted.
+
+        This has time complexity:
+
+        .. math:: O(n \log{n})
+
+        Parameters
+        ----------
+        by
+            Column used to determine the smallest elements.
+            Accepts expression input. Strings are parsed as column names.
+        k
+            Number of elements to return.
+        reverse
+            Consider the `k` largest elements of the `by` column( (instead of the `k`
+            smallest). This can be specified per column by passing a sequence of
+            booleans.
+
+        See Also
+        --------
+        top_k
+        top_k_by
+        bottom_k
+
+        Examples
+        --------
+        >>> s = pl.Series("a", [2, 5, 1, 4, 3])
+        >>> s.bottom_k_by("a", 3)
+        shape: (3,)
+        Series: 'a' [i64]
+        [
+            1
+            2
+            3
+        ]
+        """
+
     def arg_sort(self, *, descending: bool = False, nulls_last: bool = False) -> Series:
         """
         Get the index values that would sort this Series.

py-polars/tests/unit/operations/test_top_k.py

@@ -544,3 +544,21 @@ def test_top_k_sorted_21260() -> None:
     assert s.bottom_k(3).sort().to_list() == [1, 2, 3]
     assert s.sort(descending=False).bottom_k(3).sort().to_list() == [1, 2, 3]
     assert s.sort(descending=True).bottom_k(3).sort().to_list() == [1, 2, 3]
+
+
+def test_top_k_by() -> None:
+    # expression
+    s = pl.Series("a", [3, 8, 1, 5, 2])
+
+    assert_series_equal(
+        s.top_k_by("a", 3), pl.Series("a", [8, 5, 3]), check_order=False
+    )
+
+
+def test_bottom_k_by() -> None:
+    # expression
+    s = pl.Series("a", [3, 8, 1, 5, 2])
+
+    assert_series_equal(
+        s.bottom_k_by("a", 4), pl.Series("a", [3, 2, 1, 5]), check_order=False
+    )