Add asynchronous load method (#10327)

* new blank whatsnew

* test async load using special zarr LatencyStore

* don't use dask

* async all the way down

* remove assert False

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add pytest-asyncio to CI envs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* assert results are identical

* implement async load for dataarray and dataset

* factor out common logic

* consolidate tests via a parametrized fixture

* async_load -> load_async

* make BackendArray an ABC

* explain how to add async support for any backend in the docs

* add new methods to api docs

* whatsnew

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix ci/minimum_versions.py

* fix formatting

* concurrently load different variables in ds.load_async using asyncio.gather

* test concurrent loading of multiple variables in one dataset

* fix non-awaited load_async

* rearrange test order

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add test for orthogonal indexing

* explicitly forbid orthogonal indexing

* support async orthogonal indexing via zarr-developers/zarr-python#3083

* add test for vectorized indexing (even if it doesn't work)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add test for basic indexing

* correct test to actually use vectorized indexing

* refactor to parametrize indexing test

* implement async vectorized indexing

* revert breaking change to BackendArray

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove indirection in _ensure_cached method

* IndexingAdapters don't need async get

* Add tests

* Add decoding test

* Add IndexingAdapter mixin

* [cherry] Making decoding arrays lazy too

* parametrized over isel and sel

* mock zarr.AsyncArray.getitem in test

* tidy up the mocking

* ensure the correct zarr class's method is patched for each test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add degenerate test case of no indexing

* factor out the Latency part of LatencyStore

* use mocks in multiple objects test

* use mocks in multiple variables test

* trim latencystore down to just what's needed to dodge zarr-developers/zarr-python#3105 (comment)

* parametrizing indexing test over xarray classes

* ensure we actually test vectorized indexing for Variable

* use create_test_data

* add @pytest.mark.asyncio

* remove outdated readonly_store

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* enable tests to run when recent version of zarr-python is not available

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* separate tests to only run on correct versions of zarr

* clear error message if async oindexing not available

* parametrize over zarr_format

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add pytest-asyncio to other test CI env

* fix some mypy errors

* use method directly when possible

* remove repeated API docs from bad merge

* fix bad merge in release note

* fix other bad merge in whatsnew

* remove prints

Co-authored-by: Deepak Cherian <[email protected]>

* remove last print statement

* test async basic indexing raises informative error before zarr-python v3.0.0

* test correct error message is raised for each indexing case

* ensure each test runs on the earliest version of xaarr it can

* remove pointless repeated getitem

* set N_LAZY_VARS correctly in test

* remove unused import

* rename flag to make it more clear its only for orthogonal and vectorized indexing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove IndexingAdapter special case

* type fixes

* return a deepcopy

* try again

* one more

* Try again

* try fixing _in_memory error by not returning the adapter class

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove scope=module from fixture for robustness

* modify test to be happy with either error message

* use Variable instead of Dataset to avoid race condition of indexing between different variables

* fix bad merge in API docs

* add test to test_backends.py

* fix bug found by new test, causing pandas indexes to be converted to numpy arrays

* add test to test_variable.py for lazy async indexing

* move async load tests from test_async.py to test_backends.py

* parametrize all tests over zarr_format

* remove test_async.py file entirely

* lint

* Stricter assertions

Co-authored-by: Deepak Cherian <[email protected]>

* Update doc/whats-new.rst

Co-authored-by: Deepak Cherian <[email protected]>

* add docstring for Variable.load_async

* make all load-related docstrings consistent

* note about users being responsible for limiting concurrency

* remove parametrization over zarr_format

* account for Dataset having multiple lazy vars

* refactor test parametrization to use pytest.param(..., id=...) syntax

* refactor TestBackendIndexing to combine sync and async checks in one function

* move test_load_async onto test base class

* should fix mypy error

* add back in the parametrize_zarr_format to avoid trying to write v3 data using zarr-python v2

* parametrize test over async

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak Cherian <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

ci/minimum_versions.py

@@ -30,6 +30,7 @@
     "coveralls",
     "pip",
     "pytest",
+    "pytest-asyncio",
     "pytest-cov",
     "pytest-env",
     "pytest-mypy-plugins",

ci/requirements/all-but-dask.yml

@@ -28,6 +28,7 @@ dependencies:
   - pip
   - pydap
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/all-but-numba.yml

@@ -41,6 +41,7 @@ dependencies:
   - pyarrow # pandas raises a deprecation warning without this, breaking doctests
   - pydap
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/bare-min-and-scipy.yml

@@ -7,6 +7,7 @@ dependencies:
   - coveralls
   - pip
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/bare-minimum.yml

@@ -7,6 +7,7 @@ dependencies:
   - coveralls
   - pip
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/environment-3.14.yml

@@ -37,6 +37,7 @@ dependencies:
   - pyarrow # pandas raises a deprecation warning without this, breaking doctests
   - pydap
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/environment-windows-3.14.yml

@@ -32,6 +32,7 @@ dependencies:
   - pyarrow # importing dask.dataframe raises an ImportError without this
   - pydap
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/environment-windows.yml

@@ -32,6 +32,7 @@ dependencies:
   - pyarrow # importing dask.dataframe raises an ImportError without this
   - pydap
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/environment.yml

@@ -39,6 +39,7 @@ dependencies:
   - pydap
   - pydap-server
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins

ci/requirements/min-all-deps.yml

@@ -40,6 +40,7 @@ dependencies:
   - pip
   - pydap=3.5.0
   - pytest
+  - pytest-asyncio
   - pytest-cov
   - pytest-env
   - pytest-mypy-plugins