Repository: dask/dask (python)

Issue:
Generically rebuild a collection with different keys
I'm coding a generic library function that takes in input an arbitrary dask collection and produces in output a variant of that same collection, which computes to the same values but with a different dask graph *and cohexists on the original collection*.

An example is https://github.com/dask/dask/pull/6804.

I can't find how to implement it in a way that supports third-party collections - not without making assumptions beyond those stipulated by https://docs.dask.org/en/latest/spec.html and https://docs.dask.org/en/latest/custom-collections.html.

I am therefore proposing to tighten the definitions in both documents to align them to the de-facto implementation of all builtin dask collections. This is also a good thing because I expect the dask codebase to already blindly assume the below in several places simply because it's true for Array, Bag, Dataframe, Series and Delayed:

#### 1
In https://docs.dask.org/en/latest/spec.html#definitions, replace

> A key is any hashable value that is not a task:

with

> A key is either a string or a tuple of 1+ elements, whose first element is a string and the remainder are arbitrary hashables:

#### 2
 In https://docs.dask.org/en/latest/custom-collections.html#__dask_keys__, add the requirement that, if more than one key is returned, all keys must be tuples with an identical first element (commonly referred to as the *name* of the collection)

My problem here is that I don't know anything about how custom collections look like in other projects. However, I also strongly suspect that they already adhere to the de facto standard above, simply because otherwise the dask codebase would fall over in multiple points.

#### 3
In https://docs.dask.org/en/latest/custom-collections.html#__dask_postpersist__, add that the returned callable must accept an optional keyword parameter, ``name``, which will be used to replace the collection's name.


## Example
Here we're introducing a 1 second delay in the compute() call of a dask bag:

```python
>>> import time
>>> import dask.bag
>>> from dask import delayed
>>> from dask.base import tokenize
>>> from dask.highlevelgraph import HighLevelGraph

>>> a = dask.bag.from_sequence([1,2,3,4], npartitions=2)
>>> dict(a.__dask_graph__())
{('from_sequence-00e91bb38e7830800935abba6e2614f8', 0): [1, 2],
 ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1): [3, 4]}
>>> s = delayed(time.sleep)(1)
>>> name = "with_sleep-" + tokenize(a, s)
>>> def _with_sleep(node, deps):
...     return node

# Simplified - actual implementation is a bit more involved to make it
# work for arbitrary collections, e.g. 2+ dimensional arrays
>>> layer = {(name, i): (_with_sleep, k, s.key) for i, k in enumerate(a.__dask_keys__())}

>>> dsk = HighLevelGraph.from_collections(name, layer, dependencies=(a, s))
>>> rebuild, args = a.__dask_postpersist__()
>>> b = rebuild(dsk, *args, name=name)
>>> dict(b.__dask_graph__())
{('from_sequence-00e91bb38e7830800935abba6e2614f8', 0): [1, 2],
 ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1): [3, 4],
 'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92': (<function time.sleep>, 1),
 ('with_sleep-0e771ca8ce0b7fd08e0b0dc7c4dde9d7', 0): (
     <function _with_sleep>, 
     ('from_sequence-00e91bb38e7830800935abba6e2614f8', 0), 
     'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92'
 ),
 ('with_sleep-0e771ca8ce0b7fd08e0b0dc7c4dde9d7', 1): (
     <function _with_sleep>,
     ('from_sequence-00e91bb38e7830800935abba6e2614f8', 1),
     'sleep-1485da81-3fc0-4a0d-b099-7dc21ccc7d92'
 ),
}
```
---
Additionally, it doesn't seem that https://docs.dask.org/en/latest/custom-collections.html has been updated since the introduction of HighLevelGraph.

Specifically:
- it states that the output of ``__dask_graph__`` is a MutableMapping; however, HighLevelGraph is a Mapping.
- it makes no mention of the ``__dask_layers__`` method. I think it would be sensible to stipulate that, if ``__dask_layers__`` is defined, then ``__dask_graph__`` must return a HighLevelGraph?

Pull request:
Generically rebuild a collection with different keys
Closes #7093

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 3c691382f..de5acd2cd 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1177,7 +1177,10 @@ class Array(DaskMethodsMixin):
         return finalize, ()
 
     def __dask_postpersist__(self):
-        return Array, (self.name, self.chunks, self.dtype, self._meta)
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        return Array(dsk, name or self.name, self.chunks, self.dtype, self._meta)
 
     def _reset_cache(self, key=None):
         """
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 6cebbbf3b..84ddfc991 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -37,10 +37,16 @@ from tlz import (
 
 from .. import config
 from .avro import to_avro
-from ..base import tokenize, dont_optimize, DaskMethodsMixin
+from ..base import tokenize, dont_optimize, replace_name_in_key, DaskMethodsMixin
 from ..bytes import open_files
 from ..context import globalmethod
-from ..core import quote, istask, get_dependencies, reverse_dict, flatten
+from ..core import (
+    quote,
+    istask,
+    get_dependencies,
+    reverse_dict,
+    flatten,
+)
 from ..sizeof import sizeof
 from ..delayed import Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
@@ -362,7 +368,11 @@ class Item(DaskMethodsMixin):
         return finalize_item, ()
 
     def __dask_postpersist__(self):
-        return Item, (self.key,)
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        key = replace_name_in_key(self.key, name) if name else self.key
+        return Item(dsk, key)
 
     @staticmethod
     def from_delayed(value):
@@ -468,7 +478,10 @@ class Bag(DaskMethodsMixin):
         return finalize, ()
 
     def __dask_postpersist__(self):
-        return type(self), (self.name, self.npartitions)
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        return type(self)(dsk, name or self.name, self.npartitions)
 
     def __str__(self):
         return "dask.bag<%s, npartitions=%d>" % (key_split(self.name), self.npartitions)
diff --git a/dask/base.py b/dask/base.py
index 1e148f23a..b39a0f563 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -32,6 +32,8 @@ __all__ = (
     "visualize",
     "tokenize",
     "normalize_token",
+    "get_collection_name",
+    "replace_name_in_key",
 )
 
 
@@ -1186,3 +1188,54 @@ def wait(x, timeout=None, return_when="ALL_COMPLETED"):
         return wait(x, timeout=timeout, return_when=return_when)
     except (ImportError, ValueError):
         return x
+
+
+def get_collection_name(collection) -> str:
+    """Infer the collection name from the dask keys, under the assumption that all keys
+    are either tuples with matching first element, and that element is a string, or
+    there is exactly one key and it is a string.
+
+    Examples
+    --------
+    >>> a.__dask_keys__()
+    ["foo-123"]
+    >>> get_collection_name(a)
+    "foo"
+    >>> b.__dask_keys__()
+    [[("foo-123", 0, 0), ("foo-123", 0, 1)], [("foo-123", 1, 0), ("foo-123", 1, 1)]]
+    >>> get_collection_name(b)
+    "foo-123"
+    """
+    if not is_dask_collection(collection):
+        raise TypeError(f"Expected Dask collection; got {type(collection)}")
+    try:
+        key = next(flatten(collection.__dask_keys__()))
+    except StopIteration:
+        # Collection with no keys; this is a legitimate use case but, at the moment of
+        # writing, can only happen with third-party collections
+        raise KeyError("Dask collection has no keys")
+
+    if isinstance(key, tuple) and key and isinstance(key[0], str):
+        return key[0]
+    if isinstance(key, str):
+        return key
+    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+
+
+def replace_name_in_key(key, name: str):
+    """Given a dask key, which must be either a single string or a tuple whose first
+    element is a string (commonly referred to as a collection's 'name'), replace the
+    name with a new one.
+
+    Examples
+    --------
+    >>> replace_name_in_key("foo", "bar")
+    "bar"
+    >>> replace_name_in_key(("foo-123", 1, 2), "bar-456")
+    ("bar-456", 1, 2)
+    """
+    if isinstance(key, tuple) and key and isinstance(key[0], str):
+        return (name,) + key[1:]
+    if isinstance(key, str):
+        return name
+    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 550d7fe4c..809681a7c 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -59,8 +59,8 @@ def blockwise(
 ):
     """Create a Blockwise symbolic mutable mapping
 
-    This is like the ``make_blockwise_graph`` function, but rather than construct a dict, it
-    returns a symbolic Blockwise object.
+    This is like the ``make_blockwise_graph`` function, but rather than construct a
+    dict, it returns a symbolic Blockwise object.
 
     See Also
     --------
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index ac0a3b25b..a8a4f43a5 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -146,7 +146,10 @@ class Scalar(DaskMethodsMixin, OperatorMethodMixin):
         return first, ()
 
     def __dask_postpersist__(self):
-        return Scalar, (self._name, self._meta, self.divisions)
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        return Scalar(dsk, name or self._name, self._meta, self.divisions)
 
     @property
     def _meta_nonempty(self):
@@ -321,7 +324,10 @@ class _Frame(DaskMethodsMixin, OperatorMethodMixin):
         return finalize, ()
 
     def __dask_postpersist__(self):
-        return type(self), (self._name, self._meta, self.divisions)
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        return type(self)(dsk, name or self._name, self._meta, self.divisions)
 
     @property
     def _constructor(self):
diff --git a/dask/delayed.py b/dask/delayed.py
index 3ad139645..2bf26cd1e 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -8,7 +8,7 @@ from tlz import curry, concat, unique, merge
 
 from . import config, threaded
 from .base import is_dask_collection, dont_optimize, DaskMethodsMixin
-from .base import tokenize as _tokenize
+from .base import replace_name_in_key, tokenize as _tokenize
 from .compatibility import is_dataclass, dataclass_fields
 
 from .core import quote
@@ -467,10 +467,6 @@ def optimize(dsk, keys, **kwargs):
     return dsk2
 
 
-def rebuild(dsk, key, length):
-    return Delayed(key, dsk, length)
-
-
 class Delayed(DaskMethodsMixin, OperatorMethodMixin):
     """Represents a value to be computed by dask.
 
@@ -513,7 +509,11 @@ class Delayed(DaskMethodsMixin, OperatorMethodMixin):
         return single_key, ()
 
     def __dask_postpersist__(self):
-        return rebuild, (self._key, getattr(self, "_length", None))
+        return self._rebuild, ()
+
+    def _rebuild(self, dsk, name=None):
+        key = replace_name_in_key(self.key, name) if name else self.key
+        return Delayed(key, dsk, getattr(self, "_length", None))
 
     def __getstate__(self):
         return tuple(getattr(self, i) for i in self.__slots__)
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index d638ed7b8..ad45a9406 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -385,15 +385,15 @@ Dataframe
 
 -   Bump up tolerance for rolling tests (:pr:`6502`) `Julia Signell`_
 -   Implement DatFrame.__len__ (:pr:`6515`) `Tom Augspurger`_
--   Infer arrow schema in to_parquet  (for ArrowEngine`) (:pr:`6490`) `Richard Zamora`_
+-   Infer arrow schema in to_parquet  (for ArrowEngine`) (:pr:`6490`) `Richard (Rick) Zamora`_
 -   Fix parquet test when no pyarrow (:pr:`6524`) `Martin Durant`_
--   Remove problematic ``filter`` arguments in ArrowEngine (:pr:`6527`) `Richard Zamora`_
--   Avoid schema validation by default in ArrowEngine (:pr:`6536`) `Richard Zamora`_
+-   Remove problematic ``filter`` arguments in ArrowEngine (:pr:`6527`) `Richard (Rick) Zamora`_
+-   Avoid schema validation by default in ArrowEngine (:pr:`6536`) `Richard (Rick) Zamora`_
 
 Core
 ++++
 
--   Use unpack_collections in make_blockwise_graph (:pr:`6517`) `Thomas Fan`_
+-   Use unpack_collections in make_blockwise_graph (:pr:`6517`) `Thomas J Fan`_
 -   Move key_split() from optimization.py to utils.py (:pr:`6529`) `Mads R. B. Kristensen`_
 -   Make tests run on moto server (:pr:`6528`) `Martin Durant`_
 
@@ -3724,7 +3724,6 @@ Other
 .. _`Stuart Berg`: https://github.com/stuarteberg
 .. _`Guillaume Eynard-Bontemps`: https://github.com/guillaumeeb
 .. _`Adam Beberg`: https://github.com/beberg
-.. _`Johnnie Gray`: https://github.com/jcmgray
 .. _`Roma Sokolov`: https://github.com/little-arhat
 .. _`Daniel Severo`: https://github.com/dsevero
 .. _`Michał Jastrzębski`: https://github.com/inc0
diff --git a/docs/source/custom-collections.rst b/docs/source/custom-collections.rst
index a7ae0d31c..8358c1878 100644
--- a/docs/source/custom-collections.rst
+++ b/docs/source/custom-collections.rst
@@ -44,11 +44,14 @@ interface is used inside Dask.
 
     Returns
     -------
-    dsk : MutableMapping, None
+    dsk : Mapping, None
         The Dask graph.  If ``None``, this instance will not be interpreted as a
         Dask collection, and none of the remaining interface methods will be
         called.
 
+    If the collection also specifies :meth:`__dask_layers__`, then ``dsk`` must be a
+    :class:`~dask.highlevelgraph.HighLevelGraph` or ``None``.
+
 
 .. method:: __dask_keys__(self)
 
@@ -61,6 +64,29 @@ interface is used inside Dask.
         After computation, the results will be returned in the same layout,
         with the keys replaced with their corresponding outputs.
 
+    If more than one key is returned, all keys must be tuples sharing the same first
+    string element, which is known as the *collection name*, followed by zero or more
+    arbitrary hashables. If the output of the graph is represented by a single key, it
+    *may* also be a bare string.
+
+    These are all valid outputs:
+
+    - ``[]``
+    - ``["x"]``
+    - ``[[("y", "a", 0), ("y", "a", 1)], [("y", "b", 0), ("y", "b", 1)]``
+
+
+.. method:: __dask_layers__(self)
+
+    This method should only be implemented if the collection uses
+    :class:`~dask.highlevelgraph.HighLevelGraph` to implement its dask graph.
+
+    Returns
+    -------
+    names : tuple
+        Tuple of names of the HighLevelGraph layers which contain all keys returned by
+        :meth:`__dask_keys__`.
+
 
 .. staticmethod:: __dask_optimize__(dsk, keys, \*\*kwargs)
 
@@ -131,23 +157,27 @@ interface is used inside Dask.
 .. method:: __dask_postpersist__(self)
 
     Return the rebuilder and (optional) extra arguments to rebuild an equivalent
-    Dask collection from a persisted graph.
+    Dask collection from a persisted or rebuilt graph.
 
-    Used to implement ``dask.persist``.
+    Used to implement :func:`dask.persist`.
 
     Returns
     -------
     rebuild : callable
-        A function with the signature ``rebuild(dsk, *extra_args)``. Called
-        with a persisted graph containing only the keys and results from
-        ``__dask_keys__``, as well as any extra arguments as specified in
-        ``extra_args``. Should return an equivalent Dask collection with the
-        same keys as ``self``, but with the results already computed. For
-        example, the ``rebuild`` function for ``dask.array.Array`` is just the
-        ``__init__`` method called with the new graph but the same metadata.
+        A function with the signature ``rebuild(dsk, *extra_args, name : str = None)``.
+        ``dsk`` is a Mapping which contains at least the output keys returned by
+        :meth:`__dask_keys__`. The callable should return an equivalent Dask collection with the
+        same keys as ``self``, but with the results that are computed through a
+        different graph. In the case of :func:`dask.persist`, the new graph will have
+        just the output keys and the values already computed.
+
+        If the optional parameter ``name`` is specified, it indicates that all output
+        keys are changing too; e.g. if the previous output of :meth:`__dask_keys__` was
+        ``[("a", 0), ("a", 1)]``, after calling ``rebuild(dsk, *extra_args, name="b")``
+        it must become ``[("b", 0), ("b", 1)]``.
     extra_args : tuple
         Any extra arguments to pass to ``rebuild`` after ``dsk``. If no extra
-        arguments should be an empty tuple.
+        arguments are necessary, it must be an empty tuple.
 
 
 .. note:: It's also recommended to define ``__dask_tokenize__``,
@@ -436,7 +466,8 @@ elements of ``dask.delayed``:
 .. code:: python
 
     # Saved as dask_tuple.py
-    from dask.base import DaskMethodsMixin
+    import dask
+    from dask.base import DaskMethodsMixin, replace_name_in_key
     from dask.optimization import cull
 
     # We subclass from DaskMethodsMixin to add common dask methods to our
@@ -472,16 +503,21 @@ elements of ``dask.delayed``:
             return tuple, ()
 
         def __dask_postpersist__(self):
-            # Since our __init__ takes a graph as its first argument, our
-            # rebuild function can just be the class itself. For extra
-            # arguments we also return a tuple containing just the keys.
-            return Tuple, (self._keys,)
+            # We need to return a callable with the signature
+            # rebuild(dsk, *extra_args, name: str = None)
+            return Tuple._rebuild, (self._keys,)
+
+        @staticmethod
+        def _rebuild(dsk, keys, name=None):
+            if name is not None:
+                keys = [replace_name_in_key(key, name) for key in keys]
+            return Tuple(dsk, keys)
 
         def __dask_tokenize__(self):
             # For tokenize to work we want to return a value that fully
             # represents this object. In this case it's the list of keys
             # to be computed.
-            return tuple(self._keys)
+            return self._keys
 
 Demonstrating this class:
 
@@ -491,14 +527,16 @@ Demonstrating this class:
     >>> from operator import add, mul
 
     # Define a dask graph
-    >>> dsk = {'a': 1,
-    ...        'b': 2,
-    ...        'c': (add, 'a', 'b'),
-    ...        'd': (mul, 'b', 2),
-    ...        'e': (add, 'b', 'c')}
+    >>> dsk = {"k0": 1,
+    ...        ("x", "k1"): 2,
+    ...        ("x", 1): (add, "k0", ("x", "k1")),
+    ...        ("x", 2): (mul, ("x", "k1"), 2),
+    ...        ("x", 3): (add, ("x", "k1"), ("x", 1))}
 
-    # The output keys for this graph
-    >>> keys = ['b', 'c', 'd', 'e']
+    # The output keys for this graph.
+    # The first element of each tuple must be the same across the whole collection;
+    # the remainder are arbitrary, unique hashables
+    >>> keys = [("x", "k1"), ("x", 1), ("x", 2), ("x", 3)]
 
     >>> x = Tuple(dsk, keys)
 
@@ -511,10 +549,7 @@ Demonstrating this class:
     >>> isinstance(x2, Tuple)
     True
     >>> x2.__dask_graph__()
-    {'b': 2,
-     'c': 3,
-     'd': 4,
-     'e': 5}
+    {('x', 'k1'): 2, ('x', 1): 3, ('x', 2): 4, ('x', 3): 5}
     >>> x2.compute()
     (2, 3, 4, 5)
 
diff --git a/docs/source/spec.rst b/docs/source/spec.rst
index 9811c6c0c..b006ec554 100644
--- a/docs/source/spec.rst
+++ b/docs/source/spec.rst
@@ -20,7 +20,12 @@ A **Dask graph** is a dictionary mapping **keys** to **computations**:
     'w': (sum, ['x', 'y', 'z']),
     'v': [(sum, ['w', 'z']), 2]}
 
-A **key** is any hashable value that is not a **task**:
+
+A **key** is either:
+
+- a string, or
+- a tuple of one or more elements, whose first element is a string and the remainder are
+  arbitrary hashables:
 
 .. code-block:: python
 


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 04fc8f084..1d349481a 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -25,6 +25,8 @@ from dask.base import (
     named_schedulers,
     get_scheduler,
     collections_to_dsk,
+    get_collection_name,
+    replace_name_in_key,
 )
 from dask.core import literal
 from dask.delayed import Delayed
@@ -279,7 +281,7 @@ def test_tokenize_pandas_mixed_unicode_bytes():
 
 @pytest.mark.skipif("not pd")
 def test_tokenize_pandas_no_pickle():
-    class NoPickle(object):
+    class NoPickle:
         # pickling not supported because it is a local class
         pass
 
@@ -330,7 +332,7 @@ def test_tokenize_kwargs():
 
 
 def test_tokenize_same_repr():
-    class Foo(object):
+    class Foo:
         def __init__(self, x):
             self.x = x
 
@@ -341,7 +343,7 @@ def test_tokenize_same_repr():
 
 
 def test_tokenize_method():
-    class Foo(object):
+    class Foo:
         def __init__(self, x):
             self.x = x
 
@@ -464,8 +466,8 @@ def test_tokenize_object_with_recursion_error_returns_uuid():
 
 
 def test_is_dask_collection():
-    class DummyCollection(object):
-        def __init__(self, dsk=None):
+    class DummyCollection:
+        def __init__(self, dsk):
             self.dask = dsk
 
         def __dask_graph__(self):
@@ -475,7 +477,7 @@ def test_is_dask_collection():
     assert is_dask_collection(x)
     assert not is_dask_collection(2)
     assert is_dask_collection(DummyCollection({}))
-    assert not is_dask_collection(DummyCollection())
+    assert not is_dask_collection(DummyCollection(None))
     assert not is_dask_collection(DummyCollection)
 
 
@@ -537,7 +539,7 @@ def test_unpack_collections():
 
     # Result that looks like a task
     def fail(*args):
-        raise ValueError("Shouldn't have been called")
+        raise ValueError("Shouldn't have been called")  # pragma: nocover
 
     collections, repack = unpack_collections(
         a, (fail, 1), [(fail, 2, 3)], traverse=False
@@ -546,6 +548,52 @@ def test_unpack_collections():
     repack([(fail, 1)])  # Smoketest results that look like tasks
 
 
+def test_get_collection_name():
+    class DummyCollection:
+        def __init__(self, dsk, keys):
+            self.dask = dsk
+            self.keys = keys
+
+        def __dask_graph__(self):
+            return self.dask
+
+        def __dask_keys__(self):
+            return self.keys
+
+    with pytest.raises(TypeError):
+        get_collection_name(object())
+    with pytest.raises(TypeError):
+        get_collection_name(DummyCollection(None, []))
+    with pytest.raises(KeyError):
+        get_collection_name(DummyCollection({}, []))
+    with pytest.raises(TypeError):
+        # Keys must either be a string or a tuple where the first element is a string
+        get_collection_name(DummyCollection({1: 2}, [1]))
+
+    assert get_collection_name(DummyCollection({"a-123": 1}, ["a-123"])) == "a-123"
+
+    # Arbitrary hashables
+    h1 = object()
+    h2 = object()
+    # __dask_keys__() returns a nested list
+    assert (
+        get_collection_name(
+            DummyCollection({("a", h1): 1, ("a", h2): 2}, [[[("a", h1), ("a", h2)]]])
+        )
+        == "a"
+    )
+
+
+def test_replace_name_in_keys():
+    assert replace_name_in_key("foo", "bar") == "bar"
+    assert replace_name_in_key("foo-123", "bar-456") == "bar-456"
+    h1 = object()  # Arbitrary hashables
+    h2 = object()
+    assert replace_name_in_key(("foo-123", h1, h2), "bar") == ("bar", h1, h2)
+    with pytest.raises(TypeError):
+        replace_name_in_key(1, "foo")
+
+
 class Tuple(DaskMethodsMixin):
     __slots__ = ("_dask", "_keys")
     __dask_scheduler__ = staticmethod(dask.threaded.get)
@@ -555,9 +603,11 @@ class Tuple(DaskMethodsMixin):
         self._keys = keys
 
     def __add__(self, other):
-        if isinstance(other, Tuple):
-            return Tuple(merge(self._dask, other._dask), self._keys + other._keys)
-        return NotImplemented
+        if not isinstance(other, Tuple):
+            return NotImplemented  # pragma: nocover
+        name = tokenize(self, other)
+        dsk = {(name, i): k for i, k in enumerate(self._keys + other._keys)}
+        return Tuple(merge(dsk, self._dask, other._dask), list(dsk))
 
     def __dask_graph__(self):
         return self._dask
@@ -572,45 +622,96 @@ class Tuple(DaskMethodsMixin):
         return tuple, ()
 
     def __dask_postpersist__(self):
-        return Tuple, (self._keys,)
+        return Tuple._rebuild, (self._keys,)
+
+    @staticmethod
+    def _rebuild(dsk, keys, name=None):
+        if name is not None:
+            keys = [replace_name_in_key(key, name) for key in keys]
+        return Tuple(dsk, keys)
 
 
 def test_custom_collection():
-    dsk = {"a": 1, "b": 2}
-    dsk2 = {"c": (add, "a", "b"), "d": (add, "c", 1)}
+    # Arbitrary hashables
+    h1 = object()
+    h2 = object()
+
+    # Collections with 2+ keys must have all keys in the format of tuples where the
+    # first element is the same string, referred to as collection name, and the rest are
+    # arbitrary hashables
+    dsk = {("x", h1): 1, ("x", h2): 2}
+    dsk2 = {("y", h1): (add, ("x", h1), ("x", h2)), ("y", h2): (add, ("y", h1), 1)}
     dsk2.update(dsk)
-    dsk3 = {"e": (add, "a", 4), "f": (inc, "e")}
-    dsk3.update(dsk)
 
-    x = Tuple(dsk, ["a", "b"])
-    y = Tuple(dsk2, ["c", "d"])
-    z = Tuple(dsk3, ["e", "f"])
+    # If and only if there is only one top-level key, it can be just a string
+    dsk3 = {"z": (add, ("y", h1), ("y", h2))}
+    dsk3.update(dsk2)
+
+    w = Tuple({}, [])  # A collection can have no keys at all
+    x = Tuple(dsk, [("x", h1), ("x", h2)])
+    y = Tuple(dsk2, [("y", h1), ("y", h2)])
+    z = Tuple(dsk3, ["z"])
 
     # __slots__ defined on base mixin class propagates
     with pytest.raises(AttributeError):
         x.foo = 1
 
     # is_dask_collection
+    assert is_dask_collection(w)
     assert is_dask_collection(x)
+    assert is_dask_collection(y)
+    assert is_dask_collection(z)
 
     # tokenize
+    assert tokenize(w) == tokenize(w)
     assert tokenize(x) == tokenize(x)
-    assert tokenize(x) != tokenize(y)
+    assert tokenize(y) == tokenize(y)
+    assert tokenize(z) == tokenize(z)
+    assert len({tokenize(coll) for coll in (w, x, y, z)}) == 4
+
+    # get_collection_name
+    with pytest.raises(KeyError):
+        get_collection_name(w)
+    assert get_collection_name(x) == "x"
+    assert get_collection_name(y) == "y"
+    assert get_collection_name(z) == "z"
 
     # compute
+    assert w.compute() == ()
     assert x.compute() == (1, 2)
-    assert dask.compute(x, [y, z]) == ((1, 2), [(3, 4), (5, 6)])
-    t = x + y + z
-    assert t.compute() == (1, 2, 3, 4, 5, 6)
+    assert y.compute() == (3, 4)
+    assert z.compute() == (7,)
+    assert dask.compute(w, [{"x": x}, y, z]) == ((), [{"x": (1, 2)}, (3, 4), (7,)])
+    t = w + x + y + z
+    assert t.compute() == (1, 2, 3, 4, 7)
 
     # persist
     t2 = t.persist()
     assert isinstance(t2, Tuple)
-    assert t2._dask == dict(zip("abcdef", range(1, 7)))
-    assert t2.compute() == (1, 2, 3, 4, 5, 6)
-    x2, y2, z2 = dask.persist(x, y, z)
+    assert t2._keys == t._keys
+    assert sorted(t2._dask.values()) == [1, 2, 3, 4, 7]
+    assert t2.compute() == (1, 2, 3, 4, 7)
+
+    w2, x2, y2, z2 = dask.persist(w, x, y, z)
+    assert y2._keys == y._keys
+    assert y2._dask == {("y", h1): 3, ("y", h2): 4}
+    assert y2.compute() == (3, 4)
+
     t3 = x2 + y2 + z2
-    assert t2._dask == t3._dask
+    assert t3.compute() == (1, 2, 3, 4, 7)
+
+    # __dask_postpersist__ with name change
+    rebuild, args = w.__dask_postpersist__()
+    w3 = rebuild({}, *args, name="w3")
+    assert w3.compute() == ()
+
+    rebuild, args = x.__dask_postpersist__()
+    x3 = rebuild({("x3", h1): 10, ("x3", h2): 20}, *args, name="x3")
+    assert x3.compute() == (10, 20)
+
+    rebuild, args = z.__dask_postpersist__()
+    z3 = rebuild({"z3": 70}, *args, name="z3")
+    assert z3.compute() == (70,)
 
 
 @pytest.mark.skipif("not db")
@@ -666,6 +767,18 @@ def test_persist_array():
     assert len(y.dask) == y.npartitions
 
 
+@pytest.mark.skipif("not da")
+def test_persist_array_change_name():
+    a = da.zeros(4, dtype=int, chunks=2)
+    rebuild, args = a.__dask_postpersist__()
+    dsk = {("b", 0): np.array([1, 2]), ("b", 1): np.array([3, 4])}
+    b = rebuild(dsk, *args, name="b")
+    assert isinstance(b, da.Array)
+    assert b.name == "b"
+    assert b.__dask_keys__() == [("b", 0), ("b", 1)]
+    da.utils.assert_eq(b, [1, 2, 3, 4])
+
+
 @pytest.mark.skipif("not dd")
 def test_compute_dataframe():
     df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 5, 3, 3]})
@@ -673,8 +786,75 @@ def test_compute_dataframe():
     ddf1 = ddf.a + 1
     ddf2 = ddf.a + ddf.b
     out1, out2 = compute(ddf1, ddf2)
-    dd._compat.tm.assert_series_equal(out1, df.a + 1)
-    dd._compat.tm.assert_series_equal(out2, df.a + df.b)
+    dd.utils.assert_eq(out1, df.a + 1)
+    dd.utils.assert_eq(out2, df.a + df.b)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_dataframe():
+    df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})
+    ddf1 = dd.from_pandas(df, npartitions=2) * 2
+    assert len(ddf1.__dask_graph__()) == 4
+    ddf2 = ddf1.persist()
+    assert isinstance(ddf2, dd.DataFrame)
+    assert len(ddf2.__dask_graph__()) == 2
+    dd.utils.assert_eq(ddf2, ddf1)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_series():
+    ds = pd.Series([1, 2, 3, 4])
+    dds1 = dd.from_pandas(ds, npartitions=2) * 2
+    assert len(dds1.__dask_graph__()) == 4
+    dds2 = dds1.persist()
+    assert isinstance(dds2, dd.Series)
+    assert len(dds2.__dask_graph__()) == 2
+    dd.utils.assert_eq(dds2, dds1)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_scalar():
+    ds = pd.Series([1, 2, 3, 4])
+    dds1 = dd.from_pandas(ds, npartitions=2).min()
+    assert len(dds1.__dask_graph__()) == 5
+    dds2 = dds1.persist()
+    assert isinstance(dds2, dd.core.Scalar)
+    assert len(dds2.__dask_graph__()) == 1
+    dd.utils.assert_eq(dds2, dds1)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_dataframe_change_name():
+    df1 = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})
+    df2 = pd.DataFrame({"a": [2, 3, 5, 6], "b": [6, 7, 9, 10]})
+    ddf1 = dd.from_pandas(df1, npartitions=2)
+    rebuild, args = ddf1.__dask_postpersist__()
+    dsk = {("x", 0): df2.iloc[:2], ("x", 1): df2.iloc[2:]}
+    ddf2 = rebuild(dsk, *args, name="x")
+    assert ddf2.__dask_keys__() == [("x", 0), ("x", 1)]
+    dd.utils.assert_eq(ddf2, df2)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_series_change_name():
+    ds1 = pd.Series([1, 2, 3, 4])
+    ds2 = pd.Series([5, 6, 7, 8])
+    dds1 = dd.from_pandas(ds1, npartitions=2)
+    rebuild, args = dds1.__dask_postpersist__()
+    dsk = {("x", 0): ds2.iloc[:2], ("x", 1): ds2.iloc[2:]}
+    dds2 = rebuild(dsk, *args, name="x")
+    assert dds2.__dask_keys__() == [("x", 0), ("x", 1)]
+    dd.utils.assert_eq(dds2, ds2)
+
+
+@pytest.mark.skipif("not dd")
+def test_persist_scalar_change_name():
+    ds1 = pd.Series([1, 2, 3, 4])
+    dds1 = dd.from_pandas(ds1, npartitions=2).min()
+    rebuild, args = dds1.__dask_postpersist__()
+    dds2 = rebuild({("x", 0): 5}, *args, name="x")
+    assert dds2.__dask_keys__() == [("x", 0)]
+    dd.utils.assert_eq(dds2, 5)
 
 
 @pytest.mark.skipif("not dd or not da")
@@ -685,7 +865,7 @@ def test_compute_array_dataframe():
     ddf = dd.from_pandas(df, npartitions=2).a + 2
     arr_out, df_out = compute(darr, ddf)
     assert np.allclose(arr_out, arr + 1)
-    dd._compat.tm.assert_series_equal(df_out, df.a + 2)
+    dd.utils.assert_eq(df_out, df.a + 2)
 
 
 @pytest.mark.skipif("not dd")
@@ -941,6 +1121,33 @@ def test_persist_delayed():
     assert x3.compute() == xx.compute()
 
 
+some_hashable = object()
+
+
+@pytest.mark.parametrize("key", ["a", ("a-123", some_hashable)])
+def test_persist_delayed_custom_key(key):
+    d = Delayed(key, {key: "b", "b": 1})
+    assert d.compute() == 1
+    dp = d.persist()
+    assert dp.compute() == 1
+    assert dp.key == key
+    assert dict(dp.dask) == {key: 1}
+
+
+@pytest.mark.parametrize(
+    "key,new_name,new_key",
+    [("a", "b", "b"), (("a-123", some_hashable), "b-123", ("b-123", some_hashable))],
+)
+def test_persist_delayed_change_name(key, new_name, new_key):
+    d = Delayed(key, {key: 1})
+    assert d.compute() == 1
+    rebuild, args = d.__dask_postpersist__()
+    dp = rebuild({new_key: 2}, *args, name=new_name)
+    assert dp.compute() == 2
+    assert dp.key == new_key
+    assert dict(dp.dask) == {new_key: 2}
+
+
 @pytest.mark.skipif("not da or not db")
 def test_persist_array_bag():
     x = da.arange(5, chunks=2) + 1
@@ -963,6 +1170,48 @@ def test_persist_array_bag():
     assert list(b) == list(bb)
 
 
+@pytest.mark.skipif("not db")
+def test_persist_bag():
+    a = db.from_sequence([1, 2, 3], npartitions=2).map(lambda x: x * 2)
+    assert len(a.__dask_graph__()) == 4
+    b = a.persist(scheduler="sync")
+    assert isinstance(b, db.Bag)
+    assert len(b.__dask_graph__()) == 2
+    db.utils.assert_eq(a, b)
+
+
+@pytest.mark.skipif("not db")
+def test_persist_item():
+    a = db.from_sequence([1, 2, 3], npartitions=2).map(lambda x: x * 2).min()
+    assert len(a.__dask_graph__()) == 7
+    b = a.persist(scheduler="sync")
+    assert isinstance(b, db.Item)
+    assert len(b.__dask_graph__()) == 1
+    db.utils.assert_eq(a, b)
+
+
+@pytest.mark.skipif("not db")
+def test_persist_bag_change_name():
+    a = db.from_sequence([1, 2, 3], npartitions=2)
+    rebuild, args = a.__dask_postpersist__()
+    dsk = {("b", 0): [4], ("b", 1): [5, 6]}
+    b = rebuild(dsk, *args, name="b")
+    assert isinstance(b, db.Bag)
+    assert b.name == "b"
+    assert b.__dask_keys__() == [("b", 0), ("b", 1)]
+    db.utils.assert_eq(b, [4, 5, 6])
+
+
+@pytest.mark.skipif("not db")
+def test_persist_item_change_name():
+    a = db.from_sequence([1, 2, 3]).min()
+    rebuild, args = a.__dask_postpersist__()
+    b = rebuild({"x": 4}, *args, name="x")
+    assert isinstance(b, db.Item)
+    assert b.__dask_keys__() == ["x"]
+    db.utils.assert_eq(b, 4)
+
+
 def test_normalize_function_limited_size():
     for i in range(1000):
         normalize_function(lambda x: x)