Add shards to valid_encodings to enable sharded Zarr writing (#9948)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

xarray/backends/zarr.py

@@ -444,6 +444,7 @@ def extract_zarr_variable_encoding(
     safe_to_drop = {"source", "original_shape", "preferred_chunks"}
     valid_encodings = {
         "chunks",
+        "shards",
         "compressor",  # TODO: delete when min zarr >=3
         "compressors",
         "filters",
@@ -825,6 +826,7 @@ def open_store_variable(self, name):
                 {
                     "compressors": zarr_array.compressors,
                     "filters": zarr_array.filters,
+                    "shards": zarr_array.shards,
                 }
             )
             if self.zarr_group.metadata.zarr_format == 3:

xarray/tests/test_backends.py

@@ -2496,6 +2496,24 @@ def test_chunk_encoding(self) -> None:
             with self.roundtrip(data) as actual:
                 pass
 
+    def test_shard_encoding(self) -> None:
+        # These datasets have no dask chunks. All chunking/sharding specified in
+        # encoding
+        if has_zarr_v3 and zarr.config.config["default_zarr_format"] == 3:
+            data = create_test_data()
+            chunks = (1, 1)
+            shards = (5, 5)
+            data["var2"].encoding.update({"chunks": chunks})
+            data["var2"].encoding.update({"shards": shards})
+            with self.roundtrip(data) as actual:
+                assert shards == actual["var2"].encoding["shards"]
+
+            # expect an error with shards not divisible by chunks
+            data["var2"].encoding.update({"chunks": (2, 2)})
+            with pytest.raises(ValueError):
+                with self.roundtrip(data) as actual:
+                    pass
+
     @requires_dask
     @pytest.mark.skipif(
         ON_WINDOWS,