CLN: overriding IndexEngine subclasses (#35498)

Author: alexhlim

pandas/_libs/index.pyx

@@ -36,6 +36,7 @@ from pandas._libs import (
 )
 from pandas._libs.missing import (
     checknull,
+    is_matching_na,
     isnaobj,
 )
 
@@ -270,6 +271,25 @@ cdef class IndexEngine:
         self._ensure_mapping_populated()
         return self.mapping.lookup(values)
 
+    def get_stargets(self, ndarray targets) -> set:
+        return set(targets)
+
+    def convert_val_if_nan(self, object val) -> object:
+        # unable to utilize val if nan when updating
+        # hashable data structures (ie. sets, dict)
+        if checknull(val):
+            return -1
+        else:
+            return val
+
+    def should_update_d(self, object target, object val) -> bool:
+        # -1 in targets could be either -1 or nan
+        # ensures values in d[-1] to be included only once
+        if target == val or is_matching_na(target, val):
+            return True
+
+        return False
+
     def get_indexer_non_unique(self, ndarray targets):
         """
         Return an indexer suitable for taking from a non unique index
@@ -293,11 +313,7 @@ cdef class IndexEngine:
 
         self._ensure_mapping_populated()
         values = np.array(self._get_index_values(), copy=False)
-        values_mask = isnaobj(values)
-        targets_mask = isnaobj(targets)
-        stargets = set(targets)
-        if -1 not in stargets and targets_mask.any():
-            stargets.add(-1)
+        stargets = self.get_stargets(targets)
         n = len(values)
         n_t = len(targets)
         if n > 10_000:
@@ -328,22 +344,15 @@ cdef class IndexEngine:
         if stargets:
             # otherwise, map by iterating through all items in the index
             for i in range(n):
-                if values_mask[i]:
-                    val = -1
-                else:
-                    val = values[i]
+                val = self.convert_val_if_nan(values[i])
 
                 if val in stargets:
                     if val not in d:
                         d[val] = []
                     d[val].append(i)
 
         for i in range(n_t):
-            nan_target = targets_mask[i]
-            if nan_target:
-                val = -1
-            else:
-                val = targets[i]
+            val = self.convert_val_if_nan(targets[i])
 
             # found
             if val in d:
@@ -354,21 +363,7 @@ cdef class IndexEngine:
                         n_alloc += 10_000
                         result = np.resize(result, n_alloc)
 
-                    # -1 in targets could be either -1 or nan
-                    # ensures values in d[-1] to be included only once
-                    if val == -1:
-                        nan_val = values_mask[j]
-                        # nan
-                        if nan_target:
-                            if nan_val:
-                                result[count] = j
-                                count += 1
-                        # -1
-                        else:
-                            if not nan_val:
-                                result[count] = j
-                                count += 1
-                    else:
+                    if self.should_update_d(targets[i], values[j]):
                         result[count] = j
                         count += 1
 
@@ -419,6 +414,13 @@ cdef class ObjectEngine(IndexEngine):
     cdef _make_hash_table(self, Py_ssize_t n):
         return _hash.PyObjectHashTable(n)
 
+    def get_stargets(self, ndarray targets) -> set:
+        stargets = set(targets)
+        # account for NA-like targets
+        if -1 not in stargets and isnaobj(targets).any():
+            stargets.add(-1)
+
+        return stargets
 
 cdef class DatetimeEngine(Int64Engine):
 
@@ -490,6 +492,14 @@ cdef class DatetimeEngine(Int64Engine):
         except KeyError:
             raise KeyError(val)
 
+    def get_stargets(self, ndarray targets) -> set:
+        stargets = set(targets)
+        # account for NaTs
+        if -1 not in stargets and isnaobj(targets).any():
+            stargets.add(-1)
+
+        return stargets
+
     def get_indexer_non_unique(self, ndarray targets):
         # we may get datetime64[ns] or timedelta64[ns], cast these to int64
         return super().get_indexer_non_unique(targets.view("i8"))

pandas/_libs/index_class_helper.pxi.in

@@ -77,4 +77,14 @@ cdef class {{name}}Engine(IndexEngine):
 
         return self._unpack_bool_indexer(indexer, val)
 
+    {{if name in {'Float64', 'Float32'} }}
+    def get_stargets(self, ndarray targets):
+        stargets = set(targets)
+        # account for nans
+        if -1 not in stargets and np.isnan(targets).any():
+            stargets.add(-1)
+
+        return stargets
+    {{endif}}
+
 {{endfor}}