Repository: zarr-developers/zarr-python

Issue:
FsspecStore directory deletion performance improvements
Hi all,

There appears to be a bottleneck in directory delete operations for `FsspecStore`, which inherits the default [delete_dir](https://github.com/zarr-developers/zarr-python/blob/71f635464d6480d28635c7c370cfda84182bae79/src/zarr/abc/store.py#L322) method from `Store` . This method lists keys and deletes them one-by-one, which is rather slow for a large number of keys. Certain fsspec implementations like [s3fs](https://github.com/fsspec/s3fs/blob/51e3c80ef380a82081a171de652e2b699753be2b/s3fs/core.py#L2042) and [gcfs](https://github.com/fsspec/gcsfs/blob/e3ebfb9531aa8c5486b468e5c90e39aca9d80f95/gcsfs/core.py#L1282) allow passing a list of paths to their `_rm` methods, and send delete requests in batches to improve performance and reduce likelihood of throttling errors. As a test, I override the `delete_dir` method in `FsspecStore` [here](https://github.com/zarr-developers/zarr-python/commit/11ac3d9993bbf520c55e93774405bb5a6092829f), and saw a big speedup in array deletion for s3fs-backed stores. I am not suggesting this is the optimal way to go about things, merely offering my observation.

Here is a MRE:
```python
import time

import zarr
from zarr.core.sync import sync
import numpy as np


def main():
    start = time.time()
    z = zarr.open("s3://test-bucket/foo.zarr", mode="w")
    end = time.time()
    print(f"Time to open store: {end - start:.2f} seconds")

    # Create an array and time it
    start = time.time()
    a = z.create_array("foo", shape=(2000, 2000), chunks=(100, 100), dtype=np.uint16)
    a[:] = np.random.randint(0, 1000, size=(2000, 2000), dtype=np.uint16)
    end = time.time()
    print(f"Time to create array: {end - start:.2f} seconds")

    # Delete the array and time it
    start = time.time()
    sync(z.store.delete_dir("foo"))
    end = time.time()
    print(f"Time to delete array: {end - start:.2f} seconds")



if __name__ == "__main__":
    main()
    
```

If using the overridden implementation, the output is:
```
Time to open store: 2.51 seconds
Time to create array: 7.86 seconds
Time to delete array: 1.02 seconds
```

With the original implementation:
```
Time to open store: 2.47 seconds
Time to create array: 7.73 seconds
Time to delete array: 106.09 seconds
```

Discussion:
carshadi: > Thanks @carshadi for continued efforts here. I'd like to see some fsspec specific tests here if possible.

@jhamman added a few test cases. Let me know if there are others that would be useful.
carshadi: Hi @martindurant , @jhamman , let me know if there's anything else I can do here. thanks
martindurant: +1

Patch:
diff --git a/changes/2661.feature.1.rst b/changes/2661.feature.1.rst
new file mode 100644
index 00000000..5d0209c5
--- /dev/null
+++ b/changes/2661.feature.1.rst
@@ -0,0 +1,1 @@
+Improves performance of FsspecStore.delete_dir for remote filesystems supporting concurrent/batched deletes, e.g., s3fs.
\ No newline at end of file
diff --git a/src/zarr/storage/_fsspec.py b/src/zarr/storage/_fsspec.py
index 92c14fcc..1cc7039e 100644
--- a/src/zarr/storage/_fsspec.py
+++ b/src/zarr/storage/_fsspec.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from contextlib import suppress
 from typing import TYPE_CHECKING, Any
 
 from zarr.abc.store import (
@@ -286,6 +287,19 @@ class FsspecStore(Store):
         except self.allowed_exceptions:
             pass
 
+    async def delete_dir(self, prefix: str) -> None:
+        # docstring inherited
+        if not self.supports_deletes:
+            raise NotImplementedError(
+                "This method is only available for stores that support deletes."
+            )
+        self._check_writable()
+
+        path_to_delete = _dereference_path(self.path, prefix)
+
+        with suppress(*self.allowed_exceptions):
+            await self.fs._rm(path_to_delete, recursive=True)
+
     async def exists(self, key: str) -> bool:
         # docstring inherited
         path = _dereference_path(self.path, key)


Test patch:
diff --git a/tests/test_store/test_fsspec.py b/tests/test_store/test_fsspec.py
index 929de378..a710b9e2 100644
--- a/tests/test_store/test_fsspec.py
+++ b/tests/test_store/test_fsspec.py
@@ -217,6 +217,14 @@ class TestFsspecStoreS3(StoreTests[FsspecStore, cpu.Buffer]):
         store = await self.store_cls.open(**store_kwargs)
         assert await store.is_empty("")
 
+    async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None:
+        store.supports_deletes = False
+        with pytest.raises(
+            NotImplementedError,
+            match="This method is only available for stores that support deletes.",
+        ):
+            await store.delete_dir("test_prefix")
+
 
 @pytest.mark.skipif(
     parse_version(fsspec.__version__) < parse_version("2024.12.0"),
@@ -244,3 +252,28 @@ def test_no_wrap_async_filesystem():
 
     assert not isinstance(store.fs, AsyncFileSystemWrapper)
     assert store.fs.async_impl
+
+
+@pytest.mark.skipif(
+    parse_version(fsspec.__version__) < parse_version("2024.12.0"),
+    reason="No AsyncFileSystemWrapper",
+)
+async def test_delete_dir_wrapped_filesystem(tmpdir) -> None:
+    from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
+    from fsspec.implementations.local import LocalFileSystem
+
+    wrapped_fs = AsyncFileSystemWrapper(LocalFileSystem(auto_mkdir=True))
+    store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmpdir}/test/path")
+
+    assert isinstance(store.fs, AsyncFileSystemWrapper)
+    assert store.fs.asynchronous
+
+    await store.set("zarr.json", cpu.Buffer.from_bytes(b"root"))
+    await store.set("foo-bar/zarr.json", cpu.Buffer.from_bytes(b"root"))
+    await store.set("foo/zarr.json", cpu.Buffer.from_bytes(b"bar"))
+    await store.set("foo/c/0", cpu.Buffer.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")