Repository: dask/dask

Issue:
merge_asof works incorrectly (with reproducible example)
I am trying to implement the following logic in dask: https://stackoverflow.com/questions/58810517/pd-merge-asof-with-multiple-matches-per-time-period. The idea is to have merge_asof with multiple matches per time period if the right table contains more than one matching record.

In the following example, I try to find all matching events for given payments within a time window before each payment.

I find that my code works correctly with Pandas, but Dask returns invalid results:

**UPDATE** I have noticed, that merge_asof with pandas dataframes returns "right_on" column. And with dask dataframes it does not. That might be the cause of the problem

```python
import pandas as pd
import dask.dataframe as dd
from datetime import datetime
from dask.distributed import Client
import numpy as np

def merge_payments_events(payments_dd, events_dd, window):
    payments_sorted = payments_dd.sort_values('BORN_DATE')
    events_sorted = events_dd.sort_values('time_event')

    left_merge = dd.merge_asof(
        payments_sorted, 
        events_sorted, 
        left_by=['CLIENT_UID'],
        right_by=['CLIENT_UID'], 
        left_on='BORN_DATE', 
        right_on='time_event',
        tolerance=pd.Timedelta(window), 
        direction='backward'
    ).reset_index()

    right_merge = dd.merge_asof(
        events_sorted, 
        payments_sorted, 
        left_by=['CLIENT_UID'],
        right_by=['CLIENT_UID'], 
        left_on='time_event', 
        right_on='BORN_DATE', 
        tolerance=pd.Timedelta(window), 
        direction='forward'
    ).reset_index()

    payments_events = left_merge.merge(right_merge, how='outer').dropna(subset=['BORN_DATE'])
    return payments_events


def random_dates(start, end, n, unit='D', seed=None):
    if not seed:  # from piR's answer
        np.random.seed(0)

    ndays = (end - start).days + 1
    return pd.to_timedelta(np.random.rand(n) * ndays, unit=unit) + start


client = Client()
print(client.dashboard_link)

n_payments = 100_000
n_events = 200_000

payments = pd.DataFrame({
    'BORN_DATE': random_dates(datetime(2021, 1, 1), datetime(2021, 1, 3), n_payments, unit='min'),
    'CLIENT_UID': np.random.randint(1, 10_000, n_payments),
    'id': np.arange(0, n_payments)
})

events = pd.DataFrame({
    'time_event': random_dates(datetime(2021, 1, 1), datetime(2021, 1, 3), n_events, unit='min'),
    'CLIENT_UID': np.random.randint(1, 10_000, n_events)
})

payments_dd = dd.from_pandas(payments, npartitions=10)
events_dd = dd.from_pandas(events, npartitions=10)

print("Computed by dask")
dask_result = merge_payments_events(payments_dd, events_dd, window='1h')\
    .dropna(subset=['BORN_DATE']).groupby(['id', 'BORN_DATE'])\
    .count()\
    .compute()\
    .value_counts()

print(dask_result)

print("Computed by pandas")
pandas_result = merge_payments_events(payments, events, window='1h')\
    .dropna(subset=['BORN_DATE'])\
    .groupby(['id', 'BORN_DATE'])\
    .count()\
    .value_counts()

print(pandas_result)
```

I see the following output:
```
Computed by dask
index  CLIENT_UID  time_event
1      1           0             100000
dtype: int64

Computed by pandas
index  CLIENT_UID  time_event
1      1           1             30423
2      2           2             23115
3      3           3             14966
4      4           4              9568
5      5           5              6114
1      1           0              4926
6      6           6              3885
7      7           7              2533
8      8           8              1668
9      9           9              1018
10     10          10              655
11     11          11              401
12     12          12              273
13     13          13              170
14     14          14              108
15     15          15               66
16     16          16               37
17     17          17               25
18     18          18               22
19     19          19                9
20     20          20                8
21     21          21                6
24     24          24                2
22     22          22                1
23     23          23                1
dtype: int64
```

Dask seems to ignore all matches

- Dask version: 2021.11.2
- Python version: 3.8
- Operating System: macos
- Install method (conda, pip, source): pip



Discussion:
Thanks for opening this @kdubovikov I took an initial look and I noticed that it seems like there are incorrect values showing up in the BORN_DATE column. When I do `merge_payments_events(payments_dd, events_dd, window='1h').head()` it includes `BORN_DATE`s that are less than the min `BORN_DATE` on the original df.

There are probably several different issues at play here, but having such a nice reproducible example will definitely help us improve these kinds of workflows, so thank you for taking the time to come up with this example!
Thanks for this question @kdubovikov! I think I have an idea as to what is happening here. The approach suggested in the [SO post](https://stackoverflow.com/questions/58810517/pd-merge-asof-with-multiple-matches-per-time-period) also doesn't work as-is in their example with Dask, and it seems this is indeed (as you noted) because the `right_on` columns are not returned in the Dask `merge_asof`, which are necessary for the final outer merge to work. The results of the individual `merge_asof` calls are the same in Pandas and Dask, excluding the indices and the dropped column.

Here is a fairly minimal example, but could probably be made simpler:

```python
import pandas as pd
import dask.dataframe as dd


def big_merge(df1, df2):
    left_merge = dd.merge_asof(df1, df2,
        tolerance=pd.Timedelta('7d'), direction='backward', 
        left_on='endofweek', right_on='timestamp', 
        left_by='GroupCol', right_by='GroupVal')

    right_merge = dd.merge_asof(df2, df1, 
        tolerance=pd.Timedelta('7d'), direction='forward', 
        left_on='timestamp', right_on='endofweek',
        left_by='GroupVal', right_by='GroupCol')

    merged = left_merge.merge(right_merge, how="outer")
    
    return left_merge, right_merge, merged


def check_equal(df1, df2):
    check_cols = set(df1.columns).intersection(set(df2.columns))
    return df1[check_cols].reset_index(drop=True).fillna(0).eq(
        df2[check_cols].reset_index(drop=True).fillna(0)).all().all()


# create pandas dataframe
df1 = pd.DataFrame({'endofweek': ['2019-08-31', '2019-08-31', '2019-09-07', '2019-09-07', '2019-09-14', '2019-09-14'], 'GroupCol': [1234,8679,1234,8679,1234,8679]})
df2 = pd.DataFrame({'timestamp': ['2019-08-30 10:00', '2019-08-30 10:30', '2019-09-07 12:00', '2019-09-08 14:00'], 'GroupVal': [1234, 1234, 8679, 1234], 'TextVal': ['1234_1', '1234_2', '8679_1', '1234_3']})
df1['endofweek'] = pd.to_datetime(df1['endofweek'])
df2['timestamp'] = pd.to_datetime(df2['timestamp'])

# pandas
pd_left, pd_right, pd_outer = big_merge(df1, df2)

# dask
dd_left, dd_right, dd_outer = big_merge(
    dd.from_pandas(df1, npartitions=2),
    dd.from_pandas(df2, npartitions=2)
)

print(check_equal(pd_left, dd_left.compute()))
print(check_equal(pd_right, dd_right.compute()))
print(check_equal(pd_outer, dd_outer.compute()))
```

I'm still looking through `merge_asof` to see where the `right_on` column is being lost, but will update here when I have more info.
Here's a better minimal example to show the `right_on` column is not returned in Dask, but it is in Pandas, while both `left_by` and `right_by` are returned in both:

```python
import pandas as pd
import dask.dataframe as dd

# create pandas dataframe
df1 = pd.DataFrame({'endofweek': [1, 1, 2, 2, 3, 4], 'GroupCol': [1234,8679,1234,8679,1234,8679]})
df2 = pd.DataFrame({'timestamp': [0, 0, 1, 3], 'GroupVal': [1234, 1234, 8679, 1234]})

# pandas
pandas_cols = dd.merge_asof(df1, df2,
    left_on='endofweek', right_on='timestamp', 
    left_by='GroupCol', right_by='GroupVal').columns

# dask
dask_cols = dd.merge_asof(
    dd.from_pandas(df1, npartitions=2),
    dd.from_pandas(df2, npartitions=2),
    left_on='endofweek', right_on='timestamp', 
    left_by='GroupCol', right_by='GroupVal').columns

# missing right_on column in Dask
set(pandas_cols) ^ set(dask_cols)
```
We've noticed this in our job, the fact that `merge_asof` for some reason drops the right timestamp. We don't know why this is (in Pandas it doesn't happen), but we solve it by this naïve procedure:

1. Assigning an additional column to `right` which is a copy of the timestamp we're gonna join on.
2. Calling `merge_asof`.
3. Renaming that additional column with the name the original right timestamp had (which has already been lost).

Of course this is just a dirty fix, since it doesn't change `merge_asof` itself.
Thanks for providing this workaround @jorloplaz! I haven't gotten around to changing `merge_asof` itself yet, would you feel comfortable looking into it? 
I'm afraid I'm not that familiar with Dask internals so as to try to change `merge_asof` myself. That's just a simple workaround to fix after invoking it that the right-hand-side dataframe's timestamp has been lost, but I really don't know why it's lost!
I've looked into this issue a bit. It seems like we are setting index on both dataframes and we are always calling `pd.merge_asof` with `right_index=True` and `left_index=True`. In this case, pandas behavior is the same. Here is an example:

```
import pandas as pd

# These are the kwargs we are passing to pd.merge_asof
kwargs = {
    'left_index': True,
    'right_index': True,
    'left_by': 'GroupCol',
    'right_by': 'GroupVal',
    'suffixes': ('_x', '_y'),
    'tolerance': None,
    'allow_exact_matches': True,
    'direction': 'backward'
}
df1 = pd.DataFrame({'endofweek': [1, 1, 2, 2, 3, 4], 'GroupCol': [1234,8679,1234,8679,1234,8679]})
df1 = df1.set_index("endofweek")
df2 = pd.DataFrame({'timestamp': [0, 0, 1, 3], 'GroupVal': [1234, 1234, 8679, 1234]})
df2 = df2.set_index("timestamp")
print(pd.merge_asof(df1, df2, **kwargs))

           GroupCol  GroupVal
endofweek                    
1              1234      1234
1              8679      8679
2              1234      1234
2              8679      8679
3              1234      1234
4              8679      8679
```

Not sure if it's possible, but maybe we can avoid setting index to produce the same output as pandas.
I guess that's related to the fact that `left_on` and `right_on` (that is, the columns being joined approximately) must be sorted for `pd.merge_asof` to work, so it makes sense to have a preference for them to be indices.

Actually, `dd.merge_asof` doesn't allow `left_by` and `right_by` (that is, the columns to be joined exactly before joining approximately on `left_on` and `right_on`) to be the indices of the dataframes to be joined. In contrast, those in `left_on` and `right_on` can indeed be the indices, and according to what you say, it seems that if they happen not to be the indices already, they are set to be so before invoking `pd.merge_asof`. 

So I guess this implies in turn that in pseudocode `dd.merge_asof` does this:

1. Check if `left_on` is the left index. If it's not, do `left = left.set_index(left_on)` to guarantee sortedness.
2. Check if `right_on` is the right index. If it's not, do `right = right.set_index(right_on)` to guarantee sortedness.
3. Somehow call `pd.merge_asof(left_part, right_part, left_index=True, right_index=True, left_by=left_by, right_by=right_by`, ...) with the rest of `kwargs` you mention for the whole dataframe. I guess inside a `map_partitions`?
> So I guess this implies in turn that in pseudocode `dd.merge_asof` does this:
> 
> 1. Check if `left_on` is the left index. If it's not, do `left = left.set_index(left_on)` to guarantee sortedness.
> 2. Check if `right_on` is the right index. If it's not, do `right = right.set_index(right_on)` to guarantee sortedness.

@jorloplaz I think that's exactly right (lines 931, 936 in particular):
https://github.com/dask/dask/blob/baa378f8a09005ce6c79a7029fc0c1c1d2430a5f/dask/dataframe/multi.py#L925-L936

> 3. Somehow call `pd.merge_asof(left_part, right_part, left_index=True, right_index=True, left_by=left_by, right_by=right_by`, ...) with the rest of `kwargs` you mention for the whole dataframe. I guess inside a `map_partitions`?

I think this is essentially correct (though maybe a bit more complicated). The task graph is created in [`merge_asof_indexed`](https://github.com/dask/dask/blob/baa378f8a09005ce6c79a7029fc0c1c1d2430a5f/dask/dataframe/multi.py#L947):
https://github.com/dask/dask/blob/baa378f8a09005ce6c79a7029fc0c1c1d2430a5f/dask/dataframe/multi.py#L848-L862

Which relies on [`merge_asof_padded`](https://github.com/dask/dask/blob/baa378f8a09005ce6c79a7029fc0c1c1d2430a5f/dask/dataframe/multi.py#L784), so now I'm stepping through to see if either of these could be changed such that the right index is preserved.


Thanks for looking into this @scharlottej13 . As I see it, this amounts to the fact that [pd.merge_asof](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.merge_asof.html) behaves like a left join in [pd.merge](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.merge.html), which literally states:

> The join is done on columns or indexes. If joining columns on columns, the DataFrame indexes will be ignored. Otherwise if joining indexes on indexes or indexes on a column or columns, the index will be passed on. 

My understanding of this, backed up with @kori73 's snippet, is that when `left_index=True` and `right_index=True`, `pd.merge_asof` only "passes on" the left index (because it's a left join), and the right index is NOT passed on. Hence, the right timestamp is lost by Pandas, and Dask consequently loses it because it invokes Pandas with this setting.

So it seems that `right` should NOT be indexed by its timestamp (that is, by `right_on`). However, `right` should be indeed sorted by that column, so perhaps the way to go is calling `right.sort_values(right_on)` instead of `right.set_index(right_on)`? And then using `pd.merge_asof(left_index=True, right_index=False, right_on=right_on)`.

Nevertheless, this may break the partitions alignment that I think is done under the hood once `left` and `right` have been indexed by their respective timestamps.... What do you think?

Another option would be the same workaround that I suggested being done outside `dd.merge_asof`, but this time harcoding it inside its code, that is:

1. Duplicating in `right` the column `right_on` with a temporary copy, let's say `right = right.assign(**{'foo': right[right_on]})`.
2. Proceeding as it's done currently, that is, setting both indexes and calling `pd.merge_asof(left_index=True, right_index=True)`.
3. Renaming the `right_on` column that is lost by Pandas with `result = result.rename(columns={'foo': right_on})`.

Resolving patch:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -946,7 +946,10 @@ def merge_asof(
     if not is_dask_collection(right):
         right = from_pandas(right, npartitions=1)
     if right_on is not None:
-        right = right.set_index(right_on, sorted=True)
+        if on:
+            right = right.set_index(right_on, sorted=True)
+        else:
+            right = right.set_index(right_on, drop=False, sorted=True)
 
     if by is not None:
         kwargs["left_by"] = kwargs["right_by"] = by


Test patch:
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -636,6 +636,31 @@ def test_merge_asof_with_empty():
     assert_eq(result_dd, result_df, check_index=False)
 
 
+def test_merge_asof_on_left_right():
+    df1 = pd.DataFrame(
+        {
+            "endofweek": [1, 1, 2, 2, 3, 4],
+            "GroupCol": [1234, 8679, 1234, 8679, 1234, 8679],
+        }
+    )
+    df2 = pd.DataFrame(
+        {"timestamp": [0, 0, 1, 3], "GroupVal": [1234, 1234, 8679, 1234]}
+    )
+
+    # pandas
+    result_df = pd.merge_asof(df1, df2, left_on="endofweek", right_on="timestamp")
+
+    # dask
+    result_dd = dd.merge_asof(
+        dd.from_pandas(df1, npartitions=2),
+        dd.from_pandas(df2, npartitions=2),
+        left_on="endofweek",
+        right_on="timestamp",
+    )
+
+    assert_eq(result_df, result_dd, check_index=False)
+
+
 @pytest.mark.parametrize("join", ["inner", "outer"])
 def test_indexed_concat(join):
     A = pd.DataFrame(