Skip to content

Commit fbb5ba9

Browse files
committed
4. Create dask coordinate arrays only once.
1 parent 0f40188 commit fbb5ba9

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

xarray/core/dataset.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2895,15 +2895,27 @@ def _validate_interp_indexer(x, new_x):
28952895
obj, newidx = missing._localize(obj, {k: v})
28962896
validated_indexers[k] = newidx[k]
28972897

2898+
# optimization: create dask coordinate arrays once per Dataset
2899+
# rather than once per Variable when dask.array.unify_chunks is called later
2900+
# GH4739
2901+
if obj.__dask_graph__():
2902+
dask_indexers = {
2903+
k: (index.to_base_variable().chunk(), dest.to_base_variable().chunk())
2904+
for k, (index, dest) in validated_indexers.items()
2905+
}
2906+
28982907
variables: Dict[Hashable, Variable] = {}
28992908
for name, var in obj._variables.items():
29002909
if name in indexers:
29012910
continue
29022911

2912+
if is_duck_dask_array(var.data):
2913+
use_indexers = dask_indexers
2914+
else:
2915+
use_indexers = validated_indexers
2916+
29032917
if var.dtype.kind in "uifc":
2904-
var_indexers = {
2905-
k: v for k, v in validated_indexers.items() if k in var.dims
2906-
}
2918+
var_indexers = {k: v for k, v in use_indexers.items() if k in var.dims}
29072919
variables[name] = missing.interp(var, var_indexers, method, **kwargs)
29082920
elif all(d not in indexers for d in var.dims):
29092921
# keep unrelated object array

0 commit comments

Comments
 (0)