Repository: dask/dask

Issue:
Use repartition to increase the number of partitions
This works fine when the user supplies divisions but doesn't work otherwise.


Discussion:

Patch:
diff --git a/dask/__init__.py b/dask/__init__.py
index f5f266a68..5da166f83 100644
--- a/dask/__init__.py
+++ b/dask/__init__.py
@@ -8,7 +8,7 @@ try:
 except ImportError:
     pass
 try:
-    from .base import visualize, compute
+    from .base import visualize, compute, persist
 except ImportError:
     pass
 
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 5daa80998..d21ae1624 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -4,6 +4,7 @@ from collections import Iterable, Iterator, defaultdict
 from functools import wraps, partial
 import itertools
 import math
+from operator import getitem
 import os
 import types
 import uuid
@@ -1117,23 +1118,40 @@ class Bag(Base):
         --------
         >>> b.repartition(5)  # set to have 5 partitions  # doctest: +SKIP
         """
-        if npartitions > self.npartitions:
-            msg = ("Repartition only supports going to fewer partitions\n"
-                   " old: %d  new: %d")
-            raise NotImplementedError(msg % (self.npartitions, npartitions))
-        npartitions_ratio = self.npartitions / npartitions
-        new_partitions_boundaries = [int(old_partition_index * npartitions_ratio)
-                                     for old_partition_index in range(npartitions + 1)]
-        new_name = 'repartition-%d-%s' % (npartitions, tokenize(self))
-
-        dsk = {}
-        for new_partition_index in range(npartitions):
-            value = (list, (toolz.concat,
-                            [(self.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
+        new_name = 'repartition-%d-%s' % (npartitions, tokenize(self, npartitions))
+        if npartitions == self.npartitions:
+            return self
+        elif npartitions < self.npartitions:
+            ratio = self.npartitions / npartitions
+            new_partitions_boundaries = [int(old_partition_index * ratio)
+                                         for old_partition_index in range(npartitions + 1)]
+
+            dsk = {}
+            for new_partition_index in range(npartitions):
+                value = (list, (toolz.concat,
+                                [(self.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
+        else:  # npartitions > self.npartitions
+            ratio = npartitions / self.npartitions
+            split_name = 'split-%s' % tokenize(self, npartitions)
+            dsk = {}
+            last = 0
+            j = 0
+            for i in range(self.npartitions):
+                new = last + ratio
+                if i == self.npartitions - 1:
+                    k = npartitions - j
+                else:
+                    k = int(new - last)
+                dsk[(split_name, i)] = (split, (self.name, i), k)
+                for jj in range(k):
+                    dsk[(new_name, j)] = (getitem, (split_name, i), jj)
+                    j += 1
+                last = new
+
         return Bag(dsk=merge(self.dask, dsk), name=new_name, npartitions=npartitions)
 
     def accumulate(self, binop, initial=no_default):
@@ -1677,3 +1695,18 @@ def random_state_data_python(n, random_state=None):
     maxuint32 = 1 << 32
     return [tuple(random_state.randint(0, maxuint32) for i in range(624))
             for i in range(n)]
+
+
+def split(seq, n):
+    """ Split apart a sequence into n equal pieces
+
+    >>> split(range(10), 3)
+    [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]]
+    """
+    if not isinstance(seq, (list, tuple)):
+        seq = list(seq)
+
+    part = len(seq) / n
+    L = [seq[int(part * i): int(part * (i + 1))] for i in range(n - 1)]
+    L.append(seq[int(part * (n - 1)):])
+    return L
diff --git a/dask/base.py b/dask/base.py
index 8729ad3ea..dc187a38c 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1,9 +1,9 @@
 from __future__ import absolute_import, division, print_function
 
 from collections import OrderedDict
+import copy
 from functools import partial
 from hashlib import md5
-from operator import attrgetter
 import pickle
 import os
 import uuid
@@ -11,8 +11,9 @@ import uuid
 from toolz import merge, groupby, curry, identity
 from toolz.functoolz import Compose
 
-from .compatibility import bind_method, unicode
+from .compatibility import bind_method, unicode, PY3
 from .context import _globals
+from .core import flatten
 from .utils import Dispatch
 
 __all__ = ("Base", "compute", "normalize_token", "tokenize", "visualize")
@@ -61,8 +62,21 @@ class Base(object):
         return visualize(self, filename=filename, format=format,
                          optimize_graph=optimize_graph, **kwargs)
 
+    def persist(self, **kwargs):
+        """ Persist this dask collection into memory
+
+        See ``dask.base.persist`` for full docstring
+        """
+        (result,) = persist(self, **kwargs)
+        return result
+
     def compute(self, **kwargs):
-        """Compute several dask collections at once.
+        """ Compute this dask collection
+
+        This turns a lazy Dask collection into its in-memory equivalent.
+        For example a Dask.array turns into a NumPy array and a Dask.dataframe
+        turns into a Pandas dataframe.  The entire dataset must fit into memory
+        before calling this operation.
 
         Parameters
         ----------
@@ -76,7 +90,8 @@ class Base(object):
         kwargs
             Extra keywords to forward to the scheduler ``get`` function.
         """
-        return compute(self, **kwargs)[0]
+        (result,) = compute(self, **kwargs)
+        return result
 
     @classmethod
     def _get(cls, dsk, keys, get=None, **kwargs):
@@ -119,24 +134,6 @@ class Base(object):
         raise NotImplementedError
 
 
-def _extract_graph_and_keys(vals):
-    """Given a list of dask vals, return a single graph and a list of keys such
-    that ``get(dsk, keys)`` is equivalent to ``[v.compute() v in vals]``."""
-    dsk = {}
-    keys = []
-    for v in vals:
-        # Optimization to avoid merging dictionaries in Delayed values. Reduces
-        # memory usage for large graphs.
-        if hasattr(v, '_dasks'):
-            for d in v._dasks:
-                dsk.update(d)
-        else:
-            dsk.update(v.dask)
-        keys.append(v._keys())
-
-    return dsk, keys
-
-
 def compute(*args, **kwargs):
     """Compute several dask collections at once.
 
@@ -165,13 +162,12 @@ def compute(*args, **kwargs):
     >>> compute(a, b)
     (45, 4.5)
     """
+    optimize_graph = kwargs.pop('optimize_graph', True)
     variables = [a for a in args if isinstance(a, Base)]
     if not variables:
         return args
 
     get = kwargs.pop('get', None) or _globals['get']
-    optimizations = (kwargs.pop('optimizations', None) or
-                     _globals.get('optimizations', []))
 
     if not get:
         get = variables[0]._default_get
@@ -181,18 +177,8 @@ def compute(*args, **kwargs):
                              "scheduler `get` function using either "
                              "the `get` kwarg or globally with `set_options`.")
 
-    if kwargs.get('optimize_graph', True):
-        groups = groupby(attrgetter('_optimize'), variables)
-        groups = {opt: _extract_graph_and_keys(val)
-                  for opt, val in groups.items()}
-        for opt in optimizations:
-            groups = {k: [opt(dsk, keys), keys]
-                      for k, (dsk, keys) in groups.items()}
-        dsk = merge([opt(dsk, keys, **kwargs)
-                    for opt, (dsk, keys) in groups.items()])
-        keys = [var._keys() for var in variables]
-    else:
-        dsk, keys = _extract_graph_and_keys(variables)
+    dsk = collections_to_dsk(variables, optimize_graph, **kwargs)
+    keys = [var._keys() for var in variables]
     results = get(dsk, keys, **kwargs)
 
     results_iter = iter(results)
@@ -388,3 +374,162 @@ def tokenize(*args, **kwargs):
     if kwargs:
         args = args + (kwargs,)
     return md5(str(tuple(map(normalize_token, args))).encode()).hexdigest()
+
+
+def collections_to_dsk(collections, optimize_graph=True, **kwargs):
+    """
+    Convert many collections into a single dask graph, after optimization
+    """
+    optimizations = (kwargs.pop('optimizations', None) or
+                     _globals.get('optimizations', []))
+    if optimize_graph:
+        groups = groupby(lambda x: x._optimize, collections)
+        groups = {opt: _extract_graph_and_keys(val)
+                  for opt, val in groups.items()}
+        for opt in optimizations:
+            groups = {k: [opt(dsk, keys), keys]
+                      for k, (dsk, keys) in groups.items()}
+        dsk = merge([opt(dsk, keys, **kwargs)
+                     for opt, (dsk, keys) in groups.items()])
+    else:
+        dsk = merge(c.dask for c in collections)
+
+    return dsk
+
+
+def _extract_graph_and_keys(vals):
+    """Given a list of dask vals, return a single graph and a list of keys such
+    that ``get(dsk, keys)`` is equivalent to ``[v.compute() v in vals]``."""
+    dsk = {}
+    keys = []
+    for v in vals:
+        # Optimization to avoid merging dictionaries in Delayed values. Reduces
+        # memory usage for large graphs.
+        if hasattr(v, '_dasks'):
+            for d in v._dasks:
+                dsk.update(d)
+        else:
+            dsk.update(v.dask)
+        keys.append(v._keys())
+
+    return dsk, keys
+
+
+def redict_collection(c, dsk):
+    from dask.delayed import Delayed
+    if isinstance(c, Delayed):
+        return Delayed(c.key, [dsk])
+    else:
+        cc = copy.copy(c)
+        cc.dask = dsk
+        return cc
+
+
+def persist(*args, **kwargs):
+    """ Persist multiple Dask collections into memory
+
+    This turns lazy Dask collections into Dask collections with the same
+    metadata, but now with their results fully computed or actively computing
+    in the background.
+
+    For example a lazy dask.array built up from many lazy calls will now be a
+    dask.array of the same shape, dtype, chunks, etc., but now with all of
+    those previously lazy tasks either computed in memory as many small NumPy
+    arrays (in the single-machine case) or asynchronously running in the
+    background on a cluster (in the distributed case).
+
+    This function operates differently if a ``dask.distributed.Client`` exists
+    and is connected to a distributed scheduler.  In this case this function
+    will return as soon as the task graph has been submitted to the cluster,
+    but before the computations have completed.  Computations will continue
+    asynchronously in the background.  When using this function with the single
+    machine scheduler it blocks until the computations have finished.
+
+    When using Dask on a single machine you should ensure that the dataset fits
+    entirely within memory.
+
+    Examples
+    --------
+    >>> df = dd.read_csv('/path/to/*.csv')  # doctest: +SKIP
+    >>> df = df[df.name == 'Alice']  # doctest: +SKIP
+    >>> df['in-debt'] = df.balance < 0  # doctest: +SKIP
+    >>> df = df.persist()  # triggers computation  # doctest: +SKIP
+
+    >>> df.value().min()  # future computations are now fast  # doctest: +SKIP
+    -10
+    >>> df.value().max()  # doctest: +SKIP
+    100
+
+    >>> from dask import persist  # use persist function on multiple collections
+    >>> a, b = persist(a, b)  # doctest: +SKIP
+
+    Parameters
+    ----------
+    *args: Dask collections
+    get : callable, optional
+        A scheduler ``get`` function to use. If not provided, the default
+        is to check the global settings first, and then fall back to
+        the collection defaults.
+    optimize_graph : bool, optional
+        If True [default], the graph is optimized before computation.
+        Otherwise the graph is run as is. This can be useful for debugging.
+    **kwargs
+        Extra keywords to forward to the scheduler ``get`` function.
+
+    Returns
+    -------
+    New dask collections backed by in-memory data
+    """
+    collections = [a for a in args if isinstance(a, Base)]
+    if not collections:
+        return args
+
+    try:
+        from distributed.client import default_client
+    except ImportError:
+        pass
+    else:
+        try:
+            client = default_client()
+        except ValueError:
+            pass
+        else:
+            collections = client.persist(collections, **kwargs)
+            if isinstance(collections, list):  # distributed is inconsistent here
+                collections = tuple(collections)
+            else:
+                collections = (collections,)
+            results_iter = iter(collections)
+            return tuple(a if not isinstance(a, Base)
+                         else next(results_iter)
+                         for a in args)
+
+    optimize_graph = kwargs.pop('optimize_graph', True)
+
+    get = kwargs.pop('get', None) or _globals['get']
+
+    if not get:
+        get = collections[0]._default_get
+        if not all(a._default_get == get for a in collections):
+            raise ValueError("Compute called on multiple collections with "
+                             "differing default schedulers. Please specify a "
+                             "scheduler `get` function using either "
+                             "the `get` kwarg or globally with `set_options`.")
+
+    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
+    keys = list(flatten([var._keys() for var in collections]))
+    results = get(dsk, keys, **kwargs)
+
+    d = dict(zip(keys, results))
+
+    result = [redict_collection(c, {k: d[k]
+                                    for k in flatten(c._keys())})
+              for c in collections]
+    results_iter = iter(result)
+    return tuple(a if not isinstance(a, Base)
+                 else next(results_iter)
+                 for a in args)
+
+
+if PY3:
+    Base.persist.__doc__ = persist.__doc__


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 27e474e2e..0d114cf07 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -915,19 +915,29 @@ def test_zip(npartitions, hi=1000):
     assert list(pairs) == list(zip(range(0, hi, 2), range(1, hi, 2)))
 
 
-def test_repartition():
-    for x, y in [(10, 5), (7, 3), (5, 1), (5, 4)]:
-        b = db.from_sequence(range(20), npartitions=x)
-        c = b.repartition(y)
+@pytest.mark.parametrize('nin', [1, 2, 7, 11, 23])
+@pytest.mark.parametrize('nout', [1, 2, 5, 12, 23])
+def test_repartition(nin, nout):
+    b = db.from_sequence(range(100), npartitions=nin)
+    c = b.repartition(npartitions=nout)
 
-        assert b.npartitions == x
-        assert c.npartitions == y
-        assert list(b) == c.compute(get=dask.get)
+    assert c.npartitions == nout
+    assert b.compute(get=dask.get) == c.compute(get=dask.get)
+    results = dask.get(c.dask, c._keys())
+    assert all(results)
 
-    try:
-        b.repartition(100)
-    except NotImplementedError as e:
-        assert '100' in str(e)
+
+def test_repartition_names():
+    b = db.from_sequence(range(100), npartitions=5)
+    c = b.repartition(2)
+    assert b.name != c.name
+
+    d = b.repartition(20)
+    assert b.name != c.name
+    assert c.name != d.name
+
+    c = b.repartition(5)
+    assert b is c
 
 
 @pytest.mark.skipif('not db.core._implement_accumulate')
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 6aa7d54f9..1d7c73103 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -7,8 +7,10 @@ import subprocess
 import sys
 
 import dask
+from dask import delayed
 from dask.base import (compute, tokenize, normalize_token, normalize_function,
-                       visualize)
+                       visualize, persist)
+from dask.delayed import Delayed
 from dask.utils import tmpdir, tmpfile, ignoring
 from dask.utils_test import inc, dec
 from dask.compatibility import unicode
@@ -263,6 +265,19 @@ def test_compute_array():
     assert np.allclose(out2, arr + 2)
 
 
+@pytest.mark.skipif('not da')
+def test_persist_array():
+    from dask.array.utils import assert_eq
+    arr = np.arange(100).reshape((10, 10))
+    x = da.from_array(arr, chunks=(5, 5))
+    x = (x + 1) - x.mean(axis=0)
+    y = x.persist()
+
+    assert_eq(x, y)
+    assert set(y.dask).issubset(x.dask)
+    assert len(y.dask) == y.npartitions
+
+
 @pytest.mark.skipif('not dd')
 def test_compute_dataframe():
     df = pd.DataFrame({'a': [1, 2, 3, 4], 'b': [5, 5, 3, 3]})
@@ -378,3 +393,42 @@ def test_default_imports():
                  'partd', 's3fs', 'distributed']
     for mod in blacklist:
         assert mod not in modules
+
+
+def test_persist_literals():
+    assert persist(1, 2, 3) == (1, 2, 3)
+
+
+def test_persist_delayed():
+    x1 = delayed(1)
+    x2 = delayed(inc)(x1)
+    x3 = delayed(inc)(x2)
+
+    xx, = persist(x3)
+    assert isinstance(xx, Delayed)
+    assert xx.key == x3.key
+    assert len(xx.dask) == 1
+
+    assert x3.compute() == xx.compute()
+
+
+@pytest.mark.skipif('not da or not db')
+def test_persist_array_bag():
+    x = da.arange(5, chunks=2) + 1
+    b = db.from_sequence([1, 2, 3]).map(inc)
+
+    with pytest.raises(ValueError):
+        persist(x, b)
+
+    xx, bb = persist(x, b, get=dask.async.get_sync)
+
+    assert isinstance(xx, da.Array)
+    assert isinstance(bb, db.Bag)
+
+    assert xx.name == x.name
+    assert bb.name == b.name
+    assert len(xx.dask) == xx.npartitions < len(x.dask)
+    assert len(bb.dask) == bb.npartitions < len(b.dask)
+
+    assert np.allclose(x, xx)
+    assert list(b) == list(bb)
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 45b00a0bb..5cd3fcc51 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -1,6 +1,25 @@
 import pytest
 pytest.importorskip('distributed')
 
+from dask import persist, delayed
+from distributed.client import _wait
+from distributed.utils_test import gen_cluster, inc
+
 
 def test_can_import_client():
     from dask.distributed import Client # noqa: F401
+
+
+@gen_cluster(client=True)
+def test_persist(c, s, a, b):
+    x = delayed(inc)(1)
+    x2, = persist(x)
+
+    yield _wait(x2)
+    assert x2.key in a.data or x2.key in b.data
+
+    y = delayed(inc)(10)
+    y2, one = persist(y, 1)
+
+    yield _wait(y2)
+    assert y2.key in a.data or y2.key in b.data