Function unpack_collections

dask/delayed.py:115–293

Normalize a python object and merge all sub-graphs. - Replace ``Delayed`` with their keys - Convert literals to things the schedulers can handle - Extract dask graphs from all enclosed values. Note, that the returned _task_ is not necessarily runnable and the caller is responsi

(expr, _return_collections=True)

def unpack_collections(expr, _return_collections=True):
    """Normalize a python object and merge all sub-graphs.

    - Replace ``Delayed`` with their keys
    - Convert literals to things the schedulers can handle
    - Extract dask graphs from all enclosed values.

    Note, that the returned _task_ is not necessarily runnable and the caller is
    responsible to deal with the output types accordingly.

    The task is one of
    - `TaskRef` as a pointer to the collection returned in collections. This is
      not callable and should not be a top-level member of a dask task graph.
    - A runnable task (i.e. subclass `GraphNode`) which can be embedded
      directly into a task graph. This indicates that a dask collection was
      encountered on a deeper nesting level and this runnable task restores the
      input nesting with the computed dask collection replaced.
    - The unaltered object as provided if no dask collections are found.

    Parameters
    ----------
    expr : object
        The object to be normalized. This function knows how to handle
        dask collections, as well as most builtin python types.

    _optimize_collections: bool, optional
        Internal use only!


    Returns
    -------
    task : object
    collections : a tuple of collections

    Examples
    --------
    >>> import dask
    >>> a = delayed(1, 'a')
    >>> b = delayed(2, 'b')
    >>> task, collections = unpack_collections([a, b, 3])
    >>> task
    List((TaskRef('a'), TaskRef('b'), 3))
    >>> collections
    (Delayed('a'), Delayed('b'))

    >>> task, collections = unpack_collections({a: 1, b: 2})
    >>> task
    Dict(a: 1, b: 2)
    >>> collections
    (Delayed('a'), Delayed('b'))
    """
    if isinstance(expr, Delayed):
        if _return_collections:
            return TaskRef(expr._key), (expr,)
        else:
            expr = collections_to_expr(expr).finalize_compute()
            (name,) = expr.__dask_keys__()
            return TaskRef(name), (expr,)