From acce63e13d2372f946ab05df5b0c739c961c2218 Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Sat, 11 Jan 2025 17:55:09 -0500 Subject: [PATCH 1/3] Update set_partial_values to accept Buffer rather than BytesLike --- src/zarr/abc/store.py | 5 +---- src/zarr/storage/_fsspec.py | 5 +---- src/zarr/storage/_local.py | 3 ++- src/zarr/storage/_logging.py | 4 +--- src/zarr/storage/_memory.py | 4 ++-- src/zarr/storage/_wrapper.py | 5 +---- src/zarr/storage/_zip.py | 2 +- src/zarr/testing/store.py | 17 +++++++++++++++++ tests/test_store/test_memory.py | 2 +- 9 files changed, 27 insertions(+), 20 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index e6a5518a4b..088fdc37f1 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -16,7 +16,6 @@ from typing import Any, Self, TypeAlias from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.common import BytesLike __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"] @@ -285,9 +284,7 @@ def supports_partial_writes(self) -> bool: ... @abstractmethod - async def set_partial_values( - self, key_start_values: Iterable[tuple[str, int, BytesLike]] - ) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: """Store values at a given key, starting at byte range_start. Parameters diff --git a/src/zarr/storage/_fsspec.py b/src/zarr/storage/_fsspec.py index 99c8c778e7..265ece6808 100644 --- a/src/zarr/storage/_fsspec.py +++ b/src/zarr/storage/_fsspec.py @@ -18,7 +18,6 @@ from fsspec.asyn import AsyncFileSystem from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.common import BytesLike ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( @@ -316,9 +315,7 @@ async def get_partial_values( return [None if isinstance(r, Exception) else prototype.buffer.from_bytes(r) for r in res] - async def set_partial_values( - self, key_start_values: Iterable[tuple[str, int, BytesLike]] - ) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: # docstring inherited raise NotImplementedError diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index 5eaa85c592..c252cb7c8e 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -174,7 +174,8 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None: await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive) async def set_partial_values( - self, key_start_values: Iterable[tuple[str, int, bytes | bytearray | memoryview]] + self, + key_start_values: Iterable[tuple[str, int, Buffer]], ) -> None: # docstring inherited self._check_writable() diff --git a/src/zarr/storage/_logging.py b/src/zarr/storage/_logging.py index 5ca716df2c..4a0fa875f4 100644 --- a/src/zarr/storage/_logging.py +++ b/src/zarr/storage/_logging.py @@ -197,9 +197,7 @@ async def delete(self, key: str) -> None: with self.log(key): return await self._store.delete(key=key) - async def set_partial_values( - self, key_start_values: Iterable[tuple[str, int, bytes | bytearray | memoryview]] - ) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: # docstring inherited keys = ",".join([k[0] for k in key_start_values]) with self.log(keys): diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index d35ecbe33d..2b85b6e00c 100644 --- a/src/zarr/storage/_memory.py +++ b/src/zarr/storage/_memory.py @@ -38,7 +38,7 @@ class MemoryStore(Store): supports_writes: bool = True supports_deletes: bool = True - supports_partial_writes: bool = True + supports_partial_writes: bool = False supports_listing: bool = True _store_dict: MutableMapping[str, Buffer] @@ -134,7 +134,7 @@ async def delete(self, key: str) -> None: except KeyError: logger.debug("Key %s does not exist.", key) - async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: # docstring inherited raise NotImplementedError diff --git a/src/zarr/storage/_wrapper.py b/src/zarr/storage/_wrapper.py index 255e965439..1b9a804ed9 100644 --- a/src/zarr/storage/_wrapper.py +++ b/src/zarr/storage/_wrapper.py @@ -9,7 +9,6 @@ from zarr.abc.store import ByteRequest from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.common import BytesLike from zarr.abc.store import Store @@ -108,9 +107,7 @@ async def delete(self, key: str) -> None: def supports_partial_writes(self) -> bool: return self._store.supports_partial_writes - async def set_partial_values( - self, key_start_values: Iterable[tuple[str, int, BytesLike]] - ) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: return await self._store.set_partial_values(key_start_values) @property diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index e808b80e4e..29f3e1ad7a 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -209,7 +209,7 @@ async def set(self, key: str, value: Buffer) -> None: with self._lock: self._set(key, value) - async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None: + async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, Buffer]]) -> None: raise NotImplementedError async def set_if_not_exists(self, key: str, value: Buffer) -> None: diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 602d001693..1a631cbc03 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -181,6 +181,23 @@ async def test_set_many(self, store: S) -> None: for k, v in store_dict.items(): assert (await self.get(store, k)).to_bytes() == v.to_bytes() + @pytest.mark.parametrize("key", ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"]) + async def test_set_partial_values(self, store: S, key: str) -> None: + """ + Ensure that data can be written to the store using the store.set method. + """ + assert not store.read_only + # Create empty key + await store.set(key, self.buffer_cls.from_bytes(b"")) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + if store.supports_partial_writes: + await store.set_partial_values([(key, 0, data_buf[:2]), (key, 2, data_buf[2:])]) + observed = await self.get(store, key) + assert_bytes_equal(observed.to_bytes(), data_buf) + else: + with pytest.raises(NotImplementedError): + await store.set_partial_values([(key, 0, data_buf[:2]), (key, 2, data_buf[2:])]) + @pytest.mark.parametrize( "key_ranges", [ diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index ba38889b52..9f34f04ca6 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -41,7 +41,7 @@ def test_store_supports_listing(self, store: MemoryStore) -> None: assert store.supports_listing def test_store_supports_partial_writes(self, store: MemoryStore) -> None: - assert store.supports_partial_writes + assert not store.supports_partial_writes def test_list_prefix(self, store: MemoryStore) -> None: assert True From 655bccdafbd5dd9b77db70bf706815feef3cdb9f Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Sat, 11 Jan 2025 18:06:02 -0500 Subject: [PATCH 2/3] Fix docstring --- src/zarr/testing/store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 1a631cbc03..c5dcab2d01 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -184,7 +184,7 @@ async def test_set_many(self, store: S) -> None: @pytest.mark.parametrize("key", ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"]) async def test_set_partial_values(self, store: S, key: str) -> None: """ - Ensure that data can be written to the store using the store.set method. + Ensure that data can be written to the store using the store.set_partial_values method. """ assert not store.read_only # Create empty key From 260cdd8328c3f59df62d28443f6d33ba064579a6 Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Sat, 11 Jan 2025 18:13:56 -0500 Subject: [PATCH 3/3] Update GPU test to match not supporting partial writes --- tests/test_store/test_memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index 9f34f04ca6..c70bd6f0b2 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -81,7 +81,7 @@ def test_store_supports_listing(self, store: GpuMemoryStore) -> None: assert store.supports_listing def test_store_supports_partial_writes(self, store: GpuMemoryStore) -> None: - assert store.supports_partial_writes + assert not store.supports_partial_writes def test_list_prefix(self, store: GpuMemoryStore) -> None: assert True