Function _decompose_vectorized_indexer

xarray/core/indexing.py:1237–1316

Decompose vectorized indexer to the successive two indexers, where the first indexer will be used to index backend arrays, while the second one is used to index loaded on-memory np.ndarray. Parameters ---------- indexer : VectorizedIndexer indexing_support : one of Inde

(
    indexer: VectorizedIndexer,
    shape: _Shape,
    indexing_support: IndexingSupport,
)

def _decompose_vectorized_indexer(
    indexer: VectorizedIndexer,
    shape: _Shape,
    indexing_support: IndexingSupport,
) -> tuple[ExplicitIndexer, ExplicitIndexer]:
    """
    Decompose vectorized indexer to the successive two indexers, where the
    first indexer will be used to index backend arrays, while the second one
    is used to index loaded on-memory np.ndarray.

    Parameters
    ----------
    indexer : VectorizedIndexer
    indexing_support : one of IndexerSupport entries

    Returns
    -------
    backend_indexer: OuterIndexer or BasicIndexer
    np_indexers: an ExplicitIndexer (VectorizedIndexer / BasicIndexer)

    Notes
    -----
    This function is used to realize the vectorized indexing for the backend
    arrays that only support basic or outer indexing.

    As an example, let us consider to index a few elements from a backend array
    with a vectorized indexer ([0, 3, 1], [2, 3, 2]).
    Even if the backend array only supports outer indexing, it is more
    efficient to load a subslice of the array than loading the entire array,

    >>> array = np.arange(36).reshape(6, 6)
    >>> backend_indexer = OuterIndexer((np.array([0, 1, 3]), np.array([2, 3])))
    >>> # load subslice of the array
    ... array = NumpyIndexingAdapter(array).oindex[backend_indexer]
    >>> np_indexer = VectorizedIndexer((np.array([0, 2, 1]), np.array([0, 1, 0])))
    >>> # vectorized indexing for on-memory np.ndarray.
    ... NumpyIndexingAdapter(array).vindex[np_indexer]
    array([ 2, 21,  8])
    """
    assert isinstance(indexer, VectorizedIndexer)

    if indexing_support is IndexingSupport.VECTORIZED:
        return indexer, BasicIndexer(())

    backend_indexer_elems: list[slice | np.ndarray[Any, np.dtype[np.generic]]] = []
    np_indexer_elems: list[slice | np.ndarray[Any, np.dtype[np.generic]]] = []
    # convert negative indices
    indexer_elems = [
        np.where(k < 0, k + s, k) if isinstance(k, np.ndarray) else k
        for k, s in zip(indexer.tuple, shape, strict=True)
    ]

    for k, s in zip(indexer_elems, shape, strict=True):
        if isinstance(k, slice):
            # If it is a slice, then we will slice it as-is
            # (but make its step positive) in the backend,
            # and then use all of it (slice(None)) for the in-memory portion.
            bk_slice, np_slice = _decompose_slice(k, s)