Repository: dask/dask

Issue:
astype(CategoricalDtype) should be known categoricals
Currently, they're unknown

```
In [108]: import pandas as pd; import dask.dataframe as dd

In [109]: df = pd.DataFrame({"A": [1, 2], "B": pd.Categorical(['a', 'b']), 'C': ['a', 'b']})

In [110]: a = dd.from_pandas(df, 2)

In [111]: df.C.astype(df.B.dtype)
Out[111]:
0    a
1    b
Name: C, dtype: category
Categories (2, object): [a, b]

In [112]: a.C.astype(a.B.dtype)
Out[112]:
Dask Series Structure:
npartitions=1
0    category[unknown]
1                  ...
Name: C, dtype: category
Dask Name: astype, 3 tasks
```


Discussion:

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1418,8 +1418,24 @@ def isin(self, values):
 
     @derived_from(pd.DataFrame)
     def astype(self, dtype):
-        meta = self._meta.astype(dtype)
-        meta = clear_known_categories(meta)
+        # XXX: Pandas will segfault for empty dataframes when setting
+        # categorical dtypes. This operation isn't allowed currently anyway. We
+        # get the metadata with a non-empty frame to throw the error instead of
+        # segfaulting.
+        if isinstance(self._meta, pd.DataFrame) and is_categorical_dtype(dtype):
+            meta = self._meta_nonempty.astype(dtype)
+        else:
+            meta = self._meta.astype(dtype)
+        if hasattr(dtype, 'items'):
+            # Pandas < 0.21.0, no `categories` attribute, so unknown
+            # Pandas >= 0.21.0, known if `categories` attribute is not None
+            set_unknown = [k for k, v in dtype.items()
+                           if (is_categorical_dtype(v) and
+                               getattr(v, 'categories', None) is None)]
+            meta = clear_known_categories(meta, cols=set_unknown)
+        elif (is_categorical_dtype(dtype) and
+              getattr(dtype, 'categories', None) is None):
+            meta = clear_known_categories(meta)
         return self.map_partitions(M.astype, dtype=dtype, meta=meta)
 
     @derived_from(pd.Series)


Test patch:
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
@@ -2075,11 +2075,15 @@ def test_astype():
 
 def test_astype_categoricals():
     df = pd.DataFrame({'x': ['a', 'b', 'c', 'b', 'c'],
-                       'y': [1, 2, 3, 4, 5]})
+                       'y': ['x', 'y', 'z', 'x', 'y'],
+                       'z': [1, 2, 3, 4, 5]})
+    df = df.astype({'y': 'category'})
     ddf = dd.from_pandas(df, 2)
+    assert ddf.y.cat.known
 
     ddf2 = ddf.astype({'x': 'category'})
     assert not ddf2.x.cat.known
+    assert ddf2.y.cat.known
     assert ddf2.x.dtype == 'category'
     assert ddf2.compute().x.dtype == 'category'
 
@@ -2089,6 +2093,36 @@ def test_astype_categoricals():
     assert dx.compute().dtype == 'category'
 
 
+@pytest.mark.skipif(PANDAS_VERSION < '0.21.0',
+                    reason="No CategoricalDtype with categories")
+def test_astype_categoricals_known():
+    df = pd.DataFrame({'x': ['a', 'b', 'c', 'b', 'c'],
+                       'y': ['x', 'y', 'z', 'y', 'z'],
+                       'z': ['b', 'b', 'b', 'c', 'b'],
+                       'other': [1, 2, 3, 4, 5]})
+    ddf = dd.from_pandas(df, 2)
+
+    abc = pd.api.types.CategoricalDtype(['a', 'b', 'c'])
+    category = pd.api.types.CategoricalDtype()
+
+    # DataFrame
+    ddf2 = ddf.astype({'x': abc,
+                       'y': category,
+                       'z': 'category',
+                       'other': 'f8'})
+
+    for col, known in [('x', True), ('y', False), ('z', False)]:
+        x = getattr(ddf2, col)
+        assert pd.api.types.is_categorical_dtype(x.dtype)
+        assert x.cat.known == known
+
+    # Series
+    for dtype, known in [('category', False), (category, False), (abc, True)]:
+        dx2 = ddf.x.astype(dtype)
+        assert pd.api.types.is_categorical_dtype(dx2.dtype)
+        assert dx2.cat.known == known
+
+
 def test_groupby_callable():
     a = pd.DataFrame({'x': [1, 2, 3, None], 'y': [10, 20, 30, 40]},
                      index=[1, 2, 3, 4])