Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

BUG-20629 allow .at accessor with CategoricalIndex #26298

Merged
merged 10 commits into from
May 20, 2019
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v0.25.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`)
-
-

Expand Down
8 changes: 7 additions & 1 deletion pandas/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
16 changes: 11 additions & 5 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
21 changes: 19 additions & 2 deletions pandas/core/indexes/category.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import operator
from typing import Any
import warnings

import numpy as np
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down
10 changes: 10 additions & 0 deletions pandas/tests/indexing/test_categorical.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down