Repository: dask/dask

Issue:
Bad performance with DataFrame.categorize() and multiprocessing scheduler
```
pdf = pd.DataFrame(dict(A = ['A', 'B'] * 1000000, B = ['A', 'B'] * 1000000, C = ['A', 'B'] * 1000000))
ddf = dd.from_pandas(pdf, npartitions = 1000)
%time ddf.categorize(get = c.get) # c = distributed.Client()
%time ddf.categorize(get = dask.multiprocessing.get)

>CPU times: user 12.9 s, sys: 1.39 s, total: 14.2 s
>Wall time: 13.7 s

>CPU times: user 1min 12s, sys: 30.3 s, total: 1min 43s
>Wall time: 1min 30s
```

The multiprocessing scheduler is extremely inefficient if used together with `categorize`. The more columns in `ddf`, the bigger is the speed difference between the distributed and the multiprocessing scheduler. The main origin of the inefficency seems to be a lot of (de)serialization, as can be seen in the following profiling results.


<details>

```
`%prun ddf.categorize(get = dask.multiprocessing.get)`

         166389141 function calls (146670102 primitive calls) in 147.347 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
19249399/3436   51.293    0.000  127.712    0.037 pickle.py:269(save)
66006/50577   20.850    0.000  121.844    0.002 pickle.py:614(_batch_appends)
 18304083   14.725    0.000   20.397    0.000 pickle.py:260(get)
 20217882   10.829    0.000   10.829    0.000 {method 'write' of 'cStringIO.StringO' objects}
 19135962    5.963    0.000    5.963    0.000 {chr}
 20350667    5.179    0.000    5.179    0.000 {method 'get' of 'dict' objects}
    68860    5.159    0.000    5.159    0.000 {method 'acquire' of 'thread.lock' objects}
 19249399    3.447    0.000    3.447    0.000 pickle.py:333(persistent_id)
       24    3.128    0.130    3.128    0.130 {posix.fork}
 20705890    2.854    0.000    2.854    0.000 {id}
 18769284    2.439    0.000    2.439    0.000 {method 'append' of 'list' objects}
   619156    1.949    0.000    3.465    0.000 pickle.py:227(memoize)
    48024    1.579    0.000    1.742    0.000 core.py:159(get_dependencies)
    27575    1.519    0.000    2.594    0.000 ioloop.py:932(add_callback)
237030/3435    1.225    0.000  127.691    0.037 pickle.py:538(save_tuple)
    97725    0.936    0.000    1.910    0.000 {method '__reduce_ex__' of 'object' objects}
```
</details>

I wrote a custom categorize method that 

 1. converts the dataframe to delayed objects 
 2. calculates the categories using pandas for each delayed object
 3. concatenates the result.

This is much more efficient with the multiprocessing scheduler but a bit slower on the distributed scheduler. This type of fallback strategy generally seems to be a good option if the multiprocessing scheduler is slow. So what kind of optimizations does dask use internally that are well suited for the distributed scheduler but too challenging for the multiprocessing scheduler?


Discussion:

Resolving patch:
diff --git a/dask/dataframe/accessor.py b/dask/dataframe/accessor.py
--- a/dask/dataframe/accessor.py
+++ b/dask/dataframe/accessor.py
@@ -4,8 +4,6 @@
 import pandas as pd
 from toolz import partial
 
-from .core import Series
-
 
 def maybe_wrap_pandas(obj, x):
     if isinstance(x, np.ndarray):
@@ -27,6 +25,7 @@ class Accessor(object):
     * _accessor_name
     """
     def __init__(self, series):
+        from .core import Series
         if not isinstance(series, Series):
             raise ValueError('Accessor cannot be initialized')
         self._validate(series)
diff --git a/dask/dataframe/categorical.py b/dask/dataframe/categorical.py
--- a/dask/dataframe/categorical.py
+++ b/dask/dataframe/categorical.py
@@ -1,9 +1,10 @@
 from __future__ import absolute_import, division, print_function
 
+from collections import defaultdict
 import pandas as pd
+from toolz import partition_all
 
-from ..base import compute
-
+from ..base import tokenize
 from .accessor import Accessor
 from .utils import (has_known_categories, clear_known_categories, is_scalar,
                     is_categorical_dtype)
@@ -31,21 +32,36 @@ def _categorize_block(df, categories, index):
     return df
 
 
-def _get_categories(x):
-    return x.cat.categories if isinstance(x, pd.Series) else x.categories
-
-
-def get_categories(x, object_only=False):
-    """Return a dask object to compute the categoricals for `x`. Returns None
-    if already a known categorical"""
-    if is_categorical_dtype(x):
-        return (None if has_known_categories(x) else
-                x.map_partitions(_get_categories).unique().values)
-    return (x.dropna().drop_duplicates()
-            if not object_only or x.dtype == object else None)
-
-
-def categorize(df, columns=None, index=None, **kwargs):
+def _get_categories(df, columns, index):
+    res = {}
+    for col in columns:
+        x = df[col]
+        if is_categorical_dtype(x):
+            res[col] = pd.Series(x.cat.categories)
+        else:
+            res[col] = x.dropna().drop_duplicates()
+    if index:
+        if is_categorical_dtype(df.index):
+            return res, df.index.categories
+        return res, df.index.dropna().drop_duplicates()
+    return res, None
+
+
+def _get_categories_agg(parts):
+    res = defaultdict(list)
+    res_ind = []
+    for p in parts:
+        for k, v in p[0].items():
+            res[k].append(v)
+        res_ind.append(p[1])
+    res = {k: pd.concat(v, ignore_index=True).drop_duplicates()
+           for k, v in res.items()}
+    if res_ind[0] is None:
+        return res, None
+    return res, res_ind[0].append(res_ind[1:]).drop_duplicates()
+
+
+def categorize(df, columns=None, index=None, split_every=None, **kwargs):
     """Convert columns of the DataFrame to category dtype.
 
     Parameters
@@ -58,30 +74,65 @@ def categorize(df, columns=None, index=None, **kwargs):
         Whether to categorize the index. By default, object indices are
         converted to categorical, and unknown categorical indices are made
         known. Set True to always categorize the index, False to never.
+    split_every : int, optional
+        Group partitions into groups of this size while performing a
+        tree-reduction. If set to False, no tree-reduction will be used.
+        Default is 16.
     kwargs
         Keyword arguments are passed on to compute.
     """
+    meta = df._meta
     if columns is None:
-        columns = list(df.select_dtypes(['object', 'category']).columns)
+        columns = list(meta.select_dtypes(['object', 'category']).columns)
     elif is_scalar(columns):
         columns = [columns]
 
-    categories = [get_categories(df[col]) for col in columns]
-    if index is False:
-        index = None
-    else:
-        index = get_categories(df.index, index is None)
+    # Filter out known categorical columns
+    columns = [c for c in columns if not (is_categorical_dtype(meta[c]) and
+                                          has_known_categories(meta[c]))]
 
-    # Compute the categories
-    values = compute(index, *categories, **kwargs)
-    categories = {c: v for (c, v) in zip(columns, values[1:]) if v is not None}
+    if index is not False:
+        if is_categorical_dtype(meta.index):
+            index = not has_known_categories(meta.index)
+        elif index is None:
+            index = meta.index.dtype == object
 
     # Nothing to do
-    if not len(categories) and index is None:
+    if not len(columns) and index is False:
         return df
 
+    if split_every is None:
+        split_every = 16
+    elif split_every is False:
+        split_every = df.npartitions
+    elif not isinstance(split_every, int) or split_every < 2:
+        raise ValueError("split_every must be an integer >= 2")
+
+    token = tokenize(df, columns, index, split_every)
+    a = 'get-categories-chunk-' + token
+    dsk = {(a, i): (_get_categories, key, columns, index)
+           for (i, key) in enumerate(df._keys())}
+
+    prefix = 'get-categories-agg-' + token
+    k = df.npartitions
+    b = a
+    depth = 0
+    while k > split_every:
+        b = prefix + str(depth)
+        for part_i, inds in enumerate(partition_all(split_every, range(k))):
+            dsk[(b, part_i)] = (_get_categories_agg, [(a, i) for i in inds])
+        k = part_i + 1
+        a = b
+        depth += 1
+
+    dsk[(prefix, 0)] = (_get_categories_agg, [(a, i) for i in range(k)])
+    dsk.update(df.dask)
+
+    # Compute the categories
+    categories, index = df._get(dsk, (prefix, 0), **kwargs)
+
     # Categorize each partition
-    return df.map_partitions(_categorize_block, categories, values[0])
+    return df.map_partitions(_categorize_block, categories, index)
 
 
 class CategoricalAccessor(Accessor):
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from collections import Iterator
+from functools import wraps
 import operator
 from operator import getitem, setitem
 from pprint import pformat
@@ -28,10 +29,12 @@
 from ..base import Base, compute, tokenize, normalize_token
 from ..async import get_sync
 from . import methods
+from .accessor import DatetimeAccessor, StringAccessor
+from .categorical import CategoricalAccessor, categorize
+from .hashing import hash_pandas_object
 from .utils import (meta_nonempty, make_meta, insert_meta_param_description,
                     raise_on_meta_error, clear_known_categories,
                     is_categorical_dtype, has_known_categories)
-from .hashing import hash_pandas_object
 
 no_default = '__no_default__'
 
@@ -1562,17 +1565,14 @@ def dtype(self):
 
     @cache_readonly
     def dt(self):
-        from .accessor import DatetimeAccessor
         return DatetimeAccessor(self)
 
     @cache_readonly
     def cat(self):
-        from .categorical import CategoricalAccessor
         return CategoricalAccessor(self)
 
     @cache_readonly
     def str(self):
-        from .accessor import StringAccessor
         return StringAccessor(self)
 
     def __dir__(self):
@@ -2249,28 +2249,10 @@ def groupby(self, by=None, **kwargs):
         from dask.dataframe.groupby import DataFrameGroupBy
         return DataFrameGroupBy(self, by=by, **kwargs)
 
-    def categorize(self, columns=None, index=None, **kwargs):
-        """Convert columns of the DataFrame to category dtype.
-
-        Parameters
-        ----------
-        columns : list, optional
-            A list of column names to convert to categoricals. By default any
-            column with an object dtype is converted to a categorical, and any
-            unknown categoricals are made known.
-        index : bool, optional
-            Whether to categorize the index. By default, object indices are
-            converted to categorical, and unknown categorical indices are made
-            known. Set True to always categorize the index, False to never.
-        kwargs
-            Keyword arguments are passed on to compute.
-
-        See also
-        --------
-        dask.dataframes.categorical.categorize
-        """
-        from .categorical import categorize
-        return categorize(self, columns=columns, index=index, **kwargs)
+    @wraps(categorize)
+    def categorize(self, columns=None, index=None, split_every=None, **kwargs):
+        return categorize(self, columns=columns, index=index,
+                          split_every=split_every, **kwargs)
 
     @derived_from(pd.DataFrame)
     def assign(self, **kwargs):


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
@@ -130,6 +130,13 @@ def test_categorize():
         assert ddf2.index.cat.known == known_index
         assert_eq(ddf2, df.astype({'v': 'category'}), check_categorical=False)
 
+        # Specifying split_every works
+        ddf2 = ddf.categorize(index=index, split_every=2)
+        assert ddf2.y.cat.known
+        assert ddf2.v.cat.known
+        assert ddf2.index.cat.known == known_index
+        assert_eq(ddf2, df.astype({'v': 'category'}), check_categorical=False)
+
         # Specifying one column doesn't affect others
         ddf2 = ddf.categorize('v', index=index)
         assert not ddf2.y.cat.known
@@ -153,6 +160,13 @@ def test_categorize():
     assert ddf_known_index.categorize(['w']) is ddf_known_index
     assert ddf_known_index.categorize([]) is ddf_known_index
 
+    # Bad split_every fails
+    with pytest.raises(ValueError):
+        ddf.categorize(split_every=1)
+
+    with pytest.raises(ValueError):
+        ddf.categorize(split_every='foo')
+
 
 def test_categorize_index():
     # Object dtype