Skip to content

Commit 3b06696

Browse files
committed
Fix performance regression in interp from pydata#9881
Closes pydata#10287
1 parent 303748c commit 3b06696

File tree

3 files changed

+37
-6
lines changed

3 files changed

+37
-6
lines changed

doc/whats-new.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ Performance
6565
in :py:class:`~xarray.indexing.VectorizedIndexer` and :py:class:`~xarray.indexing.OuterIndexer`
6666
(:issue:`10316`).
6767
By `Jesse Rusak <https://github.com/jder>`_.
68+
- Fix performance regression in interp where more data was loaded than was necessary. (:issue:`10287`).
69+
By `Deepak Cherian <https://github.com/dcherian>`_.
6870

6971

7072
Documentation

xarray/core/dataset.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3805,15 +3805,16 @@ def _validate_interp_indexer(x, new_x):
38053805
for k, v in indexers.items()
38063806
}
38073807

3808+
# optimization: subset to coordinate range of the target index
3809+
if method in ["linear", "nearest"]:
3810+
for k, v in validated_indexers.items():
3811+
obj, newidx = missing._localize(obj, {k: v})
3812+
validated_indexers[k] = newidx[k]
3813+
38083814
has_chunked_array = bool(
38093815
any(is_chunked_array(v._data) for v in obj._variables.values())
38103816
)
38113817
if has_chunked_array:
3812-
# optimization: subset to coordinate range of the target index
3813-
if method in ["linear", "nearest"]:
3814-
for k, v in validated_indexers.items():
3815-
obj, newidx = missing._localize(obj, {k: v})
3816-
validated_indexers[k] = newidx[k]
38173818
# optimization: create dask coordinate arrays once per Dataset
38183819
# rather than once per Variable when dask.array.unify_chunks is called later
38193820
# GH4739
@@ -3829,7 +3830,7 @@ def _validate_interp_indexer(x, new_x):
38293830
continue
38303831

38313832
use_indexers = (
3832-
dask_indexers if is_duck_dask_array(var.data) else validated_indexers
3833+
dask_indexers if is_duck_dask_array(var._data) else validated_indexers
38333834
)
38343835

38353836
dtype_kind = var.dtype.kind

xarray/tests/test_missing.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
from __future__ import annotations
22

33
import itertools
4+
from unittest import mock
45

56
import numpy as np
67
import pandas as pd
78
import pytest
89

910
import xarray as xr
11+
from xarray.core import indexing
1012
from xarray.core.missing import (
1113
NumpyInterpolator,
1214
ScipyInterpolator,
@@ -772,3 +774,29 @@ def test_interpolators_complex_out_of_bounds():
772774
f = interpolator(xi, yi, method=method)
773775
actual = f(x)
774776
assert_array_equal(actual, expected)
777+
778+
779+
@requires_scipy
780+
def test_indexing_localize():
781+
# regression test for GH10287
782+
ds = xr.Dataset(
783+
{
784+
"sigma_a": xr.DataArray(
785+
data=np.ones((16, 8, 36811)),
786+
dims=["p", "t", "w"],
787+
coords={"w": np.linspace(0, 30000, 36811)},
788+
)
789+
}
790+
)
791+
792+
original_func = indexing.NumpyIndexingAdapter.__getitem__
793+
794+
def wrapper(self, indexer):
795+
return original_func(self, indexer)
796+
797+
with mock.patch.object(
798+
indexing.NumpyIndexingAdapter, "__getitem__", side_effect=wrapper, autospec=True
799+
) as mock_func:
800+
ds["sigma_a"].interp(w=15000.5)
801+
actual_indexer = mock_func.mock_calls[0].args[1]._key
802+
assert actual_indexer == (slice(None), slice(None), slice(18404, 18408))

0 commit comments

Comments
 (0)