From c2ee1a9aa308fb8f7a85ae106028d8624d888d70 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sat, 4 Mar 2023 21:53:29 -0700 Subject: [PATCH 1/8] Fix lazy slice rewriting. There was a bug in estimating the last index of the slice. Index a range object instead. Closes #6560 --- xarray/core/indexing.py | 20 ++++++++++++++++---- xarray/tests/test_backends.py | 21 ++++++++++++++++++--- 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index becf1554453..782dc0312ee 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -829,9 +829,21 @@ def decompose_indexer( raise TypeError(f"unexpected key type: {indexer}") -def _decompose_slice(key, size): +def _decompose_slice(key: slice, size: int) -> tuple[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: @@ -839,10 +851,10 @@ def _decompose_slice(key, size): 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( diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index a156b864315..7659094cf9e 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -711,9 +711,6 @@ def multiple_indexing(indexers): ] multiple_indexing(indexers5) - @pytest.mark.xfail( - reason="zarr without dask handles negative steps in slices incorrectly", - ) def test_vectorized_indexing_negative_step(self) -> None: # use dask explicitly when present open_kwargs: dict[str, Any] | None @@ -753,6 +750,24 @@ def multiple_indexing(indexers): ] multiple_indexing(indexers) + def test_outer_indexing_reversed(self) -> None: + # regression test for GH6560 + ds = xr.Dataset( + { + "z": ( + ("time", "isoBaricInhPa", "latitude", "longitude"), + np.ones((1, 5, 721, 1440)), + ) + }, + coords={"latitude": np.linspace(-90, 90, 721)}, + ) + + with self.roundtrip(ds) as on_disk: + subset = on_disk.isel(time=[0], isoBaricInhPa=1).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() From 2c75a10d1fc68299abccebc77f71f07278a1a3d7 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sat, 4 Mar 2023 22:30:33 -0700 Subject: [PATCH 2/8] Support reversed slices with Zarr --- xarray/backends/zarr.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 6686d67ed4d..62c3e4f95a8 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -77,7 +77,13 @@ def __getitem__(self, key): ] else: assert isinstance(key, indexing.OuterIndexer) - return array.oindex[key.tuple] + # decompose so that negative slices are now positive slices for the zarr library + # Then reverse the in-memory data. + backend_indexer, numpy_indexer = indexing._decompose_outer_indexer( + key, self.shape, indexing.IndexingSupport.OUTER + ) + part = array.oindex[backend_indexer.tuple] + return part[numpy_indexer.tuple] # if self.ndim == 0: # could possibly have a work-around for 0d data here From 483bead634b746e9759d0126385e1c533ccb4fdf Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Sat, 4 Mar 2023 22:46:35 -0700 Subject: [PATCH 3/8] Update xarray/core/indexing.py --- xarray/core/indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 782dc0312ee..9009318f012 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -829,7 +829,7 @@ def decompose_indexer( raise TypeError(f"unexpected key type: {indexer}") -def _decompose_slice(key: slice, size: int) -> tuple[slice, ...]: +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. From e0a46b4d74efb4af00b4fd8018a067d5b9d13ec1 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sun, 5 Mar 2023 20:52:25 -0700 Subject: [PATCH 4/8] Fix test --- xarray/backends/zarr.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index f104414c10a..4ef30e5e668 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -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): @@ -77,13 +80,12 @@ def __getitem__(self, key): ] else: assert isinstance(key, indexing.OuterIndexer) - # decompose so that negative slices are now positive slices for the zarr library - # Then reverse the in-memory data. - backend_indexer, numpy_indexer = indexing._decompose_outer_indexer( - key, self.shape, indexing.IndexingSupport.OUTER + # Lie about IndexingSupport so that we do + # rewrite negative slices to positive slices + return indexing.explicit_indexing_adapter( + key, array.shape, indexing.IndexingSupport.OUTER, self._oindex ) - part = array.oindex[backend_indexer.tuple] - return part[numpy_indexer.tuple] + # if self.ndim == 0: # could possibly have a work-around for 0d data here From a217485eb3c706183e309ecf9efe8355d1b7ef96 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sun, 5 Mar 2023 21:14:36 -0700 Subject: [PATCH 5/8] bring back xfail --- xarray/tests/test_backends.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 7d59e4cabfd..2bea337c359 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -711,6 +711,9 @@ def multiple_indexing(indexers): ] multiple_indexing(indexers5) + @pytest.mark.xfail( + reason="zarr without dask handles negative steps in slices incorrectly", + ) def test_vectorized_indexing_negative_step(self) -> None: # use dask explicitly when present open_kwargs: dict[str, Any] | None From e911e12857db8b10bdaab568341240f2327b434b Mon Sep 17 00:00:00 2001 From: dcherian Date: Mon, 13 Mar 2023 21:07:42 -0600 Subject: [PATCH 6/8] Smaller test --- xarray/tests/test_backends.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 2bea337c359..4d1d1f39fe5 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -756,19 +756,11 @@ def multiple_indexing(indexers): def test_outer_indexing_reversed(self) -> None: # regression test for GH6560 ds = xr.Dataset( - { - "z": ( - ("time", "isoBaricInhPa", "latitude", "longitude"), - np.ones((1, 5, 721, 1440)), - ) - }, - coords={"latitude": np.linspace(-90, 90, 721)}, + {"z": (("t", "p", "y", "x"), np.ones((1, 1, 31, 40)))}, ) with self.roundtrip(ds) as on_disk: - subset = on_disk.isel(time=[0], isoBaricInhPa=1).z[:, ::10, ::10][ - :, ::-1, : - ] + subset = on_disk.isel(t=[0], p=0).z[:, ::10, ::10][:, ::-1, :] assert subset.sizes == subset.load().sizes def test_isel_dataarray(self) -> None: From fb7ccc9ed93e243beae67b03336f55528eb63e3b Mon Sep 17 00:00:00 2001 From: dcherian Date: Thu, 16 Mar 2023 15:50:47 -0600 Subject: [PATCH 7/8] Better? --- xarray/backends/zarr.py | 4 +--- xarray/core/indexing.py | 22 ++++++++++++++++++---- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 4ef30e5e668..bc251d05631 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -80,10 +80,8 @@ def __getitem__(self, key): ] else: assert isinstance(key, indexing.OuterIndexer) - # Lie about IndexingSupport so that we do - # rewrite negative slices to positive slices return indexing.explicit_indexing_adapter( - key, array.shape, indexing.IndexingSupport.OUTER, self._oindex + key, array.shape, indexing.IndexingSupport.VECTORIZED, self._oindex ) # if self.ndim == 0: diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 9009318f012..0ba716fbf7d 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -23,6 +23,7 @@ NDArrayMixin, either_dict_or_kwargs, get_valid_numpy_dtype, + is_scalar, to_0d_array, ) @@ -980,12 +981,25 @@ def _decompose_outer_indexer( [14, 15, 14], [ 8, 9, 8]]) """ - if indexing_support == IndexingSupport.VECTORIZED: - return indexer, BasicIndexer(()) - assert isinstance(indexer, (OuterIndexer, BasicIndexer)) - backend_indexer: list[Any] = [] np_indexer = [] + + assert isinstance(indexer, (OuterIndexer, BasicIndexer)) + + 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): From de10523d8d0f9426f9d8cebfe60f08802b942cfd Mon Sep 17 00:00:00 2001 From: dcherian Date: Fri, 17 Mar 2023 09:17:05 -0600 Subject: [PATCH 8/8] fix typing --- xarray/core/indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 0ba716fbf7d..46d128b9e9d 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -982,7 +982,7 @@ def _decompose_outer_indexer( [ 8, 9, 8]]) """ backend_indexer: list[Any] = [] - np_indexer = [] + np_indexer: list[Any] = [] assert isinstance(indexer, (OuterIndexer, BasicIndexer))