Repository: dask/dask

Issue:
groupby.apply + multiprocessing causes error
Running the following code causes an RemoteError for me:

import dask
import pandas as pd
import dask.dataframe as dd
pdf = pd.DataFrame({'A':[1,2,3,4,5], 'B':['1','1','a','a','a']})
ddf = dd.from_pandas(pdf, npartitions = 3)
groups = ddf.groupby('B').apply(lambda x: x).compute(get = dask.multiprocessing.get)

Using the threaded scheduler, this problem does not occur.

I get the following error stack:
  File "<stdin>", line 1, in <module>
  File "/[path_to_anaconda]/lib/python2.7/site-packages/dask/base.py", line 86, in compute
    return compute(self, *_kwargs)[0]
  File "/[path_to_anaconda]/lib/python2.7/site-packages/dask/base.py", line 179, in compute
    results = get(dsk, keys, *_kwargs)
  File "[path_to_anaconda]a/lib/python2.7/site-packages/dask/multiprocessing.py", line 83, in get
    queue=queue, get_id=_process_get_id, **kwargs)
  File "/[path_to_anaconda]/lib/python2.7/site-packages/dask/async.py", line 470, in get_async
    key, res, tb, worker_id = queue.get()
  File "<string>", line 2, in get
  File "/[path_to_anaconda]/lib/python2.7/multiprocessing/managers.py", line 774, in _callmethod
    raise convert_to_error(kind, result)
multiprocessing.managers.RemoteError: 

---

Unserializable message: ('#RETURN', (('zpartd-7ec42c1484d311e6b85f782bcb709c3f',), <partd.encode.Encode object at 0x7fe80bd8a250>, None, 1007))

---



Discussion:
Thanks for the bug report - I'm able to reproduce on master, although I get an empty dataframe instead of an error.

The issue is due to a bug in partd (a composable key-value store that we use for shuffle operations). We're writing to a buffered file, but not flushing the buffer (a python dictionary) afterwards. For larger datasets the buffer is usually empty, so it just happened to work before. Since the test set is small, the buffer has all the data and the result is an empty dataframe.

So does this mean, that this error should not occur on larger datasets? Because for me, as far as i can tell, the size of the dataframe doesn't matter: using the multiprocessing scheduler, I always get an error.

> Am 27.09.2016 um 23:30 schrieb Jim Crist notifications@github.com:
> 
> Thanks for the bug report - I'm able to reproduce on master, although I get an empty dataframe instead of an error.
> 
> The issue is due to a bug in partd (a composable key-value store that we use for shuffle operations). We're writing to a buffered file, but not flushing the buffer (a python dictionary) afterwards. For larger datasets the buffer is usually empty, so it just happened to work before. Since the test set is small, the buffer has all the data and the result is an empty dataframe.
> 
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub, or mute the thread.

It appears you're using an old version of dask (from a comment in a different issue) - try upgrading to the latest version (0.11), or master (`pip install -U git+https://github.com/dask/dask.git`) and running again. You'll also need to upgrade `partd`. I get a wrong answer, but not an error, so there's still a bug here.

Note that if you're loading all the data into memory already, then the task-based shuffle in the distributed scheduler might be more performant here. You can specify this by passing in `shuffle=task` to the groupby call. I recommend using this with the distributed scheduler.


Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -164,6 +164,23 @@ def rearrange_by_column(df, col, npartitions=None, max_branch=None,
         raise NotImplementedError("Unknown shuffle method %s" % shuffle)
 
 
+class maybe_buffered_partd(object):
+    """If serialized, will return non-buffered partd. Otherwise returns a
+    buffered partd"""
+    def __init__(self, buffer=True):
+        self.buffer = buffer
+
+    def __reduce__(self):
+        return (maybe_buffered_partd, (False,))
+
+    def __call__(self, *args, **kwargs):
+        import partd
+        if self.buffer:
+            return partd.PandasBlocks(partd.Buffer(partd.Dict(), partd.File()))
+        else:
+            return partd.PandasBlocks(partd.File())
+
+
 def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
     """ Shuffle using local disk """
     if npartitions is None:
@@ -172,10 +189,8 @@ def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
     token = tokenize(df, column, npartitions)
     always_new_token = uuid.uuid1().hex
 
-    import partd
     p = ('zpartd-' + always_new_token,)
-    dsk1 = {p: (partd.PandasBlocks, (partd.Buffer, (partd.Dict,),
-                                                   (partd.File,)))}
+    dsk1 = {p: (maybe_buffered_partd(),)}
 
     # Partition data on disk
     name = 'shuffle-partition-' + always_new_token


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -484,3 +484,13 @@ def test_groupby_apply_tasks():
             b = ddf.groupby(ind(ddf)).B.apply(len)
             assert eq(a, b.compute())
             assert not any('partd' in k[0] for k in b.dask)
+
+
+def test_groupby_multiprocessing():
+    from dask.multiprocessing import get
+    df = pd.DataFrame({'A': [1, 2, 3, 4, 5],
+                       'B': ['1','1','a','a','a']})
+    ddf = dd.from_pandas(df, npartitions=3)
+    with dask.set_options(get=get):
+        assert eq(ddf.groupby('B').apply(lambda x: x),
+                  df.groupby('B').apply(lambda x: x))
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1,13 +1,17 @@
 import pandas as pd
 import pandas.util.testing as tm
 import pytest
+import pickle
 import numpy as np
 
 import dask.dataframe as dd
+from dask.threaded import get as threaded_get
+from dask.multiprocessing import get as mp_get
 from dask.dataframe.shuffle import (shuffle, hash_series,
                                     partitioning_index,
                                     rearrange_by_column,
-                                    rearrange_by_divisions)
+                                    rearrange_by_divisions,
+                                    maybe_buffered_partd)
 from dask.async import get_sync
 from dask.dataframe.utils import eq, make_meta
 
@@ -216,7 +220,8 @@ def test_shuffle_sort(shuffle):
 
 
 @pytest.mark.parametrize('shuffle', ['tasks', 'disk'])
-def test_rearrange(shuffle):
+@pytest.mark.parametrize('get', [threaded_get, mp_get])
+def test_rearrange(shuffle, get):
     df = pd.DataFrame({'x': range(10)})
     ddf = dd.from_pandas(df, npartitions=4)
     ddf2 = ddf.assign(y=ddf.x % 4)
@@ -226,8 +231,8 @@ def test_rearrange(shuffle):
     assert set(ddf.dask).issubset(result.dask)
 
     # Every value in exactly one partition
-    a = result.compute()
-    parts = get_sync(result.dask, result._keys())
+    a = result.compute(get=get)
+    parts = get(result.dask, result._keys())
     for i in a.y.drop_duplicates():
         assert sum(i in part.y for part in parts) == 1
 
@@ -239,3 +244,14 @@ def test_rearrange_by_column_with_narrow_divisions():
 
     df = rearrange_by_divisions(a, 'x', (0, 2, 5))
     list_eq(df, a)
+
+
+def test_maybe_buffered_partd():
+    import partd
+    f = maybe_buffered_partd()
+    p1 = f()
+    assert isinstance(p1.partd, partd.Buffer)
+    f2 = pickle.loads(pickle.dumps(f))
+    assert not f2.buffer
+    p2 = f2()
+    assert isinstance(p2.partd, partd.File)