Merge branch 'v3' into fix/modes

Author: jhamman

src/zarr/abc/store.py

@@ -197,6 +197,22 @@ async def set(self, key: str, value: Buffer) -> None:
         """
         ...
 
+    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+        """
+        Store a key to ``value`` if the key is not already present.
+
+        Parameters
+        -----------
+        key : str
+        value : Buffer
+        """
+        # Note for implementers: the default implementation provided here
+        # is not safe for concurrent writers. There's a race condition between
+        # the `exists` check and the `set` where another writer could set some
+        # value at `key` or delete `key`.
+        if not await self.exists(key):
+            await self.set(key, value)
+
     async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
         """
         Insert multiple (key, value) pairs into storage.
@@ -322,6 +338,8 @@ async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -
 
     async def delete(self) -> None: ...
 
+    async def set_if_not_exists(self, default: Buffer) -> None: ...
+
 
 async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None:
     if value is None:

src/zarr/codecs/sharding.py

@@ -97,6 +97,9 @@ async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -
     async def delete(self) -> None:
         del self.shard_dict[self.chunk_coords]
 
+    async def set_if_not_exists(self, default: Buffer) -> None:
+        self.shard_dict.setdefault(self.chunk_coords, default)
+
 
 class _ShardIndex(NamedTuple):
     # dtype uint64, shape (chunks_per_shard_0, chunks_per_shard_1, ..., 2)

src/zarr/core/array.py

@@ -13,7 +13,12 @@
 from zarr.codecs import BytesCodec
 from zarr.codecs._v2 import V2Compressor, V2Filters
 from zarr.core.attributes import Attributes
-from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer, default_buffer_prototype
+from zarr.core.buffer import (
+    BufferPrototype,
+    NDArrayLike,
+    NDBuffer,
+    default_buffer_prototype,
+)
 from zarr.core.chunk_grids import RegularChunkGrid, _guess_chunks
 from zarr.core.chunk_key_encodings import (
     ChunkKeyEncoding,
@@ -71,6 +76,7 @@
     from collections.abc import Iterable, Iterator, Sequence
 
     from zarr.abc.codec import Codec, CodecPipeline
+    from zarr.core.group import AsyncGroup
     from zarr.core.metadata.common import ArrayMetadata
 
 # Array and AsyncArray are defined in the base ``zarr`` namespace
@@ -337,7 +343,7 @@ async def _create_v3(
         )
 
         array = cls(metadata=metadata, store_path=store_path)
-        await array._save_metadata(metadata)
+        await array._save_metadata(metadata, ensure_parents=True)
         return array
 
     @classmethod
@@ -376,7 +382,7 @@ async def _create_v2(
             attributes=attributes,
         )
         array = cls(metadata=metadata, store_path=store_path)
-        await array._save_metadata(metadata)
+        await array._save_metadata(metadata, ensure_parents=True)
         return array
 
     @classmethod
@@ -621,9 +627,24 @@ async def getitem(
         )
         return await self._get_selection(indexer, prototype=prototype)
 
-    async def _save_metadata(self, metadata: ArrayMetadata) -> None:
+    async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None:
         to_save = metadata.to_buffer_dict(default_buffer_prototype())
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
+
+        if ensure_parents:
+            # To enable zarr.create(store, path="a/b/c"), we need to create all the intermediate groups.
+            parents = _build_parents(self)
+
+            for parent in parents:
+                awaitables.extend(
+                    [
+                        (parent.store_path / key).set_if_not_exists(value)
+                        for key, value in parent.metadata.to_buffer_dict(
+                            default_buffer_prototype()
+                        ).items()
+                    ]
+                )
+
         await gather(*awaitables)
 
     async def _set_selection(
@@ -2354,3 +2375,21 @@ def chunks_initialized(array: Array | AsyncArray) -> tuple[str, ...]:
             out.append(chunk_key)
 
     return tuple(out)
+
+
+def _build_parents(node: AsyncArray | AsyncGroup) -> list[AsyncGroup]:
+    from zarr.core.group import AsyncGroup, GroupMetadata
+
+    required_parts = node.store_path.path.split("/")[:-1]
+    parents = []
+
+    for i, part in enumerate(required_parts):
+        path = "/".join(required_parts[:i] + [part])
+        parents.append(
+            AsyncGroup(
+                metadata=GroupMetadata(zarr_format=node.metadata.zarr_format),
+                store_path=StorePath(store=node.store_path.store, path=path),
+            )
+        )
+
+    return parents

src/zarr/core/group.py

@@ -13,7 +13,7 @@
 import zarr.api.asynchronous as async_api
 from zarr.abc.metadata import Metadata
 from zarr.abc.store import Store, set_or_delete
-from zarr.core.array import Array, AsyncArray
+from zarr.core.array import Array, AsyncArray, _build_parents
 from zarr.core.attributes import Attributes
 from zarr.core.buffer import default_buffer_prototype
 from zarr.core.common import (
@@ -144,7 +144,7 @@ async def from_store(
             metadata=GroupMetadata(attributes=attributes, zarr_format=zarr_format),
             store_path=store_path,
         )
-        await group._save_metadata()
+        await group._save_metadata(ensure_parents=True)
         return group
 
     @classmethod
@@ -279,9 +279,22 @@ async def delitem(self, key: str) -> None:
         else:
             raise ValueError(f"unexpected zarr_format: {self.metadata.zarr_format}")
 
-    async def _save_metadata(self) -> None:
+    async def _save_metadata(self, ensure_parents: bool = False) -> None:
         to_save = self.metadata.to_buffer_dict(default_buffer_prototype())
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
+
+        if ensure_parents:
+            parents = _build_parents(self)
+            for parent in parents:
+                awaitables.extend(
+                    [
+                        (parent.store_path / key).set_if_not_exists(value)
+                        for key, value in parent.metadata.to_buffer_dict(
+                            default_buffer_prototype()
+                        ).items()
+                    ]
+                )
+
         await asyncio.gather(*awaitables)
 
     @property

src/zarr/store/common.py

@@ -51,6 +51,9 @@ async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -
     async def delete(self) -> None:
         await self.store.delete(self.path)
 
+    async def set_if_not_exists(self, default: Buffer) -> None:
+        await self.store.set_if_not_exists(self.path, default)
+
     async def exists(self) -> bool:
         return await self.store.exists(self.path)
 

src/zarr/store/local.py

@@ -60,6 +60,7 @@ def _put(
     path: Path,
     value: Buffer,
     start: int | None = None,
+    exclusive: bool = False,
 ) -> int | None:
     path.parent.mkdir(parents=True, exist_ok=True)
     if start is not None:
@@ -68,7 +69,13 @@ def _put(
             f.write(value.as_numpy_array().tobytes())
         return None
     else:
-        return path.write_bytes(value.as_numpy_array().tobytes())
+        view = memoryview(value.as_numpy_array().tobytes())
+        if exclusive:
+            mode = "xb"
+        else:
+            mode = "wb"
+        with path.open(mode=mode) as f:
+            return f.write(view)
 
 
 class LocalStore(Store):
@@ -155,14 +162,23 @@ async def get_partial_values(
         return await concurrent_map(args, to_thread, limit=None)  # TODO: fix limit
 
     async def set(self, key: str, value: Buffer) -> None:
+        return await self._set(key, value)
+
+    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+        try:
+            return await self._set(key, value, exclusive=True)
+        except FileExistsError:
+            pass
+
+    async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         if not self._is_open:
             await self._open()
         self._check_writable()
         assert isinstance(key, str)
         if not isinstance(value, Buffer):
             raise TypeError("LocalStore.set(): `value` must a Buffer instance")
         path = self.root / key
-        await to_thread(_put, path, value)
+        await 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]]

src/zarr/store/logging.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Self
 
 from zarr.abc.store import AccessMode, ByteRangeRequest, Store
+from zarr.core.buffer import Buffer
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator, Iterable
@@ -149,6 +150,10 @@ async def set(self, key: str, value: Buffer) -> None:
         with self.log():
             return await self._store.set(key=key, value=value)
 
+    async def set_if_not_exists(self, key: str, default: Buffer) -> None:
+        with self.log():
+            return await self._store.set_if_not_exists(key=key, value=default)
+
     async def delete(self, key: str) -> None:
         with self.log():
             return await self._store.delete(key=key)

src/zarr/store/memory.py

@@ -88,9 +88,8 @@ async def exists(self, key: str) -> bool:
         return key in self._store_dict
 
     async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:
-        if not self._is_open:
-            await self._open()
         self._check_writable()
+        await self._ensure_open()
         assert isinstance(key, str)
         if not isinstance(value, Buffer):
             raise TypeError(f"Expected Buffer. Got {type(value)}.")
@@ -102,6 +101,11 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
         else:
             self._store_dict[key] = value
 
+    async def set_if_not_exists(self, key: str, default: Buffer) -> None:
+        self._check_writable()
+        await self._ensure_open()
+        self._store_dict.setdefault(key, default)
+
     async def delete(self, key: str) -> None:
         self._check_writable()
         try:

src/zarr/store/remote.py

@@ -5,6 +5,7 @@
 import fsspec
 
 from zarr.abc.store import ByteRangeRequest, Store
+from zarr.core.buffer import Buffer
 from zarr.store.common import _dereference_path
 
 if TYPE_CHECKING:

src/zarr/store/zip.py

@@ -191,6 +191,13 @@ async def set(self, key: str, value: Buffer) -> None:
     async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
         raise NotImplementedError
 
+    async def set_if_not_exists(self, key: str, default: Buffer) -> None:
+        self._check_writable()
+        with self._lock:
+            members = self._zf.namelist()
+            if key not in members:
+                self._set(key, default)
+
     async def delete(self, key: str) -> None:
         raise NotImplementedError