tests

Author: normanrz

src/zarr/core/codec_pipeline.py

@@ -332,12 +332,21 @@ async def write_batch(
         drop_axes: tuple[int, ...] = (),
     ) -> None:
         if self.supports_partial_encode:
-            await self.encode_partial_batch(
-                [
-                    (byte_setter, value[out_selection], chunk_selection, chunk_spec)
-                    for byte_setter, chunk_spec, chunk_selection, out_selection in batch_info
-                ],
-            )
+            # Pass scalar values as is
+            if len(value.shape) == 0:
+                await self.encode_partial_batch(
+                    [
+                        (byte_setter, value, chunk_selection, chunk_spec)
+                        for byte_setter, chunk_spec, chunk_selection, out_selection in batch_info
+                    ],
+                )
+            else:
+                await self.encode_partial_batch(
+                    [
+                        (byte_setter, value[out_selection], chunk_selection, chunk_spec)
+                        for byte_setter, chunk_spec, chunk_selection, out_selection in batch_info
+                    ],
+                )
 
         else:
             # Read existing bytes if not total slice

tests/test_array.py

@@ -20,6 +20,7 @@
     VLenUTF8Codec,
     ZstdCodec,
 )
+from zarr.codecs.sharding import ShardingCodec
 from zarr.core._info import ArrayInfo
 from zarr.core.array import (
     CompressorsLike,
@@ -479,88 +480,87 @@ def test_update_attrs(zarr_format: ZarrFormat) -> None:
     assert arr2.attrs["foo"] == "bar"
 
 
-@pytest.mark.parametrize(("chunks", "shards"), [((2, 2), None), (2, 2), (4, 4)])
+@pytest.mark.parametrize(("chunks", "shards"), [((2, 2), None), ((2, 2), (4, 4))])
 class TestInfo:
-    chunks: tuple[int, int]
-    shards: tuple[int, int] | None
-
-    def __init__(self, chunks: tuple[int, int], shards: tuple[int, int] | None) -> None:
-        self.chunks = chunks
-        self.shards = shards
-
-    def test_info_v2(self) -> None:
-        arr = zarr.create_array(
-            store={}, shape=(8, 8), dtype="f8", chunks=self.chunks, zarr_format=2
-        )
+    def test_info_v2(self, chunks: tuple[int, int], shards: tuple[int, int] | None) -> None:
+        arr = zarr.create_array(store={}, shape=(8, 8), dtype="f8", chunks=chunks, zarr_format=2)
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=2,
             _data_type=np.dtype("float64"),
             _shape=(8, 8),
-            _chunk_shape=self.chunks,
+            _chunk_shape=chunks,
             _shard_shape=None,
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _count_bytes=128,
+            _count_bytes=512,
             _compressor=numcodecs.Zstd(),
         )
         assert result == expected
 
-    def test_info_v3(self) -> None:
-        arr = zarr.create_array(
-            store={}, shape=(8, 8), dtype="f8", chunks=self.chunks, shards=self.shards
-        )
+    def test_info_v3(self, chunks: tuple[int, int], shards: tuple[int, int] | None) -> None:
+        arr = zarr.create_array(store={}, shape=(8, 8), dtype="f8", chunks=chunks, shards=shards)
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=3,
             _data_type=DataType.parse("float64"),
             _shape=(8, 8),
-            _chunk_shape=(2, 2),
+            _chunk_shape=chunks,
+            _shard_shape=shards,
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=[BytesCodec(), ZstdCodec()],
-            _count_bytes=128,
+            _codecs=[BytesCodec(), ZstdCodec()]
+            if shards is None
+            else [ShardingCodec(chunk_shape=chunks, codecs=[BytesCodec(), ZstdCodec()])],
+            _count_bytes=512,
         )
         assert result == expected
 
-    def test_info_complete(self) -> None:
+    def test_info_complete(self, chunks: tuple[int, int], shards: tuple[int, int] | None) -> None:
         arr = zarr.create_array(
             store={},
             shape=(8, 8),
             dtype="f8",
-            chunks=self.chunks,
-            shards=self.shards,
-            compressors=None,
+            chunks=chunks,
+            shards=shards,
+            compressors=(),
         )
         result = arr.info_complete()
         expected = ArrayInfo(
             _zarr_format=3,
             _data_type=DataType.parse("float64"),
             _shape=(8, 8),
-            _chunk_shape=self.chunks,
-            _shard_shape=self.shards,
+            _chunk_shape=chunks,
+            _shard_shape=shards,
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=[BytesCodec()],
-            _count_bytes=128,
+            _codecs=[BytesCodec()] if shards is None else [ShardingCodec(chunk_shape=chunks)],
+            _count_bytes=512,
             _count_chunks_initialized=0,
-            _count_bytes_stored=373,  # the metadata?
+            _count_bytes_stored=373 if shards is None else 578,  # the metadata?
         )
         assert result == expected
 
-        arr[:2, :2] = 10
+        arr[:4, :4] = 10
         result = arr.info_complete()
-        expected = dataclasses.replace(
-            expected, _count_chunks_initialized=1, _count_bytes_stored=405
-        )
+        if shards is None:
+            expected = dataclasses.replace(
+                expected, _count_chunks_initialized=4, _count_bytes_stored=501
+            )
+        else:
+            expected = dataclasses.replace(
+                expected, _count_chunks_initialized=1, _count_bytes_stored=774
+            )
         assert result == expected
 
-    async def test_info_v2_async(self) -> None:
+    async def test_info_v2_async(
+        self, chunks: tuple[int, int], shards: tuple[int, int] | None
+    ) -> None:
         arr = await zarr.api.asynchronous.create_array(
-            store={}, shape=(8, 8), dtype="f8", chunks=self.chunks, zarr_format=2
+            store={}, shape=(8, 8), dtype="f8", chunks=chunks, zarr_format=2
         )
         result = arr.info
         expected = ArrayInfo(
@@ -572,65 +572,76 @@ async def test_info_v2_async(self) -> None:
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _count_bytes=128,
+            _count_bytes=512,
             _compressor=numcodecs.Zstd(),
         )
         assert result == expected
 
-    async def test_info_v3_async(self) -> None:
+    async def test_info_v3_async(
+        self, chunks: tuple[int, int], shards: tuple[int, int] | None
+    ) -> None:
         arr = await zarr.api.asynchronous.create_array(
             store={},
             shape=(8, 8),
             dtype="f8",
-            chunks=self.chunks,
-            shards=self.shards,
+            chunks=chunks,
+            shards=shards,
         )
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=3,
             _data_type=DataType.parse("float64"),
             _shape=(8, 8),
-            _chunk_shape=self.chunks,
-            _shard_shape=self.shards,
+            _chunk_shape=chunks,
+            _shard_shape=shards,
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=[BytesCodec(), ZstdCodec()],
-            _count_bytes=128,
+            _codecs=[BytesCodec(), ZstdCodec()]
+            if shards is None
+            else [ShardingCodec(chunk_shape=chunks, codecs=[BytesCodec(), ZstdCodec()])],
+            _count_bytes=512,
         )
         assert result == expected
 
-    async def test_info_complete_async(self) -> None:
+    async def test_info_complete_async(
+        self, chunks: tuple[int, int], shards: tuple[int, int] | None
+    ) -> None:
         arr = await zarr.api.asynchronous.create_array(
             store={},
             dtype="f8",
             shape=(8, 8),
-            chunks=self.chunks,
-            shards=self.shards,
+            chunks=chunks,
+            shards=shards,
             compressors=None,
         )
         result = await arr.info_complete()
         expected = ArrayInfo(
             _zarr_format=3,
             _data_type=DataType.parse("float64"),
             _shape=(8, 8),
-            _chunk_shape=self.chunks,
-            _shard_shape=self.shards,
+            _chunk_shape=chunks,
+            _shard_shape=shards,
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=[BytesCodec()],
-            _count_bytes=128,
+            _codecs=[BytesCodec()] if shards is None else [ShardingCodec(chunk_shape=chunks)],
+            _count_bytes=512,
             _count_chunks_initialized=0,
-            _count_bytes_stored=373,  # the metadata?
+            _count_bytes_stored=373 if shards is None else 578,  # the metadata?
         )
         assert result == expected
 
-        await arr.setitem((slice(2), slice(2)), 10)
+        await arr.setitem((slice(4), slice(4)), 10)
         result = await arr.info_complete()
-        expected = dataclasses.replace(
-            expected, _count_chunks_initialized=1, _count_bytes_stored=405
-        )
+        if shards is None:
+            expected = dataclasses.replace(
+                expected, _count_chunks_initialized=4, _count_bytes_stored=501
+            )
+        else:
+            expected = dataclasses.replace(
+                expected, _count_chunks_initialized=1, _count_bytes_stored=774
+            )
         assert result == expected
 
 

tests/test_codecs/test_sharding.py

@@ -70,6 +70,35 @@ def test_sharding(
     assert np.array_equal(data, read_data)
 
 
+@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
+@pytest.mark.parametrize("index_location", ["start", "end"])
+@pytest.mark.parametrize("offset", [0, 10])
+def test_sharding_scalar(
+    store: Store,
+    index_location: ShardingCodecIndexLocation,
+    offset: int,
+) -> None:
+    """
+    Test that we can create an array with a sharding codec, write data to that array, and get
+    the same data out via indexing.
+    """
+    spath = StorePath(store)
+
+    arr = zarr.create_array(
+        spath,
+        shape=(128, 128),
+        chunks=(32, 32),
+        shards={"shape": (64, 64), "index_location": index_location},
+        dtype="uint8",
+        fill_value=6,
+        filters=[TransposeCodec(order=order_from_dim("F", 2))],
+        compressors=BloscCodec(cname="lz4"),
+    )
+    arr[:16, :16] = 10  # intentionally write partial chunks
+    read_data = arr[:16, :16]
+    np.testing.assert_array_equal(read_data, 10)
+
+
 @pytest.mark.parametrize("index_location", ["start", "end"])
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
 @pytest.mark.parametrize(