Repository: dask/dask

Issue:
read_parquet not passing through categories
Although fastparquet correctly loads a categorical, and it shows up in the underlying pandas dataframe, this information does not appear in the dask dataframe.

Failing case:
```
def test_cats_in_parquet():
    s = pd.Series([b'a', b'b', b'c']*20, dtype='category')
    d = pd.DataFrame({'b': s})
    df = dd.from_pandas(d, npartitions=1)
    to_parquet('temp.parq', df, write_index=True)
    df2 = read_parquet('temp.parq', index='index', categories=['b'])
    assert df2.compute().b.cat,categories == df.b.cat.categories
    assert df2.b.cat.categories == df.b.cat.categories
```



Discussion:
@TomAugspurger what is the right way to specify categories on a dask.dataframe object?

If you know which columns are categoricals, and the set of categories per column, then
Using `df = df.assign(b=lambda x: x.b.cat.set_categories(categories))`

Or if you have a dictionary of `{col: [categories]}` then I think `df = df.assign(**{k: lambda x: x[k].cat.set_categories(cats) for k, cats in cat_map.items()}` would work. I haven't followed the parquet stuff closely, but [this](http://fastparquet.readthedocs.io/en/latest/details.html#categoricals) suggests that you should know the set of categories?

It seems to me that the right thing to do is to provide the categories in the `meta` value on dataframe creation. The trouble is that we don't know the cat values before load. I notice by looking at the from_bcolz function that loading the values and calling unique is sometimes used. 
In the parquet case, the cat values are stored separately and could be loaded beforehand fairly quickly. In fact, we could use this to enforce the same cat values for every piece by passing them around at load time - at some cost in complexity.

FWIW right now, using `dd.read_csv(..., dtype={'col: 'category'})` (which works in pandas 0.19+) leaves you in an inconsistent state where `df.col.cat.categories` won't be correct since we haven't done the full scan of the data.

I should just read the code, but will different "blocks" (not sure if that's the right term) in the parquet file all have the same dictionary encoding of `{val: int}`? Or is that what you mean by

> In fact, we could use this to enforce the same cat values for every piece by passing them around at load time

> The trouble is that we don't know the cat values before load. 

With the current design we _must_ know the categories before load.  Otherwise we just can't faithfully support categoricals.  

This can change if we start making use of newer `union_categories` functionality in Pandas in the future, but that's not implemented yet.

They don't necessarily have the same dictionary encoding, but this is opt-in for the case that the user is sure that they do. If fastparquet wrote the file, then the encoding will be consistent (if indeed they were consistent in the original data-frame)

Solved by this? https://github.com/dask/fastparquet/pull/28/


Resolving patch:
diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -1,5 +1,6 @@
 import struct
 
+import pandas as pd
 from toolz import first, partial
 
 from ..core import DataFrame, Series
@@ -71,6 +72,11 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None):
            not(fastparquet.api.filter_out_stats(rg, filters, pf.helper)) and
            not(fastparquet.api.filter_out_cats(rg, filters))]
 
+    # get category values from first row-group
+    categories = categories or []
+    cats = pf.grab_cats(categories)
+    categories = [cat for cat in categories if cats.get(cat, None) is not None]
+
     # Find an index among the partially sorted columns
     minmax = fastparquet.api.sorted_partitioned_columns(pf)
 
@@ -105,6 +111,9 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None):
               pf.dtypes.items() if k in all_columns}
 
     meta = make_meta(dtypes)
+    for cat in categories:
+        meta[cat] = pd.Series(pd.Categorical([], categories=cats[cat]))
+
     if index_col:
         meta = meta.set_index(index_col)
 


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -111,3 +111,18 @@ def test_optimize(fn, c):
     dsk = x._optimize(x.dask, x._keys())
     assert len(dsk) == x.npartitions
     assert all(v[4] == c for v in dsk.values())
+
+
+def test_categorical():
+    with tmpdir() as tmp:
+        df = pd.DataFrame({'x': ['a', 'b', 'c'] * 100},
+                          dtype='category')
+        ddf = dd.from_pandas(df, npartitions=3)
+        to_parquet(tmp, ddf)
+
+        ddf2 = read_parquet(tmp, categories=['x'])
+
+        assert ddf2.x.cat.categories.tolist() == ['a', 'b', 'c']
+        ddf2.loc[:1000].compute()
+        df.index.name = 'index'  # defaults to 'index' in this case
+        assert assert_eq(df, ddf2)