Repository: dask/dask

Issue:
Add persist function
The distributed scheduler makes heavy use of `persist`, which creates a new dask collection with a collection that points to more evaluated data.  

It might be convenient to add this verb to dask generally.

```python
df = dd.read_csv(...)  # lazy loading
df = dask.persist(df)  # df is dask.dataframe that points to pandas dataframes
```

Persist would be a blocking call (to support general dask schedulers) that in the distributed case would be the same as 

```python
def persist(...)
    collection = global_client.persist(collection)
    wait(collection)
    return collection
```

In the single machine case it would probably be more like the following:

```python
def persist(...):
    keys = collection._keys()
    evaluated = get(collection.dask, keys)
    dsk = zip(keys, evaluated)
    return redict_collection(collection, dsk)

def redict_collection(c, dsk):  # this was taken from distributed/client.py
    if isinstance(c, Delayed):
        return Delayed(c.key, [dsk])
    else:
        cc = copy.copy(c)
        cc.dask = dsk
        return cc
```

Having a general `persist` function would make it easier to build algorithms that were consistent between single and distributed schedulers, such as occurs in dask-glm (cc @moody-marlin).  It would also help users who want to use dask for parallelism needs rather than out-of-core computation (cc @jbednar)

@jcrist @shoyer thoughts?


Discussion:
I like the idea. It may be nice to also support a general mutable-mapping to persist to, to allow for spilling to disk. I could see a few options for how to do this:

- Set persist mapping globally using `set_options`
- Pass a mapping to the persist call `persist(collection, cache=cache)`

The graph returned from `persist` might then look like `{new_key: (getitem, cache, key), ...}` instead of just `{new_key: result_of_old_key}`.

---

Alternatively, persist could be a method on a context instead of a function. Something like:
    
```python
# Needs a better name than `Context`
context = Context(cache=my_spill_to_disk_dictionary_thing)
result = context.persist(collection)
# The compute method would then pull results from the local
# cache before computing
context.compute(other_collection_maybe_based_off_this_one)
```

One perk of this instead of a global/parameter method is that it mirrors the distributed client context, and would allow for a mimicking of the incremental computation api distributed allows. The global compute/persist could then just be

```python
global_context = Context(cache={})
compute = global_context.compute
persist = global_context.persist
```

This isn't as simple as a global thing though, and I'm unsure that it's needed. Just a thought.
I guess this would help to organize the current set of globals that we keep around in `set_options`.  

@jcrist do you have any desire/time to experiment with implementing this?
Yes, I think this is a good idea. With xarray, I found myself writing things like `.compute().chunk(...)` with some frequency, both for parallelism and to solve scheduler fail cases like https://github.com/dask/dask/issues/874. Persisting to a dictionary with automatic spilling to disk could make this very convenient.
Yeah, I'd be interested in working on this. I'm trying to finish up the categorical work today, then should have some time to work on this.

Resolving patch:
diff --git a/dask/__init__.py b/dask/__init__.py
--- a/dask/__init__.py
+++ b/dask/__init__.py
@@ -8,7 +8,7 @@
 except ImportError:
     pass
 try:
-    from .base import visualize, compute
+    from .base import visualize, compute, persist
 except ImportError:
     pass
 
diff --git a/dask/base.py b/dask/base.py
--- 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 @@
 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 @@ def visualize(self, filename='mydask', format=None, optimize_graph=False,
         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 @@ def compute(self, **kwargs):
         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 @@ def _get_binary_operator(cls, op, inv=False):
         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)
@@ -383,3 +369,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/tests/test_base.py b/dask/tests/test_base.py
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -7,8 +7,10 @@
 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
@@ -259,6 +261,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]})
@@ -374,3 +389,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
--- 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