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,43 @@ def __iter__(self):
         """Returns an Iterator over the values of this Categorical."""
         return iter(self.get_values().tolist())
 
+    @staticmethod
+    def _contains(key, categories, container):
+        """Returns True if `key` is in `categories` and the
+        location of `key` in `categories` is in `container`.
+
+        This is a helper method used in :method:`Categorical.__contains__`
+        and in :class:`CategoricalIndex.__contains__`.
+        """
+
+        # is key in categories? Then get its location in categories.
+        # If not (i.e. KeyError), its location logically can't be in
+        # container either.
+        try:
+            loc = categories.get_loc(key)
+        except KeyError:
+            return False
+
+        # loc is the location of key in categories, but also the value
+        # for key in container. So, key may be in categories,
+        # but still not in container, this must be checked. Example:
+        # 'b' in Categorical(['a'], categories=['a', 'b']) #  False
+        if is_scalar(loc):
+            return loc in container
+        else:
+            # if categories is an IntervalIndex, loc is an array.
+            # Check if any scalar of the array is in the container
+            return any(loc_ in container for loc_ in loc)
+
+    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()
+
+        return self._contains(key, self.categories, container=self._codes)
+
     def _tidy_repr(self, max_vals=10, footer=True):
         """ a short repr displaying only max_vals and an optional (but default
         footer)

pandas/core/indexes/category.py

@@ -24,6 +24,7 @@
 import pandas.core.common as com
 import pandas.core.missing as missing
 import pandas.core.indexes.base as ibase
+from pandas.core.arrays import Categorical
 
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 _index_doc_kwargs.update(dict(target_klass='CategoricalIndex'))
@@ -125,7 +126,6 @@ def _create_from_codes(self, codes, categories=None, ordered=None,
         CategoricalIndex
         """
 
-        from pandas.core.arrays import Categorical
         if categories is None:
             categories = self.categories
         if ordered is None:
@@ -162,7 +162,6 @@ def _create_categorical(self, data, categories=None, ordered=None,
         if not isinstance(data, ABCCategorical):
             if ordered is None and dtype is None:
                 ordered = False
-            from pandas.core.arrays import Categorical
             data = Categorical(data, categories=categories, ordered=ordered,
                                dtype=dtype)
         else:
@@ -328,23 +327,8 @@ def __contains__(self, key):
         if isna(key):  # if key is a NaN, check if any NaN is in self.
             return self.hasnans
 
-        # 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 and in self._engine. key may be in categories,
-        # but still not in self, check this. Example:
-        # 'b' in CategoricalIndex(['a'], categories=['a', 'b']) #  False
-        if is_scalar(loc):
-            return loc in self._engine
-        else:
-            # if self.categories is IntervalIndex, loc is an array
-            # check if any scalar of the array is in self._engine
-            return any(loc_ in self._engine for loc_ in loc)
+        return Categorical._contains(key, categories=self.categories,
+                                     container=self._engine)
 
     @Appender(_index_shared_docs['contains'] % _index_doc_kwargs)
     def contains(self, key):
@@ -479,7 +463,6 @@ def where(self, cond, other=None):
             other = self._na_value
         values = np.where(cond, self.values, other)
 
-        from pandas.core.arrays import Categorical
         cat = Categorical(values,
                           categories=self.categories,
                           ordered=self.ordered)
@@ -862,7 +845,6 @@ def _delegate_method(self, name, *args, **kwargs):
     def _add_accessors(cls):
         """ add in Categorical accessor methods """
 
-        from pandas.core.arrays import Categorical
         CategoricalIndex._add_delegate_accessors(
             delegate=Categorical, accessors=["rename_categories",
                                              "reorder_categories",

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():
+
+    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