Repository: dask/dask

Issue:
Rows can be dropped when setting index on a categorical with non-lexical order.
When setting the index on a partitioned dataframe to a categorical whose order is not lexical, rows can be lost. That's a mouthful for sure, and I'm not certain that it would be reasonable to expect otherwise, but I have produced an MRP.

```python
import os
import string, random

import pandas as pd
import numpy as np
import dask
import dask.dataframe as dd

print(dask.__version__)

def random_word(len=10):
    return ''.join(np.random.choice(list(string.ascii_letters), len))

rows = 5000
labels = [random_word(np.random.randint(15, 30)) for i in range(rows)]
data = [i for i in range(len(labels))]

df = pd.DataFrame({'ID': labels, 'data': data})


ddf = dd.from_pandas(df, npartitions=10)
print(ddf.shape[0].compute())

# Drop rows
ddf = ddf.categorize(columns=['ID'])
ddf['ID'] = ddf['ID'].cat.as_ordered()
ddf = ddf.set_index('ID')
print(ddf.shape[0].compute())

# Does not drop rows
ddf = dd.from_pandas(df, npartitions=10)
ddf = ddf.categorize(columns=['ID'])

# Order category lexically
cats = ddf['ID'].cat.categories.values
cats = sorted(cats)
ddf = ddf
ddf['ID'] = ddf['ID'].cat.reorder_categories(cats, ordered=True)

ddf.set_index('ID')
print(ddf.shape[0].compute())
```

In my case, the output generated is:

```
2.8.1
5000
1223
5000
```

This doesn't occur when the dataframe is a single partition. Digging around in the documentation makes me think there might be a scalar lexical comparison instead of a categorical compare somewhere. That or I'm not expecting the right behavior!


Discussion:
Thanks for the excellent issue report, I'm able to reproduce the issue. This is definitely a bug. I'd guess that it's an issue in our `set_index` code, which does sorting along the index to repartition. If you feel like debugging deeper, I'd start looking there.
I've spent some more time with this and I think I've narrowed it down to `dataframe.shuffle:rearrange_by_column_disk()` but wow it is pretty gnarly in there. I'm not familiar enough with dask's internals yet to figure out exactly what is going on. It looks like there are some `as_ordered` references in the shuffle-collect stage of the graph, `dsk4`, and I think they might be related.
Ok I think I've found the defect, but I'm unsure of the best way to patch it. In `dask.dataframe.shuffle.set_partitions_pre` the `divisions.searchsorted` operation returns invalid results since Series `s` has dtype `pandas.CategoricalDtype` while `divisions` is of type `object`. If I override the dtype of `divisions` to match `s` then the results are correct. Performing that operation in `set_partitions_pre` seems like a bad way to do this.

I think a better fix would be changing can be  `dask.dataframe.shuffle.set_partitions`:

`    divisions = df._meta._constructor_sliced(divisions)`

to

`    divisions = df._meta._constructor_sliced(divisions).astype(df[index].dtype)`

Making sure the types match seems like a good idea, but I note that in `dask.dataframe.shuffle.set_index` we have already explicitly changed them from a series of the correct type to a list(!):

`divisions = divisions.tolist()`

Swapping the type of divisions back and forth seems sub-optimal.

I'm going through the process testing a commit and all that, but wanted to post just incase this is a ridiculous proposal, and because I'm going to be traveling starting tomorrow and I don't want to forget the details.
Thanks for continuing to update. Could you share snippets of input /outputs in `set_partitions` / `set_partitions_pre` and `searchsorted`? That'd make it easier for us to gauge things.

I'll also take a look locally now.
A  bit easier of a reproducer since it doesn't modify `ddf`


```python
import string

import pandas as pd
import numpy as np
import dask
import dask.dataframe as dd

print(dask.__version__)


def random_word(len=10):
    return ''.join(np.random.choice(list(string.ascii_letters), len))


rows = 5000
labels = [random_word(np.random.randint(15, 30)) for i in range(rows)]
labels = pd.Categorical(labels, categories=pd.unique(labels), ordered=True)
data = np.arange(len(labels))

df = pd.DataFrame({'ID': labels, 'data': data})
ddf = dd.from_pandas(df, npartitions=10)
r = ddf.set_index("ID").shape[0].compute()
assert r == 5000, r
```
Unfortunately I don't have a nice and tidy debug session since this involved stepping through the scheduler and such to find exactly where things broke. Digging through what's left in my console history here are some interesting bits:

Broken:
```
> dask/dataframe/shuffle.py(560)set_partitions_pre()                                                                                                                                                                                                   
-> partitions = divisions.searchsorted(s, side="right") - 1                                                                                                                                                                                                                                                                  
(Pdb) !s                                                                                                                                                                                                                                                                                                                     
0       HLKqjkquDubWNrpAVdUZZiNMjzal                                                                                                                                                                                                                                                                                         
1      SOYeMUTdJUzwZliNjYBHEWWvpvTmk                                                                                                                                                                                                                                                                                         
2             focbliZQYtABAjrVdJZCEN                                                                                                                                                                                                                                                                                         
3                  DunGqQmGuJlbjCMoq                                                                                                                                                                                                                                                                                         
4                  wlErPVmTibdjGrEiZ                                                                                                                                                                                                                                                                                         
                   ...                                                                                                                                                                                                                                                                                                       
495           nPhmEiMMEcmJTzJYVAIOFQ                                                                                                                                                                                                                                                                                         
496    aWhjcqaiAMCfiypafmzTvGxfqGxzw                                                                                                                                                                                                                                                                                         
497     OcXGYFNaXuhtEjtikCAtkFEiCxKK                                                                                                                                                                                                                                                                                         
498                kbGMqzwPEWrSbiwxC                                                                                                                                                                                                                                                                                         
499     xjKOAmQADwMarfmqFYsBifdXxtxR                                                                                                                                                                                                                                                                                         
Name: ID, Length: 500, dtype: category                                                                                                                                                                                                                                                                                       
Categories (1000, object): [HLKqjkquDubWNrpAVdUZZiNMjzal < SOYeMUTdJUzwZliNjYBHEWWvpvTmk <                                                                                                                                                                                                                                   
                            focbliZQYtABAjrVdJZCEN < DunGqQmGuJlbjCMoq ... iqDRQaMxpCOuNaNCOOObJasR <                                                                                                                                                                                                                        
                            zCVyliMRpVtQAWvmFMZvqcpcPa < jexaZlCqfSEoIQAsfYivprkhtJG <                                                                                                                                                                                                                                       
                            hkzuLcghsBzShCwxHyTDPtzwBE]                                                                                                                                                                                                                                                                      
(Pdb) divisions                                                                                                                                                                                                                                                                                                              
0    HLKqjkquDubWNrpAVdUZZiNMjzal                                                                                                                                                                                                                                                                                            
1    xjKOAmQADwMarfmqFYsBifdXxtxR                                                                                                                                                                                                                                                                                            
2      hkzuLcghsBzShCwxHyTDPtzwBE                                                                                                                                                                                                                                                                                            
dtype: object                                                                                                                                                                                                                                                                                                                
(Pdb) divisions.searchsorted(s, side="right") - 1                                                                                                                                                                                                                                                                            
array([ 0,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  2,  0,  0,  0,  2,  2,                                                                                                                                                                                                                                                   
        2,  0, -1,  0,  0,  0,  0,  2,  0,  2,  0,  0,  2,  0,  0,  2, -1,                                                                                                                                                                                                                                                   
        0,  0,  2,  0, -1,  0,  0,  0, -1,  0,  2,  0,  2,  0,  2,  0,  0,                                                                                                                                                                                                                                                   
       -1,  0,  0,  0,  2,  0,  0,  2, -1,  0,  0,  0, -1,  0,  0,  0, -1,                                                                                                                                                                                                                                                   
        0,  0, -1,  0,  0,  2,  0,  2,  2,  0,  2,  0,  2,  0,  2,  2,  2,                                                                                                                                                                                                                                                   
        0,  0, -1,  0,  0,  2, -1,  0,  2,  0,  0,  2,  0, -1,  0, -1,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  0,  0,  0,  2, -1,  0,  2,  0,  0, -1, -1,  0,  0,  0,                                                                                                                                                                                                                                                   
        0,  0, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0,  2,  0,  2,  0,  0,                                                                                                                                                                                                                                                   
       -1,  0,  0,  0,  0,  0,  2,  0,  2,  0,  0,  0,  0, -1,  0,  2,  0,                                                                                                                                                                                                                                                   
        0,  0,  2, -1, -1, -1,  0,  0,  2,  0,  0,  2,  0,  2,  0,  0,  0,                                                                                                                                                                                                                                                   
        2,  0,  0,  0,  2,  2, -1,  0,  0,  0,  0,  2,  0,  0, -1,  0,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  2,  0,  2,  0,  0, -1,  0,  0,  0,  2,  0,  0,  0, -1,                                                                                                                                                                                                                                                   
        0, -1, -1,  0, -1,  0,  0,  2,  0,  0,  0, -1,  0,  0,  2,  0,  0,                                                                                                                                                                                                                                                   
        2,  0,  0, -1,  0,  0,  0,  0,  0,  0,  2,  0,  2, -1,  0,  0,  2,                                                                                                                                                                                                                                                   
        0,  2, -1,  0, -1, -1,  0, -1,  0,  0,  2,  0,  0,  0,  0,  2, -1,                                                                                                                                                                                                                                                   
        0,  0,  2,  0,  0,  2,  0,  0,  0,  2,  0,  0, -1,  0,  2, -1,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  0,  2,  0,  0,  0,  0,  2,  2, -1,  0,  0,  0,  2, -1,                                                                                                                                                                                                                                                   
       -1, -1,  0, -1,  0,  0,  0,  2,  0,  0,  2,  0,  2,  0,  2,  0,  2,                                                                                                                                                                                                                                                   
        0, -1,  0, -1,  0,  2,  0,  2, -1,  0,  2,  2, -1,  0,  2,  0,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  2,  0, -1,  0,  2,  0,  2,  0,  0,  0,  2,  2,  2,  0,                                                                                                                                                                                                                                                   
        0,  0,  2,  0,  0,  2,  0,  0,  2,  0,  0,  2,  0,  0,  2, -1,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  2,  0,  2,  2,  0, -1,  0,  2,  0,  0,  2,  2,  0,  2,                                                                                                                                                                                                                                                   
        0,  0, -1,  0,  0, -1,  0,  0,  0,  2, -1,  0,  0, -1,  0,  0, -1,                                                                                                                                                                                                                                                   
        0,  0, -1,  0,  0,  0,  0,  0,  0,  2,  0,  0,  2,  0,  2,  0,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  0,  0,  0, -1,  0,  0,  2, -1,  0,  0,  0,  2,  2,  0,                                                                                                                                                                                                                                                   
        0,  2,  0, -1,  0, -1, -1, -1,  0,  2,  0,  0,  0,  2, -1,  0,  0,                                                                                                                                                                                                                                                   
        0,  2,  0,  2,  0,  0,  2,  0,  0,  0, -1, -1,  2, -1,  0,  0,  0,                                                                                                                                                                                                                                                   
        2,  2,  0,  0,  0,  2,  0,  2,  0,  0,  0,  0, -1, -1,  0,  0,  2,                                                                                                                                                                                                                                                   
        0,  2,  0,  2,  0,  0,  0,  2,  0,  0,  0,  2,  2,  0,  0,  2,  2,                                                                                                                                                                                                                                                   
        0,  0,  0,  0,  0,  2,  2])
```

Not broken:

```
> dask/local.py(224)execute_task()                                                                                                                                                                                                                     
-> result = dumps((result, id))                                                                                                                                                                                                                                                                                              
(Pdb) result                                                                                                                                                                                                                                                                                                                 
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,                                                                                                                                                                                                                                                     
       1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,                                                                                                                                                                                                                                                     
       1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
```

Although I'm definitely concerned about the labels being clumped together like that, so maybe it's just broken in a different way. All the rows are at least in the dataframe at the end of the set_index operation this way!

Since I don't have a proper development environment I used a patch:

```
def set_partitions_pre(s, divisions):
    from pandas import CategoricalDtype
    if isinstance(s.dtype, CategoricalDtype):
        divisions = divisions.astype(s.dtype)
    partitions = divisions.searchsorted(s, side="right") - 1
    partitions[(s >= divisions.iloc[-1]).values] = len(divisions) - 2
    return partitions

dask.dataframe.shuffle.set_partitions_pre = set_partitions_pre
```

Oh, one other hint for reproducing - this still fails with npartitions = 2, so for debug I was using 1000 rows with two partitions.

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -90,6 +90,10 @@ def set_index(
 
         mins = remove_nans(mins)
         maxes = remove_nans(maxes)
+        if pd.api.types.is_categorical_dtype(index2.dtype):
+            dtype = index2.dtype
+            mins = pd.Categorical(mins, dtype=dtype).codes.tolist()
+            maxes = pd.Categorical(maxes, dtype=dtype).codes.tolist()
 
         if (
             mins == sorted(mins)
@@ -170,8 +174,22 @@ def set_partition(
     shuffle
     partd
     """
-    divisions = df._meta._constructor_sliced(divisions)
     meta = df._meta._constructor_sliced([0])
+    if isinstance(divisions, tuple):
+        # pd.isna considers tuples to be scalars. Convert to a list.
+        divisions = list(divisions)
+
+    if np.isscalar(index):
+        dtype = df[index].dtype
+    else:
+        dtype = index.dtype
+
+    if pd.isna(divisions).any() and pd.api.types.is_integer_dtype(dtype):
+        # Can't construct a Series[int64] when any / all of the divisions are NaN.
+        divisions = df._meta._constructor_sliced(divisions)
+    else:
+        divisions = df._meta._constructor_sliced(divisions, dtype=dtype)
+
     if np.isscalar(index):
         partitions = df[index].map_partitions(
             set_partitions_pre, divisions=divisions, meta=meta


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
@@ -1,5 +1,7 @@
 import itertools
 import os
+import random
+
 import pandas as pd
 import pytest
 import pickle
@@ -724,6 +726,22 @@ def test_set_index_on_empty():
         assert ddf.npartitions == 1
 
 
+def test_set_index_categorical():
+    # https://github.com/dask/dask/issues/5671
+    order = list(reversed(string.ascii_letters))
+    values = list(string.ascii_letters)
+    random.shuffle(values)
+    dtype = pd.api.types.CategoricalDtype(order, ordered=True)
+    df = pd.DataFrame({"A": pd.Categorical(values, dtype=dtype), "B": 1})
+
+    result = dd.from_pandas(df, npartitions=2).set_index("A")
+    assert len(result) == len(df)
+
+    # sorted with the metric defined by the Categorical
+    divisions = pd.Categorical(result.divisions, dtype=dtype)
+    tm.assert_categorical_equal(divisions, divisions.sort_values())
+
+
 def test_compute_divisions():
     from dask.dataframe.shuffle import compute_divisions