Skip to content

Commit

Permalink
chore(python): split Expr.top_k and Expr.top_k_by into separate
Browse files Browse the repository at this point in the history 
functions
  • Loading branch information
@MarcoGorelli
MarcoGorelli committed May 3, 2024
1 parent 19c46f5 commit 9629096
Show file tree
Hide file tree
Showing 5 changed files with 166 additions and 128 deletions.
2 changes: 2 additions & 0 deletions py-polars/docs/source/reference/expressions/modify_select.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ Manipulation/selection
Expr.arg_true
Expr.backward_fill
Expr.bottom_k
Expr.bottom_k_by
Expr.cast
Expr.ceil
Expr.clip
Expand Down Expand Up @@ -61,5 +62,6 @@ Manipulation/selection
Expr.take_every
Expr.to_physical
Expr.top_k
Expr.top_k_by
Expr.upper_bound
Expr.where
212 changes: 136 additions & 76 deletions py-polars/polars/expr/expr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2035,8 +2035,6 @@ def top_k(
self,
k: int | IntoExprColumn = 5,
*,
by: IntoExpr | Iterable[IntoExpr] | None = None,
descending: bool | Sequence[bool] = False,
nulls_last: bool = False,
maintain_order: bool = False,
multithreaded: bool = True,
Expand All @@ -2046,19 +2044,12 @@ def top_k(
This has time complexity:
.. math:: O(n + k \log{}n - \frac{k}{2})
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
k
Number of elements to return.
by
Column(s) included in sort order. Accepts expression input.
Strings are parsed as column names.
If not provided, each column will be treated induvidually.
descending
Return the k smallest. Top-k by multiple columns can be specified per
column by passing a sequence of booleans.
nulls_last
Place null values last.
maintain_order
Expand All @@ -2068,7 +2059,9 @@ def top_k(
See Also
--------
top_k_by
bottom_k
bottom_k_by
Examples
--------
Expand Down Expand Up @@ -2097,15 +2090,60 @@ def top_k(
│ 3 ┆ 4 │
│ 2 ┆ 98 │
└───────┴──────────┘
"""
k = parse_as_expression(k)
return self._from_pyexpr(self._pyexpr.top_k(k, nulls_last, multithreaded))

def top_k_by(
self,
by: IntoExpr | Iterable[IntoExpr],
k: int | IntoExprColumn = 5,
*,
descending: bool | Sequence[bool] = False,
nulls_last: bool = False,
maintain_order: bool = False,
multithreaded: bool = True,
) -> Self:
r"""
Return elements corresponding to the `k` largest elements of the `by` column(s).
>>> df2 = pl.DataFrame(
This has time complexity:
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
by
Column(s) included in sort order. Accepts expression input.
Strings are parsed as column names.
k
Number of elements to return.
descending
Return the k smallest. Top-k by multiple columns can be specified per
column by passing a sequence of booleans.
nulls_last
Place null values last.
maintain_order
Whether the order should be maintained if elements are equal.
multithreaded
Sort using multiple threads.
See Also
--------
top_k
bottom_k
bottom_k_by
Examples
--------
>>> df = pl.DataFrame(
... {
... "a": [1, 2, 3, 4, 5, 6],
... "b": [6, 5, 4, 3, 2, 1],
... "c": ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"],
... }
... )
>>> df2
>>> df
shape: (6, 3)
┌─────┬─────┬────────┐
│ a ┆ b ┆ c │
Expand All @@ -2122,9 +2160,9 @@ def top_k(
Get the top 2 rows by column `a` or `b`.
>>> df2.select(
... pl.all().top_k(2, by="a").name.suffix("_top_by_a"),
... pl.all().top_k(2, by="b").name.suffix("_top_by_b"),
>>> df.select(
... pl.all().top_k_by("a", 2).name.suffix("_top_by_a"),
... pl.all().top_k_by("b", 2).name.suffix("_top_by_b"),
... )
shape: (2, 6)
┌────────────┬────────────┬────────────┬────────────┬────────────┬────────────┐
Expand All @@ -2138,12 +2176,12 @@ def top_k(
Get the top 2 rows by multiple columns with given order.
>>> df2.select(
>>> df.select(
... pl.all()
... .top_k(2, by=["c", "a"], descending=[False, True])
... .top_k_by(["c", "a"], 2, descending=[False, True])
... .name.suffix("_by_ca"),
... pl.all()
... .top_k(2, by=["c", "b"], descending=[False, True])
... .top_k_by(["c", "b"], 2, descending=[False, True])
... .name.suffix("_by_cb"),
... )
shape: (2, 6)
Expand All @@ -2159,8 +2197,8 @@ def top_k(
Get the top 2 rows by column `a` in each group.
>>> (
... df2.group_by("c", maintain_order=True)
... .agg(pl.all().top_k(2, by="a"))
... df.group_by("c", maintain_order=True)
... .agg(pl.all().top_k_by("a", 2))
... .explode(pl.all().exclude("c"))
... )
shape: (5, 3)
Expand All @@ -2177,32 +2215,22 @@ def top_k(
└────────┴─────┴─────┘
"""
k = parse_as_expression(k)
if by is not None:
by = parse_as_list_of_expressions(by)
if isinstance(descending, bool):
descending = [descending]
elif len(by) != len(descending):
msg = f"the length of `descending` ({len(descending)}) does not match the length of `by` ({len(by)})"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.top_k_by(
k, by, descending, nulls_last, maintain_order, multithreaded
)
)
else:
if not isinstance(descending, bool):
msg = "`descending` should be a boolean if no `by` is provided"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.top_k(k, descending, nulls_last, multithreaded)
by = parse_as_list_of_expressions(by)
if isinstance(descending, bool):
descending = [descending]
elif len(by) != len(descending):
msg = f"the length of `descending` ({len(descending)}) does not match the length of `by` ({len(by)})"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.top_k_by(
k, by, descending, nulls_last, maintain_order, multithreaded
)
)

def bottom_k(
self,
k: int | IntoExprColumn = 5,
*,
by: IntoExpr | Iterable[IntoExpr] | None = None,
descending: bool | Sequence[bool] = False,
nulls_last: bool = False,
maintain_order: bool = False,
multithreaded: bool = True,
Expand All @@ -2212,19 +2240,12 @@ def bottom_k(
This has time complexity:
.. math:: O(n + k \log{}n - \frac{k}{2})
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
k
Number of elements to return.
by
Column(s) included in sort order.
Accepts expression input. Strings are parsed as column names.
If not provided, each column will be treated induvidually.
descending
Return the k largest. Bottom-k by multiple columns can be specified per
column by passing a sequence of booleans.
nulls_last
Place null values last.
maintain_order
Expand All @@ -2235,6 +2256,8 @@ def bottom_k(
See Also
--------
top_k
top_k_by
bottom_k_by
Examples
--------
Expand All @@ -2261,15 +2284,60 @@ def bottom_k(
│ 3 ┆ 4 │
│ 2 ┆ 98 │
└───────┴──────────┘
"""
k = parse_as_expression(k)
return self._from_pyexpr(self._pyexpr.bottom_k(k, nulls_last, multithreaded))

def bottom_k_by(
self,
by: IntoExpr | Iterable[IntoExpr],
k: int | IntoExprColumn = 5,
*,
descending: bool | Sequence[bool] = False,
nulls_last: bool = False,
maintain_order: bool = False,
multithreaded: bool = True,
) -> Self:
r"""
Return elements corresponding to the `k` smallest elements of `by` column(s).
>>> df2 = pl.DataFrame(
This has time complexity:
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
by
Column(s) included in sort order.
Accepts expression input. Strings are parsed as column names.
k
Number of elements to return.
descending
Return the k largest. Bottom-k by multiple columns can be specified per
column by passing a sequence of booleans.
nulls_last
Place null values last.
maintain_order
Whether the order should be maintained if elements are equal.
multithreaded
Sort using multiple threads.
See Also
--------
top_k
top_k_by
bottom_k
Examples
--------
>>> df = pl.DataFrame(
... {
... "a": [1, 2, 3, 4, 5, 6],
... "b": [6, 5, 4, 3, 2, 1],
... "c": ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"],
... }
... )
>>> df2
>>> df
shape: (6, 3)
┌─────┬─────┬────────┐
│ a ┆ b ┆ c │
Expand All @@ -2286,9 +2354,9 @@ def bottom_k(
Get the bottom 2 rows by column `a` or `b`.
>>> df2.select(
... pl.all().bottom_k(2, by="a").name.suffix("_btm_by_a"),
... pl.all().bottom_k(2, by="b").name.suffix("_btm_by_b"),
>>> df.select(
... pl.all().bottom_k_by("a", 2).name.suffix("_btm_by_a"),
... pl.all().bottom_k_by("b", 2).name.suffix("_btm_by_b"),
... )
shape: (2, 6)
┌────────────┬────────────┬────────────┬────────────┬────────────┬────────────┐
Expand All @@ -2302,12 +2370,12 @@ def bottom_k(
Get the bottom 2 rows by multiple columns with given order.
>>> df2.select(
>>> df.select(
... pl.all()
... .bottom_k(2, by=["c", "a"], descending=[False, True])
... .bottom_k_by(["c", "a"], 2, descending=[False, True])
... .name.suffix("_by_ca"),
... pl.all()
... .bottom_k(2, by=["c", "b"], descending=[False, True])
... .bottom_k_by(["c", "b"], 2, descending=[False, True])
... .name.suffix("_by_cb"),
... )
shape: (2, 6)
Expand All @@ -2323,8 +2391,8 @@ def bottom_k(
Get the bottom 2 rows by column `a` in each group.
>>> (
... df2.group_by("c", maintain_order=True)
... .agg(pl.all().bottom_k(2, by="a"))
... df.group_by("c", maintain_order=True)
... .agg(pl.all().bottom_k_by("a", 2))
... .explode(pl.all().exclude("c"))
... )
shape: (5, 3)
Expand All @@ -2341,25 +2409,17 @@ def bottom_k(
└────────┴─────┴─────┘
"""
k = parse_as_expression(k)
if by is not None:
by = parse_as_list_of_expressions(by)
if isinstance(descending, bool):
descending = [descending]
elif len(by) != len(descending):
msg = f"the length of `descending` ({len(descending)}) does not match the length of `by` ({len(by)})"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.bottom_k_by(
k, by, descending, nulls_last, maintain_order, multithreaded
)
)
else:
if not isinstance(descending, bool):
msg = "`descending` should be a boolean if no `by` is provided"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.bottom_k(k, descending, nulls_last, multithreaded)
by = parse_as_list_of_expressions(by)
if isinstance(descending, bool):
descending = [descending]
elif len(by) != len(descending):
msg = f"the length of `descending` ({len(descending)}) does not match the length of `by` ({len(by)})"
raise ValueError(msg)
return self._from_pyexpr(
self._pyexpr.bottom_k_by(
k, by, descending, nulls_last, maintain_order, multithreaded
)
)

def arg_sort(self, *, descending: bool = False, nulls_last: bool = False) -> Self:
"""
Expand Down
4 changes: 2 additions & 2 deletions py-polars/polars/series/series.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3402,7 +3402,7 @@ def top_k(self, k: int | IntoExprColumn = 5) -> Series:
This has time complexity:
.. math:: O(n + k \log{}n - \frac{k}{2})
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
Expand Down Expand Up @@ -3432,7 +3432,7 @@ def bottom_k(self, k: int | IntoExprColumn = 5) -> Series:
This has time complexity:
.. math:: O(n + k \log{}n - \frac{k}{2})
.. math:: O(n + k \log{n} - \frac{k}{2})
Parameters
----------
Expand Down
Loading

0 comments on commit 9629096

Please sign in to comment.