fix distributed writes

doc/whats-new.rst

@@ -112,6 +112,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/common.py

@@ -178,11 +178,7 @@ def add(self, source, target):
             self.sources.append(source)
             self.targets.append(target)
         else:
-            try:
-                target[...] = source
-            except TypeError:
-                # workaround for GH: scipy/scipy#6880
-                target[:] = source
+            target[...] = source
 
     def sync(self):
         if self.sources:

xarray/backends/h5netcdf_.py

@@ -166,7 +166,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

@@ -43,6 +43,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]
@@ -398,7 +403,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,18 @@ 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()
+            try:
+                data[key] = value
+            except TypeError:
+                if key is Ellipsis:
+                    # workaround for GH: scipy/scipy#6880
+                    data[:] = value
+                else:
+                    raise
+
 
 def _open_scipy_netcdf(filename, mode, mmap, version):
     import scipy.io
@@ -201,7 +213,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

@@ -6,6 +6,7 @@
 distributed = pytest.importorskip('distributed')
 da = pytest.importorskip('dask.array')
 import dask
+from dask.distributed import Client
 from distributed.utils_test import cluster, gen_cluster
 from distributed.utils_test import loop  # flake8: noqa
 from distributed.client import futures_of
@@ -17,36 +18,52 @@
                requires_zarr)
 
 
-ENGINES = []
-if has_scipy:
-    ENGINES.append('scipy')
-if has_netCDF4:
-    ENGINES.append('netcdf4')
-if has_h5netcdf:
-    ENGINES.append('h5netcdf')
-
-
-@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)
+@gen_cluster(client=True, timeout=None)
+def test_dask_distributed_netcdf_integration_test_scipy(c, s, a, b):
+    chunks = {'dim1': 4, 'dim2': 3, 'dim3': 6}
+    original = create_test_data().chunk(chunks)
+    with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
+        original.to_netcdf(filename, engine='scipy')
+        with xr.open_dataset(filename, chunks=chunks,
+                             engine='scipy') as restored:
+            assert isinstance(restored.var1.data, da.Array)
+            computed = restored.compute()
+            assert_allclose(original, computed)
+
+@gen_cluster(client=True, timeout=None)
+def test_dask_distributed_netcdf_integration_test_netcdf4(c, s, a, b):
+    chunks = {'dim1': 4, 'dim2': 3, 'dim3': 6}
+    original = create_test_data().chunk(chunks)
+    with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
+        original.to_netcdf(filename, engine='netcdf4')
+        with xr.open_dataset(filename, chunks=chunks,
+                             engine='netcdf4') as restored:
+            assert isinstance(restored.var1.data, da.Array)
+            computed = restored.compute()
+            assert_allclose(original, computed)
+
+
+@gen_cluster(client=True, timeout=None)
+def test_dask_distributed_netcdf_integration_test_h5netcdf(c, s, a, b):
+    chunks = {'dim1': 4, 'dim2': 3, 'dim3': 6}
+    original = create_test_data().chunk(chunks)
+    with create_tmp_file(allow_cleanup_failure=ON_WINDOWS) as filename:
+        original.to_netcdf(filename, engine='h5netcdf')
+        with xr.open_dataset(filename, chunks=chunks,
+                             engine='h5netcdf') 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)