Skip to content

get_numeric_columns

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

Get list of numeric column names from a DataFrame.

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 numeric column names.

Examples

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

>>> get_numeric_columns(df, exclude=["time"])
['y']

See Also

  • cast : Cast DataFrame columns according to a schema.

Source Code

View on GitHub

Show/Hide source 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def get_numeric_columns(df: pl.DataFrame, exclude: list[str] | None = None) -> list[str]:
    """Get list of numeric column names from a DataFrame.

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

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

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

    >>> get_numeric_columns(df, exclude=["time"])
    ['y']

    See Also
    --------
    - [`cast`][yohou.utils.polars.cast] : Cast DataFrame columns according to a schema.

    """
    exclude = exclude or []
    numeric_types = [
        pl.Int8,
        pl.Int16,
        pl.Int32,
        pl.Int64,
        pl.UInt8,
        pl.UInt16,
        pl.UInt32,
        pl.UInt64,
        pl.Float32,
        pl.Float64,
    ]
    return [col for col in df.columns if any(df[col].dtype == dtype for dtype in numeric_types) and col not in exclude]