diff --git a/doc/source/whatsnew/v0.20.0.txt b/doc/source/whatsnew/v0.20.0.txt index 10684021d5599..63c2a93368de1 100644 --- a/doc/source/whatsnew/v0.20.0.txt +++ b/doc/source/whatsnew/v0.20.0.txt @@ -311,6 +311,7 @@ Bug Fixes +- Bug in ``Series.iloc`` where a ``Categorical`` object for list-like indexes input was returned, where a ``Series`` was expected. (:issue:`14580`) diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index 107d68c192ead..dad5bf5bc70ba 100755 --- a/pandas/core/indexing.py +++ b/pandas/core/indexing.py @@ -1596,6 +1596,27 @@ def _get_slice_axis(self, slice_obj, axis=0): else: return self.obj.take(slice_obj, axis=axis, convert=False) + def _get_list_axis(self, key_list, axis=0): + """ + Return Series values by list or array of integers + + Parameters + ---------- + key_list : list-like positional indexer + axis : int (can only be zero) + + Returns + ------- + Series object + """ + + # validate list bounds + self._is_valid_list_like(key_list, axis) + + # force an actual list + key_list = list(key_list) + return self.obj.take(key_list, axis=axis, convert=False) + def _getitem_axis(self, key, axis=0): if isinstance(key, slice): @@ -1606,26 +1627,20 @@ def _getitem_axis(self, key, axis=0): self._has_valid_type(key, axis) return self._getbool_axis(key, axis=axis) - # a single integer or a list of integers - else: - - if is_list_like_indexer(key): - - # validate list bounds - self._is_valid_list_like(key, axis) - - # force an actual list - key = list(key) + # a list of integers + elif is_list_like_indexer(key): + return self._get_list_axis(key, axis=axis) - else: - key = self._convert_scalar_indexer(key, axis) + # a single integer + else: + key = self._convert_scalar_indexer(key, axis) - if not is_integer(key): - raise TypeError("Cannot index by location index with a " - "non-integer key") + if not is_integer(key): + raise TypeError("Cannot index by location index with a " + "non-integer key") - # validate the location - self._is_valid_integer(key, axis) + # validate the location + self._is_valid_integer(key, axis) return self._get_loc(key, axis=axis) diff --git a/pandas/tests/test_categorical.py b/pandas/tests/test_categorical.py index 5d2c317cc0f81..fe61f16cf7eb1 100644 --- a/pandas/tests/test_categorical.py +++ b/pandas/tests/test_categorical.py @@ -54,6 +54,27 @@ def test_getitem_listlike(self): expected = c[np.array([100000]).astype(np.int64)].codes self.assert_numpy_array_equal(result, expected) + def test_getitem_category_type(self): + # GH 14580 + # test iloc() on Series with Categorical data + + s = pd.Series([1, 2, 3]).astype('category') + + # get slice + result = s.iloc[0:2] + expected = pd.Series([1, 2]).astype('category', categories=[1, 2, 3]) + tm.assert_series_equal(result, expected) + + # get list of indexes + result = s.iloc[[0, 1]] + expected = pd.Series([1, 2]).astype('category', categories=[1, 2, 3]) + tm.assert_series_equal(result, expected) + + # get boolean array + result = s.iloc[[True, False, False]] + expected = pd.Series([1]).astype('category', categories=[1, 2, 3]) + tm.assert_series_equal(result, expected) + def test_setitem(self): # int/positional