Merge branch 'feat/chunks-shards' into feat/read-funcs

Author: normanrz

src/zarr/core/array.py

@@ -803,6 +803,7 @@ def shape(self) -> ChunkCoords:
     @property
     def chunks(self) -> ChunkCoords:
         """Returns the chunk shape of the Array.
+        If sharding is used the inner chunk shape is returned.
 
         Only defined for arrays using using `RegularChunkGrid`.
         If array doesn't use `RegularChunkGrid`, `NotImplementedError` is raised.
@@ -812,14 +813,22 @@ def chunks(self) -> ChunkCoords:
         ChunkCoords:
             The chunk shape of the Array.
         """
-        if isinstance(self.metadata.chunk_grid, RegularChunkGrid):
-            return self.metadata.chunk_grid.chunk_shape
+        return self.metadata.chunks
 
-        msg = (
-            f"The `chunks` attribute is only defined for arrays using `RegularChunkGrid`."
-            f"This array has a {self.metadata.chunk_grid} instead."
-        )
-        raise NotImplementedError(msg)
+    @property
+    def shards(self) -> ChunkCoords | None:
+        """Returns the shard shape of the Array.
+        Returns None if sharding is not used.
+
+        Only defined for arrays using using `RegularChunkGrid`.
+        If array doesn't use `RegularChunkGrid`, `NotImplementedError` is raised.
+
+        Returns
+        -------
+        ChunkCoords:
+            The shard shape of the Array.
+        """
+        return self.metadata.shards
 
     @property
     def size(self) -> int:
@@ -1733,6 +1742,10 @@ def shape(self, value: ChunkCoords) -> None:
     @property
     def chunks(self) -> ChunkCoords:
         """Returns a tuple of integers describing the length of each dimension of a chunk of the array.
+        If sharding is used the inner chunk shape is returned.
+
+        Only defined for arrays using using `RegularChunkGrid`.
+        If array doesn't use `RegularChunkGrid`, `NotImplementedError` is raised.
 
         Returns
         -------
@@ -1741,6 +1754,21 @@ def chunks(self) -> ChunkCoords:
         """
         return self._async_array.chunks
 
+    @property
+    def shards(self) -> ChunkCoords | None:
+        """Returns a tuple of integers describing the length of each dimension of a shard of the array.
+        Returns None if sharding is not used.
+
+        Only defined for arrays using using `RegularChunkGrid`.
+        If array doesn't use `RegularChunkGrid`, `NotImplementedError` is raised.
+
+        Returns
+        -------
+        tuple | None
+            A tuple of integers representing the length of each dimension of a shard or None if sharding is not used.
+        """
+        return self._async_array.shards
+
     @property
     def size(self) -> int:
         """Returns the total number of elements in the array.

src/zarr/core/metadata/v2.py

@@ -44,7 +44,7 @@ class ArrayV2MetadataDict(TypedDict):
 @dataclass(frozen=True, kw_only=True)
 class ArrayV2Metadata(Metadata):
     shape: ChunkCoords
-    chunks: tuple[int, ...]
+    chunks: ChunkCoords
     dtype: np.dtype[Any]
     fill_value: int | float | str | bytes | None = 0
     order: MemoryOrder = "C"
@@ -102,6 +102,10 @@ def ndim(self) -> int:
     def chunk_grid(self) -> RegularChunkGrid:
         return RegularChunkGrid(chunk_shape=self.chunks)
 
+    @property
+    def shards(self) -> ChunkCoords | None:
+        return None
+
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         def _json_convert(
             o: Any,

src/zarr/core/metadata/v3.py

@@ -296,6 +296,40 @@ def dtype(self) -> np.dtype[Any]:
     def ndim(self) -> int:
         return len(self.shape)
 
+    @property
+    def chunks(self) -> ChunkCoords:
+        if isinstance(self.chunk_grid, RegularChunkGrid):
+            from zarr.codecs.sharding import ShardingCodec
+
+            if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec):
+                sharding_codec = self.codecs[0]
+                assert isinstance(sharding_codec, ShardingCodec)  # for mypy
+                return sharding_codec.chunk_shape
+            else:
+                return self.chunk_grid.chunk_shape
+
+        msg = (
+            f"The `chunks` attribute is only defined for arrays using `RegularChunkGrid`."
+            f"This array has a {self.chunk_grid} instead."
+        )
+        raise NotImplementedError(msg)
+
+    @property
+    def shards(self) -> ChunkCoords | None:
+        if isinstance(self.chunk_grid, RegularChunkGrid):
+            from zarr.codecs.sharding import ShardingCodec
+
+            if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec):
+                return self.chunk_grid.chunk_shape
+            else:
+                return None
+
+        msg = (
+            f"The `shards` attribute is only defined for arrays using `RegularChunkGrid`."
+            f"This array has a {self.chunk_grid} instead."
+        )
+        raise NotImplementedError(msg)
+
     def get_chunk_spec(
         self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype
     ) -> ArraySpec:

tests/test_array.py

@@ -928,3 +928,30 @@ def test_auto_partition_auto_shards(
         array_shape=array_shape, chunks=chunk_shape, shards="auto", dtype=dtype
     )
     assert auto_shards == expected_shards
+
+
+def test_chunks_and_shards() -> None:
+    store = StorePath(MemoryStore())
+    shape = (100, 100)
+    chunks = (5, 5)
+    shards = (10, 10)
+
+    arr_v3 = zarr.create_array(store=store / "v3", shape=shape, chunks=chunks, dtype="i4")
+    assert arr_v3.chunks == chunks
+    assert arr_v3.shards is None
+
+    arr_v3_sharding = zarr.create_array(
+        store=store / "v3_sharding",
+        shape=shape,
+        chunks=chunks,
+        shards=shards,
+        dtype="i4",
+    )
+    assert arr_v3_sharding.chunks == chunks
+    assert arr_v3_sharding.shards == shards
+
+    arr_v2 = zarr.create_array(
+        store=store / "v2", shape=shape, chunks=chunks, zarr_format=2, dtype="i4"
+    )
+    assert arr_v2.chunks == chunks
+    assert arr_v2.shards is None