Repository: dask/dask

Issue:
Different categories in different partitions
Categoricals are currently difficult in dask.dataframe because all partitions must have the same set of categories.  This synchronization can be very expensive, often requiring full table scans.  Operations like `.astype('category')` or `read_csv(..., dtypes={'x': 'category'})` don't work as expected.

Recently pandas started to support union-ing categoricals on concatenation.  This ability might allow dask.dataframe to instead maintain categories per-partition, which removes the burden of synchronization.  We would only combine categoricals when different partitions need to interact.  This would help us in data loading (csv, parquet) and in aggregations (groupbys, etc.) but not for operations like `pivot_table` which expect to know the categories ahead of time.

This might be as easy as adding the right call to `union_categoricals` in `dask.dataframe.core._concat`, but there are probably other nuances.

What is the best way to do this?

cc @TomAugspurger @martindurant @sinhrks 


Discussion:
Quick note: getting the full set of categories is very fast if the data is in memory (df.col.cat.categories as opposed to unique).
Also, specifically for parquet: if the data was written by fastparquet (via dask or not), then we can guarantee that the categories will be the same for each chunk.
If we remove the requirement that Dask.dataframe partitions all agree on categories then we might want to change the dask.dataframe parquet writer so that it writes individual partitions without synchronizing categories beforehand.
I think this is a good idea... This should be transparent to the user, correct? It'll be up to dask to know when to union categoricals in the background?

> Operations like .astype('category') or read_csv(..., dtypes={'x': 'category'}) don't work as expected.

This should be less painful once I finish the CategoricalDtype refactoring in pandas, for the cases where you know the set of categories ahead of time. The proposal here would still be useful for the unknown categories case.
"CategoricalDtype refactoring" - is there a link for this?
@martindurant https://github.com/pandas-dev/pandas/issues/14711

short version: move `categories` and `ordered` into the dtype itself. Then add that to the public API and let users say `pd.read_csv('...', dtype={'a': pd.CategoricalDtype(['a', 'b', 'c'], ordered=True)})`. Should be entirely backwards compatible.
@TomAugspurger can you help me create a list of operations that this change will affect?

### Operations that become harder

*  `pivot_table`
*  `df.x.cat.codes`
*  `df.x.cat.categories`

### Operations that requires intelligent merging

*  Anything relying on apply-concat-apply (groupby-aggregations, reductions, drop_duplicates, etc.)
*  merge?
*  groupby-apply?

### What is the right way to merge two dataframes with categories?

```
In [1]: import pandas as pd

In [2]: a = pd.DataFrame({'x': [1, 2, 3], 'y': ['a', 'b', 'a']})

In [3]: b = pd.DataFrame({'x': [4, 5, 6], 'y': ['a', 'c', 'c']})

In [4]: a['y'] = a.y.astype('category')

In [5]: b['y'] = b.y.astype('category')

In [6]: c = ... # concatenate a and b, respecting categories
```
http://pandas.pydata.org/pandas-docs/stable/categorical.html#unioning
The function, not yet in concat, is
```
In []: union_categoricals([a['y'], b['y']])
Out[]:
[a, b, a, a, c, c]
Categories (3, object): [a, b, c]
```
(series by series)
@martindurant this change might affect dask's parquet reading and writing, which no longer needs to have a consistent set of categories in the metadata of all partitions.
Writing will not be affected, since each partition is treated separately.
On reading, this should actually simplify things a bit, I think. Am I right in thinking that the change here just means that the labels are not the same everywhere, but we do guarantee that the given column is categorical everywhere?
> Am I right in thinking that the change here just means that the labels are not the same everywhere, but we do guarantee that the given column is categorical everywhere?

Correct
I have a good feeling that removing L80-84, 123-124 in dask.io.parquet will do the right thing. I'll test that (and make a specific test) when I get to it.
Does the metadata now not need to know about category labels at all?

I believe that @jcrist established a convention around what to specify if categories are unknown.
Unknown categoricals should include `dd.utils.UNKNOWN_CATEGORIES` in their categories. You could either build each categorical yourself, or use `dd.utils.clear_known_categories`, which can be used to set known categorical columns to unknown.
The writing *and* reading seem to magically just work right now, although the lines above cause an unnecessary preemptive read.
I realise, though, that reading categoricals in fastparquet itself does require the same labels in every row-group, which I think is documented, but may be surprising. 
Here is a fail case:

```python
In [1]: import pandas as pd
In [2]: import dask.dataframe as dd

In [3]: df = pd.DataFrame({'x': [1, 2, 3, 4, 5], 'y': list('caaab')})

In [4]: ddf = dd.from_pandas(df, npartitions=2)

In [5]: ddf['y'] = ddf.y.astype('category')

In [6]: ddf.y.cat.categories
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-6-5e223c5a5c45> in <module>()
----> 1 ddf.y.cat.categories

/home/mrocklin/workspace/dask/dask/dataframe/categorical.py in categories(self)
    159                    "supported.  Please use `column.cat.as_known()` or "
    160                    "`df.categorize()` beforehand to ensure known categories")
--> 161             raise NotImplementedError(msg)
    162         return self._delegate_property(self._series._meta, 'cat', 'categories')
    163 

NotImplementedError: `df.column.cat.categories` with unknown categories is not supported.  Please use `column.cat.as_known()` or `df.categorize()` beforehand to ensure known categories

In [7]: ddf.to_parquet('foo')

In [8]: ddf2 = dd.read_parquet('foo', categories=['y'])

In [9]: ddf2.y.cat.categories
Out[9]: Index(['a', 'c'], dtype='object')

In [10]: ddf2.y.compute().cat.categories
Out[10]: Index(['a', 'c', 'b'], dtype='object')
```

I suggest that we don't use the categories of the first partition in read_parquet, but instead use the unknown categories convention.
OK, so read_parquet should call dd.utils.clear_known_categories, and then will have the same behaviour.
In these lines:

```python
    for cat in categories:
        meta[cat] = pd.Series(pd.Categorical([], categories=cats[cat]))
```

We probably need to replace `cats[cat]` with `[dd.utils.UNKNOWN_CATEGORIES]`.

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
@@ -2,6 +2,7 @@
 from toolz import first, partial
 
 from ..core import DataFrame, Series
+from ..utils import UNKNOWN_CATEGORIES
 from ...base import tokenize, normalize_token
 from ...compatibility import PY3
 from ...delayed import delayed
@@ -12,6 +13,7 @@
     from fastparquet import parquet_thrift
     from fastparquet.core import read_row_group_file
     from fastparquet.api import _pre_allocate
+    from fastparquet.util import check_column_names
     default_encoding = parquet_thrift.Encoding.PLAIN
 except:
     fastparquet = False
@@ -71,17 +73,14 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None,
     except:
         pf = fastparquet.ParquetFile(path, open_with=myopen, sep=myopen.fs.sep)
 
+    check_column_names(pf.columns, categories)
+    categories = categories or []
     name = 'read-parquet-' + tokenize(pf, columns, categories)
 
     rgs = [rg for rg in pf.row_groups if
            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)
 
@@ -113,7 +112,7 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None,
     if index_col and index_col not in all_columns:
         all_columns = all_columns + (index_col,)
 
-    dtypes = {k: ('category' if k in (categories or []) else v) for k, v in
+    dtypes = {k: ('category' if k in categories else v) for k, v in
               pf.dtypes.items() if k in all_columns}
 
     meta = pd.DataFrame({c: pd.Series([], dtype=d)
@@ -121,7 +120,8 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None,
                         columns=[c for c in pf.columns if c in dtypes])
 
     for cat in categories:
-        meta[cat] = pd.Series(pd.Categorical([], categories=cats[cat]))
+        meta[cat] = pd.Series(pd.Categorical([],
+                              categories=[UNKNOWN_CATEGORIES]))
 
     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
@@ -4,6 +4,7 @@
 import pytest
 
 import dask
+import dask.async
 from dask.utils import tmpdir, tmpfile
 import dask.dataframe as dd
 from dask.dataframe.io.parquet import read_parquet, to_parquet
@@ -130,7 +131,7 @@ def test_categorical():
 
         ddf2 = read_parquet(tmp, categories=['x'])
 
-        assert ddf2.x.cat.categories.tolist() == ['a', 'b', 'c']
+        assert ddf2.compute().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)
@@ -192,3 +193,25 @@ def test_roundtrip(df, write_kwargs, read_kwargs):
         to_parquet(tmp, ddf, **write_kwargs)
         ddf2 = read_parquet(tmp, index=df.index.name, **read_kwargs)
         assert_eq(ddf, ddf2)
+
+
+def test_categories(fn):
+    df = pd.DataFrame({'x': [1, 2, 3, 4, 5],
+                       'y': list('caaab')})
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf['y'] = ddf.y.astype('category')
+    ddf.to_parquet(fn)
+    ddf2 = dd.read_parquet(fn, categories=['y'])
+    with pytest.raises(NotImplementedError):
+        ddf2.y.cat.categories
+    assert set(ddf2.y.compute().cat.categories) == {'a', 'b', 'c'}
+    cats_set = ddf2.map_partitions(lambda x: x.y.cat.categories).compute()
+    assert cats_set.tolist() == ['a', 'c', 'a', 'b']
+    assert_eq(ddf.y, ddf2.y, check_names=False)
+    with pytest.raises(dask.async.RemoteException):
+        # attempt to load as category that which is not so encoded
+        ddf2 = dd.read_parquet(fn, categories=['x']).compute()
+
+    with pytest.raises(ValueError):
+        # attempt to load as category unknown column
+        ddf2 = dd.read_parquet(fn, categories=['foo'])