Ensure parents are created when creating a node

This updates our Array and Group creation methods to
ensure that parents implicitly defined through a nested path are also
created.

To accomplish this semi-safely and efficiently, we require a new
setdefulat method on the Store class.

Author: TomAugspurger

src/zarr/abc/store.py

@@ -173,6 +173,16 @@ async def set(self, key: str, value: Buffer) -> None:
         """
         ...
 
+    @abstractmethod
+    async def setdefault(self, key: str, default: Buffer) -> None:
+        """
+        Store a key with a value of ``default`` if the key is not already present.
+
+        Unlike MutableMapping.default, this method does not provide any way to
+        know whether ``default`` was actually set.
+        """
+        ...
+
     async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
         """
         Insert multiple (key, value) pairs into storage.
@@ -298,9 +308,18 @@ async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -
 
     async def delete(self) -> None: ...
 
+    async def setdefault(self, default: Buffer) -> None: ...
+
 
 async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None:
     if value is None:
         await byte_setter.delete()
     else:
         await byte_setter.set(value)
+
+
+async def setdefault(byte_setter: ByteSetter, value: Buffer | None) -> None:
+    if value is None:
+        await byte_setter.delete()
+    else:
+        await byte_setter.set(value)

src/zarr/codecs/sharding.py

@@ -98,6 +98,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 setdefault(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
@@ -276,7 +282,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
@@ -315,7 +321,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
@@ -603,9 +609,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 intermediates.
+            parents = _build_parents(self)
+
+            for parent in parents:
+                awaitables.extend(
+                    [
+                        (parent.store_path / key).setdefault(value)
+                        for key, value in parent.metadata.to_buffer_dict(
+                            default_buffer_prototype()
+                        ).items()
+                    ]
+                )
+
         await gather(*awaitables)
 
     async def _set_selection(
@@ -2336,3 +2357,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).setdefault(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 setdefault(self, default: Buffer) -> None:
+        await self.store.setdefault(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):
@@ -152,14 +159,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 setdefault(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/memory.py

@@ -85,9 +85,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)}.")
@@ -99,6 +98,11 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
         else:
             self._store_dict[key] = value
 
+    async def setdefault(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:
@@ -208,6 +209,11 @@ async def set_partial_values(
     ) -> None:
         raise NotImplementedError
 
+    async def setdefault(self, key: str, default: Buffer) -> None:
+        # this isn't safe for concurrent writers, but that's probably unavoidable.
+        if not await self.fs._exists(_dereference_path(self.path, key)):
+            await self.set(key, default)
+
     async def list(self) -> AsyncGenerator[str, None]:
         allfiles = await self.fs._find(self.path, detail=False, withdirs=False)
         for onefile in (a.replace(self.path + "/", "") for a in allfiles):

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 setdefault(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
 

src/zarr/testing/store.py

@@ -273,3 +273,20 @@ async def test_list_dir(self, store: S) -> None:
 
         keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
         assert sorted(keys_expected) == sorted(keys_observed)
+
+    async def test_setdefault(self, store: S) -> None:
+        key = "k"
+        data_buf = self.buffer_cls.from_bytes(b"0000")
+        self.set(store, key, data_buf)
+
+        new = self.buffer_cls.from_bytes(b"1111")
+        await store.setdefault("k", new)  # no error
+
+        result = await store.get(key, default_buffer_prototype())
+        assert result == data_buf
+
+        await store.setdefault("k2", new)  # no error
+        await store.get("k2", default_buffer_prototype())
+
+        result = await store.get("k2", default_buffer_prototype())
+        assert result == new