make Int8/16/32Engine work with Int64HashTable

Author: tp

pandas/_libs/index.pyx

@@ -6,7 +6,9 @@ import cython
 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)
@@ -242,6 +244,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

@@ -448,18 +448,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):
         """