`open_dataset` creates default indexes sequentially, causing significant latency in cloud high-latency stores

[aladinor]

xarray.open_dataset() and xarray.open_datatree() create default indexes for each coordinate sequentially via the _maybe_create_default_indexes() function. When reading from a cloud-based Zarr store, this results in one request per coordinate, which compounds into significant total latency.

This issue becomes especially noticeable when:

• There are many coordinate variables
• You are working with DataTrees containing many groups
• The backing store has even modest per-request latency (e.g., 100ms)

Where this happens

The issue originates in the _maybe_create_default_indexes() function:
source: xarray/backends/api.py#L383-L389

Each coordinate is processed in a Python loop, and coord.variable triggers .data access on Zarr-backed arrays. This causes a separate .get() request for every coordinate.

In high-latency stores, the result is a linear slowdown proportional to the number of coordinates.

Reproducible Test Case

To reproduce this issue, I created:

• A single Dataset with multiple coordinates
• A DataTree with multiple groups, each containing several coordinates
• A LatencyStore wrapper around MemoryStore that introduces 1s artificial latency per .get() call

Click to expand Python test code

from zarr.storage import MemoryStore
import asyncio
import time
from collections.abc import Iterable
from typing import TypeVar

import numpy as np

import xarray as xr
from xarray.tests import has_zarr_v3

if has_zarr_v3:
    from zarr.abc.store import ByteRequest, Store
    from zarr.core.buffer import Buffer, BufferPrototype
    from zarr.storage import MemoryStore
    from zarr.storage._wrapper import WrapperStore

    T_Store = TypeVar("T_Store", bound=Store)


class LatencyStore(WrapperStore[T_Store]):
    """Works the same way as the zarr LoggingStore"""
    latency: float

    __read_only: bool

    def __init__(
        self,
        store: T_Store,
        read_only: bool,
        latency: float = 0.0,
    ) -> None:
        """
        Store wrapper that adds artificial latency to each get call.

        Parameters
        ----------
        store : Store
            Store to wrap
        latency : float
            Amount of artificial latency to add to each get call, in seconds.
        """
        super().__init__(store)
        self.__read_only = read_only
        self.latency = latency

    def __str__(self) -> str:
        return f"latency-{self._store}"

    def __repr__(self) -> str:
        return f"LatencyStore({self._store.__class__.__name__}, '{self._store}', latency={self.latency})"

    def with_read_only(self, read_only=False) -> "LatencyStore":
        return LatencyStore(self._store, read_only=True, latency=self.latency)

    async def get(
        self,
        key: str,
        prototype: BufferPrototype,
        byte_range: ByteRequest | None = None,
    ) -> Buffer | None:
        await asyncio.sleep(self.latency)
        return await self._store.get(
            key=key, prototype=prototype, byte_range=byte_range
        )

    async def get_partial_values(
        self,
        prototype: BufferPrototype,
        key_ranges: Iterable[tuple[str, ByteRequest | None]],
    ) -> list[Buffer | None]:
        await asyncio.sleep(self.latency)
        return await self._store.get_partial_values(
            prototype=prototype, key_ranges=key_ranges
        )



def create_dataset_with_many_coordinates(n_coords=20):
    """Create a dataset with many tiny coordinate variables"""
    coords = {}
    for i in range(n_coords):
        coords[f'coord_{i}'] = (f'coord_{i}', np.arange(2))
    ds = xr.Dataset(
        coords=coords
    )
    return ds


def create_datatree_with_many_groups(n_groups=30, coords_per_group=5):
    """Create a DataTree with many groups, each having coordinate variables"""
    print(f"Creating DataTree with {n_groups} groups, {coords_per_group} coords each...")

    root_ds = create_dataset_with_many_coordinates(5)

    tree_dict = {'/': root_ds}

    for i in range(n_groups):
        group_name = f'/group_{i:03d}'
        group_ds = create_dataset_with_many_coordinates(n_coords=coords_per_group)
        tree_dict[group_name] = group_ds

    tree = xr.DataTree.from_dict(tree_dict)
    return tree


if __name__ == '__main__':
    print("testing dataset")
    ds = create_dataset_with_many_coordinates(10)

    store = MemoryStore()
    latency_store = LatencyStore(store, latency=1, read_only=False)
    print("storing dataset")
    ds.to_zarr(latency_store, mode="w", consolidated=False, zarr_format=3)

    print("opening dataset")
    start_time = time.time()
    ds_index = xr.open_zarr(
        latency_store,
        consolidated=False,
        zarr_format=3,
        create_default_indexes = True,

    )
    with_index_time = time.time() - start_time
    print(f"open dataset with create_default_indexes=True:  {with_index_time}")

    start_time = time.time()
    ds_no_index = xr.open_zarr(
        latency_store,
        consolidated=False,
        zarr_format=3,
        create_default_indexes=False
    )
    no_index = time.time() - start_time

    print(f"open dataset with create_default_indexes=False:  {no_index}")
    print(f"creating indexes from cloud environment might take {with_index_time / no_index:.1f}x slower")

    print("==="* 20)

    print("Now testing datatrees")

    dtree = create_datatree_with_many_groups(n_groups=5, coords_per_group=10)

    dtree.to_zarr(store, mode="w", consolidated=False, zarr_format=3)

    start_time = time.time()
    dtree_index = xr.open_datatree(
        latency_store,
        consolidated=False,
        zarr_format=3,
        create_default_indexes = True,
        engine="zarr"
    )
    with_index = time.time() - start_time

    print(f"open dtree with create_default_indexes=True:  {with_index}")

    start_time = time.time()
    dtree_no_index = xr.open_datatree(
        latency_store,
        consolidated=False,
        zarr_format=3,
        create_default_indexes = False,
        engine="zarr"
    )
    no_index = time.time() - start_time

    print(f"open dtree with create_default_indexes=False:  {no_index}")
    print(f"without Index creation is {with_index / no_index:.1f}x faster")

Results

Case	With Indexes	Without Indexes	Slowdown
Dataset (10 coordinates)	12.6 s	2.0 s	6.3×
DataTree (5 groups × 10 coords)	48.1 s	18.0 s	2.7×

This confirms that most of the time spent opening the dataset is used to build default indexes sequentially.

---

Additional Results (Larger Scale)

The 50-coordinate dataset took nearly 57 seconds to open, which was expected, as the latency added by each GET request was set to 1 second.

CC: @dcherian @TomNicholas

[keewis]

with #8051, you will be able to skip the creation of default indexes upon opening, load the variables into memory using compute (which should fetch the coordinate variables at the same time), and then create indexes from memory.

Obviously, that's a manual workaround, but #8051 also adds an obvious place on where to put this kind of optimizing (for this to work with DataTree, we'd have to refactor a bit, though)

[TomNicholas]

Obviously, that's a manual workaround, but #8051 also adds an obvious place on where to put this kind of optimizing (for this to work with DataTree, we'd have to refactor a bit, though)

The hard part of this is that the optimized way to load the data is to do it asynchronously, i.e. load all coordinate variables needed concurrently. Xarray doesn't have API for that. The Zarr backend does, but at a level lower than where the indexes are set. So I'm not totally sure yet how to achieve this refactor without breaking those abstractions.

[dcherian]

We could start passing a list of "variables to eagerly load" down to the backend, so that Zarr (and perhaps pydap in the future) could async load those. This is backwards-incompatible though I think. For now, we could just hard-code that to the Zarr backend and pass it in backend_kwargs.

This serial loop (and many others in Xarray) should be trivial to run concurrently (#8965)

xarray/xarray/backends/zarr.py

Line 1700 in 5ce69b2

for path_group, store in stores.items():

[keewis]

is there a way we can push this to the backends? Such that in load we'd be able to:

1. group by all variables of a common backend (determined by a equal close method or something similar)
2. call a method on the backend that does the loading concurrently and pass it all variables to load at once
3. assemble the resulting dataset

Either way, am I correct in assuming that this is only an issue for lazy loaded variables, but not for those that are managed by a compute / chunk manager?

[TomNicholas]

group by all variables of a common backend (determined by a equal close method or something similar)

Not sure what you mean by this.

call a method on the backend that does the loading concurrently and pass it all variables to load at once

I think that's similar to what Deepak is suggesting, though with a slightly different API.

Either way, am I correct in assuming that this is only an issue for lazy loaded variables, but not for those that are managed by a compute / chunk manager?

I believe that is true - dask/cubed would be able to load the variables concurrently. But we should check that the computemanager is actually called in the necessary way for that first.

[keewis]

Not sure what you mean by this.

I was thinking that merged datasets could be opened using different backends (e.g. if you use open_mfdataset without specifying the backend) so you'd have to map variables to the corresponding backend, but not sure if that's easy to do.

[TomNicholas]

I was thinking that merged datasets could be opened using different backends (e.g. if you use open_mfdataset without specifying the backend) so you'd have to map variables to the corresponding backend, but not sure if that's easy to do.

Is that a thing that's supported today?? I hope not, that sounds complicated and uneccessary.

[keewis]

I wouldn't say "supported", but it does work as a side-effect of the backend auto-detection:

import xarray as xr

ds = xr.tutorial.open_dataset("air_temperature")
ds2 = (ds > 273.15).rename_vars({"air": "mask"})

ds.to_netcdf("ds1.nc", mode="w")
ds2.to_zarr("ds2.zarr", mode="w")

xr.open_mfdataset("ds*")

[TomNicholas]

Summarizing a discussion that @dcherian, @aladinor, and I had about this:

• @dcherian noticed that actually creating the pandas.Index objects from loaded numpy arrays in memory doesn't release the GIL, so can't easily be parallelized. That means that even if we load the required zarr arrays concurrently, there will still effectively be a serial loop over the total number of coordinates to create. But unless the individual indexes are massive, this step doesn't take much time compared to loading from a remote store, so it's still worthwhile to make that loading concurrent.
• We can't know which variables to load until the backend store has been opened.
• We also don't know whether or not to load any variables unless the backend knows the value of the create_default_indexes kwarg (or similar).
• There are therefore two basic ways to make this work:
  1. Pass information about which variables to load down into the zarr backend, and load them concurrently there.
  2. Have the backend create lazy representations of all variables as it currently does, then use the upcoming Variable.load_async() method to concurrently load all the variables needed for the indexes outside of the backend code.
• The second approach can only be done for a backend which supports asynchronous loading, which (after Add asynchronous load method #10327) will initially only be Zarr, but could in theory be others too (perhaps PyDAP).
• But then the non-backend code needs to know whether or not to call the sync or async .load method, which means it needs to know which backends support async. Perhaps via some new .supports_async property on all backends.
• So either way we need to reach across abstractions: either we need to pass information about which variables to load down into the backend, or have backends pass information up by stating what loading methods they support.
• The .load_async method seems architecturally a lot cleaner, but then xarray has to decide how exactly to run those async functions. I'm planning to raise another issue about why this is not trivial.

[dcherian]

actually creating the pandas.Index objects from loaded numpy arrays in memory doesn't release the GIL, so can't easily be parallelized.

This is so fast, it doesn't matter

%load_ext ptime
import pandas as pd
import numpy as np

%timeit pd.Index(np.arange(100_000))

17.4 μs ± 230 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

This is what I benchmarked during our meeting

%ptime pd.Index(np.arange(1000))

Total serial time:   0.00 s
Total parallel time: 0.00 s
For a 0.26X speedup across 2 threads

But for large enough indexes, we do see speedup with concurrency

%ptime pd.Index(np.arange(10_000_000))

Total serial time:   0.03 s
Total parallel time: 0.01 s
For a 2.36X speedup across 2 threads

So I concluded wrong during our meeting.

But unless the individual indexes are massive, this step doesn't take much time compared to loading from a remote store, so it's still worthwhile to make that loading concurrent.

Yes the serial data loading is certainly the bottleneck, particularly for large hierarchies.