Skip to content

Commit

Permalink
Fix lazy negative slice rewriting. (#7586)
Browse files Browse the repository at this point in the history
* Fix lazy slice rewriting.

There was a bug in estimating the last index of the slice.
Index a range object instead.

Closes #6560

* Support reversed slices with Zarr

* Update xarray/core/indexing.py

* Fix test

* bring back xfail

* Smaller test

* Better?

* fix typing
  • Loading branch information
dcherian authored Mar 27, 2023
1 parent de05403 commit 020b4c0
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 9 deletions.
8 changes: 7 additions & 1 deletion xarray/backends/zarr.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ def __init__(self, variable_name, datastore):
def get_array(self):
return self.datastore.zarr_group[self.variable_name]

def _oindex(self, key):
return self.get_array().oindex[key]

def __getitem__(self, key):
array = self.get_array()
if isinstance(key, indexing.BasicIndexer):
Expand All @@ -77,7 +80,10 @@ def __getitem__(self, key):
]
else:
assert isinstance(key, indexing.OuterIndexer)
return array.oindex[key.tuple]
return indexing.explicit_indexing_adapter(
key, array.shape, indexing.IndexingSupport.VECTORIZED, self._oindex
)

# if self.ndim == 0:
# could possibly have a work-around for 0d data here

Expand Down
42 changes: 34 additions & 8 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
NDArrayMixin,
either_dict_or_kwargs,
get_valid_numpy_dtype,
is_scalar,
to_0d_array,
)

Expand Down Expand Up @@ -840,20 +841,32 @@ def decompose_indexer(
raise TypeError(f"unexpected key type: {indexer}")


def _decompose_slice(key, size):
def _decompose_slice(key: slice, size: int) -> tuple[slice, slice]:
"""convert a slice to successive two slices. The first slice always has
a positive step.
>>> _decompose_slice(slice(2, 98, 2), 99)
(slice(2, 98, 2), slice(None, None, None))
>>> _decompose_slice(slice(98, 2, -2), 99)
(slice(4, 99, 2), slice(None, None, -1))
>>> _decompose_slice(slice(98, 2, -2), 98)
(slice(3, 98, 2), slice(None, None, -1))
>>> _decompose_slice(slice(360, None, -10), 361)
(slice(0, 361, 10), slice(None, None, -1))
"""
start, stop, step = key.indices(size)
if step > 0:
# If key already has a positive step, use it as is in the backend
return key, slice(None)
else:
# determine stop precisely for step > 1 case
# Use the range object to do the calculation
# e.g. [98:2:-2] -> [98:3:-2]
stop = start + int((stop - start - 1) / step) * step + 1
start, stop = stop + 1, start + 1
return slice(start, stop, -step), slice(None, None, -1)
exact_stop = range(start, stop, step)[-1]
return slice(exact_stop, start + 1, -step), slice(None, None, -1)


def _decompose_vectorized_indexer(
Expand Down Expand Up @@ -979,12 +992,25 @@ def _decompose_outer_indexer(
[14, 15, 14],
[ 8, 9, 8]])
"""
if indexing_support == IndexingSupport.VECTORIZED:
return indexer, BasicIndexer(())
backend_indexer: list[Any] = []
np_indexer: list[Any] = []

assert isinstance(indexer, (OuterIndexer, BasicIndexer))

backend_indexer: list[Any] = []
np_indexer = []
if indexing_support == IndexingSupport.VECTORIZED:
for k, s in zip(indexer.tuple, shape):
if isinstance(k, slice):
# If it is a slice, then we will slice it as-is
# (but make its step positive) in the backend,
bk_slice, np_slice = _decompose_slice(k, s)
backend_indexer.append(bk_slice)
np_indexer.append(np_slice)
else:
backend_indexer.append(k)
if not is_scalar(k):
np_indexer.append(slice(None))
return type(indexer)(tuple(backend_indexer)), BasicIndexer(tuple(np_indexer))

# make indexer positive
pos_indexer: list[np.ndarray | int | np.number] = []
for k, s in zip(indexer.tuple, shape):
Expand Down
10 changes: 10 additions & 0 deletions xarray/tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -753,6 +753,16 @@ def multiple_indexing(indexers):
]
multiple_indexing(indexers)

def test_outer_indexing_reversed(self) -> None:
# regression test for GH6560
ds = xr.Dataset(
{"z": (("t", "p", "y", "x"), np.ones((1, 1, 31, 40)))},
)

with self.roundtrip(ds) as on_disk:
subset = on_disk.isel(t=[0], p=0).z[:, ::10, ::10][:, ::-1, :]
assert subset.sizes == subset.load().sizes

def test_isel_dataarray(self) -> None:
# Make sure isel works lazily. GH:issue:1688
in_memory = create_test_data()
Expand Down

0 comments on commit 020b4c0

Please sign in to comment.