diff options
Diffstat (limited to 'venv/lib/python3.8/site-packages/narwhals/expr_cat.py')
-rw-r--r-- | venv/lib/python3.8/site-packages/narwhals/expr_cat.py | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/venv/lib/python3.8/site-packages/narwhals/expr_cat.py b/venv/lib/python3.8/site-packages/narwhals/expr_cat.py new file mode 100644 index 0000000..7a0edc2 --- /dev/null +++ b/venv/lib/python3.8/site-packages/narwhals/expr_cat.py @@ -0,0 +1,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() + ) |