Function optimize

dask/base.py:539–598

Optimize several dask collections at once. Returns equivalent dask collections that all share the same merged and optimized underlying graph. This can be useful if converting multiple collections to delayed objects, or to manually apply the optimizations at strategic points. No

(*args, traverse=True, **kwargs)

def optimize(*args, traverse=True, **kwargs):
    """Optimize several dask collections at once.

    Returns equivalent dask collections that all share the same merged and
    optimized underlying graph. This can be useful if converting multiple
    collections to delayed objects, or to manually apply the optimizations at
    strategic points.

    Note that in most cases you shouldn't need to call this function directly.

    Warning::

        This function triggers a materialization of the collections and looses
        any annotations attached to HLG layers.

    Parameters
    ----------
    *args : objects
        Any number of objects. If a dask object, its graph is optimized and
        merged with all those of all other dask objects before returning an
        equivalent dask collection. Non-dask arguments are passed through
        unchanged.
    traverse : bool, optional
        By default dask traverses builtin python collections looking for dask
        objects passed to ``optimize``. For large collections this can be
        expensive. If none of the arguments contain any dask objects, set
        ``traverse=False`` to avoid doing this traversal.
    optimizations : list of callables, optional
        Additional optimization passes to perform.
    **kwargs
        Extra keyword arguments to forward to the optimization passes.

    Examples
    --------
    >>> import dask
    >>> import dask.array as da
    >>> a = da.arange(10, chunks=2).sum()
    >>> b = da.arange(10, chunks=2).mean()
    >>> a2, b2 = dask.optimize(a, b)

    >>> a2.compute() == a.compute()
    np.True_
    >>> b2.compute() == b.compute()
    np.True_
    """
    # TODO: This API is problematic. The approach to using postpersist forces us
    # to materialize the graph. Most low level optimizations will materialize as
    # well
    collections, repack = unpack_collections(*args, traverse=traverse)
    if not collections:
        return args

    dsk = collections_to_expr(collections)

    postpersists = []
    for a in collections:
        r, s = a.__dask_postpersist__()
        postpersists.append(r(dsk.__dask_graph__(), *s))