Function strip_unknown_categories

dask/dataframe/utils.py:191–217

Replace any unknown categoricals with empty categoricals. Useful for preventing ``UNKNOWN_CATEGORIES`` from leaking into results.

(x, just_drop_unknown=False)

def strip_unknown_categories(x, just_drop_unknown=False):
    """Replace any unknown categoricals with empty categoricals.

    Useful for preventing ``UNKNOWN_CATEGORIES`` from leaking into results.
    """
    if isinstance(x, (pd.Series, pd.DataFrame)):
        x = x.copy()
        if isinstance(x, pd.DataFrame):
            cat_mask = x.dtypes == "category"
            if cat_mask.any():
                cats = cat_mask[cat_mask].index
                for c in cats:
                    if not has_known_categories(x[c]):
                        if just_drop_unknown:
                            x[c].cat.remove_categories(UNKNOWN_CATEGORIES, inplace=True)
                        else:
                            x[c] = x[c].cat.set_categories([])
        elif isinstance(x, pd.Series):
            if isinstance(x.dtype, pd.CategoricalDtype) and not has_known_categories(x):
                x = x.cat.set_categories([])
        if isinstance(x.index, pd.CategoricalIndex) and not has_known_categories(
            x.index
        ):
            x.index = x.index.set_categories([])
    elif isinstance(x, pd.CategoricalIndex) and not has_known_categories(x):
        x = x.set_categories([])
    return x


def clear_known_categories(x, cols=None, index=True, dtype_backend=None):