Repository: dask/dask

Issue:
Behaviour of map_partitions with a Pandas dataframe as argument - In this case uncorrect merge
I want to perform a merge between a Dask dataframe and a Pandas Dataframe and it seems like the pandas dataframe I pass as argument is partitionned in a way that it does not give the result we expect in the end.
I wanted to be able to able to keep the index on the Dask Dataframe, that's why I tried to use a map_partitions instead of a merge on the Dask Dataframe directly.

I think the reason of this behaviour is explained by "DataFrame-like args (both dask and pandas) will be repartitioned to align (if necessary) before applying the function." in the API documentation.

```python
food_stores = pd.DataFrame({'foodstore_id':[1,2,3], 
                            'seller_firstname':['Bob', 'Alice', 'Monique'], 
                            'seller_lastname':['Marchand', 'Boul', 'Mul'],  
                            'online_status':[1,1,0]})
records = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10], 
                        'name':['product1', 'product2', 'product3', 'product4', 'product5', 
                                'product6', 'product7', 'product8', 'product9', 'product10'],
                        'seller_firstname':['Bob', 'Alice', 'Monique', 'Alice', 'Bob', 
                                            'Alice', 'Monique', 'Bob', 'Alice', 'Bob'],
                        'seller_lastname':['Marchand', 'Boul', 'Mul', 'Boul', 'Marchand', 
                                           'Boul', 'Mul', 'Marchand', 'Boul', 'Marchand'],
                        })
records = records.set_index('id')
records = dd.from_pandas(records, npartitions=3)

def mapping(part, pandas_df):
    pandas_df = pandas_df.set_index(['seller_firstname', 'seller_lastname'])    
    return part.merge(pandas_df, how='left', 
                      left_on=['seller_firstname', 'seller_lastname'], 
                      right_on=['seller_firstname', 'seller_lastname'], right_index=True)

print(records.map_partitions(mapping, food_stores).compute())
```
|   id | name      | seller_firstname   | seller_lastname   |   foodstore_id |   online_status |
|-----:|:----------|:-------------------|:------------------|---------------:|----------------:|
|    1 | product1  | Bob                | Marchand          |            nan |             nan |
|    2 | product2  | Alice              | Boul              |            nan |             nan |
|    3 | product3  | Monique            | Mul               |              3 |               0 |
|    4 | product4  | Alice              | Boul              |            nan |             nan |
|    5 | product5  | Bob                | Marchand          |            nan |             nan |
|    6 | product6  | Alice              | Boul              |            nan |             nan |
|    7 | product7  | Monique            | Mul               |            nan |             nan |
|    8 | product8  | Bob                | Marchand          |            nan |             nan |
|    9 | product9  | Alice              | Boul              |            nan |             nan |
|   10 | product10 | Bob                | Marchand          |            nan |             nan |

I thought I would get the result I expect with the following : 

```python
food_stores = pd.DataFrame({'foodstore_id':[1,2,3], 
                            'seller_firstname':['Bob', 'Alice', 'Monique'], 
                            'seller_lastname':['Marchand', 'Boul', 'Mul'],  
                            'online_status':[1,1,0]})
records = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10], 
                        'name':['product1', 'product2', 'product3', 'product4', 'product5', 
                                'product6', 'product7', 'product8', 'product9', 'product10'],
                        'seller_firstname':['Bob', 'Alice', 'Monique', 'Alice', 'Bob', 'Alice', 
                                            'Monique', 'Bob', 'Alice', 'Bob'],
                        'seller_lastname':['Marchand', 'Boul', 'Mul', 'Boul', 'Marchand', 'Boul', 
                                           'Mul', 'Marchand', 'Boul', 'Marchand'],
                        })
records = records.set_index('id')
records = dd.from_pandas(records, npartitions=3)

def mapping(part, pandas_df):
    pandas_df = pandas_df.set_index(['seller_firstname', 'seller_lastname'])    
    return part.merge(pandas_df, how='left', left_on=['seller_firstname', 'seller_lastname'], right_on=['seller_firstname', 'seller_lastname'], right_index=True)

records.map_partitions(lambda part: mapping(part, food_stores)).compute()
```
|   id | name      | seller_firstname   | seller_lastname   |   foodstore_id |   online_status |
|-----:|:----------|:-------------------|:------------------|---------------:|----------------:|
|    1 | product1  | Bob                | Marchand          |              1 |               1 |
|    2 | product2  | Alice              | Boul              |              2 |               1 |
|    3 | product3  | Monique            | Mul               |              3 |               0 |
|    4 | product4  | Alice              | Boul              |              2 |               1 |
|    5 | product5  | Bob                | Marchand          |              1 |               1 |
|    6 | product6  | Alice              | Boul              |              2 |               1 |
|    7 | product7  | Monique            | Mul               |              3 |               0 |
|    8 | product8  | Bob                | Marchand          |              1 |               1 |
|    9 | product9  | Alice              | Boul              |              2 |               1 |
|   10 | product10 | Bob                | Marchand          |              1 |               1 |

**Environment**:

- Dask version: 2.16.0 (same with latest - 2.19)
- Python version: 3.7.3
- Operating System: Ubuntu 18.04.3 LTS
- Install method (conda, pip, source): image daskdev/dask-notebook 2.16.0



Discussion:
I'm not sure about the expected behavior here, but right now your bet is to disable the alignment of the pandas DataFrame by passing it as a keyword

```
res = records.map_partitions(mapping, pandas_df=food_stores).compute()
```

FYI, since you're adjusting the index you should also call `clear_divisions()` on the result.
Passing in an argument should be fine in map_partitions.  I would take one
partition out of the dask dataframe and run your function on that to see
what happens.

    mapping(records.partitions[0].compute(), pandas_df)

On Tue, Jun 30, 2020 at 6:46 AM odovad <notifications@github.com> wrote:

> I want to perform a merge between a Dask dataframe and a Pandas Dataframe
> and it seems like the pandas dataframe I pass as argument is partitionned
> in a way that it does give the result we expect in the end.
> I wanted to be able to able to keep the index on the Dask Dataframe,
> that's why I tried to use a map_partitions instead of a merge on the Dask
> Dataframe directly.
>
> I think the reason of this behaviour is explained by "DataFrame-like args
> (both dask and pandas) will be repartitioned to align (if necessary) before
> applying the function."
>
> food_stores = pd.DataFrame({'foodstore_id':[1,2,3],
>                             'seller_firstname':['Bob', 'Alice', 'Monique'],
>                             'seller_lastname':['Marchand', 'Boul', 'Mul'],
>                             'online_status':[1,1,0]})records = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10],
>                         'name':['product1', 'product2', 'product3', 'product4', 'product5',
>                                 'product6', 'product7', 'product8', 'product9', 'product10'],
>                         'seller_firstname':['Bob', 'Alice', 'Monique', 'Alice', 'Bob',
>                                             'Alice', 'Monique', 'Bob', 'Alice', 'Bob'],
>                         'seller_lastname':['Marchand', 'Boul', 'Mul', 'Boul', 'Marchand',
>                                            'Boul', 'Mul', 'Marchand', 'Boul', 'Marchand'],
>                         })records = records.set_index('id')records = dd.from_pandas(records, npartitions=3)
> def mapping(part, pandas_df):
>     pandas_df = pandas_df.set_index(['seller_firstname', 'seller_lastname'])
>     return part.merge(pandas_df, how='left',
>                       left_on=['seller_firstname', 'seller_lastname'],
>                       right_on=['seller_firstname', 'seller_lastname'], right_index=True)
> print(records.map_partitions(mapping, food_stores).compute())
>
> id name seller_firstname seller_lastname foodstore_id online_status
> 1 product1 Bob Marchand nan nan
> 2 product2 Alice Boul nan nan
> 3 product3 Monique Mul 3 0
> 4 product4 Alice Boul nan nan
> 5 product5 Bob Marchand nan nan
> 6 product6 Alice Boul nan nan
> 7 product7 Monique Mul nan nan
> 8 product8 Bob Marchand nan nan
> 9 product9 Alice Boul nan nan
> 10 product10 Bob Marchand nan nan
>
> I thought I would get the result I expect with the following :
>
> food_stores = pd.DataFrame({'foodstore_id':[1,2,3],
>                             'seller_firstname':['Bob', 'Alice', 'Monique'],
>                             'seller_lastname':['Marchand', 'Boul', 'Mul'],
>                             'online_status':[1,1,0]})records = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10],
>                         'name':['product1', 'product2', 'product3', 'product4', 'product5',
>                                 'product6', 'product7', 'product8', 'product9', 'product10'],
>                         'seller_firstname':['Bob', 'Alice', 'Monique', 'Alice', 'Bob', 'Alice',
>                                             'Monique', 'Bob', 'Alice', 'Bob'],
>                         'seller_lastname':['Marchand', 'Boul', 'Mul', 'Boul', 'Marchand', 'Boul',
>                                            'Mul', 'Marchand', 'Boul', 'Marchand'],
>                         })records = records.set_index('id')records = dd.from_pandas(records, npartitions=3)
> def mapping(part, pandas_df):
>     pandas_df = pandas_df.set_index(['seller_firstname', 'seller_lastname'])
>     return part.merge(pandas_df, how='left', left_on=['seller_firstname', 'seller_lastname'], right_on=['seller_firstname', 'seller_lastname'], right_index=True)
> records.map_partitions(lambda part: mapping(part, food_stores)).compute()
>
> id name seller_firstname seller_lastname foodstore_id online_status
> 1 product1 Bob Marchand 1 1
> 2 product2 Alice Boul 2 1
> 3 product3 Monique Mul 3 0
> 4 product4 Alice Boul 2 1
> 5 product5 Bob Marchand 1 1
> 6 product6 Alice Boul 2 1
> 7 product7 Monique Mul 3 0
> 8 product8 Bob Marchand 1 1
> 9 product9 Alice Boul 2 1
> 10 product10 Bob Marchand 1 1
>
> *Environment*:
>
>    - Dask version: 2.16.0 (same with latest - 2.19)
>    - Python version: 3.7.3
>    - Operating System: Ubuntu 18.04.3 LTS
>    - Install method (conda, pip, source): image daskdev/dask-notebook
>    2.16.0
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6357>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTAPIFDG7RLNSPBE7LDRZHUB7ANCNFSM4OMIB3BQ>
> .
>

Thanks @TomAugspurger, I did not try this way. I think the behaviour is not easy to understand. 
Why is the pandas DataFrame aligned ? 

What do you mean when you say I'm adjusting the index ? 
The number of lines did not change, only two columns are added. 
It's not safe to say that we can trust the index? 
We need to "reindex" if we want to be able to use the index effectively ? 
@mrocklin On one partition, it seems to work well : 
Do you think the map_partitions as defined here can break the index ? 

```python
print(mapping(records.partitions[0].compute(), food_stores).to_markdown())
```

|   id | name     | seller_firstname   | seller_lastname   |   foodstore_id |   online_status |
|-----:|:---------|:-------------------|:------------------|---------------:|----------------:|
|    1 | product1 | Bob                | Marchand          |              1 |               1 |
|    2 | product2 | Alice              | Boul              |              2 |               1 |
|    3 | product3 | Monique            | Mul               |              3 |               0 |
|    4 | product4 | Alice              | Boul              |              2 |               1 |
I think this demonstrates the issue a bit more clearly

```python
import pandas as pd
import dask.dataframe as dd

df1 = pd.DataFrame({"A": [1, 2, 3, 3, 2, 3], "B": [1, 2, 3, 4, 5, 6]})
df2 = pd.DataFrame({"A": [3, 1, 2], "C": [1, 2, 3]})


def merge(a, b):
    res = pd.merge(a, b, left_on="A", right_on="A", how="left")
    return res


res_pandas = merge(df1, df2)
res_dask = dd.from_pandas(df1, npartitions=2).map_partitions(merge, df2)

```

```python

In [18]: res_pandas
Out[18]:
   A  B  C
0  1  1  2
1  2  2  3
2  3  3  1
3  3  4  1
4  2  5  3
5  3  6  1

In [19]: res_dask.compute()
Out[19]:
   A  B    C
0  1  1  2.0
1  2  2  NaN
0  3  3  NaN
0  3  4  NaN
1  2  5  NaN
2  3  6  NaN

```
@TomAugspurger and the result is not the same depending on the number of partitions you choose

Can you give more detail about the indexation ? 
Why a ```clear_divisions()``` is needed ? 
(I wanted to be able not to break the index ... maybe there is a way ?)
You can ignore my clear divisions comment.

On Tue, Jun 30, 2020 at 9:52 AM odovad <notifications@github.com> wrote:

> @TomAugspurger <https://github.com/TomAugspurger> and the result is not
> the same depending on the number of partitions you choose
>
> Can you give more detail about the indexation ?
> Why a clear_divisions() is needed ?
> (I wanted to be able not to break the index ... maybe there is a way ?)
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6357#issuecomment-651844880>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAKAOIUNLOTWDT5CFCE7QB3RZH3ZBANCNFSM4OMIB3BQ>
> .
>

On Tue, Jun 30, 2020 at 9:52 AM odovad <notifications@github.com> wrote:

> @TomAugspurger <https://github.com/TomAugspurger> and the result is not
> the same depending on the number of partitions you choose
>
> Can you give more detail about the indexation ?
> Why a clear_divisions() is needed ?
> (I wanted to be able not to break the index ... maybe there is a way ?)
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6357#issuecomment-651844880>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAKAOIUNLOTWDT5CFCE7QB3RZH3ZBANCNFSM4OMIB3BQ>
> .
>

Ok @TomAugspurger. Thanks. I was interested; I thought you said I managed to be in a situation where the dataframe has an index with known divisions that we cannot trust. Interesting to know since a set_index on a large dataframe is something we want to avoid.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -627,8 +627,25 @@ def map_partitions(self, func, *args, **kwargs):
             be the first argument, and these will be passed *after*. Arguments
             and keywords may contain ``Scalar``, ``Delayed``, ``partition_info``
             or regular python objects. DataFrame-like args (both dask and
-            pandas) will be repartitioned to align (if necessary) before
-            applying the function.
+            pandas) will be repartitioned to align  (if necessary) before
+            applying the function (see ``align_dataframes`` to control).
+        enforce_metadata : bool, default True
+            Whether to enforce at runtime that the structure of the DataFrame
+            produced by ``func`` actually matches the structure of ``meta``.
+            This will rename and reorder columns for each partition,
+            and will raise an error if this doesn't work or types don't match.
+        transform_divisions : bool, default True
+            Whether to apply the function onto the divisions and apply those
+            transformed divisions to the output.
+        align_dataframes : bool, default True
+            Whether to repartition DataFrame- or Series-like args
+            (both dask and pandas) so their divisions align before applying
+            the function. This requires all inputs to have known divisions.
+            Single-partition inputs will be split into multiple partitions.
+
+            If False, all inputs must have either the same number of partitions
+            or a single partition. Single-partition inputs will be broadcast to
+            every partition of multi-partition inputs.
         $META
 
         Examples
@@ -5794,6 +5811,7 @@ def map_partitions(
     meta=no_default,
     enforce_metadata=True,
     transform_divisions=True,
+    align_dataframes=True,
     **kwargs,
 ):
     """Apply Python function on each DataFrame partition.
@@ -5807,11 +5825,24 @@ def map_partitions(
         args should be a Dask.dataframe. Arguments and keywords may contain
         ``Scalar``, ``Delayed`` or regular python objects. DataFrame-like args
         (both dask and pandas) will be repartitioned to align (if necessary)
-        before applying the function.
-    enforce_metadata : bool
-        Whether or not to enforce the structure of the metadata at runtime.
+        before applying the function (see ``align_dataframes`` to control).
+    enforce_metadata : bool, default True
+        Whether to enforce at runtime that the structure of the DataFrame
+        produced by ``func`` actually matches the structure of ``meta``.
         This will rename and reorder columns for each partition,
         and will raise an error if this doesn't work or types don't match.
+    transform_divisions : bool, default True
+        Whether to apply the function onto the divisions and apply those
+        transformed divisions to the output.
+    align_dataframes : bool, default True
+        Whether to repartition DataFrame- or Series-like args
+        (both dask and pandas) so their divisions align before applying
+        the function. This requires all inputs to have known divisions.
+        Single-partition inputs will be split into multiple partitions.
+
+        If False, all inputs must have either the same number of partitions
+        or a single partition. Single-partition inputs will be broadcast to
+        every partition of multi-partition inputs.
     $META
     """
     name = kwargs.pop("token", None)
@@ -5830,8 +5861,10 @@ def map_partitions(
 
     from .multi import _maybe_align_partitions
 
-    args = _maybe_from_pandas(args)
-    args = _maybe_align_partitions(args)
+    if align_dataframes:
+        args = _maybe_from_pandas(args)
+        args = _maybe_align_partitions(args)
+
     dfs = [df for df in args if isinstance(df, _Frame)]
     meta_index = getattr(make_meta(dfs[0]), "index", None) if dfs else None
     if parent_meta is None and dfs:


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
@@ -1060,6 +1060,22 @@ def check(df, df_x):
     assert_eq(res, ddf)
 
 
+def test_align_dataframes():
+    df1 = pd.DataFrame({"A": [1, 2, 3, 3, 2, 3], "B": [1, 2, 3, 4, 5, 6]})
+    df2 = pd.DataFrame({"A": [3, 1, 2], "C": [1, 2, 3]})
+
+    def merge(a, b):
+        res = pd.merge(a, b, left_on="A", right_on="A", how="left")
+        return res
+
+    expected = merge(df1, df2)
+
+    ddf1 = dd.from_pandas(df1, npartitions=2)
+    actual = ddf1.map_partitions(merge, df2, align_dataframes=False)
+
+    assert_eq(actual, expected, check_index=False, check_divisions=False)
+
+
 def test_drop_duplicates():
     res = d.drop_duplicates()
     res2 = d.drop_duplicates(split_every=2)