Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update set_partial_values to accept Buffer rather than BytesLike #2688

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions src/zarr/abc/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down Expand Up @@ -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
Expand Down
5 changes: 1 addition & 4 deletions src/zarr/storage/_fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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], ...] = (
Expand Down Expand Up @@ -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

Expand Down
3 changes: 2 additions & 1 deletion src/zarr/storage/_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
4 changes: 1 addition & 3 deletions src/zarr/storage/_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions src/zarr/storage/_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member Author

@maxrjones maxrjones Jan 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm realizing this may be wrong. I interpreted supports_partial_writes to refer to whether set_partial_values is implemented whereas it may actually refer to whether set includes a byte_range parameter. I am not sure which, if either, is correct because MemoryStore had supports_partial_writes=True with byte_range implemented for set but no implemented set_partial_values whereas LocalStore had supports_partial_writes=True with no byte_range implemented for set but included a set_partial_values implementation.

supports_listing: bool = True

_store_dict: MutableMapping[str, Buffer]
Expand Down Expand Up @@ -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

Expand Down
5 changes: 1 addition & 4 deletions src/zarr/storage/_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/zarr/storage/_zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
17 changes: 17 additions & 0 deletions src/zarr/testing/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_partial_values method.
"""
assert not store.read_only
# Create empty key
await store.set(key, self.buffer_cls.from_bytes(b""))
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

set_partial_values for a LocalStore requires that a key exist. This test assumes this is the intended behavior and creates a key containing no data prior to testing set_partial_values.

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",
[
Expand Down
4 changes: 2 additions & 2 deletions tests/test_store/test_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading