Address PR review: TaskGroup, max_concurrency, and open_dataset_async

- Replace asyncio.gather with asyncio.TaskGroup for better error handling
  (cancels outstanding tasks on error)
- Add max_concurrency parameter to open_datatree for controlling parallel
  I/O operations (defaults to 10)
- Add StoreBackendEntrypoint.open_dataset_async method
- Add test for open_dataset_async equivalence

Author: aladinor

doc/whats-new.rst

@@ -16,6 +16,9 @@ New Features
 
 - Added ``inherit='all_coords'`` option to :py:meth:`DataTree.to_dataset` to inherit
   all parent coordinates, not just indexed ones (:issue:`10812`, :pull:`11230`).
+- Added ``max_concurrency`` parameter to :py:func:`open_datatree` to control
+  the maximum number of concurrent I/O operations when opening groups in parallel
+  with the Zarr backend (:pull:`10742`).
   By `Alfonso Ladino <https://github.com/aladinor>`_.
 
 Breaking Changes

xarray/backends/api.py

@@ -943,6 +943,7 @@ def open_datatree(
     chunked_array_type: str | None = None,
     from_array_kwargs: dict[str, Any] | None = None,
     backend_kwargs: dict[str, Any] | None = None,
+    max_concurrency: int | None = None,
     **kwargs,
 ) -> DataTree:
     """
@@ -1075,6 +1076,13 @@ def open_datatree(
         chunked arrays, via whichever chunk manager is specified through the `chunked_array_type` kwarg.
         For example if :py:func:`dask.array.Array` objects are used for chunking, additional kwargs will be passed
         to :py:func:`dask.array.from_array`. Experimental API that should not be relied upon.
+    max_concurrency : int, optional
+        Maximum number of concurrent I/O operations when opening groups in
+        parallel. This limits the number of groups that are loaded simultaneously.
+        Useful for controlling resource usage with large datatrees or stores
+        that may have limitations on concurrent access (e.g., icechunk).
+        Only used by backends that support parallel loading (currently Zarr v3).
+        If None (default), the backend uses its default value (typically 10).
     backend_kwargs: dict
         Additional keyword arguments passed on to the engine open function,
         equivalent to `**kwargs`.
@@ -1135,6 +1143,9 @@ def open_datatree(
     )
     overwrite_encoded_chunks = kwargs.pop("overwrite_encoded_chunks", None)
 
+    if max_concurrency is not None:
+        kwargs["max_concurrency"] = max_concurrency
+
     backend_tree = backend.open_datatree(
         filename_or_obj,
         drop_variables=drop_variables,

xarray/backends/store.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
@@ -72,5 +73,37 @@ def open_dataset(
 
         return ds
 
+    async def open_dataset_async(
+        self,
+        filename_or_obj: T_PathFileOrDataStore,
+        *,
+        mask_and_scale=True,
+        decode_times=True,
+        concat_characters=True,
+        decode_coords=True,
+        drop_variables: str | Iterable[str] | None = None,
+        set_indexes: bool = True,
+        use_cftime=None,
+        decode_timedelta=None,
+    ) -> Dataset:
+        """Async version of open_dataset.
+
+        Offloads the entire open_dataset operation to a thread to avoid blocking
+        the event loop. This is necessary because decode_cf_variables can trigger
+        data reads (e.g., for time decoding) which may use synchronous I/O.
+        """
+        return await asyncio.to_thread(
+            self.open_dataset,
+            filename_or_obj,
+            mask_and_scale=mask_and_scale,
+            decode_times=decode_times,
+            concat_characters=concat_characters,
+            decode_coords=decode_coords,
+            drop_variables=drop_variables,
+            set_indexes=set_indexes,
+            use_cftime=use_cftime,
+            decode_timedelta=decode_timedelta,
+        )
+
 
 BACKEND_ENTRYPOINTS["store"] = (None, StoreBackendEntrypoint)

xarray/backends/zarr.py

@@ -1890,6 +1890,7 @@ def open_datatree(
         storage_options=None,
         zarr_version=None,
         zarr_format=None,
+        max_concurrency: int | None = None,
     ) -> DataTree:
         filename_or_obj = _normalize_path(filename_or_obj)
 
@@ -1927,6 +1928,7 @@ def open_datatree(
                     drop_variables=drop_variables,
                     use_cftime=use_cftime,
                     decode_timedelta=decode_timedelta,
+                    max_concurrency=max_concurrency,
                 )
             )
         else:
@@ -1965,6 +1967,7 @@ async def _open_datatree_from_stores_async(
         drop_variables: str | Iterable[str] | None = None,
         use_cftime=None,
         decode_timedelta=None,
+        max_concurrency: int | None = None,
     ) -> DataTree:
         """Async helper to open datasets from pre-opened stores and create indexes.
 
@@ -1973,28 +1976,24 @@ async def _open_datatree_from_stores_async(
         """
         from xarray.backends.api import _maybe_create_default_indexes_async
 
-        # Limit concurrent to_thread calls to avoid deadlocks with some stores
-        # (e.g., icechunk can deadlock when too many threads access it simultaneously)
-        sem = asyncio.Semaphore(10)
+        if max_concurrency is None:
+            max_concurrency = 10
+        sem = asyncio.Semaphore(max_concurrency)
 
         async def open_one(path_group: str, store) -> tuple[str, Dataset]:
             async with sem:
                 store_entrypoint = StoreBackendEntrypoint()
-
-                def _load_sync():
-                    with close_on_error(store):
-                        return store_entrypoint.open_dataset(
-                            store,
-                            mask_and_scale=mask_and_scale,
-                            decode_times=decode_times,
-                            concat_characters=concat_characters,
-                            decode_coords=decode_coords,
-                            drop_variables=drop_variables,
-                            use_cftime=use_cftime,
-                            decode_timedelta=decode_timedelta,
-                        )
-
-                ds = await asyncio.to_thread(_load_sync)
+                with close_on_error(store):
+                    ds = await store_entrypoint.open_dataset_async(
+                        store,
+                        mask_and_scale=mask_and_scale,
+                        decode_times=decode_times,
+                        concat_characters=concat_characters,
+                        decode_coords=decode_coords,
+                        drop_variables=drop_variables,
+                        use_cftime=use_cftime,
+                        decode_timedelta=decode_timedelta,
+                    )
                 # Create indexes in parallel (within this group)
                 ds = await _maybe_create_default_indexes_async(ds)
                 if group:
@@ -2003,10 +2002,15 @@ def _load_sync():
                     group_name = str(NodePath(path_group))
                 return group_name, ds
 
-        # Open all datasets and create indexes concurrently
-        tasks = [open_one(path_group, store) for path_group, store in stores.items()]
-        results = await asyncio.gather(*tasks)
-        groups_dict = dict(results)
+        groups_dict: dict[str, Dataset] = {}
+
+        async def collect_result(path_group: str, store) -> None:
+            group_name, ds = await open_one(path_group, store)
+            groups_dict[group_name] = ds
+
+        async with asyncio.TaskGroup() as tg:
+            for path_group, store in stores.items():
+                tg.create_task(collect_result(path_group, store))
 
         return datatree_from_dict_with_io_cleanup(groups_dict)
 

xarray/tests/test_backends_zarr_async.py

@@ -9,6 +9,7 @@
 
 import xarray as xr
 from xarray.backends.api import _maybe_create_default_indexes_async
+from xarray.backends.store import StoreBackendEntrypoint
 from xarray.backends.zarr import ZarrBackendEntrypoint
 from xarray.testing import assert_equal
 from xarray.tests import (
@@ -232,3 +233,27 @@ def test_sync_open_datatree_uses_async_internally(self, zarr_format):
                 # For zarr v3, the async function should be called
                 assert mock_async.call_count > 0
                 assert_equal(dtree, dtree_loaded)
+
+    @pytest.mark.asyncio
+    @requires_zarr_v3
+    @parametrize_zarr_format
+    async def test_store_backend_open_dataset_async_equivalence(self, zarr_format):
+        """Test that StoreBackendEntrypoint.open_dataset_async returns same result as sync."""
+        from xarray.backends.zarr import ZarrStore
+
+        ds = create_dataset_with_coordinates(n_coords=3)
+
+        with self.create_zarr_store() as store:
+            ds.to_zarr(store, mode="w", consolidated=False, zarr_format=zarr_format)
+
+            zarr_store = ZarrStore.open_group(
+                store,
+                consolidated=False,
+                zarr_format=zarr_format,
+            )
+
+            store_entrypoint = StoreBackendEntrypoint()
+            ds_sync = store_entrypoint.open_dataset(zarr_store)
+            ds_async = await store_entrypoint.open_dataset_async(zarr_store)
+
+            assert_equal(ds_sync, ds_async)