Repository: dask/dask
Pull request: Harden DataFrame merge between clumn-selection and index
This PR is intended to fix both of the merge issues discussed in #7558.  The first problem is that `searchsorted` fails for object columns when there are null values in the data.  The `TypeError` can be avoided by explicitly mapping nulls to the 0th partition.  The second problem is that the hash values of an index can be different than the hashed values of a Series (since the Series is converted to a DataFrame).  The index problem can be avoided by calling `to_frame` whenever the attribute is available.

- [x] Closes #7558
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask` / `isort dask`


Problem statement:
Merge produces incorrect results with partitioned data
**What happened**:

Pandas result works, Dask result produces an error and does not work with example above.

Changing  in the example below: `df2_d = dd.from_pandas(df2, npartitions=3).set_index('b')` to 
`df2_d = dd.from_pandas(df2, npartitions=3).set_index('b').repartition(npartitions=5)`
causes dask to run, but produce incorrect results in the output merged dataframe.

**What you expected to happen**:
Dask and pandas agree and produce the same result when joining tables together.

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
df1 = pd.DataFrame({'a': ['0', '0', None, None, None, None, '5', '7', '15', '33']})
df2 = pd.DataFrame({'b': ['0', '5', '7', '15'], 'c': ['1', '2', '3', '4']})

pandas_result = df1.merge(df2.set_index('b'), how='left', left_on='a', right_index=True)

df1_d = dd.from_pandas(df1, npartitions=4)
df2_d = dd.from_pandas(df2, npartitions=3).set_index('b')

dask_result = df1_d.merge(df2_d, how='left', left_on='a', right_index=True).compute()
```

Hints from issue discussion:
When I ran the example above I get type error:

> TypeError: '<' not supported between instances of 'str' and 'NoneType'
Ah sorry my comment might not have been clear:
```python
import pandas as pd
df1 = pd.DataFrame({'a': ['0', '0', None, None, None, None, '5', '7', '15', '33']})
df2 = pd.DataFrame({'b': ['0', '5', '7', '15'], 'c': ['1', '2', '3', '4']})

pandas_result = df1.merge(df2.set_index('b'), how='left', left_on='a', right_index=True)

df1_d = dd.from_pandas(df1, npartitions=4)
df2_d = dd.from_pandas(df2, npartitions=3).set_index('b')

dask_result = df1_d.merge(df2_d, how='left', left_on='a', right_index=True).compute()
```
produces 
`TypeError: '<' not supported between instances of 'str' and 'NoneType'`


Adding `.repartition(npartitions=5)` to the `df2_d` line, is what produces the incorrect result. 

```python
import pandas as pd
import dask.dataframe as dd

df1 = pd.DataFrame({'a': ['0', '0', None, None, None, None, '5', '7', '15', '33']})
df2 = pd.DataFrame({'b': ['0', '5', '7', '15'], 'c': ['1', '2', '3', '4']})

pandas_result = df1.merge(df2.set_index('b'), how='left', left_on='a', right_index=True)

df1_d = dd.from_pandas(df1, npartitions=4)
df2_d = dd.from_pandas(df2, npartitions=3).set_index('b').repartition(npartitions=5)

dask_result = df1_d.merge(df2_d, how='left', left_on='a', right_index=True).compute()
```

Output:
Pandas:
```
        a	c
0	0	1
1	0	1
2	None	NaN
3	None	NaN
4	None	NaN
5	None	NaN
6	5	2
7	7	3
8	15	4
9	33	NaN
```

Output Dask:
```
	a	c
2	None	None
3	None	None
4	None	None
5	None	None
0	0	NaN
1	0	NaN
7	7	NaN
6	5	NaN
8	15	NaN
9	33	NaN
```

I think that the Error that happens and also the incorrect result are related, and are related to `.division` logic
Yeah I can reproduce and I noticed that even if you explicitly clear the divisions on `df1_d` you still get the incorrect result. 
Thank you for raising this @bluecoconut - This is a good find.  It looks like we are using a hash merge here, and the left-hand join index (a DataFrame column) is getting hashed differently than the right-hand join index (a DataFrame index).  I will definitely look into a solution for this.

Commit message: preserve divisions in Index -> to_frame

Patch to review:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 5ac2bbedf9c..e738ec36370 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -3710,7 +3710,11 @@ def to_frame(self, index=True, name=None):
             raise NotImplementedError()
 
         return self.map_partitions(
-            M.to_frame, index, name, meta=self._meta.to_frame(index, name)
+            M.to_frame,
+            index,
+            name,
+            meta=self._meta.to_frame(index, name),
+            transform_divisions=False,
         )
 
     @insert_meta_param_description(pad=12)
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 707a7b6ff7c..c3734d7a0a0 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -364,7 +364,10 @@ def hash_join(
     )
 
     # dummy result
-    meta = lhs._meta_nonempty.merge(rhs._meta_nonempty, **kwargs)
+    # Avoid using dummy data for a collection it is empty
+    _lhs_meta = lhs._meta_nonempty if len(lhs.columns) else lhs._meta
+    _rhs_meta = rhs._meta_nonempty if len(rhs.columns) else rhs._meta
+    meta = _lhs_meta.merge(_rhs_meta, **kwargs)
 
     if isinstance(left_on, list):
         left_on = (list, tuple(left_on))
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index 87036fec380..694fe741f93 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -346,6 +346,11 @@ def shuffle(
 
     if not isinstance(index, _Frame):
         index = df._select_columns_or_index(index)
+    elif hasattr(index, "to_frame"):
+        # If this is an index, we should still convert to a
+        # DataFrame. Otherwise, the hashed values of a column
+        # selection will not match (important when merging).
+        index = index.to_frame()
 
     partitions = index.map_partitions(
         partitioning_index,
@@ -732,7 +737,16 @@ def collect(p, part, meta, barrier_token):
 
 
 def set_partitions_pre(s, divisions):
-    partitions = divisions.searchsorted(s, side="right") - 1
+    try:
+        partitions = divisions.searchsorted(s, side="right") - 1
+    except TypeError:
+        # When `searchsorted` fails with `TypeError`, it may be
+        # caused by nulls in `s`. Try again with the null-values
+        # explicitly mapped to the first partition.
+        partitions = np.empty(len(s), dtype="int32")
+        partitions[s.isna()] = 0
+        not_null = s.notna()
+        partitions[not_null] = divisions.searchsorted(s[not_null], side="right") - 1
     partitions[(s >= divisions.iloc[-1]).values] = len(divisions) - 2
     return partitions
 
diff --git a/dask/dataframe/tests/test_merge_column_and_index.py b/dask/dataframe/tests/test_merge_column_and_index.py
index 5e03035e033..7974f6c63f0 100644
--- a/dask/dataframe/tests/test_merge_column_and_index.py
+++ b/dask/dataframe/tests/test_merge_column_and_index.py
@@ -209,3 +209,25 @@ def test_merge_known_to_double_bcast_left(
     assert_eq(result, expected)
     assert_eq(result.divisions, ddf_right.divisions)
     assert len(result.__dask_graph__()) < 90
+
+
+@pytest.mark.parametrize("repartition", [None, 4])
+def test_merge_column_with_nulls(repartition):
+    # See: https://github.com/dask/dask/issues/7558
+
+    df1 = pd.DataFrame({"a": ["0", "0", None, None, None, None, "5", "7", "15", "33"]})
+    df2 = pd.DataFrame({"c": ["1", "2", "3", "4"], "b": ["0", "5", "7", "15"]})
+    df1_d = dd.from_pandas(df1, npartitions=4)
+    df2_d = dd.from_pandas(df2, npartitions=3).set_index("b")
+    if repartition:
+        df2_d = df2_d.repartition(repartition)
+
+    pandas_result = df1.merge(
+        df2.set_index("b"), how="left", left_on="a", right_index=True
+    )
+    dask_result = df1_d.merge(df2_d, how="left", left_on="a", right_index=True)
+
+    assert_eq(
+        dask_result.compute().sort_index(),
+        pandas_result.sort_index(),
+    )
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
index 4618f3fdaf4..facba89c5d2 100644
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -2077,9 +2077,10 @@ def test_merge_outer_empty():
 
     for x in range(0, k_clusters + 1):
         assert_eq(
-            dd.merge(empty_df, df[df.cluster == x], how="outer").compute(),
-            df[df.cluster == x].compute(),
+            dd.merge(empty_df, df[df.cluster == x], how="outer"),
+            df[df.cluster == x],
             check_index=False,
+            check_divisions=False,
         )
 
 


Code review comments:
diff_hunk: @@ -209,3 +209,25 @@ def test_merge_known_to_double_bcast_left(
     assert_eq(result, expected)
     assert_eq(result.divisions, ddf_right.divisions)
     assert len(result.__dask_graph__()) < 90
+
+
+@pytest.mark.parametrize("repartition", [None, 4])
+def test_merge_column_with_nulls(repartition):
+    # See: https://github.com/dask/dask/issues/7558
+
+    df1 = pd.DataFrame({"a": ["0", "0", None, None, None, None, "5", "7", "15", "33"]})
+    df2 = pd.DataFrame({"c": ["1", "2", "3", "4"], "b": ["0", "5", "7", "15"]})
+    df1_d = dd.from_pandas(df1, npartitions=4)
+    df2_d = dd.from_pandas(df2, npartitions=3).set_index("b")
+    if repartition:
+        df2_d = df2_d.repartition(repartition)
+
+    pandas_result = df1.merge(
+        df2.set_index("b"), how="left", left_on="a", right_index=True
+    )
+    dask_result = df1_d.merge(df2_d, how="left", left_on="a", right_index=True)
+
+    assert_eq(
+        dask_result.compute().sort_index(),
+        pandas_result.sort_index(),
+    )
path: dask/dataframe/tests/test_merge_column_and_index.py
text: @user1:
I think this accomplishes the same thing?
```suggestion
    assert_eq(dask_result, pandas_result, check_order=False)
```

@author:
Hmm.. `TypeError: assert_frame_equal() got an unexpected keyword argument 'check_order'`

@author:
However, if seems that `assert_eq(dask_result, pandas_result)` actually works in this case