Merge branch 'main' into gh-complete

Author: normanrz

src/zarr/core/array.py

@@ -1346,18 +1346,53 @@ def info(self) -> Any:
         AsyncArray.info_complete
             All information about a group, including dynamic information
             like the number of bytes and chunks written.
+
+        Examples
+        --------
+
+        >>> arr = await zarr.api.asynchronous.create(
+        ...     path="array", shape=(3, 4, 5), chunks=(2, 2, 2))
+        ... )
+        >>> arr.info
+        Type               : Array
+        Zarr format        : 3
+        Data type          : DataType.float64
+        Shape              : (3, 4, 5)
+        Chunk shape        : (2, 2, 2)
+        Order              : C
+        Read-only          : False
+        Store type         : MemoryStore
+        Codecs             : [{'endian': <Endian.little: 'little'>}]
+        No. bytes          : 480
         """
         return self._info()
 
     async def info_complete(self) -> Any:
-        # TODO: get the size of the object from the store.
-        extra = {
-            "count_chunks_initialized": await self.nchunks_initialized(),
-            # count_bytes_stored isn't yet implemented.
-        }
-        return self._info(extra=extra)
-
-    def _info(self, extra: dict[str, int] | None = None) -> Any:
+        """
+        Return all the information for an array, including dynamic information like a storage size.
+
+        In addition to the static information, this provides
+
+        - The count of chunks initialized
+        - The sum of the bytes written
+
+        Returns
+        -------
+        ArrayInfo
+
+        See Also
+        --------
+        AsyncArray.info
+            A property giving just the statically known information about an array.
+        """
+        return self._info(
+            await self.nchunks_initialized(),
+            await self.store_path.store.getsize_prefix(self.store_path.path),
+        )
+
+    def _info(
+        self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None
+    ) -> Any:
         kwargs: dict[str, Any] = {}
         if self.metadata.zarr_format == 2:
             assert isinstance(self.metadata, ArrayV2Metadata)
@@ -1386,6 +1421,8 @@ def _info(self, extra: dict[str, int] | None = None) -> Any:
             _read_only=self.read_only,
             _store_type=type(self.store_path.store).__name__,
             _count_bytes=self.dtype.itemsize * self.size,
+            _count_bytes_stored=count_bytes_stored,
+            _count_chunks_initialized=count_chunks_initialized,
             **kwargs,
         )
 

src/zarr/core/indexing.py

@@ -1346,8 +1346,15 @@ def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords:
 
 
 def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]:
-    for i in range(product(chunk_shape)):
-        yield decode_morton(i, chunk_shape)
+    i = 0
+    order: list[ChunkCoords] = []
+    while len(order) < product(chunk_shape):
+        m = decode_morton(i, chunk_shape)
+        if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)):
+            order.append(m)
+        i += 1
+    for j in range(product(chunk_shape)):
+        yield order[j]
 
 
 def c_order_iter(chunks_per_shard: ChunkCoords) -> Iterator[ChunkCoords]:

tests/test_array.py

@@ -1,3 +1,4 @@
+import dataclasses
 import json
 import math
 import pickle
@@ -474,6 +475,87 @@ def test_info_v3(self) -> None:
         )
         assert result == expected
 
+    def test_info_complete(self) -> None:
+        arr = zarr.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = arr.info_complete()
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+            _count_chunks_initialized=0,
+            _count_bytes_stored=373,  # the metadata?
+        )
+        assert result == expected
+
+        arr[:2, :2] = 10
+        result = arr.info_complete()
+        expected = dataclasses.replace(
+            expected, _count_chunks_initialized=1, _count_bytes_stored=405
+        )
+        assert result == expected
+
+    async def test_info_v2_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=2)
+        result = arr.info
+        expected = ArrayInfo(
+            _zarr_format=2,
+            _data_type=np.dtype("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _count_bytes=128,
+        )
+        assert result == expected
+
+    async def test_info_v3_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = arr.info
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+        )
+        assert result == expected
+
+    async def test_info_complete_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = await arr.info_complete()
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+            _count_chunks_initialized=0,
+            _count_bytes_stored=373,  # the metadata?
+        )
+        assert result == expected
+
+        await arr.setitem((slice(2), slice(2)), 10)
+        result = await arr.info_complete()
+        expected = dataclasses.replace(
+            expected, _count_chunks_initialized=1, _count_bytes_stored=405
+        )
+        assert result == expected
+
 
 @pytest.mark.parametrize("store", ["memory"], indirect=True)
 @pytest.mark.parametrize("zarr_format", [2, 3])

tests/test_codecs/test_codecs.py

@@ -204,6 +204,26 @@ def test_morton() -> None:
     ]
 
 
+@pytest.mark.parametrize(
+    "shape",
+    [
+        [2, 2, 2],
+        [5, 2],
+        [2, 5],
+        [2, 9, 2],
+        [3, 2, 12],
+        [2, 5, 1],
+        [4, 3, 6, 2, 7],
+        [3, 2, 1, 6, 4, 5, 2],
+    ],
+)
+def test_morton2(shape) -> None:
+    order = list(morton_order_iter(shape))
+    for i, x in enumerate(order):
+        assert x not in order[:i]  # no duplicates
+        assert all(x[j] < shape[j] for j in range(len(shape)))  # all indices are within bounds
+
+
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 def test_write_partial_chunks(store: Store) -> None:
     data = np.arange(0, 256, dtype="uint16").reshape((16, 16))

tests/test_codecs/test_sharding.py

@@ -393,3 +393,32 @@ async def test_sharding_with_empty_inner_chunk(
     print("read data")
     data_read = await a.getitem(...)
     assert np.array_equal(data_read, data)
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+@pytest.mark.parametrize(
+    "index_location",
+    [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end],
+)
+@pytest.mark.parametrize("chunks_per_shard", [(5, 2), (2, 5), (5, 5)])
+async def test_sharding_with_chunks_per_shard(
+    store: Store, index_location: ShardingCodecIndexLocation, chunks_per_shard: tuple[int]
+) -> None:
+    chunk_shape = (2, 1)
+    shape = [x * y for x, y in zip(chunks_per_shard, chunk_shape, strict=False)]
+    data = np.ones(np.prod(shape), dtype="int32").reshape(shape)
+    fill_value = 42
+
+    path = f"test_sharding_with_chunks_per_shard_{index_location}"
+    spath = StorePath(store, path)
+    a = Array.create(
+        spath,
+        shape=shape,
+        chunk_shape=shape,
+        dtype="int32",
+        fill_value=fill_value,
+        codecs=[ShardingCodec(chunk_shape=chunk_shape, index_location=index_location)],
+    )
+    a[...] = data
+    data_read = a[...]
+    assert np.array_equal(data_read, data)

tests/test_store/test_logging.py

@@ -5,7 +5,6 @@
 import pytest
 
 import zarr
-import zarr.storage
 from zarr.core.buffer import default_buffer_prototype
 from zarr.storage.logging import LoggingStore