get_categorical_columns

yohou.utils.polars.get_categorical_columns(df, exclude=None)

Get list of categorical column names from a DataFrame.

Categorical columns are those with pl.String, pl.Categorical, or pl.Enum dtype.

Parameters

Name	Type	Description	Default
df	DataFrame	Input DataFrame.	required
exclude	list[str] | None	Column names to exclude from the result.	None

Returns

Type	Description
list[str]	List of categorical column names.

Examples

>>> import polars as pl
>>> from yohou.utils.polars import get_categorical_columns
>>> df = pl.DataFrame({"time": [1, 2, 3], "y": [10.0, 20.0, 30.0], "cat": ["A", "B", "C"]})
>>> get_categorical_columns(df)
['cat']

>>> get_categorical_columns(df, exclude=["cat"])
[]

See Also

• get_numeric_columns : List numeric column names from a DataFrame.
• is_categorical_dtype : Check if a dtype is categorical.

Source Code

View on GitHub

Show/Hide source

def get_categorical_columns(df: pl.DataFrame, exclude: list[str] | None = None) -> list[str]:
    """Get list of categorical column names from a DataFrame.

    Categorical columns are those with ``pl.String``, ``pl.Categorical``,
    or ``pl.Enum`` dtype.

    Parameters
    ----------
    df : pl.DataFrame
        Input DataFrame.
    exclude : list[str] | None, default=None
        Column names to exclude from the result.

    Returns
    -------
    list[str]
        List of categorical column names.

    Examples
    --------
    >>> import polars as pl
    >>> from yohou.utils.polars import get_categorical_columns
    >>> df = pl.DataFrame({"time": [1, 2, 3], "y": [10.0, 20.0, 30.0], "cat": ["A", "B", "C"]})
    >>> get_categorical_columns(df)
    ['cat']

    >>> get_categorical_columns(df, exclude=["cat"])
    []

    See Also
    --------
    - [`get_numeric_columns`][yohou.utils.polars.get_numeric_columns] : List numeric column names from a DataFrame.
    - [`is_categorical_dtype`][yohou.utils.polars.is_categorical_dtype] : Check if a dtype is categorical.

    """
    exclude = exclude or []
    return [col for col in df.columns if is_categorical_dtype(df[col].dtype) and col not in exclude]