Fix DataArray.__dask_scheduler__ to point to dask.threaded.get

Previously this erroneously pointed to an optimize function, likely a
copy-paste error.

For testing this also redirects the .compute methods to use the
dask.compute function directly *if* dask.__version__ >= '0.16.0'.

Author: mrocklin

xarray/core/dataarray.py

@@ -594,7 +594,7 @@ def __dask_optimize__(self):
 
     @property
     def __dask_scheduler__(self):
-        return self._to_temp_dataset().__dask_optimize__
+        return self._to_temp_dataset().__dask_scheduler__
 
     def __dask_postcompute__(self):
         func, args = self._to_temp_dataset().__dask_postcompute__()
@@ -654,8 +654,12 @@ def compute(self, **kwargs):
         --------
         dask.array.compute
         """
-        new = self.copy(deep=False)
-        return new.load(**kwargs)
+        import dask
+        if dask.__version__ >= '0.16.0':
+            return dask.compute(self, **kwargs)[0]
+        else:
+            new = self.copy(deep=False)
+            return new.load(**kwargs)
 
     def persist(self, **kwargs):
         """ Trigger computation in constituent dask arrays
@@ -673,8 +677,12 @@ def persist(self, **kwargs):
         --------
         dask.persist
         """
-        ds = self._to_temp_dataset().persist(**kwargs)
-        return self._from_temp_dataset(ds)
+        import dask
+        if dask.__version__ >= '0.16.0':
+            return dask.persist(self, **kwargs)[0]
+        else:
+            ds = self._to_temp_dataset().persist(**kwargs)
+            return self._from_temp_dataset(ds)
 
     def copy(self, deep=True):
         """Returns a copy of this array.

xarray/core/dataset.py

@@ -590,8 +590,12 @@ def compute(self, **kwargs):
         --------
         dask.array.compute
         """
-        new = self.copy(deep=False)
-        return new.load(**kwargs)
+        import dask
+        if dask.__version__ >= '0.16.0':
+            return dask.compute(self, **kwargs)[0]
+        else:
+            new = self.copy(deep=False)
+            return new.load(**kwargs)
 
     def _persist_inplace(self, **kwargs):
         """ Persist all Dask arrays in memory """
@@ -627,8 +631,12 @@ def persist(self, **kwargs):
         --------
         dask.persist
         """
-        new = self.copy(deep=False)
-        return new._persist_inplace(**kwargs)
+        import dask
+        if dask.__version__ >= '0.16.0':
+            return dask.persist(self, **kwargs)[0]
+        else:
+            new = self.copy(deep=False)
+            return new._persist_inplace(**kwargs)
 
     @classmethod
     def _construct_direct(cls, variables, coord_names, dims=None, attrs=None,

xarray/core/variable.py

@@ -352,8 +352,12 @@ def compute(self, **kwargs):
         --------
         dask.array.compute
         """
-        new = self.copy(deep=False)
-        return new.load(**kwargs)
+        import dask
+        if dask.__version__ >= '0.16.0':
+            return dask.compute(self, **kwargs)[0]
+        else:
+            new = self.copy(deep=False)
+            return new.load(**kwargs)
 
     def __dask_graph__(self):
         if isinstance(self._data, dask_array_type):