Repository: dask/dask

Issue:
Results of dask.multi.merge_asof depends on npartitions
**What happened**: Results of `merge_asof` differ depending on npartitions of the dataframe. In particular, it seems the result misses rows when the number of partitions is equal to or greater than the number of rows in the dataframe. (I am aware that setting the number of partitions higher than row count does not make sense, yet I did not expect it to lead to this behavior). 

**What you expected to happen**: Identical results of `merge_asof` independent from npartitions of the dataframe. Or, if this issue is restricted to the case where npartitions>number of rows, an exception or warning telling the user that it may lead to undefined behavior.

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame(dict(ts=[pd.to_datetime('1-1-2020')]*3, foo=[1,2,3]))
    
for npartitions in range(1,5):
    ddf = dd.from_pandas(df, npartitions=npartitions)
    result = dd.multi.merge_asof(left=ddf, right=ddf, on='ts').compute()
    print(result)

# Output
#          ts  foo_x  foo_y
# 0 2020-01-01      1      3
# 1 2020-01-01      2      3
# 2 2020-01-01      3      3
#           ts  foo_x  foo_y
# 0 2020-01-01      1      3
# 1 2020-01-01      2      3
# 2 2020-01-01      3      3
#           ts  foo_x  foo_y
# 0 2020-01-01      2      3
# 1 2020-01-01      3      3
#           ts  foo_x  foo_y
# 0 2020-01-01      2      3
# 1 2020-01-01      3      3
```
**What else do we need to know**: As far as I could verify, the issue only appears when the timestamps in the `on` column are identical.

**Environment**:

- Dask version: 2022.04.1
- Python version: 3.9.1
- Operating System: Debian GNU/Linux 11 (bullseye)
- Install method (conda, pip, source): pip



Discussion:
@JnsLns Welcome to Dask, and thanks for reporting this! 

So, Dask wouldn't assign `npartitions` equal to or greater than the number of rows, even if we explicitly say so. It can have at most *number_of_rows - 1* partitions (see example below). But looks like when the `on` column has all the same values, `merge_asof` works only on the values of the last partition. This is inconsistent with pandas. :/

I'll keep looking into this!

<details>
<summary> Minimal Example: </summary>


```python
import pandas as pd
import dask.dataframe as dd


# --- pandas ---

df = pd.DataFrame(dict(ts=[pd.to_datetime('1-1-2020')]*6, foo=list(range(6))))

pd.merge_asof(left=df, right=df, on='ts') 
#           ts  foo_x  foo_y
# 0 2020-01-01      0      5
# 1 2020-01-01      1      5
# 2 2020-01-01      2      5
# 3 2020-01-01      3      5
# 4 2020-01-01      4      5
# 5 2020-01-01      5      5


# --- Dask ---

ddf = dd.from_pandas(df, npartitions=6)
ddf.npartitions # 5

ddf.partitions[-1].compute()
#           ts  foo
# 4 2020-01-01    4
# 5 2020-01-01    5

dd.merge_asof(left=ddf, right=ddf, on='ts').compute()
#           ts  foo_x  foo_y
# 0 2020-01-01      4      5
# 1 2020-01-01      5      5


ddf1 = dd.from_pandas(df, npartitions=2)
ddf1.npartitions # 2

ddf1.partitions[-1].compute()
#           ts  foo
# 3 2020-01-01    3
# 4 2020-01-01    4
# 5 2020-01-01    5

dd.merge_asof(left=ddf1, right=ddf1, on='ts').compute()
#           ts  foo_x  foo_y
# 0 2020-01-01      3      5
# 1 2020-01-01      4      5
# 2 2020-01-01      5      5
```
</details>

Hi @pavithraes 
Thanks for taking care of it! Yes, that's what I thought. 
@ian-r-rose and I looked into this. We think the issues is that `merge_asof` sets the `on` column as index with `sorted=True`, here:
https://github.com/dask/dask/blob/6369cdb06813cf9d35f7d19f855ebabcde26d0a2/dask/dataframe/multi.py#L950 and, here:
https://github.com/dask/dask/blob/6369cdb06813cf9d35f7d19f855ebabcde26d0a2/dask/dataframe/multi.py#L955

`sorted=True` results in the same number of partitions as `ddf`, and if the `on` column doesn't have enough richness (as is that case in this issue), the resulting divisions will not be correct.

One quick solution is to set `sorted=False`, which would result in all equal index values in one partition, however that'll be an expensive operation. It may also cause a memory error if the dataset is super large.

A nicer solution could be to check for column richness (which can be tricky to infer, see next paragraph) in `set_index`. This would mean we'll raise a helpful error message anytime we set an index on an unsuitable-for-proper-divisions column.

An open question is how do we check if a column is rich enough, that is, it has enough unique values? Another way to put it -- how many unique values in a column, compared to number of partitions, would be enough to get the right divisions? And, if/how this changes if the column is sorted?

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -1029,10 +1029,12 @@ def fix_overlap(ddf, mins, maxes, lens):
         # this partition (i) if the data from this partition will need to be moved
         # to the next partition (i+1) anyway.  If we concatenate data too early,
         # we may lose rows (https://github.com/dask/dask/issues/6972).
-        if i == len(mins) - 2 or divisions[i] != divisions[i + 1]:
-            frames.append(ddf_keys[i])
-            dsk[(name, i)] = (methods.concat, frames)
-            frames = []
+        if divisions[i] == divisions[i + 1] and i + 1 in overlap:
+            continue
+
+        frames.append(ddf_keys[i])
+        dsk[(name, i)] = (methods.concat, frames)
+        frames = []
 
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
     return new_dd_object(graph, name, ddf._meta, divisions)


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
@@ -712,6 +712,18 @@ def test_merge_asof_on_left_right(left_col, right_col):
     assert_eq(result_df, result_dd, check_index=False)
 
 
+def test_merge_asof_with_various_npartitions():
+    # https://github.com/dask/dask/issues/8999
+    df = pd.DataFrame(dict(ts=[pd.to_datetime("1-1-2020")] * 3, foo=[1, 2, 3]))
+    expected = pd.merge_asof(left=df, right=df, on="ts")
+
+    for npartitions in range(1, 5):
+        ddf = dd.from_pandas(df, npartitions=npartitions)
+
+        result = dd.merge_asof(left=ddf, right=ddf, on="ts")
+        assert_eq(expected, result)
+
+
 @pytest.mark.parametrize("join", ["inner", "outer"])
 def test_indexed_concat(join):
     A = pd.DataFrame(
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
@@ -1182,6 +1182,17 @@ def test_set_index_overlap_2():
     assert ddf2.npartitions == 8
 
 
+def test_set_index_overlap_does_not_drop_rows_when_divisions_overlap():
+    # https://github.com/dask/dask/issues/9339
+    df = pd.DataFrame({"ts": [1, 1, 2, 2, 3, 3, 3, 3], "value": "abc"})
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    expected = df.set_index("ts")
+    actual = ddf.set_index("ts", sorted=True)
+
+    assert_eq(expected, actual)
+
+
 def test_compute_current_divisions_nan_partition():
     # Compute divisions 1 null partition
     a = d[d.a > 3].sort_values("a")