Repository: dask/dask

Issue:
Bug in set_index/partition on categorical index?
I would expect these two operations to result in the same divisions, but they don't.

```
    import dask as da
    import dask.dataframe as dd
    from dask.dataframe.shuffle import set_partition

    da.set_options(get=da.async.get_sync)

    df = pd.DataFrame({'x': [1, 2, 2, 2, 4], 'y': ['a', 'b', 'b', 'b', 'c']})
    df['y'] = df.y.astype('category')
    a = dd.from_pandas(df, npartitions=2)

    npartitions = 2
    divisions_y = a.y.quantile(np.linspace(0, 1, npartitions + 1)).compute().tolist()
    aa = set_partition(a, 'y', divisions_y, compute=True, drop=True)
    divisions_x = a.x.quantile(np.linspace(0, 1, npartitions + 1)).compute().tolist()
    bb = set_partition(a, 'x', divisions_x, compute=True, drop=True)
```

It looks like you're calling `strip_categories` on the DataFrames, but not applying this to the divisions. So you're doing something like

```
pd.Index([1, 2]).searchsorted(np.array(['b']))
```

Python 3.5.1, pandas 0.18.0, dask master.



Discussion:

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -121,7 +121,7 @@ def set_partition(df, index, divisions, compute=False, drop=True, **kwargs):
         dsk = merge(df.dask, dsk1, dsk2, dsk3)
         if isinstance(index, _Frame):
             dsk.update(index.dask)
-        p, barrier_token, categories = df._get(dsk, [p, barrier_token, catname], **kwargs)
+        p, barrier_token, categories = df._get(dsk, [p, barrier_token, catname2], **kwargs)
         dsk4 = {catname2: categories}
     else:
         dsk4 = {}
@@ -156,9 +156,9 @@ def barrier(args):
 def _set_partition(df, index, divisions, p, drop=True):
     """ Shard partition and dump into partd """
     df = df.set_index(index, drop=drop)
-    df = strip_categories(df)
     divisions = list(divisions)
     shards = shard_df_on_index(df, divisions[1:-1])
+    shards = list(map(strip_categories, shards))
     p.append(dict(enumerate(shards)))
 
 


Test patch:
diff --git a/dask/dataframe/tests/test_categorical.py b/dask/dataframe/tests/test_categorical.py
--- a/dask/dataframe/tests/test_categorical.py
+++ b/dask/dataframe/tests/test_categorical.py
@@ -34,11 +34,21 @@ def test_categorical_set_index():
     with dask.set_options(get=get_sync):
         b = a.set_index('y')
         df2 = df.set_index('y')
-        assert list(b.index.compute()), list(df2.index)
+        d1, d2 = b.get_division(0), b.get_division(1)
+        assert list(d1.index.compute()) == ['a']
+        assert list(sorted(d2.index.compute())) == ['b', 'b', 'c']
 
         b = a.set_index(a.y)
         df2 = df.set_index(df.y)
-        assert list(b.index.compute()), list(df2.index)
+        d1, d2 = b.get_division(0), b.get_division(1)
+        assert list(d1.index.compute()) == ['a']
+        assert list(sorted(d2.index.compute())) == ['b', 'b', 'c']
+
+        b = a.set_partition('y', ['a', 'b', 'c'])
+        df2 = df.set_index(df.y)
+        d1, d2 = b.get_division(0), b.get_division(1)
+        assert list(d1.index.compute()) == ['a']
+        assert list(sorted(d2.index.compute())) == ['b', 'b', 'c']
 
 
 def test_dataframe_categoricals():