blob: 7a0edc2449732704027d2d58d605f9639e34848e (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
|
from __future__ import annotations
from typing import TYPE_CHECKING, Generic, TypeVar
if TYPE_CHECKING:
from narwhals.expr import Expr
ExprT = TypeVar("ExprT", bound="Expr")
class ExprCatNamespace(Generic[ExprT]):
def __init__(self, expr: ExprT) -> None:
self._expr = expr
def get_categories(self) -> ExprT:
"""Get unique categories from column.
Returns:
A new expression.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame(
... {"fruits": ["apple", "mango", "mango"]},
... schema={"fruits": pl.Categorical},
... )
>>> df = nw.from_native(df_native)
>>> df.select(nw.col("fruits").cat.get_categories()).to_native()
shape: (2, 1)
┌────────┐
│ fruits │
│ --- │
│ str │
╞════════╡
│ apple │
│ mango │
└────────┘
"""
return self._expr._with_elementwise_op(
lambda plx: self._expr._to_compliant_expr(plx).cat.get_categories()
)
|