Skip to content

Commit 6302e83

Browse files
committed
Fix performance regression in interp from #9881
Closes #10287
1 parent 303748c commit 6302e83

File tree

3 files changed

+36
-6
lines changed

3 files changed

+36
-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: 27 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,28 @@ 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+
def test_indexing_localize():
780+
# regression test for GH10287
781+
ds = xr.Dataset(
782+
{
783+
"sigma_a": xr.DataArray(
784+
data=np.ones((16, 8, 36811)),
785+
dims=["p", "t", "w"],
786+
coords={"w": np.linspace(0, 30000, 36811)},
787+
)
788+
}
789+
)
790+
791+
original_func = indexing.NumpyIndexingAdapter.__getitem__
792+
793+
def wrapper(self, indexer):
794+
return original_func(self, indexer)
795+
796+
with mock.patch.object(
797+
indexing.NumpyIndexingAdapter, "__getitem__", side_effect=wrapper, autospec=True
798+
) as mock_func:
799+
ds["sigma_a"].interp(w=15000.5)
800+
actual_indexer = mock_func.mock_calls[0].args[1]._key
801+
assert actual_indexer == (slice(None), slice(None), slice(18404, 18408))

0 commit comments

Comments
 (0)