Allow setting (or skipping) new indexes in open_dataset

[benbovy]

• Closes Opening dataset without loading any indexes? #6633
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

This PR introduces a new boolean parameter set_indexes=True to xr.open_dataset(), which may be used to skip the creation of default (pandas) indexes when opening a dataset.

Currently works with the Zarr backend:

import numpy as np
import xarray as xr

# example dataset (real dataset may be much larger)
arr = np.random.random(size=1_000_000)
xr.Dataset({"x": arr}).to_zarr("dataset.zarr")


xr.open_dataset("dataset.zarr", set_indexes=False, engine="zarr")
# <xarray.Dataset>
# Dimensions:  (x: 1000000)
# Coordinates:
#     x        (x) float64 ...
# Data variables:
#     *empty*


xr.open_zarr("dataset.zarr", set_indexes=False)
# <xarray.Dataset>
# Dimensions:  (x: 1000000)
# Coordinates:
#     x        (x) float64 ...
# Data variables:
#     *empty*

I'll add it to the other Xarray backends as well, but I'd like to get your thoughts about the API first.

1. Do we want to add yet another keyword parameter to xr.open_dataset()? There are already many...
2. Do we want to add this parameter to the BackendEntrypoint.open_dataset() API?

• I'm afraid we must do it if we want this parameter in xr.open_dataset()
• this would also make it possible skipping the creation of custom indexes (if any) in custom IO backends
• con: if we require set_indexes in the signature in addition to the drop_variables parameter, this is a breaking change for all existing 3rd-party backends. Or should we group set_indexes with the other xarray decoder kwargs? This would feel a bit odd to me as setting indexes is different from decoding data.

3. Or should we leave this up to the backends?

• pros: no breaking change, more flexible (3rd party backends may want to offer more control like choosing between custom indexes and default pandas indexes or skipping the creation of indexes by default)
• cons: less discoverable, consistency is not enforced across 3rd party backends (although for such advanced case this is probably OK), not available by default in every backend.

Currently 1 and 2 are implemented in this PR, although as I write this comment I think that I would prefer 3. I guess this depends on whether we prefer open_*** vs. xr.open_dataset(engine="***") and unless I missed something there is still no real consensus about that? (e.g., #7496).

[benbovy]

@pydata/xarray any thoughts on which option among those above (top comment) would be best?

[rabernat]

1. Do we want to add yet another keyword parameter to xr.open_dataset()?

I vote for this.

[dcherian]

Is there any way for packages to say they support V1 or V2 of an entrypoint?

I just realized that if we turned off indexes by default, it would be a big win for the open_mfdataset case in many cases.

[max-sixty]

• con: if we require set_indexes in the signature in addition to the drop_variables parameter, this is a breaking change for all existing 3rd-party backends

Doesn't help with the problem at this moment, but could we add having **kwargs to the standard, so we can add parameters in future without breaking existing backends?

[benbovy]

Agreed, adding **kwargs to the standard would help! However, to be honest I find it already a bit confusing how kwargs are handled in xarray.backends.api. #8447 may eventually help having a clearer separation between common and backend-specific options.

[shoyer]

This looks great to me!

I agree with adding this into xarray.open_dataset() and BackendEntrypoint.open_dataset().

For what it's worth, I think it's OK to require backend developers to update their code more frequently -- we don't need the same level of stabily that we need for user level APIs.

[TomNicholas]

I just realized that if we turned off indexes by default, it would be a big win for the open_mfdataset case in many cases.

I would love to see this merged so that I can try this out!

[keewis]

be aware that merging now will break compatibility with any 3rd party backend, which I believe is not something we should do, even if we think that the transition window can be shorter than usual.

I my eyes the easiest way forward would be:

1. extend Backend.open_dataset with **kwargs as suggested by @max-sixty
2. wait a couple of versions (2?) to give backends the time to update and release
3. add the new option

We don't have an easy way to contact all backend developers, unfortunately.

Edit: let's discuss in the meeting today

[keewis]

In the meeting just now we decided to inspect the signature of the backend's open_dataset method and not pass the new option if it doesn't support it nor accepts **kwargs.

We should still change the spec to require **kwargs (in a new PR), and maybe emit a deprecation warning for all backends that don't have it already.

[keewis]

we discussed this in the raster index meeting a bit with @dcherian and @benbovy, and if I remember correctly we agreed that instead of passing through an option to the backend it might be cleaner to always create the backend dataset without indexes, and create the indexes in open_dataset:

Suggested change

	if set_indexes:
	coords = coord_vars
	else:
	# explicit Coordinates object with no index passed
	coords = Coordinates(coord_vars)
	# explicit Coordinates object with no index passed
	coords = Coordinates(coord_vars)

Most known backends use the BackendStoreEntrypoint to assemble the backend dataset, so these backends will work automatically, and for all other backends we'll have to figure out a way to warn.

If we also change the data model to cleanly separate I/O from decoding variables (what zarr calls "ArrayToArray" codecs, and what our current decoders do) and datasets (this would be where you'd create indexes), then we could warn about backend datasets with indexes. I think this would be a breaking change, although we could add a deprecation cycle to make the transition smoother.

[keewis]

I've taken the liberty of pushing directly to this PR (the failing flaky test is real but unrelated, and pre-commit also fails on main), so the state in this PR should be closer to what I'm proposing above.

cc @dcherian, @benbovy

[TomNicholas]

cleaner to always create the backend dataset without indexes

I fully support this. I think it's the natural way to generalize. Creating an in-memory index from a lazily-loaded array is totally unrelated from creating the lazy arrays in the first place.

[dcherian]

I've changed my mind. In reality, we are seeing that people are writing backends not just for "reading arrays", they are creating datasets with extra metadata (e.g. rioxarray might choose to add RasterIndex, xwrf might do something else). So perhaps we add a create_default_indexes that sets the default PandasIndex for any dimension coordinate that does not have an index, when returned from the backend?

EDIT: I see that's what you've done. nice!

[keewis]

if you look at the current state of the PR, that's exactly my intention

Edit: although in an ideal world we'd add support for "dataset decoders", and then people would use that to create the indexes

[TomNicholas]

So perhaps we add a create_default_indexes that sets the default PandasIndex for any dimension coordinate that does not have an index, when returned from the backend?

The general version of this would be to make the index-adding step independently pluggable, i.e. pass in a callable, which defaults to your create_default_indexes

[TomNicholas]

Edit: although in an ideal world we'd add support for "dataset decoders", and then people would use that to create the indexes

This is basically what I'm suggesting - I think you should make the API something like

@runtime_checkable
class IndexSetter(Protocol):
    def __call__(
        self,
        backend_ds: xr.Dataset,
    ) -> xr.Dataset: ...


def open_dataset(
    ...,
    index_setter: IndexSetter = create_default_indexes,
) -> xr.Dataset:
    
    if not isinstance(index_setter, IndexSetter):
        raise TypeError

    # all the normal backend logic
    backend_ds = ...

    return index_setter(backend_ds)

either adding a kwarg to open_dataset like that or alternatively you could allow people to use the backends to be the place they choose which IndexSetter callable they want to use (e.g. to always add domain-specific indexes at read time).

[keewis]

that would be an option, but instead of a separate parameter I think it might be better to figure out how a dataset decoder API could look, then reuse that for the default indexes

[TomNicholas]

That makes sense, but then would we deprecate the kwarg you're about to add?

[benbovy]

Pluggable and composable Dataset (de)coders would be nice indeed for setting indexes among other things. I like the proposal made in #8548.

Even with a decoder API it might still make sense to have a special case (kwarg) for default indexes, though, since those are also created in Dataset.__init__(), DataArray.__init__(), Coordinates.__init__(), assign_coords(), etc.

[dcherian]

I think this is fine given our rules for propagating coordinate variables when extracting DataArrays but it is potentially confusing with create_default_indexes=False, decode_coordinates=False.

[dcherian]

Very nice!