Repository: dask/dask

Issue:
Categorical data not surviving a repartition
When creating categorical columns with  
```
df[column] = df[column].astype('category')
df = df.repartition(npartitions=n)
```
  seems to result in partitions that have the old datatype for the column and one partition that has the correct category datatype.

This is with dask 15.2, I have not tested older versions.



Discussion:
This boils down to `pd.concat` not propagating categories when they don't match

```python
In [1]: import pandas as pd

In [2]: df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': ['a', 'b', 'a', 'c']})

In [3]: df2 = pd.DataFrame({'x': [1, 2, 3, 4], 'y': ['a', 'b', 'a', 'e']})

In [4]: df['y'] = df['y'].astype('category')

In [5]: df2['y'] = df2['y'].astype('category')

In [6]: df.dtypes
Out[6]: 
x       int64
y    category
dtype: object

In [7]: df2.dtypes
Out[7]: 
x       int64
y    category
dtype: object

In [8]: pd.concat([df, df2]).dtypes
Out[8]: 
x     int64
y    object
dtype: object
```

@TomAugspurger is there a way around this?

(my test case for future use)

```python
@pytest.mark.parametrize('npartitions', [1, 4])
def test_repartition_on_categoricals(npartitions):
    df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': ['a', 'b', 'a', 'c']})
    ddf = dd.from_pandas(df, npartitions=2)
    ddf['y'] = ddf['y'].astype('category')
    ddf2 = ddf.repartition(npartitions=npartitions)
    
    df = df.copy() 
    df['y'] = df['y'].astype('category')
    assert_eq(df, ddf)
    assert_eq(df, ddf2)
```
Isn't that surprising from `repartition`, where the partitions *were* known beforehand and do not get recoded?
This might be fixed by replacing the calls to `pd.concat` with `dask.dataframe.methods.concat` (which handles mixed categories).
Yep, PR coming momentarily.
Thank you @jcrist 
For reference, here's the pandas issue about concat and categoricals https://github.com/pandas-dev/pandas/issues/14177 (we weren't sure whether unioning, converting to object, or raising was "correct").

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -3642,7 +3642,7 @@ def _is_single_last_div(x):
                 raise ValueError('check for duplicate partitions\nold:\n%s\n\n'
                                  'new:\n%s\n\ncombined:\n%s'
                                  % (pformat(a), pformat(b), pformat(c)))
-            d[(out2, j - 1)] = (pd.concat, tmp)
+            d[(out2, j - 1)] = (methods.concat, tmp)
         j += 1
     return d
 
@@ -3674,10 +3674,10 @@ def repartition_npartitions(df, npartitions):
                                      for new_partition_index in range(npartitions + 1)]
         dsk = {}
         for new_partition_index in range(npartitions):
-            value = (pd.concat, [(df._name, old_partition_index)
-                                 for old_partition_index in
-                                 range(new_partitions_boundaries[new_partition_index],
-                                       new_partitions_boundaries[new_partition_index + 1])])
+            value = (methods.concat,
+                     [(df._name, old_partition_index) for old_partition_index in
+                      range(new_partitions_boundaries[new_partition_index],
+                            new_partitions_boundaries[new_partition_index + 1])])
             dsk[new_name, new_partition_index] = value
         divisions = [df.divisions[new_partition_index]
                      for new_partition_index in new_partitions_boundaries]


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
@@ -223,6 +223,19 @@ def test_categorical_set_index(shuffle):
         assert list(sorted(d2.index.compute())) == ['b', 'b', 'c']
 
 
+@pytest.mark.parametrize('npartitions', [1, 4])
+def test_repartition_on_categoricals(npartitions):
+    df = pd.DataFrame({'x': range(10), 'y': list('abababcbcb')})
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf['y'] = ddf['y'].astype('category')
+    ddf2 = ddf.repartition(npartitions=npartitions)
+
+    df = df.copy()
+    df['y'] = df['y'].astype('category')
+    assert_eq(df, ddf)
+    assert_eq(df, ddf2)
+
+
 def test_categorical_accessor_presence():
     df = pd.DataFrame({'x': list('a' * 5 + 'b' * 5 + 'c' * 5), 'y': range(15)})
     df.x = df.x.astype('category')
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -12,7 +12,7 @@
 import dask.dataframe as dd
 
 from dask.dataframe.core import repartition_divisions, aca, _concat, Scalar
-from dask.dataframe.methods import boundary_slice
+from dask.dataframe import methods
 from dask.dataframe.utils import (assert_eq, make_meta, assert_max_deps,
                                   PANDAS_VERSION)
 
@@ -1103,17 +1103,17 @@ def _check_split_data(orig, d):
 
 def test_repartition_divisions():
     result = repartition_divisions([0, 6], [0, 6, 6], 'a', 'b', 'c')
-    assert result == {('b', 0): (boundary_slice, ('a', 0), 0, 6, False),
-                      ('b', 1): (boundary_slice, ('a', 0), 6, 6, True),
+    assert result == {('b', 0): (methods.boundary_slice, ('a', 0), 0, 6, False),
+                      ('b', 1): (methods.boundary_slice, ('a', 0), 6, 6, True),
                       ('c', 0): ('b', 0),
                       ('c', 1): ('b', 1)}
 
     result = repartition_divisions([1, 3, 7], [1, 4, 6, 7], 'a', 'b', 'c')
-    assert result == {('b', 0): (boundary_slice, ('a', 0), 1, 3, False),
-                      ('b', 1): (boundary_slice, ('a', 1), 3, 4, False),
-                      ('b', 2): (boundary_slice, ('a', 1), 4, 6, False),
-                      ('b', 3): (boundary_slice, ('a', 1), 6, 7, True),
-                      ('c', 0): (pd.concat, [('b', 0), ('b', 1)]),
+    assert result == {('b', 0): (methods.boundary_slice, ('a', 0), 1, 3, False),
+                      ('b', 1): (methods.boundary_slice, ('a', 1), 3, 4, False),
+                      ('b', 2): (methods.boundary_slice, ('a', 1), 4, 6, False),
+                      ('b', 3): (methods.boundary_slice, ('a', 1), 6, 7, True),
+                      ('c', 0): (methods.concat, [('b', 0), ('b', 1)]),
                       ('c', 1): ('b', 2),
                       ('c', 2): ('b', 3)}
 
@@ -2637,23 +2637,23 @@ def test_slice_on_filtered_boundary(drop):
 def test_boundary_slice_nonmonotonic():
     x = np.array([-1, -2, 2, 4, 3])
     df = pd.DataFrame({"B": range(len(x))}, index=x)
-    result = boundary_slice(df, 0, 4)
+    result = methods.boundary_slice(df, 0, 4)
     expected = df.iloc[2:]
     tm.assert_frame_equal(result, expected)
 
-    result = boundary_slice(df, -1, 4)
+    result = methods.boundary_slice(df, -1, 4)
     expected = df.drop(-2)
     tm.assert_frame_equal(result, expected)
 
-    result = boundary_slice(df, -2, 3)
+    result = methods.boundary_slice(df, -2, 3)
     expected = df.drop(4)
     tm.assert_frame_equal(result, expected)
 
-    result = boundary_slice(df, -2, 3.5)
+    result = methods.boundary_slice(df, -2, 3.5)
     expected = df.drop(4)
     tm.assert_frame_equal(result, expected)
 
-    result = boundary_slice(df, -2, 4)
+    result = methods.boundary_slice(df, -2, 4)
     expected = df
     tm.assert_frame_equal(result, expected)
 
@@ -2674,7 +2674,7 @@ def test_boundary_slice_nonmonotonic():
 def test_with_boundary(start, stop, right_boundary, left_boundary, drop):
     x = np.array([-1, -2, 2, 4, 3])
     df = pd.DataFrame({"B": range(len(x))}, index=x)
-    result = boundary_slice(df, start, stop, right_boundary, left_boundary)
+    result = methods.boundary_slice(df, start, stop, right_boundary, left_boundary)
     expected = df.drop(drop)
     tm.assert_frame_equal(result, expected)
 
@@ -2695,7 +2695,7 @@ def test_with_boundary(start, stop, right_boundary, left_boundary, drop):
 ])
 def test_boundary_slice_same(index, left, right):
     df = pd.DataFrame({"A": range(len(index))}, index=index)
-    result = boundary_slice(df, left, right)
+    result = methods.boundary_slice(df, left, right)
     tm.assert_frame_equal(result, df)