PERF: speed up CategoricalIndex.get_loc

Author: topper-123

asv_bench/benchmarks/indexing_engines.py

@@ -1,19 +1,31 @@
 import numpy as np
 
-from pandas._libs.index import (Int64Engine, UInt64Engine, Float64Engine,
-                                ObjectEngine)
+from pandas._libs import index as li
+
+
+def _get_numeric_engines():
+    engine_names = [
+        ('Int64Engine', np.int64), ('Int32Engine', np.int32),
+        ('Int16Engine', np.int16), ('Int8Engine', np.int8),
+        ('UInt64Engine', np.uint64), ('UInt32Engine', np.uint32),
+        ('UInt16engine', np.uint16), ('UInt8Engine', np.uint8),
+        ('Float64Engine', np.float64), ('Float32Engine', np.float32),
+    ]
+    return [(getattr(li, engine_name), dtype)
+            for engine_name, dtype in engine_names if hasattr(li, engine_name)]
 
 
 class NumericEngineIndexing(object):
 
     goal_time = 0.2
-    params = [[Int64Engine, UInt64Engine, Float64Engine],
-              [np.int64, np.uint64, np.float64],
+
+    params = [_get_numeric_engines(),
               ['monotonic_incr', 'monotonic_decr', 'non_monotonic'],
               ]
-    param_names = ['engine', 'dtype', 'index_type']
+    param_names = ['engine_and_dtype', 'index_type']
 
-    def setup(self, engine, dtype, index_type):
+    def setup(self, engine_and_dtype, index_type):
+        engine, dtype = engine_and_dtype
         N = 10**5
         values = list([1] * N + [2] * N + [3] * N)
         arr = {
@@ -27,7 +39,7 @@ def setup(self, engine, dtype, index_type):
         # code belows avoids populating the mapping etc. while timing.
         self.data.get_loc(2)
 
-    def time_get_loc(self, engine, dtype, index_type):
+    def time_get_loc(self, engine_and_dtype, index_type):
         self.data.get_loc(2)
 
 
@@ -46,7 +58,7 @@ def setup(self, index_type):
             'non_monotonic': np.array(list('abc') * N, dtype=object),
         }[index_type]
 
-        self.data = ObjectEngine(lambda: arr, len(arr))
+        self.data = li.ObjectEngine(lambda: arr, len(arr))
         # code belows avoids populating the mapping etc. while timing.
         self.data.get_loc('b')
 

doc/source/whatsnew/v0.24.0.txt

@@ -759,9 +759,10 @@ Removal of prior version deprecations/changes
 Performance Improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
-- Very large improvement in performance of slicing when the index is a :class:`CategoricalIndex`,
-  both when indexing by label (using .loc) and position(.iloc).
-  Likewise, slicing a ``CategoricalIndex`` itself (i.e. ``ci[100:200]``) shows similar speed improvements (:issue:`21659`)
+- Slicing Series and Dataframes with an monotonically increasing :class:`CategoricalIndex`
+  is now very fast and has speed comparable to slicing with an ``Int64Index``.
+  The speed increase is both when indexing by label (using .loc) and position(.iloc) (:issue:`20395`)
+- Slicing a :class:`CategoricalIndex` itself (i.e. ``ci[1000:2000]``) shows similar speed improvements as above (:issue:`21659`)
 - Improved performance of :func:`Series.describe` in case of numeric dtpyes (:issue:`21274`)
 - Improved performance of :func:`pandas.core.groupby.GroupBy.rank` when dealing with tied rankings (:issue:`21237`)
 - Improved performance of :func:`DataFrame.set_index` with columns consisting of :class:`Period` objects (:issue:`21582`, :issue:`21606`)

pandas/_libs/algos.pyx

@@ -10,7 +10,8 @@ from libc.math cimport fabs, sqrt
 import numpy as np
 cimport numpy as cnp
 from numpy cimport (ndarray,
-                    NPY_INT64, NPY_UINT64, NPY_INT32, NPY_INT16, NPY_INT8,
+                    NPY_INT64, NPY_INT32, NPY_INT16, NPY_INT8,
+                    NPY_UINT64, NPY_UINT32, NPY_UINT16, NPY_UINT8,
                     NPY_FLOAT32, NPY_FLOAT64,
                     NPY_OBJECT,
                     int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t,
@@ -359,8 +360,10 @@ ctypedef fused algos_t:
     float64_t
     float32_t
     object
-    int32_t
     int64_t
+    int32_t
+    int16_t
+    int8_t
     uint64_t
     uint8_t
 
@@ -866,6 +869,8 @@ is_monotonic_float32 = is_monotonic["float32_t"]
 is_monotonic_object = is_monotonic["object"]
 is_monotonic_int64 = is_monotonic["int64_t"]
 is_monotonic_int32 = is_monotonic["int32_t"]
+is_monotonic_int16 = is_monotonic["int16_t"]
+is_monotonic_int8 = is_monotonic["int8_t"]
 is_monotonic_uint64 = is_monotonic["uint64_t"]
 is_monotonic_bool = is_monotonic["uint8_t"]
 

pandas/_libs/index.pyx

@@ -5,8 +5,10 @@ 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,
+from numpy cimport (ndarray, intp_t,
+                    float64_t, float32_t,
+                    int64_t, int32_t, int16_t, int8_t,
+                    uint64_t, uint32_t, uint16_t, uint8_t,
                     # Note: NPY_DATETIME, NPY_TIMEDELTA are only available
                     # for cimport in cython>=0.27.3
                     NPY_DATETIME, NPY_TIMEDELTA)

pandas/_libs/index_class_helper.pxi.in

@@ -10,14 +10,22 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 
 {{py:
 
-# name, dtype, ctype
-dtypes = [('Float64', 'float64', 'float64_t'),
-          ('UInt64', 'uint64', 'uint64_t'),
-          ('Int64', 'int64', 'int64_t'),
-          ('Object', 'object', 'object')]
+# name, dtype, ctype, hashtable_name, hashtable_dtype
+dtypes = [('Float64', 'float64', 'float64_t', 'Float64', 'float64'),
+          ('Float32', 'float32', 'float32_t', 'Float64', 'float64'),
+          ('Int64', 'int64', 'int64_t', 'Int64', 'int64'),
+          ('Int32', 'int32', 'int32_t', 'Int64', 'int64'),
+          ('Int16', 'int16', 'int16_t', 'Int64', 'int64'),
+          ('Int8', 'int8', 'int8_t', 'Int64', 'int64'),
+          ('UInt64', 'uint64', 'uint64_t', 'UInt64', 'uint64'),
+          ('UInt32', 'uint32', 'uint32_t', 'UInt64', 'uint64'),
+          ('UInt16', 'uint16', 'uint16_t', 'UInt64', 'uint64'),
+          ('UInt8', 'uint8', 'uint8_t', 'UInt64', 'uint64'),
+          ('Object', 'object', 'object', 'PyObject', 'object'),
+          ]
 }}
 
-{{for name, dtype, ctype in dtypes}}
+{{for name, dtype, ctype, hashtable_name, hashtable_dtype in dtypes}}
 
 
 cdef class {{name}}Engine(IndexEngine):
@@ -34,13 +42,9 @@ cdef class {{name}}Engine(IndexEngine):
                                    other, limit=limit)
 
     cdef _make_hash_table(self, n):
-        {{if name == 'Object'}}
-        return _hash.PyObjectHashTable(n)
-        {{else}}
-        return _hash.{{name}}HashTable(n)
-        {{endif}}
+        return _hash.{{hashtable_name}}HashTable(n)
 
-    {{if name != 'Float64' and name != 'Object'}}
+    {{if name not in {'Float64', 'Float32', 'Object'} }}
     cdef _check_type(self, object val):
         hash(val)
         if util.is_bool_object(val):
@@ -50,6 +54,11 @@ cdef class {{name}}Engine(IndexEngine):
     {{endif}}
 
     {{if name != 'Object'}}
+    cpdef _call_map_locations(self, values):
+        # self.mapping is of type {{hashtable_name}}HashTable,
+        # so convert dtype of values
+        self.mapping.map_locations(algos.ensure_{{hashtable_dtype}}(values))
+
     cdef _get_index_values(self):
         return algos.ensure_{{dtype}}(self.vgetter())
 
@@ -60,7 +69,7 @@ cdef class {{name}}Engine(IndexEngine):
             ndarray[{{ctype}}] values
             int count = 0
 
-        {{if name != 'Float64'}}
+        {{if name not in {'Float64', 'Float32'} }}
         if not util.is_integer_object(val):
             raise KeyError(val)
         {{endif}}

pandas/core/indexes/category.py

@@ -84,7 +84,17 @@ class CategoricalIndex(Index, accessor.PandasDelegate):
     """
 
     _typ = 'categoricalindex'
-    _engine_type = libindex.Int64Engine
+
+    @property
+    def _engine_type(self):
+        # self.codes can have dtype int8, int16, int32 or int64, so we need
+        # to return the corresponding engine type (libindex.Int8Engine, etc.).
+        return {np.int8: libindex.Int8Engine,
+                np.int16: libindex.Int16Engine,
+                np.int32: libindex.Int32Engine,
+                np.int64: libindex.Int64Engine,
+                }[self.codes.dtype.type]
+
     _attributes = ['name']
 
     def __new__(cls, data=None, categories=None, ordered=None, dtype=None,
@@ -382,7 +392,7 @@ def argsort(self, *args, **kwargs):
     def _engine(self):
 
         # we are going to look things up with the codes themselves
-        return self._engine_type(lambda: self.codes.astype('i8'), len(self))
+        return self._engine_type(lambda: self.codes, len(self))
 
     # introspection
     @cache_readonly
@@ -450,6 +460,7 @@ def get_loc(self, key, method=None):
         array([False,  True, False,  True], dtype=bool)
         """
         code = self.categories.get_loc(key)
+        code = self.codes.dtype.type(code)
         try:
             return self._engine.get_loc(code)
         except KeyError: