ENH: Add masked engine

asv_bench/benchmarks/indexing.py

@@ -10,9 +10,11 @@
 import numpy as np
 
 from pandas import (
+    NA,
     CategoricalIndex,
     DataFrame,
     Float64Index,
+    Index,
     Int64Index,
     IntervalIndex,
     MultiIndex,
@@ -87,6 +89,32 @@ def time_loc_slice(self, index, index_structure):
         self.data.loc[:800000]
 
 
+class NumericMaskedIndexing:
+    params = [
+        ("Int64", "UInt64", "Float64"),
+        (True, False),
+    ]
+    param_names = ["dtype", "monotonic"]
+
+    def setup(self, dtype, monotonic):
+        N = 10**6
+        indices = {
+            True: Index(range(N), dtype=dtype),
+            False: Index(
+                list(range(50)) + [54, 53, 52, 51] + list(range(55, N - 1)), dtype=dtype
+            ).append(Index([NA], dtype=dtype)),
+        }
+        self.data = indices[monotonic]
+        self.indexer = np.arange(300, 1_000)
+        self.data_dups = self.data.append(self.data)
+
+    def time_get_indexer(self, dtype, monotonic):
+        self.data.get_indexer(self.indexer)
+
+    def time_get_indexer_dups(self, dtype, monotonic):
+        self.data.get_indexer_for(self.indexer)
+
+
 class NonNumericSeriesIndexing:
 
     params = [

asv_bench/benchmarks/indexing_engines.py

@@ -30,6 +30,26 @@ def _get_numeric_engines():
     ]
 
 
+def _get_masked_engines():
+    engine_names = [
+        ("MaskedInt64Engine", "Int64"),
+        ("MaskedInt32Engine", "Int32"),
+        ("MaskedInt16Engine", "Int16"),
+        ("MaskedInt8Engine", "Int8"),
+        ("MaskedUInt64Engine", "UInt64"),
+        ("MaskedUInt32Engine", "UInt32"),
+        ("MaskedUInt16engine", "UInt16"),
+        ("MaskedUInt8Engine", "UInt8"),
+        ("MaskedFloat64Engine", "Float64"),
+        ("MaskedFloat32Engine", "Float32"),
+    ]
+    return [
+        (getattr(libindex, engine_name), dtype)
+        for engine_name, dtype in engine_names
+        if hasattr(libindex, engine_name)
+    ]
+
+
 class NumericEngineIndexing:
 
     params = [
@@ -80,6 +100,61 @@ def time_get_loc_near_middle(self, engine_and_dtype, index_type, unique, N):
         self.data.get_loc(self.key_middle)
 
 
+class MaskedNumericEngineIndexing:
+
+    params = [
+        _get_masked_engines(),
+        ["monotonic_incr", "monotonic_decr", "non_monotonic"],
+        [True, False],
+        [10**5, 2 * 10**6],  # 2e6 is above SIZE_CUTOFF
+    ]
+    param_names = ["engine_and_dtype", "index_type", "unique", "N"]
+
+    def setup(self, engine_and_dtype, index_type, unique, N):
+        engine, dtype = engine_and_dtype
+
+        if index_type == "monotonic_incr":
+            if unique:
+                arr = np.arange(N * 3, dtype=dtype.lower())
+            else:
+                values = list([1] * N + [2] * N + [3] * N)
+                arr = np.array(values, dtype=dtype.lower())
+            mask = np.zeros(N * 3, dtype="uint8")
+        elif index_type == "monotonic_decr":
+            if unique:
+                arr = np.arange(N * 3, dtype=dtype.lower())[::-1]
+            else:
+                values = list([1] * N + [2] * N + [3] * N)
+                arr = np.array(values, dtype=dtype.lower())[::-1]
+            mask = np.zeros(N * 3, dtype="uint8")
+        else:
+            assert index_type == "non_monotonic"
+            if unique:
+                arr = np.zeros(N * 3, dtype=dtype.lower())
+                arr[:N] = np.arange(N * 2, N * 3, dtype=dtype.lower())
+                arr[N:] = np.arange(N * 2, dtype=dtype.lower())
+
+            else:
+                arr = np.array([1, 2, 3] * N, dtype=dtype.lower())
+            mask = np.zeros(N * 3, dtype="uint8")
+            mask[-1] = True
+
+        self.data = engine(arr, mask)
+        # code belows avoids populating the mapping etc. while timing.
+        self.data.get_loc(2)
+
+        self.key_middle = arr[len(arr) // 2]
+        self.key_early = arr[2]
+
+    def time_get_loc(self, engine_and_dtype, index_type, unique, N):
+        self.data.get_loc(self.key_early)
+
+    def time_get_loc_near_middle(self, engine_and_dtype, index_type, unique, N):
+        # searchsorted performance may be different near the middle of a range
+        #  vs near an endpoint
+        self.data.get_loc(self.key_middle)
+
+
 class ObjectEngineIndexing:
 
     params = [("monotonic_incr", "monotonic_decr", "non_monotonic")]

pandas/_libs/hashtable.pyi

@@ -125,10 +125,12 @@ class HashTable:
     def map_locations(
         self,
         values: np.ndarray,  # np.ndarray[subclass-specific]
+        mask: np.ndarray | None = ...,
     ) -> None: ...
     def lookup(
         self,
         values: np.ndarray,  # np.ndarray[subclass-specific]
+        mask: np.ndarray | None = ...,
     ) -> npt.NDArray[np.intp]: ...
     def get_labels(
         self,

pandas/_libs/hashtable_class_helper.pxi.in

@@ -966,8 +966,9 @@ cdef class StringHashTable(HashTable):
         return labels
 
     @cython.boundscheck(False)
-    def lookup(self, ndarray[object] values) -> ndarray:
+    def lookup(self, ndarray[object] values, object mask = None) -> ndarray:
         # -> np.ndarray[np.intp]
+        # mask not yet implemented
         cdef:
             Py_ssize_t i, n = len(values)
             int ret = 0
@@ -1002,7 +1003,8 @@ cdef class StringHashTable(HashTable):
         return np.asarray(locs)
 
     @cython.boundscheck(False)
-    def map_locations(self, ndarray[object] values) -> None:
+    def map_locations(self, ndarray[object] values, object mask = None) -> None:
+        # mask not yet implemented
         cdef:
             Py_ssize_t i, n = len(values)
             int ret = 0
@@ -1275,7 +1277,8 @@ cdef class PyObjectHashTable(HashTable):
         else:
             raise KeyError(key)
 
-    def map_locations(self, ndarray[object] values) -> None:
+    def map_locations(self, ndarray[object] values, object mask = None) -> None:
+        # mask not yet implemented
         cdef:
             Py_ssize_t i, n = len(values)
             int ret = 0
@@ -1289,8 +1292,9 @@ cdef class PyObjectHashTable(HashTable):
             k = kh_put_pymap(self.table, <PyObject*>val, &ret)
             self.table.vals[k] = i
 
-    def lookup(self, ndarray[object] values) -> ndarray:
+    def lookup(self, ndarray[object] values, object mask = None) -> ndarray:
         # -> np.ndarray[np.intp]
+        # mask not yet implemented
         cdef:
             Py_ssize_t i, n = len(values)
             int ret = 0

pandas/_libs/index.pyi

@@ -7,7 +7,7 @@ from pandas.core.arrays import ExtensionArray
 
 class IndexEngine:
     over_size_threshold: bool
-    def __init__(self, values: np.ndarray) -> None: ...
+    def __init__(self, values: np.ndarray, mask: np.ndarray | None = ...) -> None: ...
     def __contains__(self, val: object) -> bool: ...
 
     # -> int | slice | np.ndarray[bool]
@@ -23,12 +23,19 @@ class IndexEngine:
     @property
     def is_mapping_populated(self) -> bool: ...
     def clear_mapping(self): ...
-    def get_indexer(self, values: np.ndarray) -> npt.NDArray[np.intp]: ...
+    def get_indexer(
+        self, values: np.ndarray, mask: np.ndarray | None = ...
+    ) -> npt.NDArray[np.intp]: ...
     def get_indexer_non_unique(
         self,
         targets: np.ndarray,
     ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
 
+class MaskedIndexEngine(IndexEngine):
+    def get_indexer_non_unique(  # type: ignore[override]
+        self, targets: np.ndarray, target_mask: np.ndarray
+    ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
+
 class Float64Engine(IndexEngine): ...
 class Float32Engine(IndexEngine): ...
 class Complex128Engine(IndexEngine): ...
@@ -46,6 +53,19 @@ class DatetimeEngine(Int64Engine): ...
 class TimedeltaEngine(DatetimeEngine): ...
 class PeriodEngine(Int64Engine): ...
 class BoolEngine(UInt8Engine): ...
+class MaskedBoolEngine(MaskedUInt8Engine): ...
+class MaskedFloat64Engine(MaskedIndexEngine): ...
+class MaskedFloat32Engine(MaskedIndexEngine): ...
+class MaskedComplex128Engine(MaskedIndexEngine): ...
+class MaskedComplex64Engine(MaskedIndexEngine): ...
+class MaskedInt64Engine(MaskedIndexEngine): ...
+class MaskedInt32Engine(MaskedIndexEngine): ...
+class MaskedInt16Engine(MaskedIndexEngine): ...
+class MaskedInt8Engine(MaskedIndexEngine): ...
+class MaskedUInt64Engine(MaskedIndexEngine): ...
+class MaskedUInt32Engine(MaskedIndexEngine): ...
+class MaskedUInt16Engine(MaskedIndexEngine): ...
+class MaskedUInt8Engine(MaskedIndexEngine): ...
 
 class BaseMultiIndexCodesEngine:
     levels: list[np.ndarray]
@@ -58,8 +78,7 @@ class BaseMultiIndexCodesEngine:
         offsets: np.ndarray,  # np.ndarray[np.uint64, ndim=1]
     ) -> None: ...
     def get_indexer(
-        self,
-        target: npt.NDArray[np.object_],
+        self, target: npt.NDArray[np.object_], mask: np.ndarray | None = ...
     ) -> npt.NDArray[np.intp]: ...
     def _extract_level_codes(self, target: MultiIndex) -> np.ndarray: ...
     def get_indexer_with_fill(