wip

Author: TomAugspurger

src/zarr/abc/store.py

@@ -401,8 +401,8 @@ async def getsize(self, key: str) -> int:
 
         Returns
         -------
-        nbytes: int
-            The size of the value in bytes.
+        nbytes : int
+            The size of the value (in bytes).
 
         Raises
         ------
@@ -417,6 +417,33 @@ async def getsize(self, key: str) -> int:
             raise FileNotFoundError(key)
         return len(value)
 
+    async def getsize_dir(self, prefix: str) -> int:
+        """
+        Return the size, in bytes, of all values in a directory.
+
+        This will include just values whose keys start with ``prefix`` and
+        do not contain the character ``/`` after the given prefix.
+
+        Parameters
+        ----------
+        prefix : str
+            The prefix of the directory to measure.
+
+        Returns
+        -------
+        nbytes : int
+            The sum of the sizes of the values in the directory (in bytes).
+
+        Notes
+        -----
+        ``getsize_dir`` is just provided as a potentially faster alternative to
+        listing all the keys in a directory and calling :meth:`Store.getsize`
+        on each.
+        """
+        keys = [x async for x in self.list_dir(prefix)]
+        sizes = await gather(*[self.getsize(key) for key in keys])
+        return sum(sizes)
+
 
 @runtime_checkable
 class ByteGetter(Protocol):

src/zarr/core/array.py

@@ -843,6 +843,9 @@ def nchunks_initialized(self) -> int:
         """
         return nchunks_initialized(self)
 
+    async def nbytes_stored(self) -> int:
+        return await self.store_path.store.getsize_dir(self.store_path.path)
+
     def _iter_chunk_coords(
         self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
     ) -> Iterator[ChunkCoords]:
@@ -1436,6 +1439,16 @@ def nchunks_initialized(self) -> int:
         """
         return self._async_array.nchunks_initialized
 
+    def nbytes_stored(self) -> int:
+        """
+        Determine the size, in bytes, of the array actually written to the store.
+
+        Returns
+        -------
+        size : int
+        """
+        return sync(self._async_array.nbytes_stored())
+
     def _iter_chunk_keys(
         self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
     ) -> Iterator[str]:

tests/test_array.py

@@ -371,6 +371,15 @@ def test_chunks_initialized(test_cls: type[Array] | type[AsyncArray[Any]]) -> No
         assert observed == expected
 
 
+def test_nbytes_stored():
+    arr = zarr.create(shape=(100,), chunks=(10,), dtype="i4")
+    result = arr.nbytes_stored()
+    assert result == 366  # the size of the metadata document. This is a fragile test
+    arr[:50] = 1
+    result = arr.nbytes_stored()
+    assert result == 366
+
+
 def test_default_fill_values() -> None:
     a = Array.create(MemoryStore({}, mode="w"), shape=5, chunk_shape=5, dtype="<U4")
     assert a.fill_value == ""