Support HighLevelGraphs (#2603)

* Support HighLevelGraphs

Fixes #4291

* test __dask_layers__

* Skip dependnecies test with old dask

* Reenable dask-dev test on Travis-CI

.travis.yml

@@ -44,7 +44,6 @@ matrix:
         - libhdf5-serial-dev
         - netcdf-bin
         - libnetcdf-dev
-  - env: CONDA_ENV=py36-dask-dev
   - env: CONDA_ENV=py36-pandas-dev
   - env: CONDA_ENV=py36-bottleneck-dev
   - env: CONDA_ENV=py36-condaforge-rc

doc/whats-new.rst

@@ -43,7 +43,8 @@ Enhancements
   By `Stephan Hoyer <https://github.com/shoyer>`_
 - Like :py:class:`pandas.DatetimeIndex`, :py:class:`CFTimeIndex` now supports
   "dayofyear" and "dayofweek" accessors (:issue:`2597`).  By `Spencer Clark
-  <https://github.com/spencerkclark>`_. 
+  <https://github.com/spencerkclark>`_.
+- Support Dask ``HighLevelGraphs`` by `Matthew Rocklin <https://matthewrocklin.com>`_.
 
 
 Bug fixes
@@ -159,9 +160,9 @@ Enhancements
   to returning (and is now deprecated). This was changed in order to facilitate
   using tutorial datasets with dask.
   By `Joe Hamman <https://github.com/jhamman>`_.
-- ``DataArray`` can now use ``xr.set_option(keep_attrs=True)`` and retain attributes in binary operations, 
+- ``DataArray`` can now use ``xr.set_option(keep_attrs=True)`` and retain attributes in binary operations,
   such as (``+, -, * ,/``). Default behaviour is unchanged (*Attributes will be dismissed*). By `Michael Blaschek <https://github.com/MBlaschek>`_
-  
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/dataarray.py

@@ -587,6 +587,9 @@ def __dask_graph__(self):
     def __dask_keys__(self):
         return self._to_temp_dataset().__dask_keys__()
 
+    def __dask_layers__(self):
+        return self._to_temp_dataset().__dask_layers__()
+
     @property
     def __dask_optimize__(self):
         return self._to_temp_dataset().__dask_optimize__

xarray/core/dataset.py

@@ -509,14 +509,24 @@ def __dask_graph__(self):
         if not graphs:
             return None
         else:
-            from dask import sharedict
-            return sharedict.merge(*graphs.values())
+            try:
+                from dask.highlevelgraph import HighLevelGraph
+                return HighLevelGraph.merge(*graphs.values())
+            except ImportError:
+                from dask import sharedict
+                return sharedict.merge(*graphs.values())
+
 
     def __dask_keys__(self):
         import dask
         return [v.__dask_keys__() for v in self.variables.values()
                 if dask.is_dask_collection(v)]
 
+    def __dask_layers__(self):
+        import dask
+        return sum([v.__dask_layers__() for v in self.variables.values() if
+                    dask.is_dask_collection(v)], ())
+
     @property
     def __dask_optimize__(self):
         import dask.array as da

xarray/core/variable.py

@@ -352,6 +352,9 @@ def __dask_graph__(self):
     def __dask_keys__(self):
         return self._data.__dask_keys__()
 
+    def __dask_layers__(self):
+        return self._data.__dask_layers__()
+
     @property
     def __dask_optimize__(self):
         return self._data.__dask_optimize__

xarray/tests/test_dask.py

@@ -843,3 +843,16 @@ def test_basic_compute():
             ds.compute()
             ds.foo.compute()
             ds.foo.variable.compute()
+
+
+@pytest.mark.skipif(LooseVersion(dask.__version__) < LooseVersion('0.20.0'),
+                    reason='needs newer dask')
+def test_dask_layers_and_dependencies():
+    ds = Dataset({'foo': ('x', range(5)),
+                  'bar': ('x', range(5))}).chunk()
+
+    x = dask.delayed(ds)
+    assert set(x.__dask_graph__().dependencies).issuperset(
+        ds.__dask_graph__().dependencies)
+    assert set(x.foo.__dask_graph__().dependencies).issuperset(
+        ds.__dask_graph__().dependencies)