Skip to content

Commit 9165c26

Browse files
authored
Use _unstack_once for valid dask and sparse versions (pydata#5315)
* optional import versions to pycompat * Update indexing.py * move dtypes to avoid circular import * Update pycompat.py * faster unstacking * Update dataset.py * Update dataset.py * have to check that the array type is in use * Update dataset.py * sparse arg requires the slow path? * cupy.__version__ * test pint_array_type * Update dataset.py * Add pint check in pycompat * Update pycompat.py * revert pint test * import whole dtypes * Test turning off the ndarray check * sparse and pint doesn't work. Switch to a restrictive if * Update dataset.py * Add back some comments and add some relevant issue links * lint * Update dataset.py
1 parent 2bb5d20 commit 9165c26

File tree

1 file changed

+17
-29
lines changed

1 file changed

+17
-29
lines changed

xarray/core/dataset.py

+17-29
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@
7878
)
7979
from .missing import get_clean_interp_index
8080
from .options import OPTIONS, _get_keep_attrs
81-
from .pycompat import is_duck_dask_array, sparse_array_type
81+
from .pycompat import dask_version, is_duck_dask_array
8282
from .utils import (
8383
Default,
8484
Frozen,
@@ -4028,36 +4028,24 @@ def unstack(
40284028

40294029
result = self.copy(deep=False)
40304030
for dim in dims:
4031-
4032-
if (
4033-
# Dask arrays don't support assignment by index, which the fast unstack
4034-
# function requires.
4035-
# https://github.com/pydata/xarray/pull/4746#issuecomment-753282125
4036-
any(is_duck_dask_array(v.data) for v in self.variables.values())
4037-
# Sparse doesn't currently support (though we could special-case
4038-
# it)
4039-
# https://github.com/pydata/sparse/issues/422
4040-
or any(
4041-
isinstance(v.data, sparse_array_type)
4042-
for v in self.variables.values()
4043-
)
4044-
or sparse
4045-
# Until https://github.com/pydata/xarray/pull/4751 is resolved,
4046-
# we check explicitly whether it's a numpy array. Once that is
4047-
# resolved, explicitly exclude pint arrays.
4048-
# # pint doesn't implement `np.full_like` in a way that's
4049-
# # currently compatible.
4050-
# # https://github.com/pydata/xarray/pull/4746#issuecomment-753425173
4051-
# # or any(
4052-
# # isinstance(v.data, pint_array_type) for v in self.variables.values()
4053-
# # )
4054-
or any(
4055-
not isinstance(v.data, np.ndarray) for v in self.variables.values()
4056-
)
4031+
if not sparse and all(
4032+
# Dask arrays recently supports assignment by index,
4033+
# https://github.com/dask/dask/pull/7393
4034+
dask_version >= "2021.04.0" and is_duck_dask_array(v.data)
4035+
# Numpy arrays handles the fast path:
4036+
or isinstance(v.data, np.ndarray)
4037+
for v in self.variables.values()
40574038
):
4058-
result = result._unstack_full_reindex(dim, fill_value, sparse)
4059-
else:
4039+
# Fast unstacking path:
40604040
result = result._unstack_once(dim, fill_value)
4041+
else:
4042+
# Slower unstacking path, examples of array types that
4043+
# currently has to use this path:
4044+
# * sparse doesn't support item assigment,
4045+
# https://github.com/pydata/sparse/issues/114
4046+
# * pint has some circular import issues,
4047+
# https://github.com/pydata/xarray/pull/4751
4048+
result = result._unstack_full_reindex(dim, fill_value, sparse)
40614049
return result
40624050

40634051
def update(self, other: "CoercibleMapping") -> "Dataset":

0 commit comments

Comments
 (0)