add storage_transformers and get/set_partial_values

[jstriebel]

This PR adds two features of the current v3 spec to the zarr-python implementation:

• storage transformers
• get_partial_values and set_partial_values

The current spec is not finalized and waiting for finishing review and approval of ZEP001. Since all v3 features are currently behind the ZARR_V3_EXPERIMENTAL_API flag, I think it might still be fine to incorporate those features to gain early feedback.

Both features are prerequisites for sharding as described here, which I'd like to add as a follow-up PR, together with a new ZEP for it.

TODO:

• [x] Add unit tests and/or doctests in docstrings
• [x] Add docstrings and API docs for any new/modified user-facing classes and functions
• [New/modified features documented in docs/tutorial.rst] (v3 features seem not to be documented yet?)
• [x] Changes documented in docs/release.rst
• [x] GitHub Actions have all passed
• [x] Test coverage is 100% (Codecov passes)
• [x] Add follow-up issue to optimize get_partial_values and set_partial_values for specific stores: #1106

[jstriebel]

The Linux tests fail due to some s3 tests, and the Windows test timeout in the s3 fixture setup, which also fail on different branches and seem to be unrelated to this PR.

[jstriebel]

@joshmoore & @grlee77 Would one of you have time to review my changes here? Asking since you authored the initial v3 additions here. From my side this PR is ready for review.

[jstriebel]

I have some minor suggestions, but this looks pretty good already.

Thank you so much for the review, @grlee77! I applied your feedback in the recent commits.

I think we will get more experience and a better idea of a real-world use case with the follow-up sharding PR.

Yes, I think the followup-PR will help to make more sense for those methods.

• does passing storage_transformers=[] work equivalently to storage_transformers=None?

Yes, it does, I added a notice for it in the docstring. I also thought about just having an empty list as the default, but lists as default arguments are somewhat deprecated due to their mutability.

• are there plans to expose storage_transformers via the higher level convenience/creation functions? (e.g. zarr/creation.py)

I think it is already exposed via create. I also thought about exposing specific shorthand parameters for single storage transformers, e.g. create(…, chunks_per_shard=(32, 32, 32)), but that's not strictly necessary. Are there any other places where storage transformers should be passed, that I'm missing?

[jakirkham]

cc @martindurant (in case this is of interest)

[martindurant]

@jakirkham , I must be missing something: allowing partial byte ranges to pass through to the storage layer and transformers seem completely orthogonal to me.

In the case of this code, on this line we get the full chunk and then slice it anyway. Is the intend for some storage classes (fsspec!) to subclass this and implement a new get_partial_values?

[jstriebel]

allowing partial byte ranges to pass through to the storage layer and transformers seem completely orthogonal to me.

Yes, it is to some extend. I can also have them as separate PR. My motivation was to implement sharding, and I tried to keep the bases for it as a separate PR, that's how they both ended up here.

In the case of this code, on this line we get the full chunk and then slice it anyway. Is the intend for some storage classes (fsspec!) to subclass this and implement a new get_partial_values?

Yes, the idea is to overwrite those functions for stores that allow partial access. Here I just implemented the API and the fallback strategy, I'll open an issue for the optimized paths as a follow-up, also mentioned in the TODOs in the description above:

• Add follow-up issue to optimize get_partial_values and set_partial_values for specific stores

Does that answer your questions, @martindurant?

[martindurant]

I didn't mean to denigrate this PR, I know very little about how the storage transformer or sharding are to be implemented. I am wondering how the new get_partial_values would get called, and then it sounds like we can agree on something that works for us all.

[jstriebel]

I also thought about just having an empty list as the default, but lists as default arguments are somewhat deprecated due to their mutability.

And a tuple?

Good idea, I changed the default to an empty tuple now :+1:

[codecov[bot]]

Codecov Report

Merging #1096 (eba9006) into main (1fd607a) will not change coverage. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##              main     #1096    +/-   ##
==========================================
  Coverage   100.00%   100.00%            
==========================================
  Files           35        35            
  Lines        14157     14388   +231     
==========================================
+ Hits         14157     14388   +231     

Impacted Files	Coverage Δ
zarr/creation.py	100.00% <ø> (ø)
zarr/_storage/store.py	100.00% <100.00%> (ø)
zarr/core.py	100.00% <100.00%> (ø)
zarr/meta.py	100.00% <100.00%> (ø)
zarr/storage.py	100.00% <100.00%> (ø)
zarr/tests/test_core.py	100.00% <100.00%> (ø)
zarr/tests/test_creation.py	100.00% <100.00%> (ø)
zarr/tests/test_storage_v3.py	100.00% <100.00%> (ø)

[jstriebel]

I opened the follow-up issue #1106, which is about implementing optimized partial read/write methods for specific stores, and using them for the partial decode routines.

[jstriebel]

@grlee77 @joshmoore Thank you very much for your reviews! I incorporated your feedback where I marked the comments as resolved, the other's are still open conversations. Could you give this PR another look? CI finally is completely green, I added some more tests for a higher coverage.

[jstriebel]

@grlee77 @joshmoore I've incorporated some more changes that were useful for the sharding PR #1111. I added them here to keep the PRs clean. Could you have a look at those as well? Thanks a ton!

[martindurant]

Above I asked "how would the new get_partial_values would get called"? At it stands, this is all unused code in this PR (from a zarr user's perspective, as opposed to tests), right?

I am still trying to wrap my head around why we need a new method here, as opposed, to just getitems(keys, ranges=None).

In fact, I have been mulling over the idea of passing a context (e.g., requested chunk key, chunk slice, versus array position) through the filters/compressors and to the storage layer, such that they can take whatever action they like with that information. For example, could have Codec.decode(buf, context=None) which applies a different scale factor as a factor of dimension 3 ( https://github.com/fsspec/kerchunk/issues/212 ) and use the same information object in getitems(keys, contexts=None) to decide what byterange to fetch.

[jstriebel]

Above I asked "how would the new get_partial_values would get called"? At it stands, this is all unused code in this PR (from a zarr user's perspective, as opposed to tests), right?

Exactly, the methods are only used for "real" code in the sharding PR #1111. From the description of this PR:

Both features are prerequisites for sharding as described here, which I'd like to add as a follow-up PR, together with a new ZEP for it.

---

I am still trying to wrap my head around why we need a new method here, as opposed, to just getitems(keys, ranges=None).

That proposal is new to me, what exactly would ranges or context be here? I wonder how the returned value should look like, especially when multiple ranges can be specified for the same key. Atm getitems returns a dict with the storage-keys as dict-keys, but it might be fine to add the ranges to the dict-keys as well, returning something like {("a_key", (0, 10)): b"…", ("a_key", (-16, None)): b"…"}. The main reason for the new methods is that I described them this way in the v3 spec, but any other API for the same functionality should be fine :+1: It's definitely useful to have some coherence when a codec allows partial reads, or could even request sequential partial reads from a store. Is there some work on this going on already?

I don't have a strong opinion here, please also feel free to directly propose code changes, e.g. as a PR against this PR.

[martindurant]

I would very much like for partial reads to be available for uncompressed (unsharded) data, that is a major concern of mine, but this PR doesn't provide it as it stands, and I don't see a path to providing it based on the proposed architecture.

The other concern is about a codec knowing which chunk its working on when en/decoding. Maybe the two concerns are not part of the same puzzle...

About contexts, clearly I don't have too concrete a picture yet and certainly not an implementation. Your point about the current return value of getitems() is well made - obviously I had only considered one range per key (a list would be fine, allowing repeated key names - this is what fsspec actually uses internally). Whether the context is passed to the codec or the storage, and whether it would look the same in both cases (and whether a codec might directly call the storage as the current blosc partial read in v2 does in a hacky way) - all to be resolved.

[jstriebel]

I would very much like for partial reads to be available for uncompressed (unsharded) data, that is a major concern of mine, but this PR doesn't provide it as it stands, and I don't see a path to providing it based on the proposed architecture.

This is also implemented in #1111, please have a look at this PR as well, I think it might clarify the usage of the proposed methods quite a bit.

The other concern is about a codec knowing which chunk its working on when en/decoding. Maybe the two concerns are not part of the same puzzle...

I agree, that's quite important to allow generic partial reads (and maybe writes?) via codecs, not just for blosc. However, I think that's a separate problem, and initially allowing partial reads from the stores will make this much easier. Changing the API slightly (e.g. merging the partial reads into getitems) might still be fine later, after all v3 is only in a preview state atm. Therefore I see this rather as a minor issue, do you agree @martindurant?

[martindurant]

(for the codecs, I was actually thinking about parameters that depend on location in the overall array, not partial reads, although that could be part of it too)

[jstriebel]

@grlee77 @joshmoore @MSanKeys963 I merged the newest changes into this branch. From my POV this seems ready to be merged. I'll also update #1111 (the PR about sharding) as discussed (putting sharding behind a zep 2 flag), which we can merge separately afterwards.

[martindurant]

Did you have any thoughts of the interaction of this with https://github.com/zarr-developers/zarr-python/pull/1131 ?

[jstriebel]

Did you have any thoughts of the interaction of this with https://github.com/zarr-developers/zarr-python/pull/1131 ?

Thanks for making me aware of that, not yet. #1131 is mostly about providing a context, which might just get passed through storage transformers. It's not clear to me why stores do handle arrays directly instead of byte-buffers, but that's a completely orthogonal discussion I'd rather like to keep out of this PR. As #1131 stands right now this would not change much wrt this PR. @madsbk, do you have any concerns regarding this PR?

[martindurant]

I am not really talking about transformers - the main thrust of this PR - but the ranges.

I raise it, because the byterange (or array slice, more likely) is something that should be included in the context, so then you would not need extra methods on the storage class. That's what I was after in my comments above. It also provides a direct way to actually activate the partial read. As opposed to in the code here, where the methods exist, but there is no obvious path to calling them.

[jstriebel]

@martindurant Thanks for the clarification.

I raise it, because the byterange (or array slice, more likely) is something that should be included in the context, so then you would not need extra methods on the storage class.

Sure, that's possible, but I don't see a benefit. The context would then need to force stores into different behavior, and the returned content would depend heavily on the context. I find it more clean to separate the different behavior into a different method. But not a super-strong opinion.

As opposed to in the code here, where the methods exist, but there is no obvious path to calling them.

The method's can just be called where appropriate, e.g. for partial reads on uncompressed chunks or sharding. Please see #1111 for usage examples.

[martindurant]

The methods can just be called where appropriate

Can I be super annoying and ask for a smallest possible example here? That other PR and thread is very long and not focussed on partial reads.

x = group.create_dataset(name="x", data=np.array([1,2,3,4]), compression=None, chunks=(4, ))
x[:2]

where the second line should read only 16 bytes from the store, not all 32, if the store supports partials.

The context would then need to force stores into different behavior, and the returned content would depend heavily on the context.

I don't follow this. For the above case,

context={"key": "data/0.0", "out_selection": [slice(0, 2)], "chunk_selection": [slice(0, 2), "meta_array": np.NDArray, "dtype": np.int64]}

where all these things are known in Array._chunk_getitem (meta_array is new and shows that we can do more with this idea, from the other PR). Return bytes for the chunk would be

b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00'

and 16 more bytes of or empty/zeros; if we can't just truncate instead.

[jstriebel]

@martindurant No worries! This is a short example for your case:

x = group.create_dataset(name="x", data=np.array([1,2,3,4]), compression=None, chunks=(4, ))
x[:2]

# with PR #1111 in place, the indexing would do:
store.get_partial_values[(store_key, (0, 2 * itemsize))]
# where store_key is probably sth like "x/0", and itemsize=8 for the int64 dtype
# returning [b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00']

(For uncompressed partial reads the implementation is here in PR #1111

Also the tests have some more examples: https://github.com/zarr-developers/zarr-python/pull/1096/files#diff-49f50043d526642ad824a8d6a944b896540238b06ec11f182476ff8fdb2b5977R257-R316

I think I finally understand where you are coming from. With the context variant the parts that are not requested would simply be filled I guess (or truncated, as you mentioned). However, I'd prefer to only have a buffer for the actual parts of the chunk, which might be spread through the chunk, so one single buffer would not be enough. Therefore the return type would need to be different for the partial reads, which fits more nicely with a new method IMO. Does that make sense?

[martindurant]

OK, I got it: so the code that calls the methods introduced here actually are in the other PR. That call is not really to do with sharding, which is why I got confused. Actually I don't see how that uncompressed-partial-buffer gets used either, there seems to be many levels here.

Are you saying that you want a single call to chunk-getitem to produce many buffers, which will then be filled into the target output array? That would work with contexts too, get([key1, key1], contexts=[context1, context2]). I think in the proposal contexts= was a dict, so it would indeed need to be a list instead. That makes the API the same as here (and we could even have both).

[jstriebel]

Are you saying that you want a single call to chunk-getitem to produce many buffers, which will then be filled into the target output array?

Exactly :+1:

That would work with contexts too, get([key1, key1], contexts=[context1, context2]). I think in the proposal contexts= was a dict, so it would indeed need to be a list instead. That makes the API the same as here (and we could even have both).

That would work, but it would also make the signature of the get call very complicated (the return type would depend on the supplied arguments). Personally I'm in favor of the current variant with the additional method. @grlee77 @joshmoore @normanrz any other opinions here? I'd like to get this merged soon-ish, so I can finalize #1111 as well.

[joshmoore]

I should add re-reviewing the interaction with #1131/meta_array I start to wonder if the definition of "context" and "partial" don't need to end up in the spec or at least a clarifying ZEP together.

[jstriebel]

@joshmoore Thanks for the review!

Primary question is if you would see still getting this into a 2.13 release (especially if one were to come out over the holidays) or would it make more sense to roll it into 2.14.

I wouldn't mind having this in the 2.13 release, so we can play around with it a bit already. I can also add the additional flag for sharding in #1111 so we are on the safe side.

I should add re-reviewing the interaction with https://github.com/zarr-developers/zarr-python/pull/1131/meta_array I start to wonder if the definition of "context" and "partial" don't need to end up in the spec or at least a clarifying ZEP together.

The partial access is part of the v3 spec already: https://zarr-specs.readthedocs.io/en/latest/core/v3.0.html#abstract-store-interface.

[martindurant]

Is there any reason I should not press the big green button? We have three approvals already. Last chance to object!

[joshmoore]

Now with the somewhat large 2.13.4 out the door, I'm going to go ahead and click the button (but would have also supported @martindurant's pushing thereof :wink: :heart:)