Feat: Add a mechanism for Serializing any index - `to_zarr` on a big RangeIndex causes a crash

[ianhi]

Is your feature request related to a problem?

As I understand it the current way that indexes are serialized is by writing the full contents of their coords to disk. This has some admirable properties in terms of interoperability with other tools and backcompatibility, but also some pitfalls.

1. Places burden on user's of xarray to remember re-install the index after loading + no mechanism for remembering what coords had set_xindex called on them.

A simple example:

import xarray as xr

ds = xr.Dataset(
    {"data": ("time", [0, 1, 2, 3])},
    coords={"time": [0.1, 0.2, 0.3, 0.4], "time_metadata": ("time", [10, 15, 20, 25])},
).set_xindex("time_metadata")
ds


ds.to_zarr("extra_index.zarr", mode="w")
roundtripped = xr.open_zarr("extra_index.zarr")

assert len(ds.xindexes) == len(roundtripped.xindexes)

fails with an assertion error.

2. For lazy indexes it's possible to crash your computer or run out of disk space with a simple to_zarr

For example, this snippet will create a 120 MiB zarr store when

import xarray as xr
step="1e-5"
idx1 = xr.indexes.RangeIndex.arange(0.0, 1000.0, float(step), dim="x")
ds1 = xr.Dataset(coords=xr.Coordinates.from_xindex(idx1))

ds1.to_zarr(f"float_range_step_{step}.zarr")

and I had to restart my laptop to get it unfrozen after running the below (offending line commented out) with the step value from the example (https://xarray-indexes.readthedocs.io/builtin/range.html)

import xarray as xr
step="1e-9"
idx1 = xr.indexes.RangeIndex.arange(0.0, 1000.0, float(step), dim="x")
ds1 = xr.Dataset(coords=xr.Coordinates.from_xindex(idx1))

# ds1.to_zarr(f"float_range_step_{step}.zarr")

Describe the solution you'd like

Allow indexes more control over how they are serialized and support a mechanism for automatically re-creating them if they are available in the current environment, with a graceful fallback.

Describe alternatives you've considered

If writing a custom index you could write your own loading function, but that would not solve the coord-bomb behavior and places more friction on the user.

Additional context

No response

[sshekhar563]

Can you please assign this issue to me?

[benbovy]

The user experience could be improved indeed! I'm not sure if those pitfalls would be best addressed at the index level, though.

1. Roundtripping indexes

IMO a better abstraction level would be the Xarray IO backend API (assuming 3rd-party backends eventually support writing datasets too) and/or Dataset "transformers" or "coders": #8548.

• For a generic backend like "zarr", I think that it is fine to leave some burden and freedom to the user. We don't always want to automatically rebuild potentially expensive indexes when roundtripping. Too much magic / complex transformation logic may have pitfalls too. With Automatically create xindex #11029 it is also less an issue now, at least for 1-dimensional coordinates.

• For a domain-specific backend it would make sense to have fancier roundtripping, i.e., apply a specific pipeline of transformers that involve index creation and custom coordinate (de)serialization. It would still be possible for different backends relying on the same Xarray indexes to implement different serialization strategies depending on their needs.

2. Lazy-coordinate loading

This is a problem related to lazy coordinates, at least not directly related to Xarray indexes. A simple ds1.x.values or ds1.x.load() will probably freeze your laptop as well. EDIT: Addressing the issue at the index level won't solve ds1.drop_indexes("x").to_zarr(...).

Alternatively, in all places triggering data loading we could first check for lazy coordinates with no chunks and maybe raise an error if the memory would exceed some threshold. And we could expose an option to change the threshold value:

with xr.set_options(lazy_coord_max_load_nbytes=...):
    ds1.x.values