Repository: dask/dask

Issue:
ShareDict seems to cause issues with HighLevelGraphs in optimize_atop
Following up on https://github.com/dask/dask/issues/4038#issuecomment-434338006, this script:

```python
from dask.sharedict import ShareDict
import numpy as np


def foo(A):
    return A[None, ...]


A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))

name = 'D-' + dask.base.tokenize(A)

dsk = da.core.top(
        foo, name, ("nsrc", "ntime", "nbl", "npol"),
        A.name, ("ntime", "nbl", "npol"),
        new_axes={"nsrc": 1},
        numblocks={a.name: a.numblocks for a in (A,)})

array_dsk = ShareDict()
array_dsk.update(dsk)
array_dsk.update(A.__dask_graph__())

chunks = ((1,),) + A.chunks

D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
print(D.compute().shape)    # Succeeds
D = D.sum(axis=0)
print(D.compute().shape)    # Fails
```

causes the following error:


```traceback
$ python dask_top_2.py 
(1, 10, 20, 4)
Traceback (most recent call last):
  File "dask_top_2.py", line 32, in <module>
    pprint(D.compute().shape)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 155, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 389, in compute
    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 193, in collections_to_dsk
    for opt, (dsk, keys) in groups.items()]))
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/array/optimization.py", line 38, in optimize
    dsk = optimize_atop(dsk, keys=keys)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/array/top.py", line 594, in optimize_atop
    if (isinstance(layers[dep], TOP) and
KeyError: 'D-122ec2f0f994205583cdc43c8fba6f6d'
```

Discussion:

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 8bfb23a28..b63d5aa69 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -348,7 +348,7 @@ def apply_infer_dtype(func, args, kwargs, funcname, suggest_dtype='dtype', nout=
         Function for which output dtype is to be determined
 
     args: List of array like
-        Arguments to the function, which qwould usually be used. Only attributes
+        Arguments to the function, which would usually be used. Only attributes
         ``ndim`` and ``dtype`` are used.
 
     kwargs: dict
diff --git a/dask/array/top.py b/dask/array/top.py
index 66526a185..4f1cad48a 100644
--- a/dask/array/top.py
+++ b/dask/array/top.py
@@ -591,7 +591,8 @@ def optimize_atop(full_graph, keys=()):
             deps = set(top_layers)
             while deps:  # we gather as many sub-layers as we can
                 dep = deps.pop()
-                if (isinstance(layers[dep], TOP) and
+                if (dep in layers and
+                        isinstance(layers[dep], TOP) and
                         not (dep != layer and dep in keep) and  # output layer
                         layers[dep].concatenate == layers[layer].concatenate):  # punt on mixed concatenate
                     top_layers.add(dep)
diff --git a/dask/base.py b/dask/base.py
index 7cb0c660f..cf2cb3cc7 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -13,7 +13,8 @@ import uuid
 from toolz import merge, groupby, curry, identity
 from toolz.functoolz import Compose
 
-from .compatibility import long, unicode, Iterator
+from .compatibility import (apply, long, unicode, Iterator, is_dataclass,
+                            dataclass_fields)
 from .context import thread_state
 from .core import flatten, quote, get as simple_get
 from .hashing import hash_buffer_hex
@@ -264,6 +265,10 @@ def unpack_collections(*args, **kwargs):
             elif typ is dict:
                 tsk = (dict, [[_unpack(k), _unpack(v)]
                               for k, v in expr.items()])
+            elif is_dataclass(expr):
+                tsk = (apply, typ, (), (dict,
+                       [[f.name, _unpack(getattr(expr, f.name))] for f in
+                        dataclass_fields(expr)]))
             else:
                 return expr
 
diff --git a/dask/compatibility.py b/dask/compatibility.py
index 25adaa1aa..9c9b9c32a 100644
--- a/dask/compatibility.py
+++ b/dask/compatibility.py
@@ -307,7 +307,7 @@ def getargspec(func):
     """Version of inspect.getargspec that works with partial and warps."""
     if isinstance(func, functools.partial):
         return getargspec(func.func)
- 
+
     func = getattr(func, '__wrapped__', func)
     if isinstance(func, type):
         return _getargspec(func.__init__)
@@ -337,3 +337,14 @@ def bind_method(cls, name, func):
         setattr(cls, name, types.MethodType(func, None, cls))
     else:
         setattr(cls, name, func)
+
+
+try:
+    from dataclasses import is_dataclass, fields as dataclass_fields
+
+except ImportError:
+    def is_dataclass(x):
+        return False
+
+    def dataclass_fields(x):
+        return []
diff --git a/dask/delayed.py b/dask/delayed.py
index ff3cbe358..aa2cde5d3 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -12,7 +12,8 @@ except ImportError:
 from . import config, threaded
 from .base import is_dask_collection, dont_optimize, DaskMethodsMixin
 from .base import tokenize as _tokenize
-from .compatibility import apply, Iterator
+from .compatibility import apply, Iterator, is_dataclass, dataclass_fields
+
 from .core import quote
 from .context import globalmethod
 from .optimization import cull
@@ -91,6 +92,12 @@ def to_task_dask(expr):
         args, dsk = to_task_dask([[k, v] for k, v in expr.items()])
         return (dict, args), dsk
 
+    if is_dataclass(expr):
+        args, dsk = to_task_dask([[f.name, getattr(expr, f.name)] for f in
+                                  dataclass_fields(expr)])
+
+        return (apply, typ, (), (dict, args)), dsk
+
     if typ is slice:
         args, dsk = to_task_dask([expr.start, expr.stop, expr.step])
         return (slice,) + tuple(args), dsk
diff --git a/docs/source/array-sparse.rst b/docs/source/array-sparse.rst
index 952c9d591..14189f6af 100644
--- a/docs/source/array-sparse.rst
+++ b/docs/source/array-sparse.rst
@@ -1,39 +1,39 @@
 Sparse Arrays
 =============
 
-By swapping out in-memory numpy arrays with in-memory sparse arrays we can
-reuse the blocked algorithms of Dask.array to achieve parallel and distributed
+By swapping out in-memory NumPy arrays with in-memory sparse arrays, we can
+reuse the blocked algorithms of Dask's Array to achieve parallel and distributed
 sparse arrays.
 
-The blocked algorithms in Dask.array normally parallelize around in-memory
-numpy arrays.  However, if another in-memory array library supports the NumPy
-interface then it too can take advantage of dask.array's parallel algorithms.
+The blocked algorithms in Dask Array normally parallelize around in-memory
+NumPy arrays.  However, if another in-memory array library supports the NumPy
+interface, then it too can take advantage of Dask Array's parallel algorithms.
 In particular the `sparse <https://github.com/mrocklin/sparse/>`_ array library
-satisfies a subset of the NumPy API and works well with, and is tested against,
-Dask.array.
+satisfies a subset of the NumPy API and works well with (and is tested against)
+Dask Array.
 
 Example
 -------
 
-Say we have a dask.array with mostly zeros
+Say we have a Dask array with mostly zeros:
 
 .. code-block:: python
 
    x = da.random.random((100000, 100000), chunks=(1000, 1000))
    x[x < 0.95] = 0
 
-We can convert each of these chunks of NumPy arrays into a sparse.COO array.
+We can convert each of these chunks of NumPy arrays into a **sparse.COO** array:
 
 .. code-block:: python
 
    import sparse
    s = x.map_blocks(sparse.COO)
 
-Now our array is composed not of many NumPy arrays, but rather of many
-sparse arrays.  Semantically this does not change anything.  Operations that
-work will work identically (assuming that the behavior of ``numpy`` and
-``sparse`` are identical) but performance characteristics and storage costs may
-change significantly
+Now, our array is not composed of many NumPy arrays, but rather of many
+sparse arrays.  Semantically, this does not change anything.  Operations that
+work will continue to work identically (assuming that the behavior of ``numpy`` 
+and ``sparse`` are identical), but performance characteristics and storage costs 
+may change significantly:
 
 .. code-block:: python
 
@@ -50,24 +50,24 @@ Requirements
 
 Any in-memory library that copies the NumPy ndarray interface should work here.
 The `sparse <https://github.com/mrocklin/sparse/>`_ library is a minimal
-example.  In particular an in-memory library should implement at least the
+example.  In particular, an in-memory library should implement at least the
 following operations:
 
 1.  Simple slicing with slices, lists, and elements (for slicing, rechunking,
-    reshaping, etc).
+    reshaping, etc)
 2.  A ``concatenate`` function matching the interface of ``np.concatenate``.
-    This must be registered in ``dask.array.core.concatenate_lookup``.
+    This must be registered in ``dask.array.core.concatenate_lookup``
 3.  All ufuncs must support the full ufunc interface, including ``dtype=`` and
     ``out=`` parameters (even if they don't function properly)
 4.  All reductions must support the full ``axis=`` and ``keepdims=`` keywords
-    and behave like numpy in this respect
+    and behave like NumPy in this respect
 5.  The array class should follow the ``__array_priority__`` protocol and be
-    prepared to respond to other arrays of lower priority.
+    prepared to respond to other arrays of lower priority
 6.  If ``dot`` support is desired, a ``tensordot`` function matching the
     interface of ``np.tensordot`` should be registered in
-    ``dask.array.core.tensordot_lookup``.
+    ``dask.array.core.tensordot_lookup``
 
-The implementation of other operations like reshape, transpose, etc.
+The implementation of other operations like reshape, transpose, etc.,
 should follow standard NumPy conventions regarding shape and dtype.  Not
 implementing these is fine; the parallel ``dask.array`` will err at runtime if
 these operations are attempted.
@@ -76,8 +76,8 @@ these operations are attempted.
 Mixed Arrays
 ------------
 
-Dask.array supports mixing different kinds of in-memory arrays.  This relies
+Dask's Array supports mixing different kinds of in-memory arrays.  This relies
 on the in-memory arrays knowing how to interact with each other when necessary.
-When two arrays interact the functions from the array with the highest
-``__array_priority__`` will take precedence (for example for concatenate,
+When two arrays interact, the functions from the array with the highest
+``__array_priority__`` will take precedence (for example, for concatenate,
 tensordot, etc.).


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index d0f2c316c..117a15b9c 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -380,3 +380,28 @@ def test_namedtuple(tup):
                 dtype=A.dtype)
 
     assert_eq(A, B)
+
+
+def test_gh_4176():
+    def foo(A):
+        return A[None, ...]
+
+    A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))
+
+    name = 'D'
+
+    dsk = da.core.top(
+        foo, name, ("nsrc", "ntime", "nbl", "npol"),
+        A.name, ("ntime", "nbl", "npol"),
+        new_axes={"nsrc": 1},
+        numblocks={a.name: a.numblocks for a in (A,)}
+    )
+
+    array_dsk = dask.sharedict.ShareDict()
+    array_dsk.update(dsk)
+    array_dsk.update(A.__dask_graph__())
+
+    chunks = ((1,),) + A.chunks
+
+    D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
+    D.sum(axis=0).compute()
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index ab981e41d..12a867206 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -346,13 +346,21 @@ def test_is_dask_collection():
     assert not is_dask_collection(DummyCollection)
 
 
+try:
+    import dataclasses
+    # Avoid @dataclass decorator as Python < 3.7 fail to interpret the type hints
+    ADataClass = dataclasses.make_dataclass('ADataClass', [('a', int)])
+except ImportError:
+    dataclasses = None
+
+
 def test_unpack_collections():
     a = delayed(1) + 5
     b = a + 1
     c = a + 2
 
     def build(a, b, c, iterator):
-        return (a, b,               # Top-level collections
+        t = (a, b,                  # Top-level collections
                 {'a': a,            # dict
                  a: b,              # collections as keys
                  'b': [1, 2, [b]],  # list
@@ -361,6 +369,11 @@ def test_unpack_collections():
                  'e': {a, 2, 3}},   # set
                 iterator)           # Iterator
 
+        if dataclasses is not None:
+            t[2]['f'] = ADataClass(a=a)
+
+        return t
+
     args = build(a, b, c, (i for i in [a, b, c]))
 
     collections, repack = unpack_collections(*args)
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 4a4c888fa..c887f29ca 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -94,6 +94,22 @@ def test_delayed():
     assert a.key in b.dask
 
 
+def test_delayed_with_dataclass():
+    dataclasses = pytest.importorskip("dataclasses")
+
+    # Avoid @dataclass decorator as Python < 3.7 fail to interpret the type hints
+    ADataClass = dataclasses.make_dataclass('ADataClass', [('a', int)])
+
+    literal = dask.delayed(3)
+    with_class = dask.delayed({"a": ADataClass(a=literal)})
+
+    def return_nested(obj):
+        return obj["a"].a
+    final = delayed(return_nested)(with_class)
+
+    assert final.compute() == 3
+
+
 def test_operators():
     a = delayed([1, 2, 3])
     assert a[0].compute() == 1