ENH: implement _should_compare/_is_comparable_dtype for all Index subclasses (#38251)

Author: jbrockmendel

pandas/core/indexes/base.py

@@ -4906,16 +4906,31 @@ def get_indexer_non_unique(self, target):
             # Treat boolean labels passed to a numeric index as not found. Without
             # this fix False and True would be treated as 0 and 1 respectively.
             # (GH #16877)
-            no_matches = -1 * np.ones(self.shape, dtype=np.intp)
-            return no_matches, no_matches
+            return self._get_indexer_non_comparable(target, method=None, unique=False)
 
         pself, ptarget = self._maybe_promote(target)
         if pself is not self or ptarget is not target:
             return pself.get_indexer_non_unique(ptarget)
 
-        if not self._is_comparable_dtype(target.dtype):
-            no_matches = -1 * np.ones(self.shape, dtype=np.intp)
-            return no_matches, no_matches
+        if not self._should_compare(target):
+            return self._get_indexer_non_comparable(target, method=None, unique=False)
+
+        if not is_dtype_equal(self.dtype, target.dtype):
+            # TODO: if object, could use infer_dtype to pre-empt costly
+            #  conversion if still non-comparable?
+            dtype = find_common_type([self.dtype, target.dtype])
+            if (
+                dtype.kind in ["i", "u"]
+                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)
+
+            this = self.astype(dtype, copy=False)
+            that = target.astype(dtype, copy=False)
+            return this.get_indexer_non_unique(that)
 
         if is_categorical_dtype(target.dtype):
             tgt_values = np.asarray(target)
@@ -4968,7 +4983,7 @@ def _get_indexer_non_comparable(self, target: "Index", method, unique: bool = Tr
             If doing an inequality check, i.e. method is not None.
         """
         if method is not None:
-            other = _unpack_nested_dtype(target)
+            other = unpack_nested_dtype(target)
             raise TypeError(f"Cannot compare dtypes {self.dtype} and {other.dtype}")
 
         no_matches = -1 * np.ones(target.shape, dtype=np.intp)
@@ -5019,7 +5034,7 @@ def _should_compare(self, other: "Index") -> bool:
         """
         Check if `self == other` can ever have non-False entries.
         """
-        other = _unpack_nested_dtype(other)
+        other = unpack_nested_dtype(other)
         dtype = other.dtype
         return self._is_comparable_dtype(dtype) or is_object_dtype(dtype)
 
@@ -6172,7 +6187,7 @@ def get_unanimous_names(*indexes: Index) -> Tuple[Label, ...]:
     return names
 
 
-def _unpack_nested_dtype(other: Index) -> Index:
+def unpack_nested_dtype(other: Index) -> Index:
     """
     When checking if our dtype is comparable with another, we need
     to unpack CategoricalDtype to look at its categories.dtype.

pandas/core/indexes/category.py

@@ -554,6 +554,9 @@ def _maybe_cast_slice_bound(self, label, side: str, kind):
 
     # --------------------------------------------------------------------
 
+    def _is_comparable_dtype(self, dtype):
+        return self.categories._is_comparable_dtype(dtype)
+
     def take_nd(self, *args, **kwargs):
         """Alias for `take`"""
         warnings.warn(

pandas/core/indexes/interval.py

@@ -11,7 +11,7 @@
 from pandas._libs import lib
 from pandas._libs.interval import Interval, IntervalMixin, IntervalTree
 from pandas._libs.tslibs import BaseOffset, Timedelta, Timestamp, to_offset
-from pandas._typing import AnyArrayLike, Label
+from pandas._typing import AnyArrayLike, DtypeObj, Label
 from pandas.errors import InvalidIndexError
 from pandas.util._decorators import Appender, Substitution, cache_readonly
 from pandas.util._exceptions import rewrite_exception
@@ -38,6 +38,7 @@
     is_object_dtype,
     is_scalar,
 )
+from pandas.core.dtypes.dtypes import IntervalDtype
 
 from pandas.core.algorithms import take_1d
 from pandas.core.arrays.interval import IntervalArray, _interval_shared_docs
@@ -50,6 +51,7 @@
     default_pprint,
     ensure_index,
     maybe_extract_name,
+    unpack_nested_dtype,
 )
 from pandas.core.indexes.datetimes import DatetimeIndex, date_range
 from pandas.core.indexes.extension import ExtensionIndex, inherit_names
@@ -803,15 +805,27 @@ def _convert_list_indexer(self, keyarr):
 
         return locs
 
+    def _is_comparable_dtype(self, dtype: DtypeObj) -> bool:
+        if not isinstance(dtype, IntervalDtype):
+            return False
+        common_subtype = find_common_type([self.dtype.subtype, dtype.subtype])
+        return not is_object_dtype(common_subtype)
+
+    def _should_compare(self, other) -> bool:
+        if not super()._should_compare(other):
+            return False
+        other = unpack_nested_dtype(other)
+        return other.closed == self.closed
+
+    # TODO: use should_compare and get rid of _is_non_comparable_own_type
     def _is_non_comparable_own_type(self, other: "IntervalIndex") -> bool:
         # different closed or incompatible subtype -> no matches
 
         # TODO: once closed is part of IntervalDtype, we can just define
         #  is_comparable_dtype GH#19371
         if self.closed != other.closed:
             return True
-        common_subtype = find_common_type([self.dtype.subtype, other.dtype.subtype])
-        return is_object_dtype(common_subtype)
+        return not self._is_comparable_dtype(other.dtype)
 
     # --------------------------------------------------------------------
 

pandas/core/indexes/multi.py

@@ -20,7 +20,7 @@
 
 from pandas._libs import algos as libalgos, index as libindex, lib
 from pandas._libs.hashtable import duplicated_int64
-from pandas._typing import AnyArrayLike, Label, Scalar, Shape
+from pandas._typing import AnyArrayLike, DtypeObj, Label, Scalar, Shape
 from pandas.compat.numpy import function as nv
 from pandas.errors import InvalidIndexError, PerformanceWarning, UnsortedIndexError
 from pandas.util._decorators import Appender, cache_readonly, doc
@@ -3583,6 +3583,9 @@ def union(self, other, sort=None):
             zip(*uniq_tuples), sortorder=0, names=result_names
         )
 
+    def _is_comparable_dtype(self, dtype: DtypeObj) -> bool:
+        return is_object_dtype(dtype)
+
     def intersection(self, other, sort=False):
         """
         Form the intersection of two MultiIndex objects.
@@ -3618,15 +3621,9 @@ def intersection(self, other, sort=False):
     def _intersection(self, other, sort=False):
         other, result_names = self._convert_can_do_setop(other)
 
-        if not is_object_dtype(other.dtype):
+        if not self._is_comparable_dtype(other.dtype):
             # The intersection is empty
-            # TODO: we have no tests that get here
-            return MultiIndex(
-                levels=self.levels,
-                codes=[[]] * self.nlevels,
-                names=result_names,
-                verify_integrity=False,
-            )
+            return self[:0].rename(result_names)
 
         lvals = self._values
         rvals = other._values

pandas/core/indexes/numeric.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from pandas._libs import index as libindex, lib
-from pandas._typing import Dtype, Label
+from pandas._typing import Dtype, DtypeObj, Label
 from pandas.util._decorators import doc
 
 from pandas.core.dtypes.cast import astype_nansafe
@@ -148,6 +148,10 @@ def _convert_tolerance(self, tolerance, target):
                 )
         return tolerance
 
+    def _is_comparable_dtype(self, dtype: DtypeObj) -> bool:
+        # If we ever have BoolIndex or ComplexIndex, this may need to be tightened
+        return is_numeric_dtype(dtype)
+
     @classmethod
     def _assert_safe_casting(cls, data, subarr):
         """

pandas/tests/indexes/test_base.py

@@ -1251,10 +1251,9 @@ def test_get_indexer_numeric_index_boolean_target(self, method, idx_class):
         if method == "get_indexer":
             tm.assert_numpy_array_equal(result, expected)
         else:
-            expected = np.array([-1, -1, -1, -1], dtype=np.intp)
-
+            missing = np.arange(3, dtype=np.intp)
             tm.assert_numpy_array_equal(result[0], expected)
-            tm.assert_numpy_array_equal(result[1], expected)
+            tm.assert_numpy_array_equal(result[1], missing)
 
     def test_get_indexer_with_NA_values(
         self, unique_nulls_fixture, unique_nulls_fixture2
@@ -2359,5 +2358,6 @@ def construct(dtype):
 
     else:
         no_matches = np.array([-1] * 6, dtype=np.intp)
+        missing = np.arange(6, dtype=np.intp)
         tm.assert_numpy_array_equal(result[0], no_matches)
-        tm.assert_numpy_array_equal(result[1], no_matches)
+        tm.assert_numpy_array_equal(result[1], missing)