Repository: pandas-dev/pandas

Issue:
nth groupby method on DataFrame
The nth groupby method on a Series takes the nth **non-NaN** value in the Series.

This means on  a DataFrame it's not going to be well defined...

Should we make this a Series only method?

```
In [1]: g = pd.DataFrame([['a'], ['a'], ['a'], ['b'], ['b'], ['a']],columns=['A']).groupby('A')

In [2]: g.nth(0)
Out[2]: 
Empty DataFrame
Columns: []
Index: []

In [3]: g.A.nth(0)
Out[3]: 
A
a    a
b    b
Name: A, dtype: object
```



Discussion:
Maybe I just don't understand what nth does, not always empty.

@hayd Maybe `nth` on a `DataFrame` should take the column as an argument and then it would then solve #5503

It's also really slow: http://stackoverflow.com/a/20087789/1240268

I'm of the opinion that nth should take the nth, regardless of NaN, or take kwarg for NaN-ness. Need to think through the logic of what it is doing atm (it's only small!)

For convenience the current impl is:

```
def nth(self, n):
    def picker(arr):
        arr = arr[notnull(arr)]
        if len(arr) >= n + 1:
            return arr.iget(n)
        else:
            return np.nan
    return self.agg(picker)
```

With a frame the `arr = arr[notnull(arr)]` does nothing, I really don't understand how the next bit even works since iget is not a DataFrame method....!

I like the idea of a kwarg to describe this, grr to current behaviour: Series/DataFrame being inconsistent. Not sure how we should play this.

...in any case we should use cumcount, much faster.

Other possibility is plonk/override this into Series Groupby "as is*", but have it in Groupby the just take the nth row regardless of NaN, at least for now.

Also atm this is not that great, as it NaNs for groups smaller than n.

\* but using cumcount.

Any thoughts on this @jreback @TomAugspurger?

Should we somehow depreciate old Series behaviour (of getting nth non-null result, I think this may be something that R does which is why we do it...)? Not sure on way forward.

I still don't understand how this even works on a DataFrame... cumcount impl much easier (confusing to be different to Series though it already is).

I think u should break the API and have a na='drop' (but default to na=None), meaning don't remove Nan's 

Should we have something similar for DataFrames? Could be drop_na (think this is inline with other args), You think this arg should only be for Series or hack something which will work for Dataframes too?

I think for frames u could do any/all
equivalent of doing a dropna beforehand and just default o None ?

sure, I think that's not so bad. I will remember to use `_selected_obj` :)

I'm not sure on the API for this to the old behaviour (i.e. not just filtering), for one thing there are two NaNs:
- NaNs because there are no more results in the group.
- NaNs values from the actual DataFrame/Series.

For another these are two very different results, you only get the groupby's by info if you use the old method.

I still think should change... just not sure on API for it.

```
In [11]: s = pd.Series([1, 2, 3, 1, np.nan])

In [12]: s.groupby([1, 2, 3, 1, 2]).nth(1)  # the NaN has meaning here (it was missing in s)
Out[12]:
3     1
4   NaN
dtype: float64

s.groupby([1, 2, 3, 1, 2]).nth(1)  # old, has index from by, the NaN means group not that big (with dropped nas), observe that index is from by
1     1
2   NaN
3   NaN
dtype: float64
```

travis is now happy...so go ahead at your leisure

@jreback still unsure about the API, any thoughts. I guess the old behaviour is an R thing?

I guess it's annoying that these are two very different things:

a. filter frame and series to take just the nth rows. result is subframe/series.
b. what it's doing before (at least in the series case, for dataframe it currently doesn't work*), which is grab the nth non-null for each group or NaN if not there, and index but the groupby by. result is not sub frame/series.

\* amusingly I think the iget call I was confused about raises but is caught so get NaN

I think you need to support both cases, but make default a)

`df.groupby(...).nth(5,drop_na=False)`

so `drop_na=False` is a), `drop_na=True` is b)


Resolving patch:
diff --git a/pandas/core/groupby.py b/pandas/core/groupby.py
--- a/pandas/core/groupby.py
+++ b/pandas/core/groupby.py
@@ -523,15 +523,75 @@ def ohlc(self):
         """
         return self._cython_agg_general('ohlc')
 
-    def nth(self, n):
-        def picker(arr):
-            arr = arr[notnull(arr)]
-            if len(arr) >= n + 1:
-                return arr.iget(n)
+    def nth(self, n, dropna=None):
+        """
+        Take the nth row from each group.
+
+        If dropna, will not show nth non-null row, dropna is either
+        Truthy (if a Series) or 'all', 'any' (if a DataFrame); this is equivalent
+        to calling dropna(how=dropna) before the groupby.
+
+        Examples
+        --------
+        >>> DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
+        >>> g = df.groupby('A')
+        >>> g.nth(0)
+           A   B
+        0  1 NaN
+        2  5   6
+        >>> g.nth(1)
+           A  B
+        1  1  4
+        >>> g.nth(-1)
+           A  B
+        1  1  4
+        2  5  6
+        >>> g.nth(0, dropna='any')
+           B
+        A
+        1  4
+        5  6
+        >>> g.nth(1, dropna='any')  # NaNs denote group exhausted when using dropna
+            B
+        A
+        1 NaN
+        5 NaN
+
+        """
+
+        if not dropna:  # good choice
+            m = self.grouper._max_groupsize
+            if n >= m or n < -m:
+                return self._selected_obj.loc[[]]
+            rng = np.zeros(m, dtype=bool)
+            if n >= 0:
+                rng[n] = True
+                is_nth = self._cumcount_array(rng)
             else:
+                rng[- n - 1] = True
+                is_nth = self._cumcount_array(rng, ascending=False)
+            return self._selected_obj[is_nth]
+
+        if (isinstance(self._selected_obj, DataFrame)
+            and dropna not in ['any', 'all']):
+            # Note: when agg-ing picker doesn't raise this, just returns NaN
+            raise ValueError("For a DataFrame groupby, dropna must be "
+                             "either None, 'any' or 'all', "
+                             "(was passed %s)." % (dropna),)
+
+        # old behaviour, but with all and any support for DataFrames.
+
+        max_len = n if n >= 0 else - 1 - n
+        def picker(x):
+            x = x.dropna(how=dropna)  # Note: how is ignored if Series
+            if len(x) <= max_len:
                 return np.nan
+            else:
+                return x.iloc[n]
+
         return self.agg(picker)
 
+
     def cumcount(self, **kwargs):
         """
         Number each item in each group from 0 to the length of that group - 1.
@@ -579,8 +639,7 @@ def cumcount(self, **kwargs):
         ascending = kwargs.pop('ascending', True)
 
         index = self.obj.index
-        rng = np.arange(self.grouper._max_groupsize, dtype='int64')
-        cumcounts = self._cumcount_array(rng, ascending=ascending)
+        cumcounts = self._cumcount_array(ascending=ascending)
         return Series(cumcounts, index)
 
     def head(self, n=5):
@@ -606,8 +665,7 @@ def head(self, n=5):
 
         """
         obj = self._selected_obj
-        rng = np.arange(self.grouper._max_groupsize, dtype='int64')
-        in_head = self._cumcount_array(rng) < n
+        in_head = self._cumcount_array() < n
         head = obj[in_head]
         return head
 
@@ -639,11 +697,17 @@ def tail(self, n=5):
         tail = obj[in_tail]
         return tail
 
-    def _cumcount_array(self, arr, **kwargs):
+    def _cumcount_array(self, arr=None, **kwargs):
+        """
+        arr is where cumcount gets it's values from
+        """
         ascending = kwargs.pop('ascending', True)
 
+        if arr is None:
+            arr = np.arange(self.grouper._max_groupsize, dtype='int64')
+
         len_index = len(self.obj.index)
-        cumcounts = np.zeros(len_index, dtype='int64')
+        cumcounts = np.empty(len_index, dtype=arr.dtype)
         if ascending:
             for v in self.indices.values():
                 cumcounts[v] = arr[:len(v)]
diff --git a/vb_suite/groupby.py b/vb_suite/groupby.py
--- a/vb_suite/groupby.py
+++ b/vb_suite/groupby.py
@@ -269,6 +269,22 @@ def f(g):
 groupby_frame_apply = Benchmark("df.groupby(['key', 'key2']).apply(f)", setup,
                                  start_date=datetime(2011, 10, 1))
 
+
+#----------------------------------------------------------------------
+# DataFrame nth
+
+setup = common_setup + """
+df = pd.DataFrame(np.random.randint(1, 100, (10000, 2)))
+"""
+
+# Not really a fair test as behaviour has changed!
+groupby_frame_nth = Benchmark("df.groupby(0).nth(0)", setup,
+                                start_date=datetime(2014, 3, 1))
+
+groupby_series_nth = Benchmark("df[1].groupby(df[0]).nth(0)", setup,
+                                 start_date=datetime(2014, 3, 1))
+
+
 #----------------------------------------------------------------------
 # Sum booleans #2692
 


Test patch:
diff --git a/pandas/tests/test_groupby.py b/pandas/tests/test_groupby.py
--- a/pandas/tests/test_groupby.py
+++ b/pandas/tests/test_groupby.py
@@ -156,8 +156,7 @@ def test_first_last_nth(self):
         assert_frame_equal(last, expected, check_names=False)
 
         nth = grouped.nth(1)
-        expected = self.df.ix[[3, 2], ['B', 'C', 'D']]
-        expected.index = ['bar', 'foo']
+        expected = self.df.iloc[[2, 3]]
         assert_frame_equal(nth, expected, check_names=False)
 
         # it works!
@@ -165,10 +164,10 @@ def test_first_last_nth(self):
         grouped['B'].last()
         grouped['B'].nth(0)
 
-        self.df['B'][self.df['A'] == 'foo'] = np.nan
+        self.df.loc[self.df['A'] == 'foo', 'B'] = np.nan
         self.assert_(com.isnull(grouped['B'].first()['foo']))
         self.assert_(com.isnull(grouped['B'].last()['foo']))
-        self.assert_(com.isnull(grouped['B'].nth(0)['foo']))
+        self.assert_(com.isnull(grouped['B'].nth(0)[0]))  # not sure what this is testing
 
     def test_first_last_nth_dtypes(self):
 
@@ -189,8 +188,7 @@ def test_first_last_nth_dtypes(self):
         assert_frame_equal(last, expected, check_names=False)
 
         nth = grouped.nth(1)
-        expected = df.ix[[3, 2], ['B', 'C', 'D', 'E', 'F']]
-        expected.index = ['bar', 'foo']
+        expected = df.iloc[[2, 3]]
         assert_frame_equal(nth, expected, check_names=False)
 
         # GH 2763, first/last shifting dtypes
@@ -201,6 +199,29 @@ def test_first_last_nth_dtypes(self):
         f = s.groupby(level=0).first()
         self.assertEqual(f.dtype, 'int64')
 
+    def test_nth(self):
+        df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
+        g = df.groupby('A')
+
+        assert_frame_equal(g.nth(0), df.iloc[[0, 2]])
+        assert_frame_equal(g.nth(1), df.iloc[[1]])
+        assert_frame_equal(g.nth(2), df.loc[[]])
+        assert_frame_equal(g.nth(-1), df.iloc[[1, 2]])
+        assert_frame_equal(g.nth(-2), df.iloc[[0]])
+        assert_frame_equal(g.nth(-3), df.loc[[]])
+        assert_series_equal(g.B.nth(0), df.B.iloc[[0, 2]])
+        assert_series_equal(g.B.nth(1), df.B.iloc[[1]])
+        assert_frame_equal(g[['B']].nth(0), df.ix[[0, 2], ['B']])
+
+        exp = df.set_index('A')
+        assert_frame_equal(g.nth(0, dropna='any'), exp.iloc[[1, 2]])
+        assert_frame_equal(g.nth(-1, dropna='any'), exp.iloc[[1, 2]])
+
+        exp['B'] = np.nan
+        assert_frame_equal(g.nth(7, dropna='any'), exp.iloc[[1, 2]])
+        assert_frame_equal(g.nth(2, dropna='any'), exp.iloc[[1, 2]])
+
+
     def test_grouper_index_types(self):
         # related GH5375
         # groupby misbehaving when using a Floatlike index