slicing a slice with an array without expanding the slice

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,
@@ -41,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:
@@ -298,17 +302,52 @@ def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice:
     return slice(start, stop, step)
 
 
-def _index_indexer_1d(old_indexer, applied_indexer, size: int):
-    if isinstance(applied_indexer, slice) and applied_indexer == slice(None):
+def slice_slice_by_array(
+    old_slice: slice,
+    array: np.ndarray[Any, np.dtype[np.integer]],
+    size: int,
+) -> 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
+
+    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)
+
+    new_indexer = array * normalized.step + normalized.start
+
+    if np.any(new_indexer >= size):
+        raise IndexError("indices out of bounds")  # TODO: more helpful error message
+
+    return new_indexer
+
+
+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
     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]
+        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
@@ -387,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}")
 
@@ -419,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}")

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)