Getitems: support `meta_array` #1131

madsbk · 2022-09-12T12:59:58Z

BaseStore now implements getitems(), which takes a meta_array argument. Besides simplifying the code, this enable stores to read directly to GPU memory: rapidsai/kvikio#131

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)

lgtm-com · 2022-09-12T13:14:28Z

This pull request introduces 1 alert when merging 26b7d2d into b677db3 - view on LGTM.com

new alerts:

1 for Signature mismatch in overriding method

codecov · 2022-09-12T13:16:25Z

Codecov Report

Merging #1131 (a1d3520) into main (4b0705c) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##              main     #1131   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           36        37    +1     
  Lines        14799     14855   +56     
=========================================
+ Hits         14799     14855   +56

Impacted Files	Coverage Δ
zarr/_storage/store.py	`100.00% <100.00%> (ø)`
zarr/context.py	`100.00% <100.00%> (ø)`
zarr/core.py	`100.00% <100.00%> (ø)`
zarr/storage.py	`100.00% <100.00%> (ø)`
zarr/tests/test_storage.py	`100.00% <100.00%> (ø)`
zarr/tests/test_storage_v3.py	`100.00% <100.00%> (ø)`
zarr/tests/test_util.py	`100.00% <100.00%> (ø)`
zarr/tests/util.py	`100.00% <100.00%> (ø)`
zarr/util.py	`100.00% <100.00%> (ø)`

... and 2 files with indirect coverage changes

…titems

quasiben · 2022-09-21T13:56:26Z

Is this the last PR needed by kvik for full Zarr support ?

madsbk · 2022-09-21T16:40:33Z

Is this the last PR needed by kvik for full Zarr support ?

Yes with this PR and rapidsai/kvikio#131, all building blocks should be in place to use Zarr backed by CuPy arrays!

jakirkham

Thanks Mads! This is very helpful 🙏

Asked one question above about API design. Curious to hear yours and others thoughts 🙂

jakirkham · 2022-09-23T18:18:31Z

zarr/_storage/store.py

+        reads of multiple keys and/or to utilize the meta_array argument.
+        """
+
+        return {k: self[k] for k in keys if k in self}


Know fsspec-based storage layers are using a dict return value, but am wondering if we should be doing something different (like returning an iterable). Asking since this would make the read blocking vs. a bit more lazy. The latter can be useful when working with operations that take a bit longer (like reading from the cloud or parallelizing several reads)

cc @martindurant (who also may have thoughts)

The point of it being blocking, is that many keys may be being fetched concurrently. If you make it an iterator, you lose that, unless you have something that can wait on an async iterator, which in turn has first-completed working.

I agree with both of you.
The API specifies a return type of Mapping[str, Any] thus it is possible to return a lazy mapping that only reads self[k] when accessed. But let's do that in a follow up PR?

Understood, and will be interested in how that looks.

By the way, I mentioned elsewhere the possibility of passing key-specific metadata to the get function; that would happen in this same signature. I wonder if you have any use for an array where only some of it is destined for the GPU (perhaps because that data already exists there and doesn't need loading at all).

Do you have an example of what they would look like, @martindurant?

As noted above, think this discussion gets hairy enough it should be broken out into an issue (likely two) and discussed separately.

As noted above, think this discussion gets hairy enough it should be broken out into an issue (likely two) and discussed separately.

If we want to go the contexts route, we don't need a meta_array argument

Does this also imply that getitem() should have a contexts parameter?

Does this also imply that getitem() should have a contexts parameter?

Good point, yes getitem() should also take a contexts parameter

Wait a second, there is no getitem() and _chunk_getitem() anymore :)

jakirkham · 2022-09-23T18:24:54Z

Is this the last PR needed by kvik for full Zarr support ?

Yes with this PR and rapidsai/kvikio#131, all building blocks should be in place to use Zarr backed by CuPy arrays!

Just to clarify, this is a quite useful optimization (and much appreciated!). Though is it a requirement (in that an error was hit somewhere)? Is there a bit more context about how this came up (particularly if an issue was encountered)?

madsbk · 2022-09-26T07:02:50Z

Is this the last PR needed by kvik for full Zarr support ?

Yes with this PR and rapidsai/kvikio#131, all building blocks should be in place to use Zarr backed by CuPy arrays!

Just to clarify, this is a quite useful optimization (and much appreciated!). Though is it a requirement (in that an error was hit somewhere)? Is there a bit more context about how this came up (particularly if an issue was encountered)?

It is required in the sense that a backend doesn’t know the memory destination of a read without this PR. In KvikIO’s backend, we must rely on key names to infer if the data goes to host or device memory. E.g. the following code check for keys that should go to host memory always:

  if os.path.basename(fn) in [
      zarr.storage.array_meta_key,
      zarr.storage.group_meta_key,
      zarr.storage.attrs_key,
  ]:

This is a problem since a Store can define its own special keys with arbitrary names.

jakirkham · 2022-09-26T08:25:21Z

Thanks for clarifying. Do we need setitems as well?

madsbk · 2022-09-26T08:37:47Z

Thanks for clarifying. Do we need setitems as well?

I don't think it is as important as the getitems since a backend can always probe the input data and determine if it is host or device memory.
However, setitems would make it possible to implement concurrent write, which could benefit backends such as KvikIO significantly.

zarr/core.py

jakirkham · 2022-10-06T07:58:09Z

@joshmoore do you have thoughts on this one? 🙂

joshmoore

Some minor comments, @jakirkham. I do get the feeling that after the refactoring that enabled V3, we likely do need some storage-extension-specific documentation notes.

zarr/_storage/store.py

joshmoore · 2022-10-07T11:14:47Z

zarr/_storage/store.py

+        reads of multiple keys and/or to utilize the meta_array argument.
+        """
+
+        return {k: self[k] for k in keys if k in self}


Do you have an example of what they would look like, @martindurant?

joshmoore · 2022-10-07T11:17:10Z

zarr/core.py

@@ -1930,76 +1923,44 @@ def _process_chunk(
        # store selected data in output
        out[out_selection] = tmp

-    def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,


Is this so private that no one could have made use of it in a subclass?

Yes, AFAICT.

Yeah that's my understanding as well. Or at least by prefixing with _ we have warned users this is an implementation detail (subject to change) that they shouldn't rely on.

@jakirkham: this pre-dates me, so I defer. But if it's not documented somewhere we might want to review if that holds across the board. For me, _x is typically valid for subclassing by developers, otherwise it would be a __x. (i.e. public, protected, private in Java-parlance)

It's a Python thing generally (not specific to Zarr).

Single underscore if by far the most common thing to do, to show intent rather than enforce any privateness (which double underscore doesn't do either).

@jakirkham: see also though https://stackoverflow.com/questions/41761645/when-to-use-one-or-two-underscore-in-python/41761857#comment99492184_41762806

joshmoore · 2022-10-07T11:17:33Z

zarr/core.py

@@ -1257,20 +1257,13 @@ def _get_selection(self, indexer, out=None, fields=None):
        else:
            check_array_shape('out', out, out_shape)

-        # iterate over chunks
-        if not hasattr(self.chunk_store, "getitems") or \


❤️ for removing hasattr hacks in general. 👍

Co-authored-by: Josh Moore <josh@openmicroscopy.org>

…titems

madsbk · 2023-03-13T16:13:45Z

Resolved conflict.
I suggest we go with this approach for now and then unify get_partial_values() and getitems() later.

rabernat · 2023-03-13T19:21:44Z

@joshmoore & @jakirkham - any objections to merging this as-is and deferring the broader discussions for future work?

joshmoore · 2023-03-15T09:41:20Z

Thanks for the ping, @rabernat. I have to say I'm liking the typed Context better. I will read through again ASAP.

martindurant · 2023-03-24T17:08:09Z

Notice the link to FITS reading: another example of per-chunk parameters that could be fixed by passing around a context, although it's not the only way to do it. I'd say that there are a number of efforts (some more and some less well developed) waiting to see what zarr will do.

madsbk · 2023-03-30T12:39:52Z

@joshmoore, how can I help progress here? Maybe an online meeting?

joshmoore · 2023-04-04T16:34:56Z

Sorry, @madsbk. My last comment was shortly before leaving on several weeks of travel and I lost of track of this. It wasn’t my intention to hold you up, and big 👍 for the ping. A few thoughts (from smallest to biggest) though I get the feeling that you are being ping-ponged back and forth between design issues:

If we're less than confident on this API, perhaps we should make this a keyword-only argument. That way the second position could potentially be used for something else and deprecation/renaming would be more straight-forward.
Unlike some of the context that was going to get passed, I have the feeling that the meta_array is more likely to be used for all chunks equally. Does that sound right? If so, I'm a bit hesitant to force someone to create dictionaries of N objects all with the same meta_array instance wrapped in new Context object. As I started to think about "wildcard" keys in the context argument or an additional default_context argument, I realized that perhaps we've not gotten this abstract right yet.
Finally, I've still not got a clear picture of what other fields on the Context object are going to help with the likes of partial read. Obviously, that's less to do your PR but just to callback to the original discussions about whether or not we knew what the API should look like.

…titems

madsbk · 2023-04-11T11:51:12Z

If we're less than confident on this API, perhaps we should make this a keyword-only argument. That way the second position could potentially be used for something else and deprecation/renaming would be more straight-forward.

Agree, fixed.

Unlike some of the context that was going to get passed, I have the feeling that the meta_array is more likely to be used for all chunks equally. Does that sound right? If so, I'm a bit hesitant to force someone to create dictionaries of N objects all with the same meta_array instance wrapped in new Context object. As I started to think about "wildcard" keys in the context argument or an additional default_context argument, I realized that perhaps we've not gotten this abstract right yet.

I have added ConstantMap, which is a read-only map that maps all keys to the same constant value. This fixes the issue of having to create a Context for each key without complicating the semantic of getitems().

Finally, I've still not got a clear picture of what other fields on the Context object are going to help with the likes of partial read. Obviously, that's less to do your PR but just to callback to the original discussions about whether or not we knew what the API should look like.

I don't think any of us have a clear answer to this. Therefore, I suggest that we merge this PR and let people explore its possibilities. As long as Context is optional and we force Stores to implement a fallback if it doesn't exist, I think we are fine.

madsbk · 2023-04-14T07:36:34Z

Thanks all!

rabernat · 2023-04-14T12:12:36Z

Thank you for your patience Mads.

By using the new API in zarr-developers/zarr-python#1131, we do not have to guess whether to read into host or device memory. That is, no more filtering of specify keys like: ```python if os.path.basename(fn) in [ zarr.storage.array_meta_key, zarr.storage.group_meta_key, zarr.storage.attrs_key, ]: ``` Notice, this PR is on hold until Zarr v2.15 is released Closes #119 UPDATE: Zarr v2.15 has been released Authors: - Mads R. B. Kristensen (https://github.com/madsbk) Approvers: - Lawrence Mitchell (https://github.com/wence-) - Jordan Jacobelli (https://github.com/jjacobelli) URL: #131

madsbk added 6 commits September 13, 2022 10:47

Use _chunk_getitems() always

c268c63

Implement getitems() always

c7023f9

FSStore.getitems(): accept meta_array and on_error

89fa599

getitems(): handle on_error="omit"

f05ee3a

Removed the on_error argument

0eed377

remove redundant check

e578051

madsbk force-pushed the getitems branch from a145abb to e578051 Compare September 13, 2022 09:19

getitems(): use Sequence instead of Iterable

03c97a8

madsbk mentioned this pull request Sep 13, 2022

Use the Zarr's new getitems() API rapidsai/kvikio#131

Merged

madsbk marked this pull request as ready for review September 13, 2022 11:16

Merge branch 'main' of github.com:zarr-developers/zarr-python into ge…

8ba463c

…titems

madsbk and others added 2 commits September 22, 2022 08:31

Merge branch 'main' into getitems

dfa731b

Merge branch 'main' into getitems

8e753d6

jakirkham requested a review from joshmoore September 23, 2022 18:14

jakirkham reviewed Sep 23, 2022

View reviewed changes

martindurant reviewed Sep 26, 2022

View reviewed changes

zarr/core.py Outdated Show resolved Hide resolved

Merge branch 'main' into getitems

f05edb1

madsbk mentioned this pull request Oct 6, 2022

Dependency clean up rapidsai/kvikio#137

Merged

joshmoore reviewed Oct 7, 2022

View reviewed changes

madsbk and others added 2 commits October 7, 2022 15:27

Typo

1d2b6ea

Co-authored-by: Josh Moore <josh@openmicroscopy.org>

Merge branch 'main' into getitems

f87cb60

madsbk added 2 commits March 13, 2023 13:25

support the new get_partial_values() method

02fc80d

Merge branch 'main' of github.com:zarr-developers/zarr-python into ge…

61981f2

…titems

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Mar 13, 2023

Resolve conflict with get_partial_values()

d0afcde

martindurant mentioned this pull request Mar 24, 2023

Enable the reading of tile compressed FITS files with kerchunk astropy/astropy#14575

Open

10 tasks

madsbk added 4 commits April 11, 2023 10:56

Merge branch 'main' of github.com:zarr-developers/zarr-python into ge…

7e01831

…titems

make contexts keyword-only

d9838ef

Introduce ConstantMap

c3ee95f

use typing.Mapping

ec5f396

test_constant_map

a1d3520

rabernat merged commit b14f15f into zarr-developers:main Apr 13, 2023

jakirkham added a commit that referenced this pull request Apr 13, 2023

Add release note for #1131

a66f40b

MSanKeys963 mentioned this pull request Apr 14, 2023

Update release.rst for the upcoming release - 2.15.0 #1378

Merged

2 tasks

madsbk mentioned this pull request Apr 18, 2023

GPUdirect kanglcn/moraine#7

Open

madsbk deleted the getitems branch April 20, 2023 09:48

rabernat mentioned this pull request May 4, 2023

How should xarray use/support sparse arrays? pydata/xarray#3213

Open

rabernat mentioned this pull request May 25, 2023

Adding sparse array support zarr-developers/zarr-specs#245

Open

jakirkham mentioned this pull request Jun 22, 2023

getitems follow up work #1442

Open

cgohlke mentioned this pull request Aug 12, 2023

tifffile.ZarrTiffStore wrapped in zarr.LRUStoreCache fails to read any chunks #1497

Closed

akshaysubr mentioned this pull request Apr 22, 2024

[V3] Support for batched Store API in v3 #1806

Open

Getitems: support meta_array #1131

Getitems: support meta_array #1131

Conversation

madsbk commented Sep 12, 2022 • edited Loading

lgtm-com bot commented Sep 12, 2022

codecov bot commented Sep 12, 2022 • edited Loading

Codecov Report

quasiben commented Sep 21, 2022

madsbk commented Sep 21, 2022 • edited Loading

jakirkham left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakirkham commented Sep 23, 2022

madsbk commented Sep 26, 2022 • edited Loading

jakirkham commented Sep 26, 2022

madsbk commented Sep 26, 2022

jakirkham commented Oct 6, 2022

joshmoore left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

madsbk commented Mar 13, 2023

rabernat commented Mar 13, 2023

joshmoore commented Mar 15, 2023

martindurant commented Mar 24, 2023

madsbk commented Mar 30, 2023

joshmoore commented Apr 4, 2023

madsbk commented Apr 11, 2023

madsbk commented Apr 14, 2023

rabernat commented Apr 14, 2023

Getitems: support `meta_array` #1131

Getitems: support `meta_array` #1131

madsbk commented Sep 12, 2022 •

edited

Loading

codecov bot commented Sep 12, 2022 •

edited

Loading

madsbk commented Sep 21, 2022 •

edited

Loading

madsbk commented Sep 26, 2022 •

edited

Loading