Skip to content

Commit

Permalink
feature(group): add group setitem api (#2393)
Browse files Browse the repository at this point in the history
* feature(group): add group setitem api

* arrays proxy

* rollback to simple version

* rollback deprecation

* rollback ...

* Update tests/test_group.py
  • Loading branch information
jhamman authored Oct 24, 2024
1 parent 109f71f commit f4af51c
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 7 deletions.
8 changes: 6 additions & 2 deletions src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -396,12 +396,16 @@ async def save_array(

mode = kwargs.pop("mode", None)
store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
if np.isscalar(arr):
arr = np.array(arr)
shape = arr.shape
chunks = getattr(arr, "chunks", None) # for array-likes with chunks attribute
new = await AsyncArray.create(
store_path,
zarr_format=zarr_format,
shape=arr.shape,
shape=shape,
dtype=arr.dtype,
chunks=arr.shape,
chunks=chunks,
**kwargs,
)
await new.setitem(slice(None), arr)
Expand Down
24 changes: 22 additions & 2 deletions src/zarr/core/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -600,6 +600,23 @@ def from_dict(
store_path=store_path,
)

async def setitem(self, key: str, value: Any) -> None:
"""Fastpath for creating a new array
New arrays will be created with default array settings for the array type.
Parameters
----------
key : str
Array name
value : array-like
Array data
"""
path = self.store_path / key
await async_api.save_array(
store=path, arr=value, zarr_format=self.metadata.zarr_format, exists_ok=True
)

async def getitem(
self,
key: str,
Expand Down Expand Up @@ -1394,8 +1411,11 @@ def __len__(self) -> int:
return self.nmembers()

def __setitem__(self, key: str, value: Any) -> None:
"""__setitem__ is not supported in v3"""
raise NotImplementedError
"""Fastpath for creating a new array.
New arrays will be created using default settings for the array type.
"""
self._sync(self._async_group.setitem(key, value))

def __repr__(self) -> str:
return f"<Group {self.store_path}>"
Expand Down
5 changes: 4 additions & 1 deletion src/zarr/storage/zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,10 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:

async def delete(self, key: str) -> None:
# docstring inherited
raise NotImplementedError
# we choose to only raise NotImplementedError here if the key exists
# this allows the array/group APIs to avoid the overhead of existence checks
if await self.exists(key):
raise NotImplementedError

async def exists(self, key: str) -> bool:
# docstring inherited
Expand Down
21 changes: 19 additions & 2 deletions tests/test_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -427,8 +427,25 @@ def test_group_setitem(store: Store, zarr_format: ZarrFormat) -> None:
Test the `Group.__setitem__` method.
"""
group = Group.from_store(store, zarr_format=zarr_format)
with pytest.raises(NotImplementedError):
group["key"] = 10
arr = np.ones((2, 4))
group["key"] = arr
assert list(group.array_keys()) == ["key"]
assert group["key"].shape == (2, 4)
np.testing.assert_array_equal(group["key"][:], arr)

if store.supports_deletes:
key = "key"
else:
# overwriting with another array requires deletes
# for stores that don't support this, we just use a new key
key = "key2"

# overwrite with another array
arr = np.zeros((3, 5))
group[key] = arr
assert key in list(group.array_keys())
assert group[key].shape == (3, 5)
np.testing.assert_array_equal(group[key], arr)


def test_group_contains(store: Store, zarr_format: ZarrFormat) -> None:
Expand Down

0 comments on commit f4af51c

Please sign in to comment.