ENH: first and nth for groupby. GH #994

Author: Chang She

pandas/core/groupby.py

@@ -350,8 +350,19 @@ def ohlc(self):
         return self._cython_agg_general('ohlc')
 
     def last(self):
+        return self.nth(-1)
+
+    def first(self):
+        return self.nth(0)
+
+    def nth(self, n):
         def picker(arr):
-            return arr[-1] if arr is not None and len(arr) else np.nan
+            if arr is not None:
+                n_ok_pos = n >= 0 and len(arr) > n
+                n_ok_neg = n < 0 and len(arr) >= n
+                if n_ok_pos or n_ok_neg:
+                    return arr.iget(n)
+            return np.nan
         return self.agg(picker)
 
     def _cython_agg_general(self, how):

pandas/tests/test_groupby.py

@@ -120,6 +120,25 @@ def test_basic(self):
         # corner cases
         self.assertRaises(Exception, grouped.aggregate, lambda x: x * 2)
 
+        grouped = self.df.groupby('A')
+        first = grouped.first()
+        expected = grouped.get_group('bar')
+        expected = expected.xs(expected.index[0])[1:]
+        expected.name ='bar'
+        assert_series_equal(first.xs('bar'), expected)
+
+        last = grouped.last()
+        expected = grouped.get_group('bar')
+        expected = expected.xs(expected.index[-1])[1:]
+        expected.name ='bar'
+        assert_series_equal(last.xs('bar'), expected)
+
+        nth = grouped.nth(1)
+        expected = grouped.get_group('foo')
+        expected = expected.xs(expected.index[1])[1:]
+        expected.name ='foo'
+        assert_series_equal(nth.xs('foo'), expected)
+
     def test_groupby_dict_mapping(self):
         # GH #679
         from pandas import Series