CLN: Using masks instead of list comprehensions (pandas-dev#35498)

alexhlim · Aug 18, 2020 · ca82d69 · ca82d69
1 parent 3367203
commit ca82d69
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/pandas/_libs/index.pyx b/pandas/_libs/index.pyx
@@ -270,13 +270,11 @@ cdef class IndexEngine:
             Py_ssize_t i, j, n, n_t, n_alloc
 
         self._ensure_mapping_populated()
-        if targets[isnaobj(targets)].size > 0:
-            new_targets = [0 if checknull(t) else t for t in targets]
-            new_values = [0 if checknull(v) else v for v in self._get_index_values()]
-            targets = np.array(new_targets, dtype=targets.dtype)
-            values = np.array(new_values, dtype=self._get_index_values().dtype)
-        else:
-            values = np.array(self._get_index_values(), copy=False)
+        values = np.array(self._get_index_values(), copy=False)
+        targets_mask = isnaobj(targets)
+        if targets_mask.any():
+            targets[targets_mask] = 0
+            values[isnaobj(values)] = 0
 
         stargets = set(targets)
         n = len(values)