From 54a5dedbc485161c2647295ca31fd49b215b2afb Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 20:17:45 +0200 Subject: [PATCH 01/16] use `is_full_slice` --- xarray/core/indexing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 8e4458fb88f..e0967391954 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -26,6 +26,7 @@ get_valid_numpy_dtype, is_duck_array, is_duck_dask_array, + is_full_slice, is_scalar, is_valid_numpy_dtype, to_0d_array, @@ -299,7 +300,7 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: def _index_indexer_1d(old_indexer, applied_indexer, size: int): - if isinstance(applied_indexer, slice) and applied_indexer == slice(None): + if is_full_slice(applied_indexer): # shortcut for the usual case return old_indexer if isinstance(old_indexer, slice): From 9cab2ace361f7f2a950600f7852b9c787c66973a Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 20:18:12 +0200 Subject: [PATCH 02/16] don't try to expand a full slice --- xarray/core/indexing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index e0967391954..ec540dffa01 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -308,6 +308,8 @@ def _index_indexer_1d(old_indexer, applied_indexer, size: int): indexer = slice_slice(old_indexer, applied_indexer, size) elif isinstance(applied_indexer, integer_types): indexer = range(*old_indexer.indices(size))[applied_indexer] # type: ignore[assignment] + elif is_full_slice(old_indexer): + indexer = applied_indexer else: indexer = _expand_slice(old_indexer, size)[applied_indexer] else: From 91e316fef6099c602e7fabca14484b0e886f7c1e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 20:24:41 +0200 Subject: [PATCH 03/16] slice a slice with an array without materializing the slice --- xarray/core/indexing.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index ec540dffa01..b0378605865 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -299,6 +299,17 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: return slice(start, stop, step) +def slice_slice_by_array(old_slice, array, size): + # _expand_slice(old_indexer, size)[applied_indexer] + normalized = normalize_slice(old_slice, size) + + new_indexer = array * normalized.step + normalized.start + if np.any(new_indexer >= normalized.stop): + raise IndexError("indices out of bounds") # TODO: more helpful error message + + return new_indexer + + def _index_indexer_1d(old_indexer, applied_indexer, size: int): if is_full_slice(applied_indexer): # shortcut for the usual case @@ -311,7 +322,7 @@ def _index_indexer_1d(old_indexer, applied_indexer, size: int): elif is_full_slice(old_indexer): indexer = applied_indexer else: - indexer = _expand_slice(old_indexer, size)[applied_indexer] + indexer = slice_slice_by_array(old_indexer, applied_indexer, size) else: indexer = old_indexer[applied_indexer] return indexer From 5c467ca714161fc619900043d45b30ea083ad612 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 20:32:55 +0200 Subject: [PATCH 04/16] formatting --- xarray/core/indexing.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index b0378605865..7860dcf6ff1 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -300,10 +300,9 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: def slice_slice_by_array(old_slice, array, size): - # _expand_slice(old_indexer, size)[applied_indexer] normalized = normalize_slice(old_slice, size) - new_indexer = array * normalized.step + normalized.start + if np.any(new_indexer >= normalized.stop): raise IndexError("indices out of bounds") # TODO: more helpful error message From 697be5d7af8f0691cbf7cc21dbc6e7c9fa6b9017 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:13:20 +0200 Subject: [PATCH 05/16] type hints --- xarray/core/indexing.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 7860dcf6ff1..4445ff057bc 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -299,8 +299,14 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: return slice(start, stop, step) -def slice_slice_by_array(old_slice, array, size): +def slice_slice_by_array( + old_slice: slice, + array: np.ndarray[Any, np.dtype[np.generic]], + size: int, +) -> np.ndarray[Any, np.dtype[np.generic]]: + # to get a concrete slice, limited to the size of the array normalized = normalize_slice(old_slice, size) + new_indexer = array * normalized.step + normalized.start if np.any(new_indexer >= normalized.stop): From e9667566a3a63ffa7012d62c562faa8c7312320a Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:27:45 +0200 Subject: [PATCH 06/16] check that the new algorithm works properly --- xarray/tests/test_indexing.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index 6dd75b58c6a..38f755c0a3e 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -305,6 +305,20 @@ def test_slice_slice(self) -> None: actual = x[new_slice] assert_array_equal(expected, actual) + @pytest.mark.parametrize( + ["old_slice", "array", "size", "expected"], + ( + (slice(None, 8), np.arange(2, 6), 10, np.arange(2, 6)), + (slice(2, None), np.arange(2, 6), 10, np.arange(4, 8)), + (slice(1, 10, 2), np.arange(1, 4), 15, np.arange(3, 8, 2)), + (slice(10, None, -1), np.array([2, 5, 7]), 12, np.array([8, 5, 3])), + ), + ) + def test_slice_slice_by_array(self, old_slice, array, size, expected): + actual = indexing.slice_slice_by_array(old_slice, array, size) + + assert_array_equal(actual, expected) + def test_lazily_indexed_array(self) -> None: original = np.random.rand(10, 20, 30) x = indexing.NumpyIndexingAdapter(original) From 4d828eda577b5d479ca8180783ea15c19f4fa80b Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:29:18 +0200 Subject: [PATCH 07/16] compare to `size` instead --- 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 4445ff057bc..e5a3a1ade1d 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -309,7 +309,7 @@ def slice_slice_by_array( new_indexer = array * normalized.step + normalized.start - if np.any(new_indexer >= normalized.stop): + if np.any(new_indexer >= size): raise IndexError("indices out of bounds") # TODO: more helpful error message return new_indexer From 78a79abb83720f86ea2e787159ab7bcd34fc131f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:43:31 +0200 Subject: [PATCH 08/16] doctests --- xarray/core/indexing.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index e5a3a1ade1d..ce5290c8da7 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -304,6 +304,19 @@ def slice_slice_by_array( array: np.ndarray[Any, np.dtype[np.generic]], size: int, ) -> np.ndarray[Any, np.dtype[np.generic]]: + """Given a slice and the size of the dimension to which it will be applied, + index it with an array to return a new array equivalent to applying + the slices sequentially + + Examples + -------- + >>> slice_slice_by_array(slice(2, 10), np.array([1, 3, 5]), 12) + array([3, 5, 7]) + >>> slice_slice_by_array(slice(1, None, 2), np.array([1, 3, 7, 8]), 20) + array([ 3, 7, 15, 17]) + >>> slice_slice_by_array(slice(None, None, -1), np.array([2, 4, 7]), 20) + array([17, 15, 12]) + """ # to get a concrete slice, limited to the size of the array normalized = normalize_slice(old_slice, size) From 8fef7733eb2fadb7d31b8e3b0c62e1491c2fabfd Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:48:09 +0200 Subject: [PATCH 09/16] type hints --- xarray/core/indexing.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index ce5290c8da7..5f2bfde829c 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -42,6 +42,9 @@ from xarray.namedarray._typing import _Shape, duckarray from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint + BasicIndexerType = int | np.integer | slice + OuterIndexerType = BasicIndexerType | np.ndarray[Any, np.dtype[np.integer]] + @dataclass class IndexSelResult: @@ -328,7 +331,11 @@ def slice_slice_by_array( return new_indexer -def _index_indexer_1d(old_indexer, applied_indexer, size: int): +def _index_indexer_1d( + old_indexer: OuterIndexerType, + applied_indexer: OuterIndexerType, + size: int, +) -> OuterIndexerType: if is_full_slice(applied_indexer): # shortcut for the usual case return old_indexer @@ -419,7 +426,7 @@ class BasicIndexer(ExplicitIndexer): __slots__ = () - def __init__(self, key: tuple[int | np.integer | slice, ...]): + def __init__(self, key: tuple[BasicIndexerType, ...]): if not isinstance(key, tuple): raise TypeError(f"key must be a tuple: {key!r}") @@ -451,9 +458,7 @@ class OuterIndexer(ExplicitIndexer): def __init__( self, - key: tuple[ - int | np.integer | slice | np.ndarray[Any, np.dtype[np.generic]], ... - ], + key: tuple[BasicIndexerType | np.ndarray[Any, np.dtype[np.generic]], ...], ): if not isinstance(key, tuple): raise TypeError(f"key must be a tuple: {key!r}") From 005d5bc4a6f8f3dfc48b0b0085fb28756dc75d1b Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 21:48:50 +0200 Subject: [PATCH 10/16] stricter type hints --- xarray/core/indexing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 5f2bfde829c..492eef0b320 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -304,9 +304,9 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: def slice_slice_by_array( old_slice: slice, - array: np.ndarray[Any, np.dtype[np.generic]], + array: np.ndarray[Any, np.dtype[np.integer]], size: int, -) -> np.ndarray[Any, np.dtype[np.generic]]: +) -> np.ndarray[Any, np.dtype[np.integer]]: """Given a slice and the size of the dimension to which it will be applied, index it with an array to return a new array equivalent to applying the slices sequentially From a217e6a71c6ed467cea7682c706943d8c4f4f233 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 22:06:20 +0200 Subject: [PATCH 11/16] check against `numpy` Co-authored-by: Deepak Cherian --- xarray/tests/test_indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index 38f755c0a3e..75566579b5e 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -316,7 +316,7 @@ def test_slice_slice(self) -> None: ) def test_slice_slice_by_array(self, old_slice, array, size, expected): actual = indexing.slice_slice_by_array(old_slice, array, size) - + expected = np.arange(size)[old_slice][array] assert_array_equal(actual, expected) def test_lazily_indexed_array(self) -> None: From 938610862c5f0b0b9fe37427306ba09a964ccb5c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 22:08:15 +0200 Subject: [PATCH 12/16] remove the old parametrized expected values --- xarray/tests/test_indexing.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index 75566579b5e..fbb6eb04e2e 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -306,17 +306,17 @@ def test_slice_slice(self) -> None: assert_array_equal(expected, actual) @pytest.mark.parametrize( - ["old_slice", "array", "size", "expected"], + ["old_slice", "array", "size"], ( - (slice(None, 8), np.arange(2, 6), 10, np.arange(2, 6)), - (slice(2, None), np.arange(2, 6), 10, np.arange(4, 8)), - (slice(1, 10, 2), np.arange(1, 4), 15, np.arange(3, 8, 2)), - (slice(10, None, -1), np.array([2, 5, 7]), 12, np.array([8, 5, 3])), + (slice(None, 8), np.arange(2, 6), 10), + (slice(2, None), np.arange(2, 6), 10), + (slice(1, 10, 2), np.arange(1, 4), 15), + (slice(10, None, -1), np.array([2, 5, 7]), 12), ), ) - def test_slice_slice_by_array(self, old_slice, array, size, expected): + def test_slice_slice_by_array(self, old_slice, array, size): actual = indexing.slice_slice_by_array(old_slice, array, size) - expected = np.arange(size)[old_slice][array] + expected = np.arange(size)[old_slice][array] assert_array_equal(actual, expected) def test_lazily_indexed_array(self) -> None: From 94577abc50cd01db903bcbb2905ce25fea82736f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 22:29:11 +0200 Subject: [PATCH 13/16] fix type hints --- xarray/core/indexing.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 492eef0b320..8753e56a22f 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -339,17 +339,23 @@ def _index_indexer_1d( if is_full_slice(applied_indexer): # shortcut for the usual case return old_indexer + + indexer: OuterIndexerType if isinstance(old_indexer, slice): if isinstance(applied_indexer, slice): indexer = slice_slice(old_indexer, applied_indexer, size) elif isinstance(applied_indexer, integer_types): - indexer = range(*old_indexer.indices(size))[applied_indexer] # type: ignore[assignment] + indexer = range(*old_indexer.indices(size))[applied_indexer] elif is_full_slice(old_indexer): indexer = applied_indexer else: indexer = slice_slice_by_array(old_indexer, applied_indexer, size) - else: + elif isinstance(old_indexer, np.ndarray): indexer = old_indexer[applied_indexer] + else: + # should be unreachable + raise ValueError("cannot index integers. Please open an issuec-") + return indexer @@ -664,7 +670,8 @@ def __init__(self, array: Any, key: ExplicitIndexer | None = None): def _updated_key(self, new_key: ExplicitIndexer) -> BasicIndexer | OuterIndexer: iter_new_key = iter(expanded_indexer(new_key.tuple, self.ndim)) - full_key = [] + + full_key: list[OuterIndexerType] = [] for size, k in zip(self.array.shape, self.key.tuple, strict=True): if isinstance(k, integer_types): full_key.append(k) @@ -673,7 +680,7 @@ def _updated_key(self, new_key: ExplicitIndexer) -> BasicIndexer | OuterIndexer: full_key_tuple = tuple(full_key) if all(isinstance(k, integer_types + (slice,)) for k in full_key_tuple): - return BasicIndexer(full_key_tuple) + return BasicIndexer(cast(tuple[BasicIndexerType, ...], full_key_tuple)) return OuterIndexer(full_key_tuple) @property From 3d316a9e298eabbd4ff7cc8eb771931649af35a5 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 28 Jul 2025 22:30:40 +0200 Subject: [PATCH 14/16] shortcut for existing full slices --- xarray/core/indexing.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 8753e56a22f..1af14a4171e 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -339,6 +339,9 @@ def _index_indexer_1d( if is_full_slice(applied_indexer): # shortcut for the usual case return old_indexer + if is_full_slice(old_indexer): + # shortcut for full slices + return applied_indexer indexer: OuterIndexerType if isinstance(old_indexer, slice): @@ -346,8 +349,6 @@ def _index_indexer_1d( indexer = slice_slice(old_indexer, applied_indexer, size) elif isinstance(applied_indexer, integer_types): indexer = range(*old_indexer.indices(size))[applied_indexer] - elif is_full_slice(old_indexer): - indexer = applied_indexer else: indexer = slice_slice_by_array(old_indexer, applied_indexer, size) elif isinstance(old_indexer, np.ndarray): From ea45bde176c71cea622992609a49b66eeceea5e7 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 29 Jul 2025 12:08:52 +0200 Subject: [PATCH 15/16] move the definition of the new types out of the type checking block --- xarray/core/indexing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 1af14a4171e..4f147a56edb 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -42,8 +42,8 @@ from xarray.namedarray._typing import _Shape, duckarray from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint - BasicIndexerType = int | np.integer | slice - OuterIndexerType = BasicIndexerType | np.ndarray[Any, np.dtype[np.integer]] +BasicIndexerType = int | np.integer | slice +OuterIndexerType = BasicIndexerType | np.ndarray[Any, np.dtype[np.integer]] @dataclass From fdc90a21dbd8a6d70d07e69137da56a618af3126 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 30 Jul 2025 22:30:42 +0200 Subject: [PATCH 16/16] also support negative array values --- xarray/core/indexing.py | 26 ++++++++++++++++++++++++-- xarray/tests/test_indexing.py | 2 ++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 4f147a56edb..1153c2e06d0 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -302,6 +302,25 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: return slice(start, stop, step) +def normalize_array( + array: np.ndarray[Any, np.dtype[np.integer]], size: int +) -> np.ndarray[Any, np.dtype[np.integer]]: + """ + Ensure that the given array only contains positive values. + + Examples + -------- + >>> normalize_array(np.array([-1, -2, -3, -4]), 10) + array([9, 8, 7, 6]) + >>> normalize_array(np.array([-5, 3, 5, -1, 8]), 12) + array([ 7, 3, 5, 11, 8]) + """ + if np.issubdtype(array.dtype, np.unsignedinteger): + return array + + return np.where(array >= 0, array, array + size) + + def slice_slice_by_array( old_slice: slice, array: np.ndarray[Any, np.dtype[np.integer]], @@ -321,9 +340,12 @@ def slice_slice_by_array( array([17, 15, 12]) """ # to get a concrete slice, limited to the size of the array - normalized = normalize_slice(old_slice, size) + normalized_slice = normalize_slice(old_slice, size) + + size_after_slice = len(range(*normalized_slice.indices(size))) + normalized_array = normalize_array(array, size_after_slice) - new_indexer = array * normalized.step + normalized.start + new_indexer = normalized_array * normalized_slice.step + normalized_slice.start if np.any(new_indexer >= size): raise IndexError("indices out of bounds") # TODO: more helpful error message diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index fbb6eb04e2e..db4f6aaf0bd 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -312,6 +312,8 @@ def test_slice_slice(self) -> None: (slice(2, None), np.arange(2, 6), 10), (slice(1, 10, 2), np.arange(1, 4), 15), (slice(10, None, -1), np.array([2, 5, 7]), 12), + (slice(2, None, 2), np.array([3, -2, 5, -1]), 13), + (slice(8, None), np.array([1, -2, 2, -1, -7]), 20), ), ) def test_slice_slice_by_array(self, old_slice, array, size):