Merge pull request #7631 from jreback/groupby_nth_fix

BUG: doc example in groupby.rst (GH7559 / GH7628)

Author: jreback

pandas/core/groupby.py

@@ -829,19 +829,29 @@ def nth(self, n, dropna=None):
         dropped = self.obj.dropna(how=dropna, axis=self.axis)
 
         # get a new grouper for our dropped obj
-        grouper, exclusions, obj = _get_grouper(dropped, key=self.keys, axis=self.axis,
-                                                level=self.level, sort=self.sort)
+        if self.keys is None and self.level is None:
 
-        sizes = obj.groupby(grouper).size()
-        result = obj.groupby(grouper).nth(n)
+            # we don't have the grouper info available (e.g. we have selected out
+            # a column that is not in the current object)
+            axis = self.grouper.axis
+            grouper = axis[axis.isin(dropped.index)]
+            keys = self.grouper.names
+        else:
+
+            # create a grouper with the original parameters, but on the dropped object
+            grouper, _, _ = _get_grouper(dropped, key=self.keys, axis=self.axis,
+                                         level=self.level, sort=self.sort)
+
+        sizes = dropped.groupby(grouper).size()
+        result = dropped.groupby(grouper).nth(n)
         mask = (sizes<max_len).values
 
         # set the results which don't meet the criteria
         if len(result) and mask.any():
             result.loc[mask] = np.nan
 
         # reset/reindex to the original groups
-        if len(self.obj) == len(dropped):
+        if len(self.obj) == len(dropped) or len(result) == len(self.grouper.result_index):
             result.index = self.grouper.result_index
         else:
             result = result.reindex(self.grouper.result_index)

pandas/tests/test_groupby.py

@@ -304,6 +304,13 @@ def test_nth(self):
         result = s.groupby(g).nth(0,dropna='all')
         assert_series_equal(result,expected)
 
+        # doc example
+        df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
+        g = df.groupby('A')
+        result = g.B.nth(0, dropna=True)
+        expected = g.B.first()
+        assert_series_equal(result,expected)
+
     def test_grouper_index_types(self):
         # related GH5375
         # groupby misbehaving when using a Floatlike index