make Int8/16/32Engine work with Int64HashTable

Author: tp

pandas/_libs/algos_common_helper.pxi.in

@@ -28,6 +28,7 @@ dtypes = [('float64', 'float64_t', 'np.float64', True, True),
           ('int64', 'int64_t', 'np.int64', False, True),
           ('uint64', 'uint64_t', 'np.uint64', False, True),
           ('int32', 'int32_t', 'np.int32', False, True),
+          ('int16', 'int16_t', 'np.int16', False, True),
           ('int8', 'int8_t', 'np.int8', False, True),
           ('bool', 'uint8_t', 'np.bool', False, True)]
 

pandas/_libs/index.pyx

@@ -9,7 +9,9 @@ from cpython.slice cimport PySlice_Check
 import numpy as np
 cimport numpy as cnp
 from numpy cimport (ndarray, float64_t, int32_t,
-                    int64_t, uint8_t, uint64_t, intp_t,
+                    int8_t, int16_t, int32_t, int64_t,
+                    uint8_t, uint64_t,
+                    intp_t,
                     # Note: NPY_DATETIME, NPY_TIMEDELTA are only available
                     # for cimport in cython>=0.27.3
                     NPY_DATETIME, NPY_TIMEDELTA)
@@ -264,6 +266,8 @@ cdef class IndexEngine:
         if not self.is_mapping_populated:
 
             values = self._get_index_values()
+            if values.dtype in {'int8', 'int16', 'int32'}:
+                values = algos.ensure_int64(values)
             self.mapping = self._make_hash_table(len(values))
             self._call_map_locations(values)
 

pandas/_libs/index_class_helper.pxi.in

@@ -40,6 +40,8 @@ cdef class {{name}}Engine(IndexEngine):
     cdef _make_hash_table(self, n):
         {{if name == 'Object'}}
         return _hash.PyObjectHashTable(n)
+        {{elif name in {'Int8', 'Int16', 'Int32'} }}
+        return _hash.Int64HashTable(n)
         {{else}}
         return _hash.{{name}}HashTable(n)
         {{endif}}

pandas/core/indexes/category.py

@@ -437,18 +437,12 @@ def get_loc(self, key, method=None):
         >>> non_monotonic_index.get_loc('b')
         array([False,  True, False,  True], dtype=bool)
         """
-        codes = self.categories.get_loc(key)
-        if (codes == -1):
-            raise KeyError(key)
-
-        if self.is_monotonic_increasing and not self.is_unique:
-            if codes not in self._engine:
-                raise KeyError(key)
-            codes = self.codes.dtype.type(codes)
-            lhs = self.codes.searchsorted(codes, side='left')
-            rhs = self.codes.searchsorted(codes, side='right')
-            return slice(lhs, rhs)
-        return self._engine.get_loc(codes)
+        code = self.categories.get_loc(key)
+
+        # dtype must be same as dtype for self.codes else searchsorted is slow
+        code = self.codes.dtype.type(code)
+
+        return self._engine.get_loc(code)
 
     def get_value(self, series, key):
         """