BUG: get_indexer_non_unique with np.datetime64("NaT") and np.timedelta64("NaT")

doc/source/whatsnew/v1.4.0.rst

@@ -447,6 +447,7 @@ Indexing
 - Bug in :meth:`DataFrame.nlargest` and :meth:`Series.nlargest` where sorted result did not count indexes containing ``np.nan`` (:issue:`28984`)
 - Bug in indexing on a non-unique object-dtype :class:`Index` with an NA scalar (e.g. ``np.nan``) (:issue:`43711`)
 - Bug in :meth:`Series.__setitem__` with object dtype when setting an array with matching size and dtype='datetime64[ns]' or dtype='timedelta64[ns]' incorrectly converting the datetime/timedeltas to integers (:issue:`43868`)
+- Bug in :meth:`Index.get_indexer_non_unique` when index contains multiple ``np.datetime64("NaT")`` and ``np.timedelta64("NaT")`` (:issue:`43869`)
 -
 
 Missing

pandas/_libs/index.pyx

@@ -23,7 +23,11 @@ cnp.import_array()
 
 from pandas._libs cimport util
 from pandas._libs.hashtable cimport HashTable
-from pandas._libs.tslibs.nattype cimport c_NaT as NaT
+from pandas._libs.tslibs.nattype cimport (
+    c_NaT as NaT,
+    is_dt64nat,
+    is_td64nat,
+)
 from pandas._libs.tslibs.period cimport is_period_object
 from pandas._libs.tslibs.timedeltas cimport _Timedelta
 from pandas._libs.tslibs.timestamps cimport _Timestamp
@@ -319,10 +323,14 @@ cdef class IndexEngine:
             ndarray[intp_t] result, missing
             set stargets, remaining_stargets
             dict d = {}
-            object val
+            object val, dt64nat, td64nat
             Py_ssize_t count = 0, count_missing = 0
             Py_ssize_t i, j, n, n_t, n_alloc, start, end
             bint d_has_nan = False, stargets_has_nan = False, need_nan_check = True
+            d_has_dt64nat = False, stargets_has_dt64nat = False,
+            need_dt64nat_check = True, d_has_td64nat = False,
+            stargets_has_td64nat = False, need_td64nat_check = True,
+            need_np_nat_check = False
 
         values = self.values
         stargets = set(targets)
@@ -357,6 +365,11 @@ cdef class IndexEngine:
         if stargets:
             # otherwise, map by iterating through all items in the index
 
+            # determine if we need to check for numpy nats
+            # ie. np.datetime64("NaT") np.timedelta64("NaT")
+            if values.dtype == object:
+                need_np_nat_check = True
+
             for i in range(n):
                 val = values[i]
                 if val in stargets:
@@ -378,12 +391,59 @@ cdef class IndexEngine:
                             d_has_nan = True
                         d[np.nan].append(i)
 
+                elif need_np_nat_check:
+                    if is_dt64nat(val):
+                        if need_dt64nat_check:
+                            # Do this check only once
+                            stargets_has_dt64nat = any(
+                                is_dt64nat(starget) for starget in stargets
+                            )
+                            need_dt64nat_check = False
+
+                        if stargets_has_dt64nat:
+                            if not d_has_dt64nat:
+                                # store to ensure future access to `d` uses same key
+                                dt64nat = np.datetime64("NaT")
+                                d[dt64nat] = []
+                                d_has_dt64nat = True
+                            d[dt64nat].append(i)
+
+                    elif is_td64nat(val):
+                        if need_td64nat_check:
+                            # Do this check only once
+                            stargets_has_td64nat = any(
+                                is_td64nat(starget) for starget in stargets
+                            )
+                            need_td64nat_check = False
+
+                        if stargets_has_td64nat:
+                            if not d_has_td64nat:
+                                # store to ensure future access to `d` uses same key
+                                td64nat = np.timedelta64("NaT")
+                                d[td64nat] = []
+                                d_has_td64nat = True
+                            d[td64nat].append(i)
+
         for i in range(n_t):
             val = targets[i]
-
             # found
-            if val in d or (d_has_nan and util.is_nan(val)):
-                key = val if not util.is_nan(val) else np.nan
+            # cannot search for nan/nat target using `in`,
+            # need to lookup key using d_has_...
+            # and confirm na type via util function
+            if (
+                val in d
+                or (d_has_nan and util.is_nan(val))
+                or (d_has_dt64nat and is_dt64nat(val))
+                or (d_has_td64nat and is_td64nat(val))
+            ):
+                key = val
+                if d_has_nan and util.is_nan(key):
+                    key = np.nan
+                elif d_has_dt64nat and is_dt64nat(key):
+                    key = dt64nat
+                elif d_has_td64nat and is_td64nat(key):
+                    key = td64nat
+
                 for j in d[key]:
 
                     # realloc if needed

pandas/_libs/missing.pyx

@@ -20,7 +20,9 @@ from pandas._libs cimport util
 from pandas._libs.tslibs.nattype cimport (
     c_NaT as NaT,
     checknull_with_nat,
+    is_dt64nat,
     is_null_datetimelike,
+    is_td64nat,
 )
 from pandas._libs.tslibs.np_datetime cimport (
     get_datetime64_value,
@@ -77,18 +79,10 @@ cpdef bint is_matching_na(object left, object right, bint nan_matches_none=False
             and util.is_complex_object(right)
             and util.is_nan(right)
         )
-    elif util.is_datetime64_object(left):
-        return (
-            get_datetime64_value(left) == NPY_NAT
-            and util.is_datetime64_object(right)
-            and get_datetime64_value(right) == NPY_NAT
-        )
-    elif util.is_timedelta64_object(left):
-        return (
-            get_timedelta64_value(left) == NPY_NAT
-            and util.is_timedelta64_object(right)
-            and get_timedelta64_value(right) == NPY_NAT
-        )
+    elif is_dt64nat(left):
+        return is_dt64nat(right)
+    elif is_td64nat(left):
+        return is_td64nat(right)
     elif is_decimal_na(left):
         return is_decimal_na(right)
     return False
@@ -345,20 +339,16 @@ def isneginf_scalar(val: object) -> bool:
 cdef inline bint is_null_datetime64(v):
     # determine if we have a null for a datetime (or integer versions),
     # excluding np.timedelta64('nat')
-    if checknull_with_nat(v):
+    if checknull_with_nat(v) or is_dt64nat(v):
         return True
-    elif util.is_datetime64_object(v):
-        return get_datetime64_value(v) == NPY_NAT
     return False
 
 
 cdef inline bint is_null_timedelta64(v):
     # determine if we have a null for a timedelta (or integer versions),
     # excluding np.datetime64('nat')
-    if checknull_with_nat(v):
+    if checknull_with_nat(v) or is_td64nat(v):
         return True
-    elif util.is_timedelta64_object(v):
-        return get_timedelta64_value(v) == NPY_NAT
     return False
 
 

pandas/_libs/tslibs/nattype.pxd

@@ -16,4 +16,6 @@ cdef _NaT c_NaT
 
 
 cdef bint checknull_with_nat(object val)
+cdef bint is_dt64nat(object val)
+cdef bint is_td64nat(object val)
 cpdef bint is_null_datetimelike(object val, bint inat_is_null=*)

pandas/_libs/tslibs/nattype.pyx

@@ -1133,6 +1133,22 @@ cdef inline bint checknull_with_nat(object val):
     """
     return val is None or util.is_nan(val) or val is c_NaT
 
+cdef inline bint is_dt64nat(object val):
+    """
+    Is this a np.datetime64 object np.datetime64("NaT").
+    """
+    if util.is_datetime64_object(val):
+        return get_datetime64_value(val) == NPY_NAT
+    return False
+
+cdef inline bint is_td64nat(object val):
+    """
+    Is this a np.timedelta64 object np.timedelta64("NaT").
+    """
+    if util.is_timedelta64_object(val):
+        return get_timedelta64_value(val) == NPY_NAT
+    return False
+
 
 cpdef bint is_null_datetimelike(object val, bint inat_is_null=True):
     """

pandas/_testing/__init__.py

@@ -157,6 +157,7 @@
 )
 
 NULL_OBJECTS = [None, np.nan, pd.NaT, float("nan"), pd.NA, Decimal("NaN")]
+NP_NAT_OBJECTS = [np.datetime64("NaT"), np.timedelta64("NaT")]
 
 EMPTY_STRING_PATTERN = re.compile("^$")
 

pandas/conftest.py

@@ -332,6 +332,19 @@ def unique_nulls_fixture(request):
 # Generate cartesian product of unique_nulls_fixture:
 unique_nulls_fixture2 = unique_nulls_fixture
 
+
+@pytest.fixture(params=tm.NP_NAT_OBJECTS, ids=lambda x: type(x).__name__)
+def np_nat_fixture(request):
+    """
+    Fixture for each NaT type in numpy.
+    """
+    return request.param
+
+
+# Generate cartesian product of np_nat_fixture:
+np_nat_fixture2 = np_nat_fixture
+
+
 # ----------------------------------------------------------------
 # Classes
 # ----------------------------------------------------------------

pandas/tests/indexes/object/test_indexing.py

@@ -96,6 +96,52 @@ def test_get_indexer_non_unique_nas(self, nulls_fixture):
             tm.assert_numpy_array_equal(indexer, expected_indexer)
             tm.assert_numpy_array_equal(missing, expected_missing)
 
+    # TODO: distinguish between different date/time units
+    # for datetime64("NaT") and timedelta64("NaT"):
+    # ie. np.datetime64("NaT") vs np.datetime64("NaT", "ns"), np.datetime64("NaT", "ms")
+    @pytest.mark.filterwarnings("ignore:elementwise comp:DeprecationWarning")
+    def test_get_indexer_non_unique_np_nats(self, np_nat_fixture, np_nat_fixture2):
+        expected_missing = np.array([], dtype=np.intp)
+        # matching-but-not-identical nats
+        if is_matching_na(np_nat_fixture, np_nat_fixture2):
+            # ensure nats are different objects
+            index = Index(
+                np.array(
+                    ["2021-10-02", np_nat_fixture.copy(), np_nat_fixture2.copy()],
+                    dtype=object,
+                ),
+                dtype=object,
+            )
+            # pass as index to prevent target from being casted to DatetimeIndex
+            indexer, missing = index.get_indexer_non_unique(
+                Index([np_nat_fixture], dtype=object)
+            )
+            expected_indexer = np.array([1, 2], dtype=np.intp)
+            tm.assert_numpy_array_equal(indexer, expected_indexer)
+            tm.assert_numpy_array_equal(missing, expected_missing)
+        # dt64nat vs td64nat
+        else:
+            index = Index(
+                np.array(
+                    [
+                        "2021-10-02",
+                        np_nat_fixture,
+                        np_nat_fixture2,
+                        np_nat_fixture,
+                        np_nat_fixture2,
+                    ],
+                    dtype=object,
+                ),
+                dtype=object,
+            )
+            # pass as index to prevent target from being casted to DatetimeIndex
+            indexer, missing = index.get_indexer_non_unique(
+                Index([np_nat_fixture], dtype=object)
+            )
+            expected_indexer = np.array([1, 3], dtype=np.intp)
+            tm.assert_numpy_array_equal(indexer, expected_indexer)
+            tm.assert_numpy_array_equal(missing, expected_missing)
+
 
 class TestSliceLocs:
     @pytest.mark.parametrize(