Repository: dask/dask

Issue:
Error when computing top 10 bigrams in corpus
When trying to find topk bigrams in my text corpus:

``` python
b.pluck('abstract').str.split().map(lambda l: list(sliding_window(2, l))).concat().frequencies().topk(10, key=1).compute()
```

I get the following error:

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-78-03445e41555d> in <module>()
----> 1 b.pluck('abstract').str.split().map(lambda l: list(sliding_window(2, l))).concat().frequencies().topk(10, key=1).compute()

/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/bag/core.pyc in compute(self, **kwargs)
    570     def compute(self, **kwargs):
    571         """ Force evaluation of bag """
--> 572         results = get(self.dask, self._keys(), **kwargs)
    573         if isinstance(results[0], Iterable):
    574             results = toolz.concat(results)

/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/bag/core.pyc in get(dsk, keys, get, **kwargs)
     82     dsk2 = optimize(dsk, keys)
     83 
---> 84     return get(dsk2, keys, **kwargs)
     85 _get = get
     86 

/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/multiprocessing.pyc in get(dsk, keys, optimizations, num_workers, func_loads, func_dumps)
     49         # Run
     50         result = get_async(apply_async, len(pool._pool), dsk2, keys,
---> 51                            queue=queue)
     52     finally:
     53         if cleanup:

/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/async.pyc in get_async(apply_async, num_workers, dsk, result, cache, debug_counts, queue, raise_on_exception, **kwargs)
    445         if isinstance(res, Exception):
    446             raise type(res)(" Exception in remote process\n\n"
--> 447                 + str(res) + "\n\nTraceback:\n" + tb)
    448         state['cache'][key] = res
    449         finish_task(dsk, key, state, results)

TypeError:  Exception in remote process



Traceback:
  File "/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/async.py", line 258, in execute_task
    result = _execute_task(task, data)
  File "/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/async.py", line 240, in _execute_task
    args2 = [_execute_task(a, cache) for a in args]
  File "/Users/cdoig/anaconda/envs/dask-tokenize/lib/python2.7/site-packages/dask/async.py", line 241, in _execute_task
    return func(*args2)
  File "cytoolz/functoolz.pyx", line 256, in cytoolz.functoolz.curry.__call__ (cytoolz/functoolz.c:3300)
    raise val
```



Discussion:
`heapq.nlargest` (used by `dask.Bag.topk`) expects a callable for the `key` parameter.  `toolz.topk` can accept a callable, index, or list of indices as `key`.  Maybe we should switch to `toolz.topk`.

Indeed.  This was actually the original use case for why I wanted `toolz.topk`.  I'll get around to this sometime today.


Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -3,7 +3,6 @@
 import itertools
 import math
 import tempfile
-import heapq
 import inspect
 import gzip
 import bz2
@@ -16,12 +15,12 @@
 
 from toolz import (merge, frequencies, merge_with, take, curry, reduce,
                    join, reduceby, valmap, count, map, partition_all, filter,
-                   pluck, groupby)
+                   pluck, groupby, topk)
 import toolz
 from ..utils import tmpfile, ignoring
 with ignoring(ImportError):
     from cytoolz import (curry, frequencies, merge_with, join, reduceby,
-                         count, pluck, groupby)
+                         count, pluck, groupby, topk)
 
 from pbag import PBag
 
@@ -351,12 +350,12 @@ def topk(self, k, key=None):
         a = next(names)
         b = next(names)
         if key:
-            topk = curry(heapq.nlargest, key=key)
+            func = curry(topk, key=key)
         else:
-            topk = heapq.nlargest
-        dsk = dict(((a, i), (list, (topk, k, (self.name, i))))
+            func = topk
+        dsk = dict(((a, i), (list, (func, k, (self.name, i))))
                         for i in range(self.npartitions))
-        dsk2 = {(b, 0): (list, (topk, k, (toolz.concat, list(dsk.keys()))))}
+        dsk2 = {(b, 0): (list, (func, k, (toolz.concat, list(dsk.keys()))))}
         return Bag(merge(self.dask, dsk, dsk2), b, 1)
 
     def distinct(self):


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -106,6 +106,12 @@ def test_topk():
     assert list(b.topk(4, key=lambda x: -x)) == [0, 0, 0, 1]
 
 
+def test_topk_with_non_callable_key():
+    b = db.from_sequence([(1, 10), (2, 9), (3, 8)], npartitions=2)
+    assert list(b.topk(2, key=1)) == [(1, 10), (2, 9)]
+    assert list(b.topk(2, key=0)) == [(3, 8), (2, 9)]
+
+
 def test_lambdas():
     assert list(b.map(lambda x: x + 1)) == list(b.map(inc))