BUG-20629 allow .at accessor with CategoricalIndex (#26298)

Author: JustinZhengBC

doc/source/whatsnew/v0.25.0.rst

@@ -301,7 +301,7 @@ Bug Fixes
 Categorical
 ^^^^^^^^^^^
 
--
+- Bug in :func:`DataFrame.at` and :func:`Series.at` that would raise exception if the index was a :class:`CategoricalIndex` (:issue:`20629`)
 -
 -
 

pandas/_typing.py

@@ -8,8 +8,14 @@
 from pandas._libs.tslibs.timedeltas import Timedelta
 
 from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.dtypes.generic import ABCExtensionArray
+from pandas.core.dtypes.generic import (
+    ABCExtensionArray, ABCIndexClass, ABCSeries, ABCSparseSeries)
 
+AnyArrayLike = Union[ABCExtensionArray,
+                     ABCIndexClass,
+                     ABCSeries,
+                     ABCSparseSeries,
+                     np.ndarray]
 ArrayLike = Union[ABCExtensionArray, np.ndarray]
 DatetimeLikeScalar = Type[Union[Period, Timestamp, Timedelta]]
 Dtype = Union[str, np.dtype, ExtensionDtype]

pandas/core/frame.py

@@ -2694,13 +2694,19 @@ def _get_value(self, index, col, takeable=False):
 
         try:
             return engine.get_value(series._values, index)
+        except KeyError:
+            # GH 20629
+            if self.index.nlevels > 1:
+                # partial indexing forbidden
+                raise
         except (TypeError, ValueError):
+            pass
 
-            # we cannot handle direct indexing
-            # use positional
-            col = self.columns.get_loc(col)
-            index = self.index.get_loc(index)
-            return self._get_value(index, col, takeable=True)
+        # we cannot handle direct indexing
+        # use positional
+        col = self.columns.get_loc(col)
+        index = self.index.get_loc(index)
+        return self._get_value(index, col, takeable=True)
     _get_value.__doc__ = get_value.__doc__
 
     def set_value(self, index, col, value, takeable=False):

pandas/core/indexes/category.py

@@ -1,4 +1,5 @@
 import operator
+from typing import Any
 import warnings
 
 import numpy as np
@@ -17,6 +18,7 @@
 from pandas.core.dtypes.generic import ABCCategorical, ABCSeries
 from pandas.core.dtypes.missing import isna
 
+from pandas._typing import AnyArrayLike
 from pandas.core import accessor
 from pandas.core.algorithms import take_1d
 from pandas.core.arrays.categorical import Categorical, contains
@@ -494,16 +496,31 @@ def get_loc(self, key, method=None):
         except KeyError:
             raise KeyError(key)
 
-    def get_value(self, series, key):
+    def get_value(self,
+                  series: AnyArrayLike,
+                  key: Any):
         """
         Fast lookup of value from 1-dimensional ndarray. Only use this if you
         know what you're doing
+
+        Parameters
+        ----------
+        series : Series, ExtensionArray, Index, or ndarray
+            1-dimensional array to take values from
+        key: : scalar
+            The value of this index at the position of the desired value,
+            otherwise the positional index of the desired value
+
+        Returns
+        -------
+        Any
+            The element of the series at the position indicated by the key
         """
         try:
             k = com.values_from_object(key)
             k = self._convert_scalar_indexer(k, kind='getitem')
             indexer = self.get_loc(k)
-            return series.iloc[indexer]
+            return series.take([indexer])[0]
         except (KeyError, TypeError):
             pass
 

pandas/tests/indexing/test_categorical.py

@@ -638,6 +638,16 @@ def test_loc_slice(self):
         # expected = df.iloc[[1,2,3,4]]
         # assert_frame_equal(result, expected)
 
+    def test_loc_and_at_with_categorical_index(self):
+        # GH 20629
+        s = Series([1, 2, 3], index=pd.CategoricalIndex(["A", "B", "C"]))
+        assert s.loc['A'] == 1
+        assert s.at['A'] == 1
+        df = DataFrame([[1, 2], [3, 4], [5, 6]],
+                       index=pd.CategoricalIndex(["A", "B", "C"]))
+        assert df.loc['B', 1] == 4
+        assert df.at['B', 1] == 4
+
     def test_boolean_selection(self):
 
         df3 = self.df3