BUG: iloc misbehavior with pd.Series: sometimes returns pd.Categorical, fixes #14580

closes #14580

Author: Nathalie Rud <[email protected]>

Closes #14642 from nathalier/nath-working and squashes the following commits:

f6b1790 [Nathalie Rud] DOC: add whatsnew entry for bug GH #14580 fix
bb9b601 [Nathalie Rud] TST: add test to cover bug GH #14580
9a69ca5 [Nathalie Rud] CLN: refactor _getitem_axis() of  _iLocIndexer class
5b5ead4 [Nathalie Rud] BUG: fix iloc with pd.Series on pd.Categorical

Author: nathalier

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`)
 
 
 

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)
 

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