Skip to content

BUG: CategoricalIndex.union with nans #45362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jan 16, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.5.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ Bug fixes

Categorical
^^^^^^^^^^^
-
- Bug in :meth:`CategoricalIndex.union` when the index's categories are integer-dtype and the index contains ``NaN`` values incorrectly raising instead of casting to ``float64`` (:issue:`45362`)
-

Datetimelike
Expand Down
41 changes: 41 additions & 0 deletions pandas/core/dtypes/cast.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@
)
from pandas.core.dtypes.generic import (
ABCExtensionArray,
ABCIndex,
ABCSeries,
)
from pandas.core.dtypes.inference import is_list_like
Expand All @@ -93,7 +94,9 @@

if TYPE_CHECKING:

from pandas import Index
from pandas.core.arrays import (
Categorical,
DatetimeArray,
ExtensionArray,
IntervalArray,
Expand Down Expand Up @@ -1470,6 +1473,44 @@ def find_result_type(left: ArrayLike, right: Any) -> DtypeObj:
return new_dtype


def common_dtype_categorical_compat(
objs: list[Index | ArrayLike], dtype: DtypeObj
) -> DtypeObj:
"""
Update the result of find_common_type to account for NAs in a Categorical.

Parameters
----------
objs : list[np.ndarray | ExtensionArray | Index]
dtype : np.dtype or ExtensionDtype

Returns
-------
np.dtype or ExtensionDtype
"""
# GH#38240

# TODO: more generally, could do `not can_hold_na(dtype)`
if isinstance(dtype, np.dtype) and dtype.kind in ["i", "u"]:

for obj in objs:
# We don't want to accientally allow e.g. "categorical" str here
obj_dtype = getattr(obj, "dtype", None)
if isinstance(obj_dtype, CategoricalDtype):
if isinstance(obj, ABCIndex):
# This check may already be cached
hasnas = obj.hasnans
else:
# Categorical
hasnas = cast("Categorical", obj)._hasnans

if hasnas:
# see test_union_int_categorical_with_nan
dtype = np.dtype(np.float64)
break
return dtype


@overload
def find_common_type(types: list[np.dtype]) -> np.dtype:
...
Expand Down
12 changes: 2 additions & 10 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@

from pandas.core.dtypes.cast import (
can_hold_element,
common_dtype_categorical_compat,
find_common_type,
infer_dtype_from,
maybe_cast_pointwise_result,
Expand Down Expand Up @@ -5976,17 +5977,8 @@ def _find_common_type_compat(self, target) -> DtypeObj:
return _dtype_obj

dtype = find_common_type([self.dtype, target_dtype])
dtype = common_dtype_categorical_compat([self, target], dtype)

if dtype.kind in ["i", "u"]:
# TODO: what about reversed with self being categorical?
if (
isinstance(target, Index)
and is_categorical_dtype(target.dtype)
and target.hasnans
):
# FIXME: find_common_type incorrect with Categorical GH#38240
# FIXME: some cases where float64 cast can be lossy?
dtype = np.dtype(np.float64)
if dtype.kind == "c":
dtype = _dtype_obj
return dtype
Expand Down
14 changes: 14 additions & 0 deletions pandas/tests/indexes/test_setops.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,6 +592,20 @@ def test_union_with_duplicate_index_not_subset_and_non_monotonic(cls):
tm.assert_index_equal(result, expected)


def test_union_int_categorical_with_nan():
ci = CategoricalIndex([1, 2, np.nan])
assert ci.categories.dtype.kind == "i"

idx = Index([1, 2])

result = idx.union(ci)
expected = Index([1, 2, np.nan], dtype=np.float64)
tm.assert_index_equal(result, expected)

result = ci.union(idx)
tm.assert_index_equal(result, expected)


class TestSetOpsUnsorted:
# These may eventually belong in a dtype-specific test_setops, or
# parametrized over a more general fixture
Expand Down