add Categorical.__contains__

Author: tp

asv_bench/benchmarks/categoricals.py

@@ -202,7 +202,11 @@ class Contains(object):
     def setup(self):
         N = 10**5
         self.ci = tm.makeCategoricalIndex(N)
-        self.cat = self.ci.categories[0]
+        self.c = self.ci.values
+        self.key = self.ci.categories[0]
 
-    def time_contains(self):
-        self.cat in self.ci
+    def time_categorical_index_contains(self):
+        self.key in self.ci
+
+    def time_categorical_contains(self):
+        self.key in self.c

doc/source/whatsnew/v0.23.2.txt

@@ -27,6 +27,9 @@ Performance Improvements
 - Improved performance of membership checks in :class:`CategoricalIndex`
   (i.e. ``x in ci``-style checks are much faster). :meth:`CategoricalIndex.contains`
   is likewise much faster (:issue:`21369`)
+- Improved performance of membership checks in :class:`Categorical`
+  (i.e. ``x in categorical``-style checks are much faster) (:issue:`21369`)
+
 -
 
 Documentation Changes

pandas/core/arrays/categorical.py

@@ -1847,6 +1847,31 @@ 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."""
+        hash(key)
+
+        if isna(key):  # if key is a NaN, check if any NaN is in self.
+            return self.isna().any()
+
+        # is key in self.categories? Then get its location.
+        # If not (i.e. KeyError), it logically can't be in self either
+        try:
+            loc = self.categories.get_loc(key)
+        except KeyError:
+            return False
+
+        # loc is the location of key in self.categories, but also the value
+        # for key in self.codes. key may be in categories,
+        # but still not in self, check this. Example:
+        # 'b' in Categorical(['a'], categories=['a', 'b']) #  False
+        if is_scalar(loc):
+            return loc in self._codes
+        else:
+            # if self.categories is IntervalIndex, loc is an array
+            # check if any scalar of the array is in self._codes
+            return any(loc_ in self._codes for loc_ in loc)
+
     def _tidy_repr(self, max_vals=10, footer=True):
         """ a short repr displaying only max_vals and an optional (but default
         footer)

pandas/tests/categorical/test_algos.py

@@ -71,6 +71,23 @@ def test_isin_empty(empty):
     tm.assert_numpy_array_equal(expected, result)
 
 
+def test_contains(self):
+
+    c = pd.Categorical(list('aabbca'), categories=list('cab'))
+
+    assert 'b' in c
+    assert 'z' not in c
+    assert np.nan not in c
+
+    # assert codes NOT in index
+    assert 0 not in c
+    assert 1 not in c
+
+    c = pd.Categorical(list('aabbca')+[np.nan], categories=list('cab'))
+
+    assert np.nan in c
+
+
 class TestTake(object):
     # https://github.com/pandas-dev/pandas/issues/20664