Added Store.getsize

docs/guide/storage.rst

@@ -97,5 +97,13 @@ Zarr-Python :class:`zarr.abc.store.Store` API is meant to be extended. The Store
 Class includes all of the methods needed to be a fully operational store in Zarr Python.
 Zarr also provides a test harness for custom stores: :class:`zarr.testing.store.StoreTests`.
 
+``Store.get``
+~~~~~~~~~~~~~
+
+The ``prototype`` keyword of :func:`zarr.abc.store.Store.get` uses a default of
+``None``.  When given ``None``, implementations should use
+:func:`zarr.buffer.default_buffer_prototype` to look up the prototype users have
+configured.
+
 .. _Zip Store Specification: https://github.com/zarr-developers/zarr-specs/pull/311
 .. _Fsspec: https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#consolidated-metadata

src/zarr/abc/store.py

@@ -5,6 +5,8 @@
 from itertools import starmap
 from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
 
+from zarr.core.buffer.core import default_buffer_prototype
+
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable
     from types import TracebackType
@@ -197,14 +199,17 @@ def __eq__(self, value: object) -> bool:
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         """Retrieve the value associated with a given key.
 
         Parameters
         ----------
         key : str
+        prototype : BufferPrototype, optional
+            The prototype giving the buffer classes to use for buffers and nbuffers.
+            By default, :func:`zarr.buffer.default_buffer_prototype` is used.
         byte_range : tuple[int | None, int | None], optional
 
         Returns
@@ -386,6 +391,60 @@ async def _get_many(
         for req in requests:
             yield (req[0], await self.get(*req))
 
+    async def getsize(self, key: str) -> int:
+        """
+        Return the size, in bytes, of a value in a Store.
+
+        Parameters
+        ----------
+        key : str
+
+        Returns
+        -------
+        nbytes : int
+            The size of the value (in bytes).
+
+        Raises
+        ------
+        FileNotFoundError
+            When the given key does not exist in the store.
+        """
+        # Note to implementers: this default implementation is very inefficient since
+        # it requires reading the entire object. Many systems will have ways to get the
+        # size of an object without reading it.
+        value = await self.get(key, prototype=default_buffer_prototype())
+        if value is None:
+            raise FileNotFoundError(key)
+        return len(value)
+
+    async def getsize_prefix(self, prefix: str) -> int:
+        """
+        Return the size, in bytes, of all values under a 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).
+
+        See Also
+        --------
+        zarr.Array.nbytes_stored
+        Store.getsize
+
+        Notes
+        -----
+        ``getsize_prefix`` is just provided as a potentially faster alternative to
+        listing all the keys under a prefix calling :meth:`Store.getsize` on each.
+        """
+        keys = [x async for x in self.list_prefix(prefix)]
+        sizes = await gather(*[self.getsize(key) for key in keys])
+        return sum(sizes)
+
 
 @runtime_checkable
 class ByteGetter(Protocol):

src/zarr/buffer.py

@@ -0,0 +1,5 @@
+from zarr.core.buffer import default_buffer_prototype
+
+__all__ = [
+    "default_buffer_prototype",
+]

src/zarr/core/array.py

@@ -842,6 +842,9 @@ def nchunks_initialized(self) -> int:
         """
         return nchunks_initialized(self)
 
+    async def nbytes_stored(self) -> int:
+        return await self.store_path.store.getsize_prefix(self.store_path.path)
+
     def _iter_chunk_coords(
         self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
     ) -> Iterator[ChunkCoords]:
@@ -1497,6 +1500,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]:

src/zarr/storage/local.py

@@ -9,6 +9,7 @@
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer
+from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import concurrent_map
 
 if TYPE_CHECKING:
@@ -143,10 +144,12 @@ def __eq__(self, other: object) -> bool:
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: tuple[int | None, int | None] | None = None,
     ) -> Buffer | None:
         # docstring inherited
+        if prototype is None:
+            prototype = default_buffer_prototype()
         if not self._is_open:
             await self._open()
         assert isinstance(key, str)
@@ -242,3 +245,6 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
                 yield str(key).replace(to_strip, "")
         except (FileNotFoundError, NotADirectoryError):
             pass
+
+    async def getsize(self, key: str) -> int:
+        return os.path.getsize(self.root / key)

src/zarr/storage/logging.py

@@ -159,7 +159,7 @@ def __eq__(self, other: object) -> bool:
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: tuple[int | None, int | None] | None = None,
     ) -> Buffer | None:
         # docstring inherited

src/zarr/storage/memory.py

@@ -4,6 +4,7 @@
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer, gpu
+from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import concurrent_map
 from zarr.storage._utils import _normalize_interval_index
 
@@ -79,10 +80,12 @@ def __eq__(self, other: object) -> bool:
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: tuple[int | None, int | None] | None = None,
     ) -> Buffer | None:
         # docstring inherited
+        if prototype is None:
+            prototype = default_buffer_prototype()
         if not self._is_open:
             await self._open()
         assert isinstance(key, str)

src/zarr/storage/remote.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, Self
 
 from zarr.abc.store import ByteRangeRequest, Store
+from zarr.core.buffer.core import default_buffer_prototype
 from zarr.storage.common import _dereference_path
 
 if TYPE_CHECKING:
@@ -217,10 +218,12 @@ def __eq__(self, other: object) -> bool:
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
+        if prototype is None:
+            prototype = default_buffer_prototype()
         if not self._is_open:
             await self._open()
         path = _dereference_path(self.path, key)
@@ -343,3 +346,16 @@ async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         find_str = f"{self.path}/{prefix}"
         for onefile in await self.fs._find(find_str, detail=False, maxdepth=None, withdirs=False):
             yield onefile.removeprefix(find_str)
+
+    async def getsize(self, key: str) -> int:
+        path = _dereference_path(self.path, key)
+        info = await self.fs._info(path)
+
+        size = info.get("size")
+
+        if size is None:
+            # Not all filesystems support size. Fall back to reading the entire object
+            return await super().getsize(key)
+        else:
+            # fsspec doesn't have typing. We'll need to assume or verify this is true
+            return int(size)

src/zarr/storage/zip.py

@@ -9,6 +9,7 @@
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer, BufferPrototype
+from zarr.core.buffer.core import default_buffer_prototype
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable
@@ -166,11 +167,13 @@ def _get(
     async def get(
         self,
         key: str,
-        prototype: BufferPrototype,
+        prototype: BufferPrototype | None = None,
         byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
         assert isinstance(key, str)
+        if prototype is None:
+            prototype = default_buffer_prototype()
 
         with self._lock:
             return self._get(key, prototype=prototype, byte_range=byte_range)

src/zarr/testing/store.py

@@ -121,6 +121,15 @@ async def test_get(
         expected = data_buf[start : start + length]
         assert_bytes_equal(observed, expected)
 
+    async def test_get_default_prototype(self, store: S) -> None:
+        key = "c/0"
+        data = b"\x01\x02\x03\x04"
+        data_buf = self.buffer_cls.from_bytes(data)
+        await self.set(store, key, data_buf)
+        observed = await store.get(key)
+        expected = data_buf[:]
+        assert_bytes_equal(observed, expected)
+
     async def test_get_many(self, store: S) -> None:
         """
         Ensure that multiple keys can be retrieved at once with the _get_many method.
@@ -338,3 +347,15 @@ async def test_set_if_not_exists(self, store: S) -> None:
 
         result = await store.get("k2", default_buffer_prototype())
         assert result == new
+
+    async def test_getsize(self, store: S) -> None:
+        key = "k"
+        data = self.buffer_cls.from_bytes(b"0" * 10)
+        await self.set(store, key, data)
+
+        result = await store.getsize(key)
+        assert result == 10
+
+    async def test_getsize_raises(self, store: S) -> None:
+        with pytest.raises(FileNotFoundError):
+            await store.getsize("not-a-real-key")

tests/test_array.py

@@ -374,6 +374,30 @@ def test_chunks_initialized(test_cls: type[Array] | type[AsyncArray[Any]]) -> No
         assert observed == expected
 
 
+def test_nbytes_stored() -> None:
+    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 == 566  # the size with 5 chunks filled.
+    arr[50:] = 2
+    result = arr.nbytes_stored()
+    assert result == 766  # the size with all chunks filled.
+
+
+async def test_nbytes_stored_async() -> None:
+    arr = await zarr.api.asynchronous.create(shape=(100,), chunks=(10,), dtype="i4")
+    result = await arr.nbytes_stored()
+    assert result == 366  # the size of the metadata document. This is a fragile test.
+    await arr.setitem(slice(50), 1)
+    result = await arr.nbytes_stored()
+    assert result == 566  # the size with 5 chunks filled.
+    await arr.setitem(slice(50, 100), 2)
+    result = await arr.nbytes_stored()
+    assert result == 766  # the size with all chunks filled.
+
+
 def test_default_fill_values() -> None:
     a = Array.create(MemoryStore({}, mode="w"), shape=5, chunk_shape=5, dtype="<U4")
     assert a.fill_value == ""

tests/test_codecs/test_blosc.py

@@ -5,8 +5,8 @@
 
 from zarr import AsyncArray
 from zarr.abc.store import Store
+from zarr.buffer import default_buffer_prototype
 from zarr.codecs import BloscCodec, BytesCodec, ShardingCodec
-from zarr.core.buffer import default_buffer_prototype
 from zarr.storage.common import StorePath
 
 

tests/test_codecs/test_codecs.py

@@ -8,13 +8,13 @@
 import pytest
 
 from zarr import Array, AsyncArray, config
+from zarr.buffer import default_buffer_prototype
 from zarr.codecs import (
     BytesCodec,
     GzipCodec,
     ShardingCodec,
     TransposeCodec,
 )
-from zarr.core.buffer import default_buffer_prototype
 from zarr.core.indexing import Selection, morton_order_iter
 from zarr.storage import StorePath
 

tests/test_codecs/test_sharding.py

@@ -7,14 +7,14 @@
 
 from zarr import Array, AsyncArray
 from zarr.abc.store import Store
+from zarr.buffer import default_buffer_prototype
 from zarr.codecs import (
     BloscCodec,
     BytesCodec,
     ShardingCodec,
     ShardingCodecIndexLocation,
     TransposeCodec,
 )
-from zarr.core.buffer import default_buffer_prototype
 from zarr.storage.common import StorePath
 
 from ..conftest import ArrayRequest

tests/test_group.py

@@ -14,7 +14,7 @@
 import zarr.api.synchronous
 from zarr import Array, AsyncArray, AsyncGroup, Group
 from zarr.abc.store import Store
-from zarr.core.buffer import default_buffer_prototype
+from zarr.buffer import default_buffer_prototype
 from zarr.core.group import ConsolidatedMetadata, GroupMetadata
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError

tests/test_indexing.py

@@ -11,7 +11,7 @@
 from numpy.testing import assert_array_equal
 
 import zarr
-from zarr.core.buffer import BufferPrototype, default_buffer_prototype
+from zarr.buffer import default_buffer_prototype
 from zarr.core.indexing import (
     BasicSelection,
     CoordinateSelection,
@@ -32,6 +32,7 @@
     from collections.abc import AsyncGenerator
 
     from zarr.core.array import Array
+    from zarr.core.buffer import BufferPrototype
     from zarr.core.buffer.core import Buffer
     from zarr.core.common import ChunkCoords
 

tests/test_metadata/test_consolidated.py

@@ -16,7 +16,7 @@
     open,
     open_consolidated,
 )
-from zarr.core.buffer.core import default_buffer_prototype
+from zarr.buffer import default_buffer_prototype
 from zarr.core.group import ConsolidatedMetadata, GroupMetadata
 from zarr.core.metadata import ArrayV3Metadata
 from zarr.core.metadata.v2 import ArrayV2Metadata