Possible to select by label with integer slices from an integer dimension coordinate?

[blakeNaccarato]

(EDIT: Here's a dataset and a notebook Gist with outputs illustrating the behavior. In summary, sel with range not slice if you want label-based selection of integer coordinates)

So it seems sel is sort of the "preferred"/best-practice and best documented method, as opposed to say loc, but maybe loc behaves differently.

I have a DataArray instance, call it da, with an indexed dimension coordinate, frame. The coordinate is a sorted, continuous, integer index starting from 0. If I da.sel(frame=slice(None, None, 10), I benefit from automatic unbounded selection and I get back frames 0, 10, 20, etc. It's doing integer-based indexing under the hood, not label-based, but it's no issue here.

If I then just save back every tenth frame (a small subset for testing/checking into version control), load it back in as da_small, then da.sel(frame=slice(None, None, 10) would return 0, 100, 200, etc. I'd like it to still give 0, 10, 20.

I could sidestep this by passing a range index instead, but need to wrap load_dataset and conditionally pass a range indexer instead when an all-integer slice is passed, e.g. range(start or 0, (stop or da.frame[-1]) + (step := step or 1), step).

Ideally, for my case, the same slice index would always return 0, 10, 20, but I understand that doing the "right" thing with dynamic selection like sel means picking a sensible default, and integer slices aren't label-based by default.

1. Is there a built-in, best-practice way to achieve this? Does loc make a different distinction? I think I could also do some isin masking, but it's a bit more verbose and I think it's not done "lazily" for an on-disk data array?
2. Could I spin up my own custom index with minor changes to the default behavior, that prefers label-based indexing even on integer slices? (essentially baking in the range logic above)?
3. Is da.frame[-1] "lazy", e.g. it's not gonna fetch all the frames from disk if I do that?

xarray is absolutely amazing, by the way, and serves my high-speed video processing data flows just as well as it does for geosciences. It's a heck of a lot easier than wrangling raw numpy arrays, for sure. Anyways, thanks!

[keewis]

can you post an example dataset just so we can more easily follow what you're doing? For example, does this exhibit the behavior you're describing?

ds = xr.Dataset(coords={"frame": np.arange(1000)})
ds.sel(frame=slice(None, None, 10))

[keewis]

after investigating a bit further, I believe this is inherited behavior from pandas.Index.slice_indexer: while start and end can be labels, step must always be an integer size. As evidence, this fails:

ds = xr.Dataset(coords={"x": np.linspace(0, 1, 100)})
ds.sel(x=slice(None, None, 0.1))

and this fails with the same error:

s = pd.Series(np.arange(100), index=np.linspace(0, 1, 100))
s.loc[slice(None, None, 0.1)]

To fix this, we'd either have to write our own index, or convince the pandas maintainers that this is worth changing.

[blakeNaccarato]

Here's a dataset and a notebook Gist with outputs illustrating the behavior.

Thanks for the detail. So it probably can't be helped if it's indexing behavior adopted from upstream. I'm not going to explore customization/extension avenues now, but am I correct in my understanding that I might be able to do it with a custom accessor or the (experimental) custom index?

[benbovy]

I haven't looked at your notebook gist yet, but from what I can read here this Range1DIndex custom index example may be relevant: #9543 (comment).

(Note: that example only works within the context the draft PR where it has been shown).

A coordinate associated with this index is lazy, so da.frames[-1] will only compute the last frame index.

I think you could pretty easily customize it so it handles slice steps the way you want. If that behavior makes sense in the general case we might also consider supporting it in Xarray.

[blakeNaccarato]

Ah, great, thanks for the detail! Yes it seems my problem would be a specific use case for something like Range1DIndex or the related machinery discussed in that PR.

The Gist goes into more detail, but basically I'm using a get_selector helper function to map integer slices to range selections by "peeking" at the first and last coordinate indices. It's relatively cheap for my small coordinate indices and works for the time being.

from more_itertools import first, last
from xarray import DataArray

def get_selector(
    da: DataArray, dim: str, sel: slice | range | Any | None
) -> slice | range | Any:
    """Get selector, preferring label-based selection even given an integer slice."""
    if sel is None:
        return slice(None)
    if isinstance(sel, slice) and all(
        (isinstance(s, int) or s is None) for s in (sel.start, sel.stop, sel.step)
    ):
        start = sel.start or first(da[dim].values)
        stop = (sel.stop or last(da[dim].values)) + 1
        step = sel.step or 1
        return range(start, stop, step)
    return sel if isinstance(sel, slice | range) else list(sel)

I'm looking forward to making use of the new indexing features as they emerge, but I'm in no rush, as I'm sure reaching a stable API will represent quite a lot of work.