fix distributed writes

doc/whats-new.rst

@@ -86,6 +86,8 @@ Bug fixes
 - Compatibility fixes to plotting module for Numpy 1.14 and Pandas 0.22
   (:issue:`1813`).
   By `Joe Hamman <https://github.com/jhamman>`_.
+- Fixed to_netcdf when using dask distributed (:issue:`1464`).
+  By `Joe Hamman <https://github.com/jhamman>`_..
 
 .. _whats-new.0.10.0:
 

xarray/backends/h5netcdf_.py

@@ -164,7 +164,10 @@ def prepare_variable(self, name, variable, check_encoding=False,
 
         for k, v in iteritems(attrs):
             nc4_var.setncattr(k, v)
-        return nc4_var, variable.data
+
+        target = H5NetCDFArrayWrapper(name, self)
+
+        return target, variable.data
 
     def sync(self):
         with self.ensure_open(autoclose=True):

xarray/backends/netCDF4_.py

@@ -41,6 +41,11 @@ def __init__(self, variable_name, datastore):
             dtype = np.dtype('O')
         self.dtype = dtype
 
+    def __setitem__(self, key, value):
+        with self.datastore.ensure_open(autoclose=True):
+            data = self.get_array()
+            data[key] = value
+
     def get_array(self):
         self.datastore.assert_open()
         return self.datastore.ds.variables[self.variable_name]
@@ -376,7 +381,9 @@ def prepare_variable(self, name, variable, check_encoding=False,
             # OrderedDict as the input to setncatts
             nc4_var.setncattr(k, v)
 
-        return nc4_var, variable.data
+        target = NetCDF4ArrayWrapper(name, self)
+
+        return target, variable.data
 
     def sync(self):
         with self.ensure_open(autoclose=True):

xarray/backends/scipy_.py

@@ -55,6 +55,11 @@ def __getitem__(self, key):
             copy = self.datastore.ds.use_mmap
             return np.array(data, dtype=self.dtype, copy=copy)
 
+    def __setitem__(self, key, value):
+        with self.datastore.ensure_open(autoclose=True):
+            data = self.get_array()
+            data[key] = value
+
 
 def _open_scipy_netcdf(filename, mode, mmap, version):
     import scipy.io
@@ -202,7 +207,10 @@ def prepare_variable(self, name, variable, check_encoding=False,
         for k, v in iteritems(variable.attrs):
             self._validate_attr_key(k)
             setattr(scipy_var, k, v)
-        return scipy_var, data
+
+        target = ScipyArrayWrapper(name, self)
+
+        return target, data
 
     def sync(self):
         with self.ensure_open(autoclose=True):

xarray/tests/test_distributed.py

@@ -2,13 +2,12 @@
 
 import pytest
 import xarray as xr
-from xarray.core.pycompat import suppress
 
 distributed = pytest.importorskip('distributed')
 da = pytest.importorskip('dask.array')
 import dask
-from distributed.utils_test import cluster, loop, gen_cluster
-from distributed.client import futures_of, wait
+from distributed.utils_test import gen_cluster, cluster, loop
+from distributed.client import futures_of, Client
 
 from xarray.tests.test_backends import create_tmp_file, ON_WINDOWS
 from xarray.tests.test_dataset import create_test_data
@@ -28,30 +27,36 @@
 
 @pytest.mark.xfail(sys.platform == 'win32',
                    reason='https://github.com/pydata/xarray/issues/1738')
-@pytest.mark.parametrize('engine', ENGINES)
-def test_dask_distributed_netcdf_integration_test(loop, engine):
-    with cluster() as (s, _):
-        with distributed.Client(s['address'], loop=loop):
-            original = create_test_data()
-            with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
-                original.to_netcdf(filename, engine=engine)
-                with xr.open_dataset(filename, chunks=3, engine=engine) as restored:
-                    assert isinstance(restored.var1.data, da.Array)
-                    computed = restored.compute()
-                    assert_allclose(original, computed)
+# @pytest.mark.parametrize('engine', ENGINES)
+@gen_cluster(client=True, timeout=None)
+def test_dask_distributed_netcdf_integration_test(c, s, a, b):
+    chunks = {'dim1': 4, 'dim2': 3, 'dim3': 6}
+    for engine in ENGINES:
+        original = create_test_data().chunk(chunks)
+        with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
+            original.to_netcdf(filename, engine=engine)
+            with xr.open_dataset(filename, chunks=3,
+                                 engine=engine) as restored:
+                assert isinstance(restored.var1.data, da.Array)
+                computed = restored.compute()
+                assert_allclose(original, computed)
+
 
 @requires_zarr
 def test_dask_distributed_zarr_integration_test(loop):
-    with cluster() as (s, _):
-        with distributed.Client(s['address'], loop=loop):
-            original = create_test_data()
-            with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
+    chunks = {'dim1': 4, 'dim2': 3, 'dim3': 5}
+    with cluster() as (s, [a, b]):
+        with Client(s['address'], loop=loop) as c:
+            original = create_test_data().chunk(chunks)
+            with create_tmp_file(allow_cleanup_failure=ON_WINDOWS,
+                                 suffix='.zarr') as filename:
                 original.to_zarr(filename)
                 with xr.open_zarr(filename) as restored:
                     assert isinstance(restored.var1.data, da.Array)
                     computed = restored.compute()
                     assert_allclose(original, computed)
 
+
 @pytest.mark.skipif(distributed.__version__ <= '1.19.3',
                     reason='Need recent distributed version to clean up get')
 @gen_cluster(client=True, timeout=None)