RFC: Revised V3 Store Design

[jhamman]

In #1583, we laid out a rough design for the Store API for Zarr-Python 3.0. This discussion proposes some changes to that design and includes a new Store ABC. After some iteration here, I plan to update the design doc (link below).

Background

In the next major version of this library, we’re leaving behind the generic mutable-mapping store interface in favor of a more opinionated store API. In the 3.0 design doc, we outlined a basic API based on the v3-spec. Key changes relative to the prior mutable mapping interface include:

1. All public methods are async
2. API for getting and setting partial values
3. API for Zarr specific operations like getsize and rename

Today we have implemented this API for the MemoryStore, LocalStore, and @kylebarron has even tried it out with the object-store-python project (#1661).

Proposal

Now, after a few weeks of prototyping with this API, I am proposing some changes. I’m going to open each of these as separate threads below to help focus discussion.

1. Expand the API to include metadata and chunk specific methods
2. Support additional metadata-only operations in the store.
3. Require stores to be opened using a read, write, or append modes.
4. Return async generators from all list methods
5. Support partial reads/writes on single and multiple keys at once

[jhamman]

1. Expand the API to include metadata and chunk specific methods

The idea would be to add a analog set of methods for getting/setting/listing metadata documents. Rather than the store always getting or setting metadata as strings of bytes (serialized JSON), we would start passing metadata objects to the store in the form of a dictionary: For example, the metadata_get method signature could look like this:

async def metadata_get(self, key: str) -> dict[str, Any]:
    ...

This is perhaps the biggest change in my proposal but there are some clear analogs to what Zarr-Python supports today with separate store and chunk_store namespaces but takes the concept further to allow for leveraging the store’s underlying way of storing data. Stores that represent dictionaries natively (like the MemoryStoreand the MongoDBStore) could take advantage of this.

I’ll also note that this change is aligned with how [spec](https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#abstract-store-interface) envisioned stores evolving:

The store interface defines a set of operations involving keys and values. In the context of this interface, a key is a Unicode string, where the final character is not a / character. In general, a value is a sequence of bytes. Specific stores may choose more specific storage formats, which must be stated in the specification of the respective store. E.g. a database store might encode values of *.json keys with a database-native json type.

The metadata_* methods would all have default implementations that live in the base store so this would be an opt-in API for stores that can take advantage of it.

[jhamman]

I've started prototyping something along these lines here: d-v-b#111

[jhamman]

2. Support additional metadata-only operations in the store.

Expanding on the ideas in (1), we could also ask the store to do more with metadata keys. The prime example here is generating the hierarchy to support Group.tree(). In this case, the store could provide a tree API that returns a dict of metadata objects:

async def metadata_tree(self, root: str) -> dict[str, dict]:
    ...

The rational for this method is that database type stores are going to be able to produce the tree without iterating through the entire store. The proposed API here would be opt-in with a default implementation in the base store.

[d-v-b]

doesn't this require that the store know the name of the zarr metadata, and if so, isn't the store modeling the zarr spec? This is a departure from the v2 store design insofar as the v2 stores don't depend on the zarr spec at al. I think this is a nice property of the store API to keep.

Maybe we can get the behavior you want by giving stores a search method, that a zarr group would call with the appropriate metadata key, e.g.

async def search(self, path: str, query_str: str) -> dict[str, dict]
    ...
# get all v3 metadata: store.search('/path/to/group', '*zarr.json')

[jhamman]

3. Require stores to be opened using a read, write, or append modes.

Enforcing the mode of operating with a store will allow us to be far more protective against bad behaviors and less universally defensive. It will also open to the door to store-native caching which I am going to save for a separate proposal.

@classmethod
async def open(cls, path='', mode='r', **kwargs) -> Store:
    ... 

[normanrz]

I am in favor of adding mode to the stores.

Which stores need to be closed?

[jhamman]

zarr.open_store(path_or_url, mode) ... Would that clarify semantics?

Yes, but I think we will also want to support something on the Store class itself -- even if it is mostly an internal tool. I suspect your example will cover most users though...

Which stores need to be closed?

ZipStore is the one that we hope to support in Zarr-Python. Other single-object stores probably have the same constraint.

[d-v-b]

I do think the access mode makes sense as an attribute of the storage classes. On first glance I don't have a problem requiring that store classes implement ContextManager, even if most stores are stateless. We would have to consider if there's any other use for ContextManager to anticipate here, e.g. batching store operations in a transaction.

store = MyStore().open(mode='r')
foo = open_group(store, path='foo')

I'm not a fan of two-stage initiation. What's the use of MyStore() when .open has not been called?

[martindurant]

I don't like MyStore().open(mode='r'), honestly. Would that make MyStore() alone not able to do anything? What happens to a store if it is re-opened in different threads with different modes at the same time?

[martindurant]

How to enforce closure? Do some stores require calling store.close().

fsspec stores generally don't need this, or might do some finalize work when dropped. So for the open_group(...) route, this would amount to cleanup happening when all the groups and arrays referencing the store clean themselves. You could make a finalize hook to have this be more explicit.

If you really want a context, then you could have group.store = None happen on exit (and raise approproate GroupIsClosed error on an attempt to access), or wrap the original store with a ContextStore whose only job is to call the inner store while int he context and raise an error after exit.

[jhamman]

4. Return async generators from all list methods

I don’t think this one is very controversial but is a change from the original proposal. Moving to an async generator will help the async group and array class better handle large listings. Practically speaking, this would change our list* APIs in the following way:

# from
async def list(self) -> List[str]:
    ...
# to
async def list(self) -> AsyncGenerator[str, None]:

[jhamman]

5. Support partial reads/writes on single and multiple keys at once

The initial design did not allow for a partial read of a single object using the Store.get() method. Here I’m proposing an API like this:

async def get(self, key: str, byte_range: Optional[Tuple[int, Optional[int]]] = None) -> Optional[BytesLike]:
    ...

In fact, this could be the only required interface and the bulk fetch method could be implemented by default as something like:

async def get_partial_values(self, key_ranges: List[Tuple[str, Tuple[int, int]]]) -> List[bytes]:
    awaitables = [self.get(key, byte_range=byte_range) for key, byte_range in key_ranges]
    return asyncio.gather(awaitables)

Stores that are able to provide optimizations beyond this brute force approach (e.g. coalescing) could implement this method.