Check that template dims and zarr_chunks match.

Otherwise, mismatched chunks can hide surprising silent errors. PiperOrigin-RevId: 638036862
google · May 28, 2024 · 347d954 · 347d954
1 parent 0d0cdf9
commit 347d954
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/xarray_beam/_src/zarr.py b/xarray_beam/_src/zarr.py
@@ -243,6 +243,11 @@ def _infer_zarr_chunks(dataset: xarray.Dataset) -> Dict[str, int]:
 def _setup_zarr(template, store, zarr_chunks):
   """Setup a Zarr store."""
   if zarr_chunks is not None:
+    if zarr_chunks.keys() != template.dims.keys():
+      raise ValueError(
+          'template dimensions do not match keys in zarr_chunks: '
+          f'{template.dims.keys()} vs. {zarr_chunks.keys()}'
+      )
     template = _override_chunks(template, zarr_chunks)
   _verify_template_is_lazy(template)
   # inconsistent chunks in encoding can lead to spurious failures in xarray:

diff --git a/xarray_beam/_src/zarr_test.py b/xarray_beam/_src/zarr_test.py
@@ -163,6 +163,13 @@ def test_chunks_to_zarr(self):
     ):
       inputs2 | xbeam.ChunksToZarr(temp_dir, template)
 
+    temp_dir = self.create_tempdir().full_path
+    with self.assertRaisesRegex(
+        ValueError,
+        'template dimensions do not match keys in zarr_chunks',
+    ):
+      xbeam.ChunksToZarr(temp_dir, chunked, zarr_chunks={})
+
   def test_multiple_vars_chunks_to_zarr(self):
     dataset = xarray.Dataset(
         {