REF: Simplify Index.union

Author: topper-123

pandas/core/dtypes/cast.py

@@ -68,6 +68,7 @@
     is_numeric_dtype,
     is_object_dtype,
     is_scalar,
+    is_signed_integer_dtype,
     is_string_dtype,
     is_timedelta64_dtype,
     is_unsigned_integer_dtype,
@@ -1779,13 +1780,14 @@ def ensure_nanosecond_dtype(dtype: DtypeObj) -> DtypeObj:
     return dtype
 
 
-def find_common_type(types: list[DtypeObj]) -> DtypeObj:
+def find_common_type(types: list[DtypeObj], *, strict_uint64: bool = False) -> DtypeObj:
     """
     Find a common data type among the given dtypes.
 
     Parameters
     ----------
     types : list of dtypes
+    strict_uint64 : if True, object dtype is returned if uint64 and signed int present.
 
     Returns
     -------
@@ -1831,6 +1833,13 @@ def find_common_type(types: list[DtypeObj]) -> DtypeObj:
             if is_integer_dtype(t) or is_float_dtype(t) or is_complex_dtype(t):
                 return np.dtype("object")
 
+    # Index.union is special: uint64 & signed int -> object
+    if strict_uint64:
+        has_uint64 = any(t == "uint64" for t in types)
+        has_signed_int = any(is_signed_integer_dtype(t) for t in types)
+        if has_uint64 and has_signed_int:
+            return np.dtype("object")
+
     # error: Argument 1 to "find_common_type" has incompatible type
     # "List[Union[dtype, ExtensionDtype]]"; expected "Sequence[Union[dtype,
     # None, type, _SupportsDtype, str, Tuple[Any, int], Tuple[Any, Union[int,

pandas/core/indexes/base.py

@@ -77,7 +77,6 @@
     is_float_dtype,
     is_hashable,
     is_integer,
-    is_integer_dtype,
     is_interval_dtype,
     is_iterator,
     is_list_like,
@@ -2963,19 +2962,13 @@ def union(self, other, sort=None):
                     stacklevel=2,
                 )
 
-            dtype = find_common_type([self.dtype, other.dtype])
-            if self._is_numeric_dtype and other._is_numeric_dtype:
-                # Right now, we treat union(int, float) a bit special.
-                # See https://github.com/pandas-dev/pandas/issues/26778 for discussion
-                # We may change union(int, float) to go to object.
-                # float | [u]int -> float  (the special case)
-                # <T>   | <T>    -> T
-                # <T>   | <U>    -> object
-                if not (is_integer_dtype(self.dtype) and is_integer_dtype(other.dtype)):
-                    dtype = np.dtype("float64")
-                else:
-                    # one is int64 other is uint64
-                    dtype = np.dtype("object")
+            dtype = find_common_type([self.dtype, other.dtype], strict_uint64=True)
+            # Right now, we treat union(float, [u]int) a bit special.
+            # See https://github.com/pandas-dev/pandas/issues/26778 for discussion
+            # Now it's:
+            # * float | [u]int -> float
+            # * uint64 | signed int  -> object
+            # We may change union(float  [u]int) to go to object.
 
             left = self.astype(dtype, copy=False)
             right = other.astype(dtype, copy=False)

pandas/tests/indexes/test_setops.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pytest
 
-from pandas.core.dtypes.common import is_dtype_equal
+from pandas.core.dtypes.cast import find_common_type
 
 from pandas import (
     CategoricalIndex,
@@ -25,6 +25,7 @@
 import pandas._testing as tm
 from pandas.api.types import (
     is_datetime64tz_dtype,
+    is_signed_integer_dtype,
     pandas_dtype,
 )
 
@@ -48,7 +49,11 @@ def test_union_different_types(index_flat, index_flat2):
     idx1 = index_flat
     idx2 = index_flat2
 
-    type_pair = tuple(sorted([idx1.dtype.type, idx2.dtype.type], key=lambda x: str(x)))
+    common_dtype = find_common_type([idx1.dtype, idx2.dtype])
+
+    any_uint64 = idx1.dtype == np.uint64 or idx2.dtype == np.uint64
+    idx1_signed = is_signed_integer_dtype(idx1.dtype)
+    idx2_signed = is_signed_integer_dtype(idx2.dtype)
 
     # Union with a non-unique, non-monotonic index raises error
     # This applies to the boolean index
@@ -58,23 +63,12 @@ def test_union_different_types(index_flat, index_flat2):
     res1 = idx1.union(idx2)
     res2 = idx2.union(idx1)
 
-    if is_dtype_equal(idx1.dtype, idx2.dtype):
-        assert res1.dtype == idx1.dtype
-        assert res2.dtype == idx1.dtype
-
-    elif type_pair not in COMPATIBLE_INCONSISTENT_PAIRS:
-        # A union with a CategoricalIndex (even as dtype('O')) and a
-        # non-CategoricalIndex can only be made if both indices are monotonic.
-        # This is true before this PR as well.
+    if any_uint64 and (idx1_signed or idx2_signed):
         assert res1.dtype == np.dtype("O")
         assert res2.dtype == np.dtype("O")
-
-    elif idx1.dtype.kind in ["f", "i", "u"] and idx2.dtype.kind in ["f", "i", "u"]:
-        assert res1.dtype == np.dtype("f8")
-        assert res2.dtype == np.dtype("f8")
-
     else:
-        raise NotImplementedError
+        assert res1.dtype == common_dtype
+        assert res2.dtype == common_dtype
 
 
 @pytest.mark.parametrize(