Disabled auto-caching on dask; new .compute() method

doc/whats-new.rst

@@ -25,6 +25,10 @@ Breaking changes
   merges will now succeed in cases that previously raised
   ``xarray.MergeError``. Set ``compat='broadcast_equals'`` to restore the
   previous default.
+- Pickling an xarray object based on the dask backend, or reading its
+  :py:meth:`values` property, won't automatically convert the array from dask
+  to numpy in the original object anymore.
+  By `Guido Imperiale <https://github.com/crusaderky>`_.
 
 Deprecations
 ~~~~~~~~~~~~
@@ -45,33 +49,33 @@ Deprecations
 Enhancements
 ~~~~~~~~~~~~
 - Add checking of ``attr`` names and values when saving to netCDF, raising useful
-error messages if they are invalid. (:issue:`911`).
-By `Robin Wilson <https://github.com/robintw>`_.
-
+  error messages if they are invalid. (:issue:`911`).
+  By `Robin Wilson <https://github.com/robintw>`_.
 - Added ability to save ``DataArray`` objects directly to netCDF files using
   :py:meth:`~xarray.DataArray.to_netcdf`, and to load directly from netCDF files
   using :py:func:`~xarray.open_dataarray` (:issue:`915`). These remove the need
   to convert a ``DataArray`` to a ``Dataset`` before saving as a netCDF file,
   and deals with names to ensure a perfect 'roundtrip' capability.
   By `Robin Wilson <https://github.com/robintw`_.
-
 - Multi-index levels are now accessible as "virtual" coordinate variables,
   e.g., ``ds['time']`` can pull out the ``'time'`` level of a multi-index
   (see :ref:`coordinates`). ``sel`` also accepts providing multi-index levels
   as keyword arguments, e.g., ``ds.sel(time='2000-01')``
   (see :ref:`multi-level indexing`).
   By `Benoit Bovy <https://github.com/benbovy>`_.
-
 - Added the ``compat`` option ``'no_conflicts'`` to ``merge``, allowing the
   combination of xarray objects with disjoint (:issue:`742`) or
   overlapping (:issue:`835`) coordinates as long as all present data agrees.
   By `Johnnie Gray <https://github.com/jcmgray>`_. See
   :ref:`combining.no_conflicts` for more details.
-
 - It is now possible to set ``concat_dim=None`` explicitly in
   :py:func:`~xarray.open_mfdataset` to disable inferring a dimension along
   which to concatenate.
   By `Stephan Hoyer <https://github.com/shoyer>`_.
+- Added methods :py:meth:`DataArray.compute`, :py:meth:`Dataset.compute`, and
+  :py:meth:`Variable.compute` as a non-mutating alternative to
+  :py:meth:`~DataArray.load`.
+  By `Guido Imperiale <https://github.com/crusaderky>`_.
 
 Bug fixes
 ~~~~~~~~~

xarray/core/dataarray.py

@@ -561,6 +561,19 @@ def load(self):
         self._coords = new._coords
         return self
 
+    def compute(self):
+        """Manually trigger loading of this array's data from disk or a
+        remote source into memory and return a new array. The original is
+        left unaltered.
+
+        Normally, it should not be necessary to call this method in user code,
+        because all xarray functions should either work on deferred data or
+        load data automatically. However, this method can be necessary when
+        working with many file objects on disk.
+        """
+        new = self.copy(deep=False)
+        return new.load()
+
     def copy(self, deep=True):
         """Returns a copy of this array.
 

xarray/core/dataset.py

@@ -255,8 +255,11 @@ def load_store(cls, store, decoder=None):
         return obj
 
     def __getstate__(self):
-        """Always load data in-memory before pickling"""
-        self.load()
+        """Load data in-memory before pickling (except for Dask data)"""
+        for v in self.variables.values():
+            if not isinstance(v.data, dask_array_type):
+                v.load()
+
         # self.__dict__ is the default pickle object, we don't need to
         # implement our own __setstate__ method to make pickle work
         state = self.__dict__.copy()
@@ -319,6 +322,19 @@ def load(self):
 
         return self
 
+    def compute(self):
+        """Manually trigger loading of this dataset's data from disk or a
+        remote source into memory and return a new dataset. The original is
+        left unaltered.
+
+        Normally, it should not be necessary to call this method in user code,
+        because all xarray functions should either work on deferred data or
+        load data automatically. However, this method can be necessary when
+        working with many file objects on disk.
+        """
+        new = self.copy(deep=False)
+        return new.load()
+
     @classmethod
     def _construct_direct(cls, variables, coord_names, dims=None, attrs=None,
                           file_obj=None):
@@ -401,14 +417,12 @@ def copy(self, deep=False):
         """Returns a copy of this dataset.
 
         If `deep=True`, a deep copy is made of each of the component variables.
-        Otherwise, a shallow copy is made, so each variable in the new dataset
-        is also a variable in the original dataset.
+        Otherwise, a shallow copy of each of the component variable is made, so
+        that the underlying memory region of the new dataset is the same as in
+        the original dataset.
         """
-        if deep:
-            variables = OrderedDict((k, v.copy(deep=True))
-                                    for k, v in iteritems(self._variables))
-        else:
-            variables = self._variables.copy()
+        variables = OrderedDict((k, v.copy(deep=deep))
+                                for k, v in iteritems(self._variables))
         # skip __init__ to avoid costly validation
         return self._construct_direct(variables, self._coord_names.copy(),
                                       self._dims.copy(), self._attrs_copy())

xarray/core/variable.py

@@ -271,10 +271,21 @@ def data(self, data):
                 "replacement data must match the Variable's shape")
         self._data = data
 
+    def _data_cast(self):
+        if isinstance(self._data, (np.ndarray, PandasIndexAdapter)):
+            return self._data
+        else:
+            return np.asarray(self._data)
+
     def _data_cached(self):
-        if not isinstance(self._data, (np.ndarray, PandasIndexAdapter)):
-            self._data = np.asarray(self._data)
-        return self._data
+        """Load data into memory and return it.
+        Do not cache dask arrays automatically; that should
+        require an explicit load() call.
+        """
+        new_data = self._data_cast()
+        if not isinstance(self._data, dask_array_type):
+            self._data = new_data
+        return new_data
 
     @property
     def _indexable_data(self):
@@ -287,12 +298,25 @@ def load(self):
         Normally, it should not be necessary to call this method in user code,
         because all xarray functions should either work on deferred data or
         load data automatically.
-        """
-        self._data_cached()
+        """        
+        self._data = self._data_cast()
         return self
 
+    def compute(self):
+        """Manually trigger loading of this variable's data from disk or a
+        remote source into memory and return a new variable. The original is
+        left unaltered.
+
+        Normally, it should not be necessary to call this method in user code,
+        because all xarray functions should either work on deferred data or
+        load data automatically.
+        """
+        new = self.copy(deep=False)
+        return new.load()
+
     def __getstate__(self):
-        """Always cache data as an in-memory array before pickling"""
+        """Always cache data as an in-memory array before pickling
+        (with the exception of dask backend)"""
         self._data_cached()
         # self.__dict__ is the default pickle object, we don't need to
         # implement our own __setstate__ method to make pickle work
@@ -1090,10 +1114,11 @@ def __init__(self, dims, data, attrs=None, encoding=None, fastpath=False):
             raise ValueError('%s objects must be 1-dimensional' %
                              type(self).__name__)
 
-    def _data_cached(self):
-        if not isinstance(self._data, PandasIndexAdapter):
-            self._data = PandasIndexAdapter(self._data)
-        return self._data
+    def _data_cast(self):
+        if isinstance(self._data, PandasIndexAdapter):
+            return self._data
+        else:
+            return PandasIndexAdapter(self._data)
 
     def __getitem__(self, key):
         key = self._item_key_to_tuple(key)

xarray/test/test_backends.py

@@ -149,6 +149,24 @@ def assert_loads(vars=None):
             actual = ds.load()
         self.assertDatasetAllClose(expected, actual)
 
+    def test_dataset_compute(self):
+        expected = create_test_data()
+
+        with self.roundtrip(expected) as actual:
+            # Test Dataset.compute()
+            for v in actual.variables.values():
+                self.assertFalse(v._in_memory)
+
+            computed = actual.compute()
+
+            for v in actual.data_vars.values():
+                self.assertFalse(v._in_memory)
+            for v in computed.variables.values():
+                self.assertTrue(v._in_memory)
+
+            self.assertDatasetAllClose(expected, actual)
+            self.assertDatasetAllClose(expected, computed)
+
     def test_roundtrip_None_variable(self):
         expected = Dataset({None: (('x', 'y'), [[0, 1], [2, 3]])})
         with self.roundtrip(expected) as actual:
@@ -230,18 +248,6 @@ def test_roundtrip_coordinates(self):
         with self.roundtrip(expected) as actual:
             self.assertDatasetIdentical(expected, actual)
 
-        expected = original.copy()
-        expected.attrs['coordinates'] = 'something random'
-        with self.assertRaisesRegexp(ValueError, 'cannot serialize'):
-            with self.roundtrip(expected):
-                pass
-
-        expected = original.copy(deep=True)
-        expected['foo'].attrs['coordinates'] = 'something random'
-        with self.assertRaisesRegexp(ValueError, 'cannot serialize'):
-            with self.roundtrip(expected):
-                pass
-
     def test_roundtrip_boolean_dtype(self):
         original = create_boolean_data()
         self.assertEqual(original['x'].dtype, 'bool')
@@ -872,7 +878,26 @@ def test_read_byte_attrs_as_unicode(self):
 @requires_dask
 @requires_scipy
 @requires_netCDF4
-class DaskTest(TestCase):
+class DaskTest(TestCase, DatasetIOTestCases):
+    @contextlib.contextmanager
+    def create_store(self):
+        yield Dataset()
+
+    @contextlib.contextmanager
+    def roundtrip(self, data, save_kwargs={}, open_kwargs={}):
+        yield data.chunk()
+
+    def test_roundtrip_datetime_data(self):
+        # Override method in DatasetIOTestCases - remove not applicable save_kwds
+        times = pd.to_datetime(['2000-01-01', '2000-01-02', 'NaT'])
+        expected = Dataset({'t': ('t', times), 't0': times[0]})
+        with self.roundtrip(expected) as actual:
+            self.assertDatasetIdentical(expected, actual)
+
+    def test_write_store(self):
+        # Override method in DatasetIOTestCases - not applicable to dask
+        pass
+
     def test_open_mfdataset(self):
         original = Dataset({'foo': ('x', np.random.randn(10))})
         with create_tmp_file() as tmp1:
@@ -992,7 +1017,16 @@ def test_deterministic_names(self):
                 self.assertIn(tmp, dask_name)
             self.assertEqual(original_names, repeat_names)
 
-
+    def test_dataarray_compute(self):
+        # Test DataArray.compute() on dask backend.
+        # The test for Dataset.compute() is already in DatasetIOTestCases;
+        # however dask is the only tested backend which supports DataArrays
+        actual = DataArray([1,2]).chunk()
+        computed = actual.compute()
+        self.assertFalse(actual._in_memory)
+        self.assertTrue(computed._in_memory)
+        self.assertDataArrayAllClose(actual, computed)
+
 @requires_scipy_or_netCDF4
 @requires_pydap
 class PydapTest(TestCase):

xarray/test/test_dask.py

@@ -1,3 +1,4 @@
+import pickle
 import numpy as np
 import pandas as pd
 
@@ -12,44 +13,22 @@
     import dask.array as da
 
 
-def _copy_at_variable_level(arg):
-    """We need to copy the argument at the level of xarray.Variable objects, so
-    that viewing its values does not trigger lazy loading.
-    """
-    if isinstance(arg, Variable):
-        return arg.copy(deep=False)
-    elif isinstance(arg, DataArray):
-        ds = arg.to_dataset(name='__copied__')
-        return _copy_at_variable_level(ds)['__copied__']
-    elif isinstance(arg, Dataset):
-        ds = arg.copy()
-        for k in list(ds):
-            ds._variables[k] = ds._variables[k].copy(deep=False)
-        return ds
-    else:
-        assert False
-
-
 class DaskTestCase(TestCase):
     def assertLazyAnd(self, expected, actual, test):
-        expected_copy = _copy_at_variable_level(expected)
-        actual_copy = _copy_at_variable_level(actual)
+        expected.name = None
+        actual.name = None
         with dask.set_options(get=dask.get):
-            test(actual_copy, expected_copy)
-        var = getattr(actual, 'variable', actual)
-        self.assertIsInstance(var.data, da.Array)
+            test(actual, expected)
+        if isinstance(actual, Dataset):
+            for var in actual.variables.values():
+                self.assertIsInstance(var.data, da.Array)            
+        else:    
+            var = getattr(actual, 'variable', actual)
+            self.assertIsInstance(var.data, da.Array)
 
 
 @requires_dask
 class TestVariable(DaskTestCase):
-    def assertLazyAnd(self, expected, actual, test):
-        expected_copy = expected.copy(deep=False)
-        actual_copy = actual.copy(deep=False)
-        with dask.set_options(get=dask.get):
-            test(actual_copy, expected_copy)
-        var = getattr(actual, 'variable', actual)
-        self.assertIsInstance(var.data, da.Array)
-
     def assertLazyAndIdentical(self, expected, actual):
         self.assertLazyAnd(expected, actual, self.assertVariableIdentical)
 
@@ -321,3 +300,29 @@ def test_dot(self):
         eager = self.eager_array.dot(self.eager_array[0])
         lazy = self.lazy_array.dot(self.lazy_array[0])
         self.assertLazyAndAllClose(eager, lazy)
+
+    def test_dataarray_pickle(self):
+        # Test that pickling/unpickling does not convert the dask
+        # backend to numpy
+        a1 = DataArray([1,2]).chunk()
+        self.assertFalse(a1._in_memory)        
+        a2 = pickle.loads(pickle.dumps(a1))
+        self.assertDataArrayIdentical(a1, a2)
+        self.assertFalse(a1._in_memory)
+        self.assertFalse(a2._in_memory)        
+
+    def test_dataset_pickle(self):
+        ds1 = Dataset({'a': [1,2]}).chunk()
+        self.assertFalse(ds1['a']._in_memory)        
+        ds2 = pickle.loads(pickle.dumps(ds1))
+        self.assertDatasetIdentical(ds1, ds2)
+        self.assertFalse(ds1['a']._in_memory)
+        self.assertFalse(ds2['a']._in_memory)        
+
+    def test_values(self):
+        # Test that invoking the values property does not convert the dask
+        # backend to numpy
+        a = DataArray([1,2]).chunk()
+        self.assertFalse(a._in_memory)
+        self.assertEquals(a.values.tolist(), [1, 2])
+        self.assertFalse(a._in_memory)        

xarray/test/test_dataset.py

@@ -1287,10 +1287,13 @@ def test_copy(self):
 
         for copied in [data.copy(deep=False), copy(data)]:
             self.assertDatasetIdentical(data, copied)
-            for k in data:
+            # Note: IndexVariable objects with string dtype are always
+            # copied because of xarray.core.util.safe_cast_to_index.
+            # Limiting the test to data variables.
+            for k in data.data_vars:
                 v0 = data.variables[k]
                 v1 = copied.variables[k]
-                self.assertIs(v0, v1)
+                assert source_ndarray(v0.data) is source_ndarray(v1.data)
             copied['foo'] = ('z', np.arange(5))
             self.assertNotIn('foo', data)