Repository: dask/dask
File: dask/dataframe/backends.py (Python)

Code before:

@group_split_dispatch.register((pd.DataFrame, pd.Series, pd.Index))
def group_split_pandas(df, c, k, ignore_index=False):
    if is_series_like(c):
        c = c.values
    indexer, locations = pd._libs.algos.groupsort_indexer(
        c.astype(np.intp, copy=False), k
    )
    df2 = df.take(indexer)
    locations = locations.cumsum()
    parts = [
        df2.iloc[a:b].reset_index(drop=True) if ignore_index else df2.iloc[a:b]
        for a, b in zip(locations[:-1], locations[1:])
    ]
    return ShuffleGroupResult(zip(range(k), parts))


def _union_categoricals_wrapper(
    dfs: list[pd.CategoricalIndex] | list[pd.Series], **kwargs
) -> pd.Categorical:
    """
    A wrapper around pandas' union_categoricals that handles some dtype issues.

    union_categoricals requires that the dtype of each array's categories match.
    So you can't union ``Categorical(['a', 'b'])`` and ``Categorical([1, 2])``
    since the dtype (str vs. int) doesn't match.

    *Somewhere* in Dask, we're possibly creating an empty ``Categorical``
    with a dtype of ``object``. In pandas 2.x, we could union that with string
    categories since they both used object dtype. But pandas 3.x uses string
    dtype for categories.

    This wrapper handles that by creating a new ``Categorical`` with the
    correct dtype.
    """
    categories_dtypes = {cat.dtype.categories.dtype.name for cat in dfs}
    if "object" in categories_dtypes and "str" in categories_dtypes:
        # TODO: don't drop name?
        dfs = [
            type(cat)(pd.Categorical(pd.Index([], dtype="str")), name=cat.name)
            if cat.dtype.categories.dtype.name == "object" and len(cat) == 0
            else cat
            for cat in dfs
        ]

    return union_categoricals(dfs, **kwargs)


@concat_dispatch.register((pd.DataFrame, pd.Series, pd.Index))
def concat_pandas(
    dfs,

Review comment (suggestion):
I'd like to avoid this if possible. I haven't found whose creating the empty categorical yet though :/

Code after:
def group_split_pandas(df, c, k, ignore_index=False):
    if is_series_like(c):
        c = c.values
    indexer, locations = pd._libs.algos.groupsort_indexer(
        c.astype(np.intp, copy=False), k
    )
    df2 = df.take(indexer)
    locations = locations.cumsum()
    parts = [
        df2.iloc[a:b].reset_index(drop=True) if ignore_index else df2.iloc[a:b]
        for a, b in zip(locations[:-1], locations[1:])
    ]
    return ShuffleGroupResult(zip(range(k), parts))


def _union_categoricals_wrapper(
    dfs: list[pd.CategoricalIndex] | list[pd.Series], **kwargs
) -> pd.Categorical:
    """
    A wrapper around pandas' union_categoricals that handles some dtype issues.

    union_categoricals requires that the dtype of each array's categories match.
    So you can't union ``Categorical(['a', 'b'])`` and ``Categorical([1, 2])``
    since the dtype (str vs. int) doesn't match.

    *Somewhere* in Dask, we're possibly creating an empty ``Categorical``
    with a dtype of ``object``. In pandas 2.x, we could union that with string
    categories since they both used object dtype. But pandas 3.x uses string
    dtype for categories.

    This wrapper handles that by creating a new ``Categorical`` with the
    correct dtype.
    """
    categories_dtypes = {cat.dtype.categories.dtype.name for cat in dfs}
    if "object" in categories_dtypes and "str" in categories_dtypes:
        dfs = [
            (
                type(cat)(pd.Categorical(pd.Index([], dtype="str")), name=cat.name)
                if cat.dtype.categories.dtype.name == "object" and len(cat) == 0
                else cat
            )
            for cat in dfs
        ]

    return union_categoricals(dfs, **kwargs)


@concat_dispatch.register((pd.DataFrame, pd.Series, pd.Index))
def concat_pandas(
    dfs,
    axis=0,