Repository: dask/dask

Issue:
dataframe merge returns incorrect result when data is from multi-partition parquet
Similar to issue #2237 , I am seeing dask dataframe merge result can be incorrect when data is from multi-partition parquet. pandas merge always produces correct results.

I have put my test data in https://github.com/ningsean/dask_merge_issue
I am running a local cluster at localhost:9898

My setup is Python 3.6.1, with the following module versions. 

```
dask.__version__
'0.15.2+25.ga1a75a7'

numpy.__version__
'1.13.1'

partd.__version__
'0.3.8'

pandas.__version__
'0.20.1'
```

Here is the code that shows the problem.
```
from dask.distributed import Client
import dask.dataframe as dd
import pandas as pd

client = Client("localhost:9898") 

ddf1 = dd.read_parquet("keys_133m.parq")
ddf1.columns
#Index(['KEY'], dtype='object')
ddf1.dtypes
#KEY    int64

ddf2 = dd.read_parquet("key_ids_133m.parq")
ddf2.columns
#Index(['KEY', 'ID'], dtype='object')
ddf2.dtypes
#KEY    int64
#ID     int64
```

### pandas inner merge result has about **500k** rows
```
df = ddf1.compute().merge(ddf2.compute(),how='inner',on='KEY')
df.shape
#(505847, 2)
```

### dask inner merge has only about **272k** rows
```
ddf = ddf1.merge(ddf2,how='inner',on='KEY')
len(ddf)
#272051
```

### One example of missing row in ddf is KEY==135337982:
```
#dask merge result does not contain the row
ddf.loc[ddf.KEY==135337982].compute()
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []

#pandas merge result contains the row
df.loc[df.KEY==135337982]
#              KEY         ID
#122842  135337982  149107799
```

### Single-partition parquet works
```
#use nlargest to write to a single-partition parquet
ddf1.nlargest(n=505850,columns='KEY').to_parquet('keys_one_partition.parq')
ddf2.nlargest(n=1960255,columns='KEY').to_parquet('key_ids_one_partition.parq')

ddf1_one = dd.read_parquet('keys_one_partition.parq')
ddf2_one = dd.read_parquet('key_ids_one_partition.parq')

ddf_one = ddf1_one.merge(ddf2_one,how='inner',on='KEY')
len(ddf_one)
#505847

df_one = ddf1_one.compute().merge(ddf2_one.compute(),how='inner',on='KEY')
df_one.shape
#(505847, 2)
```


Discussion:
Interestingly, dask returns the correct value on the threaded scheduler.
Hmmm, I get incorrect result with different schedulers, even with a single-threaded scheduler (get=dask.get). The only way that I can get the correct result is to do the merge with pandas dataframes.

To summarize it again, here are the different merges I have tried, and only pandas merge worked.

```
ddf = ddf1.merge(ddf2,how='inner',on='KEY')

#single-threaded scheduler returns incorrect result
ddf.loc[ddf.KEY==135337982].compute(get=dask.get)
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []

#multi-process scheduler returns incorrect result
#ddf.loc[ddf.KEY==135337982].compute(get=dask.multiprocessing.get)
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []

#multi-threaded scheduler returns incorrect result
ddf.loc[ddf.KEY==135337982].compute(get=dask.threaded.get)
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []

#distributed scheduler returns incorrect result
ddf.loc[ddf.KEY==135337982].compute(get=client.get)
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []


#pandas dataframe merge returns correct result
df = ddf1.compute().merge(ddf2.compute(),how='inner',on='KEY')
df.loc[df.KEY==135337982]
#              KEY         ID
#122842  135337982  149107799

```
Strange, if I init a distributed client, the single-threaded scheduler gives the wrong result.
```
from dask.distributed import Client
import dask.dataframe as dd
import dask
client = Client("localhost:9898")
ddf1 = dd.read_parquet("keys_133m.parq")
ddf2 = dd.read_parquet("key_ids_133m.parq")
ddf = ddf1.merge(ddf2,how='inner',on='KEY')
df = ddf.compute(get=dask.get)
print(df.loc[df.KEY==135337982])

#result:
#Empty DataFrame
#Columns: [KEY, ID]
#Index: []
```

But if I don't init the client at all, the single-threaded scheduler returns the correct result (and a few times slower):
```
from dask.distributed import Client
import dask.dataframe as dd
import dask
############# do not init the distributed client ############### client = Client("localhost:9898")
ddf1 = dd.read_parquet("keys_133m.parq")
ddf2 = dd.read_parquet("key_ids_133m.parq")
ddf = ddf1.merge(ddf2,how='inner',on='KEY')
df = ddf.compute(get=dask.get)
print(df.loc[df.KEY==135337982])

#result:
#          KEY         ID
#62  135337982  149107799
```

In both cases, I checked the status bokeh app and scheduler and worker log, there didn't seem to be any activities with the local cluster. 

Another that I noticed is that, with the client init'ed, the ddf.dask has 28044 tasks with len(ddf.dask), while without it the ddf.dask as only 4779 tasks.

So, the results seem to indicate that even though not used, a client to the local cluster affects the task graph, the correctness of the merge results as well as the speed. This is totally unexpected to me.
When dask.distributed starts up it changes the global shuffle mechanism from "disk" to "tasks".  You can also control this with the `dask.set_options(shuffle='tasks')` context manager.  It might be useful to try that out on your problem to see if you can recreate your issue without starting a client.

```python
with dask.set_options(shuffle='tasks'):
    ...
```

If so then that will help to isolate the error.
Yes, `dask.set_options(shuffle='tasks')` has the same effect as starting a client. Now the error happens even without starting the client.
Interesting. If I make the "KEY" column as the index for both dask dataframes, then use join, the result would come out correct. It is slightly slower than doing merge directly due to the overhead of set_index. 

If I persist the indexed dataframes and the non-indexed dataframes, the join on the indexed would be much faster than the merge on the non-indexed, **with the additional benefit of the result being correct!**

I tried to debug and pin-point the issue with the merge but haven't been able to get anywhere. 

```
client = Client('localhost:9898')
keys = ddf1.set_index('KEY', drop=False) 
keyids = ddf2.set_index('KEY', drop=False)
ddf = keys.join(keyids, rsuffix='_y') 
ddf.loc[ddf.KEY==135337982].compute()

# result:
#                 KEY      KEY_y         ID
#KEY                                       
#135337982  135337982  135337982  149107799

len(ddf)

# result:
# 505850

```

bump - this is something that i am seeing as well.
@mrocklin @jcrist Anybody actively working on this one?  Just ran into this too.
@brendancol I'm not working on it currently, nor is anyone that I know of.  Generally people write on an issue when they start work.  I agree that this is important.  I don't personally plan to work on it any time in the next week or two.  If anyone wants to take it on that would be great.
@mrocklin i am working on creating a reproducible example. I haven't been able to reproduce this issue outside of the production environment though. Will update this thread.

bump - Any news on this? I also find this issue very important because merge is a very common operation and forced to set the index first slows me down quite a bit.
It looks like this is waiting on a reproducible example that demonstrates the failure in a simple way.  Perhaps you are able to provide such an example @dennis-ec ?
I haven't been able to reproduce this using synthetic data. However, the issue persists with the parquet files that @ningsean uploaded.
```python

import requests
import dask
import dask.dataframe as dd
import zipfile
r = requests.get("https://github.com/ningsean/dask_merge_issue/raw/master/data.zip" )
with open('data.zip', 'wb') as f:
    f.write(r.content)

zipfile.ZipFile('data.zip', 'r').extractall('.')

dd_right = dd.read_parquet('keys_133m.parq/', index=False)
dd_left = dd.read_parquet('key_ids_133m.parq/', columns = ['KEY', 'ID'], index=False)
with dask.set_options(shuffle='tasks'):
    dd_merged = dd_left.merge(dd_right, how='inner', on='KEY')
    df_computed = dd_merged.compute()

computed_from_pandas = dd_left.compute().merge(dd_right.compute(), how='inner', on='KEY')
assert len(df_computed) == len(computed_from_pandas)
```
Output
```python
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-41-d28889968375> in <module>()
----> 1 assert len(df_computed) == len(computed_from_pandas)

AssertionError: 
```
Hope this is helpful. 

```
In [7]: dask.__version__
Out[7]: '0.17.0'

In [8]: pd.__version__
Out[8]: '0.22.0'
```
Unfortunately we won't be able to include these files in tests.  We really do need something that can be generated automatically.
yea i will play around some more to reproduce this with synthetic data
This is a really weird one. It appears that the task shuffle breaks down when number of tasks increase to a big number. here is a reproducible example (perhaps not minimal):
```python
In [14]: import numpy as np
    ...: import pandas as pd
    ...: import dask
    ...: import dask.dataframe as dd
    ...: large = pd.DataFrame({'KEY': np.arange(0,50000)})  
    ...: small = pd.DataFrame({'KEY': np.arange(2500,50000)})
    ...: dd_large = dd.from_pandas(large,npartitions=257)
    ...: dd_small = dd.from_pandas(small, npartitions=3)
    ...: with dask.set_options(shuffle='tasks'):
    ...:     dd_merged = dd_large.merge(dd_small, how='inner', on='KEY')
    ...:     result = dd_merged.compute()
    ...:     
    ...: expected = large.merge(small, how='inner', on='KEY')
    ...: pd.testing.assert_frame_equal(result.sort_values('KEY').reset_index(drop=True), expected)
    ...: 
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-14-eca716bc6f44> in <module>()
     12 
     13 expected = large.merge(small, how='inner', on='KEY')
---> 14 pd.testing.assert_frame_equal(result.sort_values('KEY').reset_index(drop=True), expected)

~/miniconda3/envs/dask_join_dev/lib/python3.6/site-packages/pandas/util/testing.py in assert_frame_equal(left, right, check_dtype, check_index_type, check_column_type, check_frame_type, check_less_precise, check_names, by_blocks, check_exact, check_datetimelike_compat, check_categorical, check_like, obj)
   1351                             'DataFrame shape mismatch',
   1352                             '{shape!r}'.format(shape=left.shape),
-> 1353                             '{shape!r}'.format(shape=right.shape))
   1354 
   1355     if check_like:

~/miniconda3/envs/dask_join_dev/lib/python3.6/site-packages/pandas/util/testing.py in raise_assert_detail(obj, message, left, right, diff)
   1091         msg += "\n[diff]: {diff}".format(diff=diff)
   1092 
-> 1093     raise AssertionError(msg)
   1094 
   1095 

AssertionError: DataFrame are different

DataFrame shape mismatch
[left]:  (47322, 1)
[right]: (47500, 1)

```

however, the following works just fine:
```python
In [12]: import numpy as np
    ...: import pandas as pd
    ...: import dask
    ...: import dask.dataframe as dd
    ...: large = pd.DataFrame({'KEY': np.arange(0,50000)})  
    ...: small = pd.DataFrame({'KEY': np.arange(2500,50000)})
    ...: dd_large = dd.from_pandas(large,npartitions=256)
    ...: dd_small = dd.from_pandas(small, npartitions=3)
    ...: with dask.set_options(shuffle='tasks'):
    ...:     dd_merged = dd_large.merge(dd_small, how='inner', on='KEY')
    ...:     result = dd_merged.compute()
    ...:     
    ...: expected = large.merge(small, how='inner', on='KEY')
    ...: 

In [13]: pd.testing.assert_frame_equal(result.sort_values('KEY').reset_index(drop=True), expected)

```
cc: @jcrist @brendancol 
Thank you for this @hussainsultan .  Most of the work to resolve this problem is probably in the reproducible example that you just made.  I really appreciate your effort here.  

I will try to take a look at this later this morning.

Resolving patch:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -283,11 +283,12 @@ def hash_join(lhs, left_on, rhs, right_on, how='inner',
                      how, npartitions, suffixes, shuffle, indicator)
     name = 'hash-join-' + token
 
-    dsk = dict(((name, i), (methods.merge, (lhs2._name, i), (rhs2._name, i),
-                            how, left_on, right_on,
-                            left_index, right_index, indicator,
-                            suffixes, lhs._meta, rhs._meta))
-               for i in range(npartitions))
+    dsk = {(name, i): (methods.merge,
+                       (lhs2._name, i), (rhs2._name, i),
+                       how, left_on, right_on,
+                       left_index, right_index, indicator,
+                       suffixes, lhs._meta, rhs._meta)
+           for i in range(npartitions)}
 
     divisions = [None] * (npartitions + 1)
     return new_dd_object(toolz.merge(lhs2.dask, rhs2.dask, dsk),
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -203,6 +203,7 @@ def shuffle(df, index, shuffle=None, npartitions=None, max_branch=32,
     """
     if not isinstance(index, _Frame):
         index = df._select_columns_or_index(index)
+
     partitions = index.map_partitions(partitioning_index,
                                       npartitions=npartitions or df.npartitions,
                                       meta=pd.Series([0]))
@@ -363,6 +364,7 @@ def rearrange_by_column_tasks(df, column, max_branch=32, npartitions=None):
     if npartitions is not None and npartitions != df.npartitions:
         parts = [i % df.npartitions for i in range(npartitions)]
         token = tokenize(df2, npartitions)
+
         dsk = {('repartition-group-' + token, i): (shuffle_group_2, k, column)
                for i, k in enumerate(df2.__dask_keys__())}
         for p in range(npartitions):
@@ -442,6 +444,11 @@ def shuffle_group_get(g_head, i):
 
 
 def shuffle_group(df, col, stage, k, npartitions):
+    """ Splits dataframe into groups
+
+    The group is determined by their final partition, and which stage we are in
+    in the shuffle
+    """
     if col == '_partitions':
         ind = df[col]
     else:
@@ -449,12 +456,11 @@ def shuffle_group(df, col, stage, k, npartitions):
 
     c = ind._values
     typ = np.min_scalar_type(npartitions * 2)
-    c = c.astype(typ)
 
     npartitions, k, stage = [np.array(x, dtype=np.min_scalar_type(x))[()]
                              for x in [npartitions, k, stage]]
 
-    c = np.mod(c, npartitions, out=c)
+    c = np.mod(c, npartitions).astype(typ, copy=False)
     c = np.floor_divide(c, k ** stage, out=c)
     c = np.mod(c, k, out=c)
 


Test patch:
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
@@ -5,6 +5,7 @@
 import numpy as np
 import string
 from copy import copy
+import pandas.util.testing as tm
 
 import dask
 import dask.dataframe as dd
@@ -710,3 +711,23 @@ def test_remove_nans():
             params = [none_val if x is None else conv(x) for x in inputs]
             expected = [conv(x) for x in expected]
             assert remove_nans(params) == expected
+
+
+@pytest.mark.slow
+def test_gh_2730():
+    large = pd.DataFrame({'KEY': np.arange(0, 50000)})
+    small = pd.DataFrame({'KEY': np.arange(25, 500)})
+
+    dd_left = dd.from_pandas(small, npartitions=3)
+    dd_right = dd.from_pandas(large, npartitions=257)
+
+    with dask.set_options(shuffle='tasks', get=dask.get):
+        dd_merged = dd_left.merge(dd_right, how='inner', on='KEY')
+        result = dd_merged.compute()
+
+    expected = large.merge(small, how='inner', on='KEY')
+
+    tm.assert_frame_equal(
+        result.sort_values('KEY').reset_index(drop=True),
+        expected
+    )