Fix performance regression in interp from pydata#9881 (pydata#10370)

Author: dcherian

doc/whats-new.rst

@@ -69,6 +69,8 @@ Performance
   in :py:class:`~xarray.indexing.VectorizedIndexer` and :py:class:`~xarray.indexing.OuterIndexer`
   (:issue:`10316`).
   By `Jesse Rusak <https://github.com/jder>`_.
+- Fix performance regression in interp where more data was loaded than was necessary. (:issue:`10287`).
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 - Speed up encoding of :py:class:`cftime.datetime` objects by roughly a factor
   of three (:pull:`8324`). By `Antoine Gibek <https://github.com/antscloud>`_.
 

xarray/core/dataset.py

@@ -3804,15 +3804,16 @@ def _validate_interp_indexer(x, new_x):
             for k, v in indexers.items()
         }
 
+        # optimization: subset to coordinate range of the target index
+        if method in ["linear", "nearest"]:
+            for k, v in validated_indexers.items():
+                obj, newidx = missing._localize(obj, {k: v})
+                validated_indexers[k] = newidx[k]
+
         has_chunked_array = bool(
             any(is_chunked_array(v._data) for v in obj._variables.values())
         )
         if has_chunked_array:
-            # optimization: subset to coordinate range of the target index
-            if method in ["linear", "nearest"]:
-                for k, v in validated_indexers.items():
-                    obj, newidx = missing._localize(obj, {k: v})
-                    validated_indexers[k] = newidx[k]
             # optimization: create dask coordinate arrays once per Dataset
             # rather than once per Variable when dask.array.unify_chunks is called later
             # GH4739
@@ -3828,7 +3829,7 @@ def _validate_interp_indexer(x, new_x):
                 continue
 
             use_indexers = (
-                dask_indexers if is_duck_dask_array(var.data) else validated_indexers
+                dask_indexers if is_duck_dask_array(var._data) else validated_indexers
             )
 
             dtype_kind = var.dtype.kind

xarray/tests/test_missing.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import itertools
+from unittest import mock
 
 import numpy as np
 import pandas as pd
 import pytest
 
 import xarray as xr
+from xarray.core import indexing
 from xarray.core.missing import (
     NumpyInterpolator,
     ScipyInterpolator,
@@ -772,3 +774,29 @@ def test_interpolators_complex_out_of_bounds():
         f = interpolator(xi, yi, method=method)
         actual = f(x)
         assert_array_equal(actual, expected)
+
+
+@requires_scipy
+def test_indexing_localize():
+    # regression test for GH10287
+    ds = xr.Dataset(
+        {
+            "sigma_a": xr.DataArray(
+                data=np.ones((16, 8, 36811)),
+                dims=["p", "t", "w"],
+                coords={"w": np.linspace(0, 30000, 36811)},
+            )
+        }
+    )
+
+    original_func = indexing.NumpyIndexingAdapter.__getitem__
+
+    def wrapper(self, indexer):
+        return original_func(self, indexer)
+
+    with mock.patch.object(
+        indexing.NumpyIndexingAdapter, "__getitem__", side_effect=wrapper, autospec=True
+    ) as mock_func:
+        ds["sigma_a"].interp(w=15000.5)
+    actual_indexer = mock_func.mock_calls[0].args[1]._key
+    assert actual_indexer == (slice(None), slice(None), slice(18404, 18408))