BUG: Groupby.nth includes group key inconsistently pandas-dev#12839

Added tests

Author: adneu

doc/source/whatsnew/v0.18.2.txt

@@ -313,7 +313,7 @@ Bug Fixes
 
 
 - Bug in ``groupby`` where ``apply`` returns different result depending on whether first result is ``None`` or not (:issue:`12824`)
-
+- Bug in ``groupby(..).nth()`` where the group key is included inconsistently (:issue:`12839`)
 
 
 

pandas/core/groupby.py

@@ -457,6 +457,11 @@ def _selected_obj(self):
         else:
             return self.obj[self._selection]
 
+    def _reset_group_selection(self):
+        if self._group_selection is not None:
+            self._group_selection = None
+            self._reset_cache('_selected_obj')
+
     def _set_selection_from_grouper(self):
         """ we may need create a selection if we have non-level groupers """
         grp = self.grouper
@@ -468,6 +473,7 @@ def _set_selection_from_grouper(self):
 
             if len(groupers):
                 self._group_selection = ax.difference(Index(groupers)).tolist()
+                self._reset_cache('_selected_obj')
 
     def _set_result_index_ordered(self, result):
         # set the result index on the passed values object and
@@ -1402,6 +1408,7 @@ def head(self, n=5):
         0  1  2
         2  5  6
         """
+        self._reset_group_selection()
         mask = self._cumcount_array() < n
         return self._selected_obj[mask]
 
@@ -1428,6 +1435,7 @@ def tail(self, n=5):
         0  a  1
         2  b  1
         """
+        self._reset_group_selection()
         mask = self._cumcount_array(ascending=False) < n
         return self._selected_obj[mask]
 

pandas/tests/test_groupby.py

@@ -354,6 +354,41 @@ def test_nth_multi_index_as_expected(self):
                                          names=['A', 'B']))
         assert_frame_equal(result, expected)
 
+    def test_group_selection_cache(self):
+        # GH 12839 nth, head, and tail should return same result consistently
+        df = DataFrame([[1, 2], [1, 4], [5, 6]], columns=['A', 'B'])
+        expected = df.iloc[[0, 2]].set_index('A')
+
+        g = df.groupby('A')
+        g.head()
+        result = g.nth(0)
+        assert_frame_equal(result, expected)
+
+        g = df.groupby('A')
+        g.tail()
+        result = g.nth(0)
+        assert_frame_equal(result, expected)
+
+        g = df.groupby('A')
+        g.nth(0)
+        result = g.head(n=2)
+        assert_frame_equal(result, df)
+
+        g = df.groupby('A')
+        g.nth(0)
+        result = g.tail(n=2)
+        assert_frame_equal(result, df)
+
+        g = df.groupby('A')
+        g.head()
+        result = g.head(n=2)
+        assert_frame_equal(result, df)
+
+        g = df.groupby('A')
+        g.tail()
+        result = g.tail(n=2)
+        assert_frame_equal(result, df)
+
     def test_grouper_index_types(self):
         # related GH5375
         # groupby misbehaving when using a Floatlike index