.. autoclass:: dask.typing.DaskCollection
   :members: __dask_graph__, __dask_keys__, __dask_postcompute__,
             __dask_postpersist__, __dask_tokenize__,
             __dask_optimize__, __dask_scheduler__, compute, persist,
             visualize

HLG Collection Protocol

The ``SchedulerGetProtocol`` defines the signature that a Dask
collection's ``__dask_scheduler__`` definition must adhere to.

.. autoclass:: dask.typing.SchedulerGetCallable
   :members: __call__

Post-persist Callable Protocol

The operation of ``compute`` can be broken into three stages:

1. **Graph Merging, finalization**

   First, the individual collections are converted to a single large expression
   and nested list of keys. This is done by
   :func:`~dask.base.collections_to_expr` and ensures that all collections are
   optimized together to eliminate common sub-expressions.

   .. note::

       At this stage, legacy HLG graphs are wrapped into a ``HLGExpr`` that
       encodes __dask_postcompute__ and the low level optimizer as determined by `__dask_optimize__` into the expression.

       The ``optimize_graph`` argument is only relevant for HLG graphs and
       controls whether low level optimizations are considered.

       - If ``optimize_graph`` is ``True`` (default), then the collections are
         first grouped by their ``__dask_optimize__`` methods.  All collections with the same ``__dask_optimize__`` method have their graphs merged and keys concatenated, and then a single call to each respective ``__dask_optimize__`` is made with the merged graphs and keys.  The resulting graphs are then merged.

       - If ``optimize_graph`` is ``False``, then all the graphs are merged and
         all the keys concatenated.

   The combined graph is _finalized_ with a ``FinalizeCompute`` expression
   which instructs the expression / graph to reduce to a single partition,
   suitable to be returned to the user after compute. This is either done by
   implemengint the ``__dask_postcompute__`` method of the collection or by
   implementing an optimization path of the expression.

   For the example of a DataFrame, the ``FinalizeCompute`` simplifies to a ``RepartitionToFewer(..., npartition=1)`` which simply concatenates all results to one ordinary DataFrame.

2. **(Expression) Optimization**

   The merged expression is optimized. This step should not be confused with
   the low level optimization that is defined by `__dask_optimize__` for legacy
   graphs. This is a step that is always performed and is a required step to
   simplify and lower expressions to their final form that can be used to
   actually generate the executable task graph. See also, :doc:`/dataframe-optimizer`.

   For legacy HLG graphs, the low level optimization step is embedded in the
   graph materialization which typically only happens after the graph has been
   passed to the scheduler (see below).


3. **Computation**

   After the graphs are merged and any optimizations performed, the resulting
   large graph and nested list of keys are passed on to the scheduler.  The
   scheduler to use is chosen as follows:

   - If a ``get`` function is specified directly as a keyword, use that
   - Otherwise, if a global scheduler is set, use that
   - Otherwise fall back to the default scheduler for the given collections.
     Note that if all collections don't share the same ``__dask_scheduler__``
     then an error will be raised.

   Once the appropriate scheduler ``get`` function is determined, it is called
   with the merged graph, keys, and extra keyword arguments.  After this stage,
   ``results`` is a nested list of values. The structure of this list mirrors
   that of ``keys``, with each key substituted with its corresponding result.


Persist

The operation of ``optimize`` can be broken into two stages:

1. **Graph Merging, *no* finalization**

   Same as in ``persist``.

2. **(Expression) Optimization**

   Same as in ``compute`` and ``persist``.

3. **Materialization and Rebuilding**

   The entire graph is materialized (which also performs low level optimization).
   Similar to ``persist``, the ``rebuild`` function and arguments from
   ``__dask_postpersist__`` are used to reconstruct equivalent collections from
   the optimized graph.


Visualize

import dask
from dask.base import DaskMethodsMixin, replace_name_in_key
from dask.expr import Expr, LLGExpr
from dask.typing import Key
from dask.task_spec import Task, DataNode, Alias


# We subclass from DaskMethodsMixin to add common dask methods to
# our class (compute, persist, and visualize). This is nice but not
# necessary for creating a Dask collection (you can define them
# yourself).
class Tuple(DaskMethodsMixin):
    def __init__(self, expr):
        self._expr = expr

    def __dask_graph__(self):
        return self._expr.__dask_graph__()

    def __dask_keys__(self):
        return self._expr.__dask_keys__()

    # Use the threaded scheduler by default.
    __dask_scheduler__ = staticmethod(dask.threaded.get)

    def __dask_postcompute__(self):
        # We want to return the results as a tuple, so our finalize
        # function is `tuple`. There are no extra arguments, so we also
        # return an empty tuple.
        return tuple, ()

    def __dask_postpersist__(self):
        return Tuple._rebuild, ("mysuffix",)

    @staticmethod
    def _rebuild(futures: dict, name: str):
        expr = LLGExpr({
            (name, i): DataNode((name, i), val)
            for i, val in  enumerate(futures.values())
        })
        return Tuple(expr)

    def __dask_tokenize__(self):
        # For tokenize to work we want to return a value that fully
        # represents this object. In this case this is done by a type
        identifier plus the (also tokenized) name of the expression
        return (type(self), self._expr._name)

class RemoteTuple(Expr):
    @property
    def npartitions(self):
        return len(self.operands)

    def __dask_keys__(self):
        return [(self._name, i) for i in range(self.npartitions)]

    def _task(self, name: Key, index: int) -> Task:
        return DataNode(name, self.operands[index])

>>> from dask_tuple import Tuple

def from_pytuple(pytup: tuple) -> Tuple:
    return Tuple(RemoteTuple(*pytup))

>>> dask_tup = from_pytuple(tuple(range(5)))

>>> dask_tup.__dask_keys__()
[('remotetuple-b7ea9a26c3ab8287c78d11fd45f26793', 0),
('remotetuple-b7ea9a26c3ab8287c78d11fd45f26793', 1),
('remotetuple-b7ea9a26c3ab8287c78d11fd45f26793', 2)]

# Compute turns Tuple into a tuple
>>> dask_tup.compute()
(0, 1, 2)

# Persist turns Tuple into a Tuple, with each task already computed
>>> dask_tup2 = dask_tup.persist()
>>> isinstance(dask_tup2, Tuple)
True

>>> dask_tup2.__dask_graph__()
{('newname', 0): DataNode(0),
('newname', 1): DataNode(1),
('newname', 2): DataNode(2)}

>>> x2.compute()
(0, 1, 2)

# Run-time typechecking
>>> from dask.typing import DaskCollection
>>> isinstance(x, DaskCollection)
True

>>> from dask.base import is_dask_collection
>>> from dask import delayed

>>> x = delayed(sum)([1, 2, 3])
>>> is_dask_collection(x)
True
>>> is_dask_collection(1)
False

.. code:: python

    >>> from dask.base import tokenize, normalize_token

    # Define a tokenize implementation using a method.
    >>> class Point:
    ...     def __init__(self, x, y):
    ...         self.x = x
    ...         self.y = y
    ...
    ...     def __dask_tokenize__(self):
    ...         # This tuple fully represents self
    ...         # Wrap non-trivial objects with normalize_token before returning them
    ...         return normalize_token(Point), self.x, self.y

    >>> x = Point(1, 2)
    >>> tokenize(x)
    '5988362b6e07087db2bc8e7c1c8cc560'
    >>> tokenize(x) == tokenize(x)  # token is idempotent
    True
    >>> tokenize(Point(1, 2)) == tokenize(Point(1, 2))  # token is deterministic
    True
    >>> tokenize(Point(1, 2)) == tokenize(Point(2, 1))  # tokens are unique
    False


    # Register an implementation with normalize_token
    >>> class Point3D:
    ...     def __init__(self, x, y, z):
    ...         self.x = x
    ...         self.y = y
    ...         self.z = z

    >>> @normalize_token.register(Point3D)
    ... def normalize_point3d(x):
    ...     return normalize_token(Point3D), x.x, x.y, x.z

    >>> y = Point3D(1, 2, 3)
    >>> tokenize(y)
    '5a7e9c3645aa44cf13d021c14452152e'


For more examples, see ``dask/base.py`` or any of the built-in Dask collections.