PERF: implement __contains__ for Categorical

Author: fjetter

asv_bench/benchmarks/categoricals.py

@@ -193,3 +193,47 @@ def time_categorical_series_is_monotonic_increasing(self):
 
     def time_categorical_series_is_monotonic_decreasing(self):
         self.s.is_monotonic_decreasing
+
+
+class Contains(object):
+
+    params = (["a", "c", "d", "z", np.nan], [True, False])
+    param_names = ["value", "has_nan"]
+
+    def setup(self, value, has_nan):
+        n = 1 * 10 ** 4
+        obj_values = list("a" * n + "b" * n + "c" * n)
+        if has_nan:
+            obj_values = [np.nan] + obj_values[:-2] + [np.nan]
+
+        self.ci = pd.CategoricalIndex(obj_values, categories=list("abcd"))
+        self.cat = pd.Categorical(obj_values, categories=list("abcd"))
+
+    def time_contains_index(self, value, has_nan):
+        value in self.ci
+
+    def time_cat_isin(self, value, has_nan):
+        value in self.cat
+
+
+class Indexing(object):
+
+    params = (["a", "c"], [True, False])
+    param_names = ["value", "has_nan"]
+
+    def setup(self, value, has_nan):
+        n = 1 * 10 ** 4
+        obj_values = list("a" * n + "b" * n + "c" * n)
+        if has_nan:
+            obj_values = [np.nan] + obj_values[:-2] + [np.nan]
+
+        self.ci = pd.CategoricalIndex(obj_values, categories=list("abcd"))
+        self.cat = pd.Categorical(obj_values, categories=list("abcd"))
+        self.df = pd.DataFrame(dict(A=range(n * 3)), index=self.ci)
+        self.ser = pd.Series(range(n * 3), index=self.ci)
+
+    def time_loc_df(self, value, has_nan):
+        self.df.loc[value]
+
+    def time_loc_ser(self, value, has_nan):
+        self.ser.loc[value]

doc/source/whatsnew/v0.24.0.txt

@@ -64,7 +64,7 @@ Performance Improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
 - Improved performance of :func:`Series.describe` in case of numeric dtpyes (:issue:`21274`)
--
+- Improved performance of indexing on a Series/DataFrame with a CategoricalIndex
 
 .. _whatsnew_0240.docs:
 

pandas/core/arrays/categorical.py

@@ -1847,6 +1847,13 @@ def __iter__(self):
         """Returns an Iterator over the values of this Categorical."""
         return iter(self.get_values().tolist())
 
+    def __contains__(self, key):
+        """Returns True if `key` is in this Categorical."""
+        if key in self.categories:
+            return self.categories.get_loc(key) in self.codes
+        else:
+            return False
+
     def _tidy_repr(self, max_vals=10, footer=True):
         """ a short repr displaying only max_vals and an optional (but default
         footer)