Avoid copying vectorized indexes (#10316)

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Author: jder

asv_bench/benchmarks/indexing.py

@@ -39,18 +39,30 @@
     "2d-1scalar": xr.DataArray(randn(100, frac_nan=0.1), dims=["x"]),
 }
 
-vectorized_indexes = {
-    "1-1d": {"x": xr.DataArray(randint(0, nx, 400), dims="a")},
-    "2-1d": {
-        "x": xr.DataArray(randint(0, nx, 400), dims="a"),
-        "y": xr.DataArray(randint(0, ny, 400), dims="a"),
-    },
-    "3-2d": {
-        "x": xr.DataArray(randint(0, nx, 400).reshape(4, 100), dims=["a", "b"]),
-        "y": xr.DataArray(randint(0, ny, 400).reshape(4, 100), dims=["a", "b"]),
-        "t": xr.DataArray(randint(0, nt, 400).reshape(4, 100), dims=["a", "b"]),
-    },
-}
+
+def make_vectorized_indexes(n_index):
+    return {
+        "1-1d": {"x": xr.DataArray(randint(0, nx, n_index), dims="a")},
+        "2-1d": {
+            "x": xr.DataArray(randint(0, nx, n_index), dims="a"),
+            "y": xr.DataArray(randint(0, ny, n_index), dims="a"),
+        },
+        "3-2d": {
+            "x": xr.DataArray(
+                randint(0, nx, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
+            ),
+            "y": xr.DataArray(
+                randint(0, ny, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
+            ),
+            "t": xr.DataArray(
+                randint(0, nt, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
+            ),
+        },
+    }
+
+
+vectorized_indexes = make_vectorized_indexes(400)
+big_vectorized_indexes = make_vectorized_indexes(400_000)
 
 vectorized_assignment_values = {
     "1-1d": xr.DataArray(randn((400, ny)), dims=["a", "y"], coords={"a": randn(400)}),
@@ -101,6 +113,20 @@ def time_indexing_basic_ds_large(self, key):
         self.ds_large.isel(**basic_indexes[key]).load()
 
 
+class IndexingOnly(Base):
+    @parameterized(["key"], [list(basic_indexes.keys())])
+    def time_indexing_basic(self, key):
+        self.ds.isel(**basic_indexes[key])
+
+    @parameterized(["key"], [list(outer_indexes.keys())])
+    def time_indexing_outer(self, key):
+        self.ds.isel(**outer_indexes[key])
+
+    @parameterized(["key"], [list(big_vectorized_indexes.keys())])
+    def time_indexing_big_vectorized(self, key):
+        self.ds.isel(**big_vectorized_indexes[key])
+
+
 class Assignment(Base):
     @parameterized(["key"], [list(basic_indexes.keys())])
     def time_assignment_basic(self, key):

doc/whats-new.rst

@@ -59,6 +59,14 @@ Bug fixes
   and prevents round-tripping them as :py:class:`numpy.datetime64` values
   (:pull:`10352`). By `Spencer Clark <https://github.com/spencerkclark>`_.
 
+Performance
+~~~~~~~~~~~
+- Lazily indexed arrays now use less memory to store keys by avoiding copies
+  in :py:class:`~xarray.indexing.VectorizedIndexer` and :py:class:`~xarray.indexing.OuterIndexer`
+  (:issue:`10316`).
+  By `Jesse Rusak <https://github.com/jder>`_.
+
+
 Documentation
 ~~~~~~~~~~~~~
 

xarray/core/indexing.py

@@ -444,7 +444,7 @@ def __init__(
                         f"invalid indexer array for {type(self).__name__}; must be scalar "
                         f"or have 1 dimension: {k!r}"
                     )
-                k = k.astype(np.int64)  # type: ignore[union-attr]
+                k = duck_array_ops.astype(k, np.int64, copy=False)
             else:
                 raise TypeError(
                     f"unexpected indexer type for {type(self).__name__}: {k!r}"
@@ -488,7 +488,7 @@ def __init__(self, key: tuple[slice | np.ndarray[Any, np.dtype[np.generic]], ...
                         "invalid indexer key: ndarray arguments "
                         f"have different numbers of dimensions: {ndims}"
                     )
-                k = k.astype(np.int64)  # type: ignore[union-attr]
+                k = duck_array_ops.astype(k, np.int64, copy=False)
             else:
                 raise TypeError(
                     f"unexpected indexer type for {type(self).__name__}: {k!r}"