Optimisation for the __contains__ method of storage.LRUStoreCache

[shikharsg]

I have a few TB large dataset with 11 parameters and about 100000 chunks, and am storing it in azure blob using the ABSStore mutable mapping. When I do zarr.open_group(store=store, mode='r') with store as zarr.LRUStoreCache(max_size=2**33, store=zarr.storage.ABSStore('testcontainer', 'mydataset', BLOB_ACCOUNT_NAME, BLOB_ACCOUNT_KEY)), it takes about 45 seconds to open the group. Without the LRU wrapper the open_group operation is instantaneous. I traced the problem to the __contains__ method in LRUStoreCache mutable mapping wrapper(open_group calls the contains_array method). The __contains__ method(here) in LRUStoreCache is implemented by listing all the keys in the mutable mapping of the underlying store, and therefore, all 100000 chunks are listed before checking for existence. In the context of cloud storage this can cause significant overhead.

This is the method of LRUStoreCache now:

    def __contains__(self, key):
        with self._mutex:
            if self._contains_cache is None:
                self._contains_cache = set(self._keys())
            return key in self._contains_cache

when I changed it to this:

    def __contains__(self, key):
        return key in self._store

the open_group operation is almost instantaneous as the __contains__ method of the underlying ABSStore class uses the exists option on azure blob and so doesn't have to list all keys.

[martindurant]

You may be interested in the consolidated metadata PR, which would grab the metadata in a single call, and thereafter access in memory (including for __contains__).

[shikharsg]

Thanks for that link @martindurant. I've been following that PR. I'm a little confused though as I'm new to zarr, so sorry if my questions sound stupid . How will the consolidated metadata feature help in the __contains__ call? Will it be possible to store all the keys of the mutable mapping in the .zmetadata key?

Also, I think the way the __contains__ method is implemented in LRUStoreCache is not efficient as it lists all the keys of the underlying store(and caches them), before checking for existence of the key. In my opinion it would be better to let the underlying store implement the __contains__ call

[martindurant]

Will it be possible to store all the keys of the mutable mapping in the .zmetadata key?

Yes, exactly - all of the non-data keys are written into the single file.

The behaviour of LRU's __contains__ may be reasonable in some cases, depending on how many calls will be made, and how the list of keys is stored. Zarr datasets tend to contain very many keys, so maybe siding with a direct lookup instead of a list of commonly the best thing.

[shikharsg]

True.

Any estimate as to when the consolidated metadata PR would be merged?

[martindurant]

The only technical missing piece in the PR is coverage (the coverage repot confuses me, seems inconsistent across runs), but @alimanfoo, the maintainer, is away at the moment, so likely nothing will happen without his sayso. As you can see in the comments, there was some considerable conversation about the right approach, but the code proposed certainly is one reasonable implementation.

[jakirkham]

For short term solutions, one might use cache_metadata on Arrays and cache_attrs on Arrays and Groups and/or specify that Groups use a different store from their chunk_store. Both of these would provide ways to keep metadata in memory via caching or by explicitly using a different store for it like a DictStore that one writes out at the end.

ref: https://zarr.readthedocs.io/en/stable/api/core.html#module-zarr.core
ref: https://zarr.readthedocs.io/en/stable/api/hierarchy.html#zarr.hierarchy.group
ref: https://zarr.readthedocs.io/en/stable/api/storage.html#zarr.storage.DictStore

[shikharsg]

Ok, I think I wasn't clear enough in explaining the issue. So here's my store object:

store = zarr.LRUStoreCache(max_size=2**33, store=zarr.storage.ABSStore('testcontainer', 'mydataset', BLOB_ACCOUNT_NAME, BLOB_ACCOUNT_KEY))

Now I call open_group:

zarr.open_group(store=store, mode='r')

Now the open_group method will call the contains_array method(see here), and if you see the contains_array method, it has the line return key in store due to which LRUStoreCache.__contains__ will be called. Now LRUStoreCache.__contains__, because of how it is implemented(see here), will go through ALL the keys of the underlying mutable mapping, not just the *.zgroup... keys, irrespective of whether metadata has been consolidated or not. So if I have a very large number of keys in my mutable mapping, say 5million, it will still go through all of them. This is what's causing the large delay in doing open_group

[martindurant]

Yes, I understand - and the consolidated metadata deals with exactly this issue. The direct exists() method instead of if in keys() for __contains__ would indeed be a better solution in the case of many many keys in general, but relatively few lookups of .z* files - but the opposite is also possible of mostly .*z files, or where the cost of an exists() is high compared to keys() (listdir).

[shikharsg]

Thanks for that @martindurant. I pulled your consolidated metadata branch and just had to figure out the relevant code(that you mentioned here) to add in the open_group function to make it work. Apologies for the hassle.

[alimanfoo]

Hi all, apologies for delayed response.

FWIW I think there's some points worth splitting apart here.

The LRUStoreCache class is intended mainly to help minimise repeated retrieval of data from slow storage. In particular, this includes caching of both metadata and chunk data, if the same store is being used for both (which is usually the case).

The goals of the consolidated metadata PR are somewhat orthogonal, because it is intended primarily to accelerate the initial retrieval of metadata, particularly for libraries like xarray which need to load all metadata up front.

So IMO it's worth considering improvements to the LRUStoreCache class, independently of whether consolidated metadata is being used.

My original implementation of __contains__ clearly isn't optimal, so it's worth considering a change. I'd be happy with the proposed change, given that the main focus is on caching store values rather than keys. I wonder if it's also worth considering a better way of caching the response to a __contains__ call. E.g., something even as simple as the following:

    @functools.lru_cache(maxsize=None)
    def __contains__(self, key):
        return key in self._store

[alimanfoo]

Btw it's also worth noting that when using LRUStoreCache, whether the metadata and data (chunks) are stored in the same store or separate stores will make a big difference, as @jakirkham alludes to.

If both are in the same store, then retrieval of chunk data is likely to push out any metadata in the cache, given that chunks are larger and more numerous than metadata. I.e., LRUStoreCache will likely help with caching chunk data, but probably not metadata.

If you care specifically about caching of chunk data or caching of metadata, it would be better to store metadata and chunk data in separate stores, each of which could be wrapped with its own LRUStoreCache. Then retrieval of chunks will not push out any cached metadata, or vice versa.

Hope I'm not overcomplicating things. Basically, it depends what you care about. Speeding up initial loading of metadata? Speeding up repeated loading of metadata? Speeding up repeated loading of chunk data? Different solutions.

[shikharsg]

I very much like this idea of yours:

    @functools.lru_cache(maxsize=None)
    def __contains__(self, key):
        return key in self._store

Separate stores could also be a solution, depending on the behaviour one wants, that's true.

As a sort of a separate issue: I don't know whether this was intentional, but open_array and open_group currently do not allow passing in a chunk_store argument right? So you have to fall back to either Group or Array, while opening groups or arrays if you want to pass in chunk_store as a separate argument.

[alimanfoo]

As a sort of a separate issue: I don't know whether this was intentional, but open_array and open_group currently do not allow passing in a chunk_store argument right? So you have to fall back to either Group or Array, while opening groups or arrays if you want to pass in chunk_store as a separate argument.

Yes currently you'd have to use Group or Array constructors to pass in chunk_store. Nothing wrong with doing that, they are part of the public API, the open...() functions are intended as convenience. The open...() functions could all support a chunk_store argument, mostly it would be just passing through to the Array() or Group() constructor, although a bit of extra logic would be needed to handle the mode='w' case to ensure that any existing data to be overwritten is cleaned up properly in the chunk store.

…

-- If I do not respond to an email within a few days, please feel free to resend your email and/or contact me by other means. Alistair Miles Head of Epidemiological Informatics Centre for Genomics and Global Health Big Data Institute Li Ka Shing Centre for Health Information and Discovery Old Road Campus Headington Oxford OX3 7LF United Kingdom Phone: +44 (0)1865 743596 or +44 (0)7866 541624 Email: alimanfoo@googlemail.com Web: http://a <http://purl.org/net/aliman>limanfoo.github.io/ Twitter: @alimanfoo <https://twitter.com/alimanfoo>

[jakirkham]

I wonder if we could come up with a single file store (e.g. JSONStore), which would allow all the metadata to live in one place to start with, or would this be problematic? It certainly would disallow multiple simultaneous writes to the metadata, but maybe that is tolerable in this case.

[alimanfoo]

Forgot to mention, there is another trick you can pull re caching of data versus metadata. If you have all data and metadata in the same store, you can still have separate caches for data and metadata. E.g.:

store = ... # some store with all data and metadata
data_cache = zarr.LRUStoreCache(store)
meta_cache = zarr.LRUStoreCache(store)
root = zarr.Group(store=meta_cache, chunk_store=data_cache)

Not sure how practically useful, but does at least provide a way to ensure data reads don't push metadata out of the cache.