Feat: improves delete_dir for s3fs-backed FsspecStore

[carshadi]

Improves performance of FsspecStore.delete_dir when underlying fs is s3fs, by passing a list of filepaths to be removed in bulk instead of one-by-one.

Resolves #2659

TODO:

• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/tutorial.rst
• Changes documented in docs/release.rst
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[martindurant]

I have reservations for this code. It seems to me, that calling ._rm() should be all that's required, and let fsspec handle everything else.

[martindurant]

This seems like a bad idea. We

• shouldn't import something the user may not need
• shouldn't special-case one backend
• make this functionality available only in one case, even though it's implemented for several.

[carshadi]

I agree on all points, the reason I didn't call ._rm(path, recursive=True) without checking the specific filesystem instance was that I wasn't sure whether the signature is the same for all fsspec backends. For example, it seems that ossfs does not support passing a top-level path str with recursive=True https://github.com/fsspec/ossfs/blob/3bac1f556fa06fdd414f6181aa153d1c93e3a460/src/ossfs/core.py#L332

Is there a cleaner way to make this functionality available for all supported backends? One thing that came to mind was subclassing FsspecStore, e.g., S3fsStore, but I figured that would be unpopular 😃

[martindurant]

It _rm is a coroutine, call it (defined as async def _rm(self, path, recursive=False, batch_size=None, **kwargs)). If not, call rm(). How this interacts with the "make it async process" method - ossfs is sync - I am not sure.

Issues and PRs should be made to increase consistency between fsspec backend implementations.

[carshadi]

Perhaps something like this?

    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('This method is only available for stores that support deletes.')
        if not self.supports_listing:
            raise NotImplementedError('This method is only available for stores that support directory listing.')
        self._check_writable()

        # Safely construct the path/prefix to delete
        # This helps avoid duplicated slashes and handles edge cases with empty prefix
        # (e.g., "path/" + "/prefix" => "path/prefix")
        def _path_join(*parts):
            return "/".join(s.strip("/") for s in parts if s)

        path_to_delete = _path_join(self.path, prefix)

        if hasattr(self.fs, "_rm") and inspect.iscoroutinefunction(self.fs._rm):
            try:
                await self.fs._rm(path_to_delete, recursive=True)
            except self.allowed_exceptions:
                pass
        else:
            try:
                loop = asyncio.get_running_loop()
                # Run the sync operation in a thread to avoid blocking the event loop
                await loop.run_in_executor(None, functools.partial(self.fs.rm, path=path_to_delete, recursive=True))
            except self.allowed_exceptions:
                pass

Although, it seems that FsspecStore only supports classes inheriting from AsyncFileSystem, so maybe handling the synchronous case isn't necessary? See

zarr-python/src/zarr/storage/_fsspec.py

Lines 89 to 90 in 12f6012

	if not self.fs.async_impl:
	raise TypeError("Filesystem needs to support async operations.")

[martindurant]

it seems that FsspecStore only supports classes inheriting from AsyncFileSystem

This is right, but you should explicitly check what happens if that async filsystem is actually a sync one wrapped with https://github.com/fsspec/filesystem_spec/blob/master/fsspec/implementations/asyn_wrapper.py#L32

[carshadi]

Wrapping a local filesystem appears to work as expected, though not sure about others

import time
import os

from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
from fsspec.implementations.local import LocalFileSystem
import zarr
from zarr.core.sync import sync
import numpy as np

from zarr.storage._fsspec import FsspecStore


def main():
    path = r"C:\Users\cameron.arshadi\Desktop\fsspec_test.zarr"
    sync_fs = LocalFileSystem(auto_mkdir=True)
    async_fs = AsyncFileSystemWrapper(sync_fs)
    store = FsspecStore(async_fs, read_only=False, path=path)

    start = time.time()
    z = zarr.create_group(store=store, overwrite=True)
    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(""))
    end = time.time()
    print(f"Time to delete array: {end - start:.2f} seconds")



if __name__ == "__main__":
    main()

Using the following in FsspecStore:

    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('This method is only available for stores that support deletes.')
        if not self.supports_listing:
            raise NotImplementedError('This method is only available for stores that support directory listing.')
        self._check_writable()

        # Safely construct the path/prefix to delete
        # This helps avoid duplicated slashes and handles edge cases with empty prefix
        # (e.g., "path/" + "/prefix" => "path/prefix")
        def _path_join(*parts):
            return "/".join(s.strip("/") for s in parts if s)

        path_to_delete = _path_join(self.path, prefix)

        if hasattr(self.fs, "_rm") and inspect.iscoroutinefunction(self.fs._rm):
            try:
                await self.fs._rm(path_to_delete, recursive=True)
            except self.allowed_exceptions:
                pass
        else:
            raise NotImplementedError("The filesystem does not support async deletes")

Looking at https://github.com/fsspec/filesystem_spec/blob/1d34249f0b043907f86064c274a33454ec670ebe/fsspec/implementations/asyn_wrapper.py#L56-L71
it seems that any wrapped filesystems would have their synchronous rm() methods converted to an async _rm(), so at least in theory the if statement should always evaluate to True. Let me know what you think.

[martindurant]

local is the one we really care about. I wonder if there are cases where _rm() does something necessary in the sync filesystem that needs wrapping - in which case, this will need a fix in async_wrapper, but shouldn't need any further work here.

[carshadi]

Got it, changes made in bfb4397

[jhamman]

Thanks @carshadi for continued efforts here. I'd like to see some fsspec specific tests here if possible.

[jhamman]

What do we think about further simplifying this to:

Suggested change

	# this is probably the same condition as `if self.fs.async_impl`
	if hasattr(self.fs, "_rm") and inspect.iscoroutinefunction(self.fs._rm):
	with suppress(*self.allowed_exceptions):
	await self.fs._rm(path_to_delete, recursive=True)
	else:
	raise NotImplementedError("The store does not support async deletes")
	with suppress(*self.allowed_exceptions):
	await self.fs._rm(path_to_delete, recursive=True)

I'll let @martindurant comment but I would think Fsspec would be responsible for raising a NotImplementedError if this isn't supported by a backend.

[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.