Repository: dask/dask

Issue:
TopK
Parallel and out-of-core sorting is hard.  Fortunately it's not often what we want.  We often want just the top k elements.  This is quite doable.

Here are some ways of spelling this that make sense to me.  Alternatives welcome.

``` python
>>> df.topk(10, 'balance')

or 

>>> df.topk(10, df.balance)

or 

>>> df.balance.topk(10)
```

This should be doable with the `apply_concat_apply` function.

After `topk` works one might extend this to also work with GroupBy objects.

``` python
>>> df.groupby(df.year).topk(10, 'balance')  # top ten accounts by balance per year
```



Discussion:
Pandas calls this [`nlargest`](http://pandas.pydata.org/pandas-docs/dev/generated/pandas.Series.nlargest.html) which is taken from the [`heapq` function of the same name](https://docs.python.org/2/library/heapq.html#heapq.nlargest).

We actually don't have anything like this yet:

``` python
df.nlargest(5, ['balance'])
```

but I think I have this in a branch that I can push up.

pr'd this in core pandas https://github.com/pydata/pandas/pull/10393

It looks like the pandas PR has gone in.  This is up for grabs in dask.dataframe


Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -885,9 +885,7 @@ def value_counts(self):
 
     @wraps(pd.Series.nlargest)
     def nlargest(self, n=5):
-        f = lambda s: s.nlargest(n)
-        token = 'series-nlargest-n={0}'.format(n)
-        return aca(self, f, f, columns=self.name, token=token)
+        return nlargest(self, n)
 
     @wraps(pd.Series.isin)
     def isin(self, other):
@@ -1102,6 +1100,13 @@ def column_info(self):
         """ Return DataFrame.columns """
         return self.columns
 
+    def nlargest(self, n=5, columns=None):
+        """
+        Return the rows which contain the largest n elements from the provided
+        column, in descending order.
+        """
+        return nlargest(self, n, columns)
+
     @wraps(pd.DataFrame.groupby)
     def groupby(self, key, **kwargs):
         return GroupBy(self, key, **kwargs)
@@ -1336,6 +1341,19 @@ def elemwise_property(attr, s):
     setattr(Index, name, property(partial(elemwise_property, name)))
 
 
+def nlargest(df, n=5, columns=None):
+    if isinstance(df, Index):
+        raise AttributeError("nlargest is not available for Index objects")
+    elif isinstance(df, Series):
+        token = 'series-nlargest-n={0}'.format(n)
+        f = lambda s: s.nlargest(n)
+    elif isinstance(df, DataFrame):
+        token = 'dataframe-nlargest-n={0}'.format(n)
+        f = lambda df: df.nlargest(n, columns)
+        columns = df.columns  # this is a hack.
+    return aca(df, f, f, columns=columns, token=token)
+
+
 def _assign(df, *pairs):
     kwargs = dict(partition(2, pairs))
     return df.assign(**kwargs)


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1,6 +1,7 @@
 from itertools import product
 from datetime import datetime
 from operator import getitem
+from distutils.version import LooseVersion
 
 import pandas as pd
 import pandas.util.testing as tm
@@ -2140,3 +2141,30 @@ def test_index_time_properties():
 
     assert (i.index.day == a.index.day.compute()).all()
     assert (i.index.month == a.index.month.compute()).all()
+
+
+@pytest.mark.skipif(LooseVersion(pd.__version__) <= '0.16.2',
+        reason="nlargest not in pandas pre 0.16.2")
+def test_nlargest():
+    from string import ascii_lowercase
+    df = pd.DataFrame({'a': np.random.permutation(10),
+                       'b': list(ascii_lowercase[:10])})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    res = ddf.nlargest(5, 'a')
+    exp = df.nlargest(5, 'a')
+    eq(res, exp)
+
+
+@pytest.mark.skipif(LooseVersion(pd.__version__) <= '0.16.2',
+        reason="nlargest not in pandas pre 0.16.2")
+def test_nlargest_multiple_columns():
+    from string import ascii_lowercase
+    df = pd.DataFrame({'a': np.random.permutation(10),
+                       'b': list(ascii_lowercase[:10]),
+                       'c': np.random.permutation(10).astype('float64')})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    result = ddf.nlargest(5, ['a', 'b'])
+    expected = df.nlargest(5, ['a', 'b'])
+    eq(result, expected)