Implement async support for open_datatree

doc/whats-new.rst

@@ -83,6 +83,9 @@ Internal Changes
 Performance
 ~~~~~~~~~~~
 
+- Improve performance of :py:func:`open_datatree` for zarr stores by using async/concurrent
+  loading of groups and indexes (:pull:`10742`).
+  By `Alfonso Ladino <https://github.com/aladinor>`_.
 - Add a fastpath to the backend plugin system for standard engines (:issue:`10178`, :pull:`10937`).
   By `Sam Levang <https://github.com/slevang>`_.
 

xarray/backends/api.py

@@ -65,6 +65,7 @@
         NestedSequence,
         T_Chunks,
     )
+    from xarray.core.variable import Variable
 
     T_NetcdfEngine = Literal["netcdf4", "scipy", "h5netcdf"]
     T_Engine = Union[
@@ -348,7 +349,37 @@ def _datatree_from_backend_datatree(
 
     _protect_datatree_variables_inplace(backend_tree, cache)
     if create_default_indexes:
-        tree = backend_tree.map_over_datasets(_maybe_create_default_indexes)
+        _use_zarr_async = False
+        if engine == "zarr":
+            from xarray.backends.zarr import _zarr_v3
+
+            _use_zarr_async = _zarr_v3()
+
+        if _use_zarr_async:
+            from zarr.core.sync import sync as zarr_sync
+
+            async def create_indexes_async() -> dict[str, Dataset]:
+                import asyncio
+
+                results: dict[str, Dataset] = {}
+                tasks = [
+                    _create_index_for_node(path, node.dataset)
+                    for path, [node] in group_subtrees(backend_tree)
+                ]
+                for fut in asyncio.as_completed(tasks):
+                    path, ds = await fut
+                    results[path] = ds
+                return results
+
+            async def _create_index_for_node(
+                path: str, ds: Dataset
+            ) -> tuple[str, Dataset]:
+                return path, await _maybe_create_default_indexes_async(ds)
+
+            results = zarr_sync(create_indexes_async())
+            tree = DataTree.from_dict(results, name=backend_tree.name)
+        else:
+            tree = backend_tree.map_over_datasets(_maybe_create_default_indexes)
     else:
         tree = backend_tree
     if chunks is not None:
@@ -385,6 +416,38 @@ def _datatree_from_backend_datatree(
     return tree
 
 
+async def _maybe_create_default_indexes_async(ds: Dataset) -> Dataset:
+    """Create default indexes for dimension coordinates asynchronously.
+
+    This function parallelizes both data loading and index creation,
+    which can significantly speed up opening datasets with many coordinates.
+    """
+    import asyncio
+
+    to_index_names = [
+        name
+        for name, coord in ds.coords.items()
+        if coord.dims == (name,) and name not in ds.xindexes
+    ]
+
+    if not to_index_names:
+        return ds
+
+    async def load_var(var: Variable) -> Variable:
+        try:
+            return await var.load_async()
+        except NotImplementedError:
+            return await asyncio.to_thread(var.load)
+
+    await asyncio.gather(
+        *[load_var(ds.coords[name].variable) for name in to_index_names]
+    )
+
+    variables = {name: ds.variables[name] for name in to_index_names}
+    new_coords = Coordinates(variables)
+    return ds.assign_coords(new_coords)
+
+
 def open_dataset(
     filename_or_obj: T_PathFileOrDataStore,
     *,