fix zarr chunking bug (#2228)

Author: Joe Hamman

doc/whats-new.rst

@@ -39,6 +39,10 @@ Enhancements
 Bug fixes
 ~~~~~~~~~
 
+- Fixed a bug in ``zarr`` backend which prevented use with datasets with
+  incomplete chunks in multiple dimensions (:issue:`2225`).
+  By `Joe Hamman <https://github.com/jhamman>`_.
+
 .. _whats-new.0.10.7:
 
 v0.10.7 (7 June 2018)
@@ -60,12 +64,13 @@ Enhancements
   See :ref:`interpolating values with interp` for the detail.
   (:issue:`2079`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.
-  
+
 Bug fixes
 ~~~~~~~~~
 
 - Fixed a bug in ``rasterio`` backend which prevented use with ``distributed``.
   The ``rasterio`` backend now returns pickleable objects (:issue:`2021`).
+  By `Joe Hamman <https://github.com/jhamman>`_.
 
 .. _whats-new.0.10.6:
 

xarray/backends/zarr.py

@@ -78,24 +78,18 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim):
     # while dask chunks can be variable sized
     # http://dask.pydata.org/en/latest/array-design.html#chunks
     if var_chunks and enc_chunks is None:
-        all_var_chunks = list(product(*var_chunks))
-        first_var_chunk = all_var_chunks[0]
-        # all but the last chunk have to match exactly
-        for this_chunk in all_var_chunks[:-1]:
-            if this_chunk != first_var_chunk:
-                raise ValueError(
-                    "Zarr requires uniform chunk sizes excpet for final chunk."
-                    " Variable %r has incompatible chunks. Consider "
-                    "rechunking using `chunk()`." % (var_chunks,))
-        # last chunk is allowed to be smaller
-        last_var_chunk = all_var_chunks[-1]
-        for len_first, len_last in zip(first_var_chunk, last_var_chunk):
-            if len_last > len_first:
-                raise ValueError(
-                    "Final chunk of Zarr array must be smaller than first. "
-                    "Variable %r has incompatible chunks. Consider rechunking "
-                    "using `chunk()`." % var_chunks)
-        return first_var_chunk
+        if any(len(set(chunks[:-1])) > 1 for chunks in var_chunks):
+            raise ValueError(
+                "Zarr requires uniform chunk sizes excpet for final chunk."
+                " Variable %r has incompatible chunks. Consider "
+                "rechunking using `chunk()`." % (var_chunks,))
+        if any((chunks[0] < chunks[-1]) for chunks in var_chunks):
+            raise ValueError(
+                "Final chunk of Zarr array must be smaller than first. "
+                "Variable %r has incompatible chunks. Consider rechunking "
+                "using `chunk()`." % var_chunks)
+        # return the first chunk for each dimension
+        return tuple(chunk[0] for chunk in var_chunks)
 
     # from here on, we are dealing with user-specified chunks in encoding
     # zarr allows chunks to be an integer, in which case it uses the same chunk

xarray/tests/test_backends.py

@@ -1330,6 +1330,17 @@ def test_auto_chunk(self):
                 # chunk size should be the same as original
                 self.assertEqual(v.chunks, original[k].chunks)
 
+    def test_write_uneven_dask_chunks(self):
+        # regression for GH#2225
+        original = create_test_data().chunk({'dim1': 3, 'dim2': 4, 'dim3': 3})
+
+        with self.roundtrip(
+                original, open_kwargs={'auto_chunk': True}) as actual:
+            for k, v in actual.data_vars.items():
+                print(k)
+                assert v.chunks == actual[k].chunks
+
+
     def test_chunk_encoding(self):
         # These datasets have no dask chunks. All chunking specified in
         # encoding