Feature: store learns to delete prefixes when overwriting/creating hierarchy nodes (#2430)

* fix(array): thread order parameter through to array __init__

* type fixes

* move defaults

* apply MemoryOrder type to ArrayV2Metadata

* more more

* more more more

* feature(store,group,array): stores learn to delete prefixes when
overwriting nodes

- add Store.delete_dir and Store.delete_prefix
- update array and group creation methods to call delete_dir
- change list_prefix to return absolue keys

* fixup

* fixup

* respond to review

* fixup

* fixup

* Update src/zarr/abc/store.py

* style: pre-commit fixes

---------

Co-authored-by: Davis Bennett <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/zarr/abc/store.py

@@ -342,8 +342,8 @@ def list(self) -> AsyncGenerator[str, None]:
     @abstractmethod
     def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         """
-        Retrieve all keys in the store that begin with a given prefix. Keys are returned with the
-        common leading prefix removed.
+        Retrieve all keys in the store that begin with a given prefix. Keys are returned relative
+        to the root of the store.
 
         Parameters
         ----------
@@ -371,6 +371,20 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
         ...
 
+    async def delete_dir(self, prefix: str) -> None:
+        """
+        Remove all keys and prefixes in the store that begin with a given prefix.
+        """
+        if not self.supports_deletes:
+            raise NotImplementedError
+        if not self.supports_listing:
+            raise NotImplementedError
+        self._check_writable()
+        if not prefix.endswith("/"):
+            prefix += "/"
+        async for key in self.list_prefix(prefix):
+            await self.delete(key)
+
     def close(self) -> None:
         """Close the store."""
         self._is_open = False

src/zarr/core/array.py

@@ -553,7 +553,12 @@ async def _create_v3(
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
     ) -> AsyncArray[ArrayV3Metadata]:
-        if not exists_ok:
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir()
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=3)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=3)
 
         shape = parse_shapelike(shape)
@@ -605,7 +610,12 @@ async def _create_v2(
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
     ) -> AsyncArray[ArrayV2Metadata]:
-        if not exists_ok:
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir()
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=2)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=2)
 
         if order is None:

src/zarr/core/group.py

@@ -404,7 +404,13 @@ async def from_store(
         zarr_format: ZarrFormat = 3,
     ) -> AsyncGroup:
         store_path = await make_store_path(store)
-        if not exists_ok:
+
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir()
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=zarr_format)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=zarr_format)
         attributes = attributes or {}
         group = cls(
@@ -727,19 +733,8 @@ def _getitem_consolidated(
 
     async def delitem(self, key: str) -> None:
         store_path = self.store_path / key
-        if self.metadata.zarr_format == 3:
-            await (store_path / ZARR_JSON).delete()
-
-        elif self.metadata.zarr_format == 2:
-            await asyncio.gather(
-                (store_path / ZGROUP_JSON).delete(),  # TODO: missing_ok=False
-                (store_path / ZARRAY_JSON).delete(),  # TODO: missing_ok=False
-                (store_path / ZATTRS_JSON).delete(),  # TODO: missing_ok=True
-            )
-
-        else:
-            raise ValueError(f"unexpected zarr_format: {self.metadata.zarr_format}")
 
+        await store_path.delete_dir()
         if self.metadata.consolidated_metadata:
             self.metadata.consolidated_metadata.metadata.pop(key, None)
             await self._save_metadata()

src/zarr/storage/common.py

@@ -101,6 +101,15 @@ async def delete(self) -> None:
         """
         await self.store.delete(self.path)
 
+    async def delete_dir(self) -> None:
+        """
+        Delete all keys with the given prefix from the store.
+        """
+        path = self.path
+        if not path.endswith("/"):
+            path += "/"
+        await self.store.delete_dir(path)
+
     async def set_if_not_exists(self, default: Buffer) -> None:
         """
         Store a key to ``value`` if the key is not already present.

src/zarr/storage/local.py

@@ -226,9 +226,8 @@ async def list(self) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        to_strip = (
-            (self.root / prefix).as_posix() + "/"
-        )  # TODO: fixme in https://github.com/zarr-developers/zarr-python/issues/2438
+        to_strip = self.root.as_posix() + "/"
+        prefix = prefix.rstrip("/")
         for p in (self.root / prefix).rglob("*"):
             if p.is_file():
                 yield p.as_posix().replace(to_strip, "")

src/zarr/storage/logging.py

@@ -222,6 +222,11 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
             async for key in self._store.list_dir(prefix=prefix):
                 yield key
 
+    async def delete_dir(self, prefix: str) -> None:
+        # docstring inherited
+        with self.log(prefix):
+            await self._store.delete_dir(prefix=prefix)
+
     def with_mode(self, mode: AccessModeLiteral) -> Self:
         # docstring inherited
         with self.log(mode):

src/zarr/storage/memory.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from logging import getLogger
 from typing import TYPE_CHECKING, Self
 
 from zarr.abc.store import ByteRangeRequest, Store
@@ -14,6 +15,9 @@
     from zarr.core.common import AccessModeLiteral
 
 
+logger = getLogger(__name__)
+
+
 class MemoryStore(Store):
     """
     In-memory store for testing purposes.
@@ -137,7 +141,7 @@ async def delete(self, key: str) -> None:
         try:
             del self._store_dict[key]
         except KeyError:
-            pass
+            logger.debug("Key %s does not exist.", key)
 
     async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
         # docstring inherited
@@ -150,9 +154,10 @@ async def list(self) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        for key in self._store_dict:
+        # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix)
+        for key in list(self._store_dict):
             if key.startswith(prefix):
-                yield key.removeprefix(prefix)
+                yield key
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited

src/zarr/storage/remote.py

@@ -340,6 +340,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        find_str = f"{self.path}/{prefix}"
-        for onefile in await self.fs._find(find_str, detail=False, maxdepth=None, withdirs=False):
-            yield onefile.removeprefix(find_str)
+        for onefile in await self.fs._find(
+            f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False
+        ):
+            yield onefile.removeprefix(f"{self.path}/")

src/zarr/storage/zip.py

@@ -244,7 +244,7 @@ async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
         async for key in self.list():
             if key.startswith(prefix):
-                yield key.removeprefix(prefix)
+                yield key
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited

src/zarr/testing/store.py

@@ -213,11 +213,26 @@ async def test_exists(self, store: S) -> None:
         assert await store.exists("foo/zarr.json")
 
     async def test_delete(self, store: S) -> None:
+        if not store.supports_deletes:
+            pytest.skip("store does not support deletes")
         await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar"))
         assert await store.exists("foo/zarr.json")
         await store.delete("foo/zarr.json")
         assert not await store.exists("foo/zarr.json")
 
+    async def test_delete_dir(self, store: S) -> None:
+        if not store.supports_deletes:
+            pytest.skip("store does not support deletes")
+        await store.set("zarr.json", self.buffer_cls.from_bytes(b"root"))
+        await store.set("foo-bar/zarr.json", self.buffer_cls.from_bytes(b"root"))
+        await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar"))
+        await store.set("foo/c/0", self.buffer_cls.from_bytes(b"chunk"))
+        await store.delete_dir("foo")
+        assert await store.exists("zarr.json")
+        assert await store.exists("foo-bar/zarr.json")
+        assert not await store.exists("foo/zarr.json")
+        assert not await store.exists("foo/c/0")
+
     async def test_empty(self, store: S) -> None:
         assert await store.empty()
         await self.set(
@@ -249,8 +264,7 @@ async def test_list(self, store: S) -> None:
     async def test_list_prefix(self, store: S) -> None:
         """
         Test that the `list_prefix` method works as intended. Given a prefix, it should return
-        all the keys in storage that start with this prefix. Keys should be returned with the shared
-        prefix removed.
+        all the keys in storage that start with this prefix.
         """
         prefixes = ("", "a/", "a/b/", "a/b/c/")
         data = self.buffer_cls.from_bytes(b"")
@@ -264,7 +278,7 @@ async def test_list_prefix(self, store: S) -> None:
             expected: tuple[str, ...] = ()
             for key in store_dict:
                 if key.startswith(prefix):
-                    expected += (key.removeprefix(prefix),)
+                    expected += (key,)
             expected = tuple(sorted(expected))
             assert observed == expected