|
78 | 78 | )
|
79 | 79 | from .missing import get_clean_interp_index
|
80 | 80 | 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 |
82 | 82 | from .utils import (
|
83 | 83 | Default,
|
84 | 84 | Frozen,
|
@@ -4028,36 +4028,24 @@ def unstack(
|
4028 | 4028 |
|
4029 | 4029 | result = self.copy(deep=False)
|
4030 | 4030 | 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() |
4057 | 4038 | ):
|
4058 |
| - result = result._unstack_full_reindex(dim, fill_value, sparse) |
4059 |
| - else: |
| 4039 | + # Fast unstacking path: |
4060 | 4040 | 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) |
4061 | 4049 | return result
|
4062 | 4050 |
|
4063 | 4051 | def update(self, other: "CoercibleMapping") -> "Dataset":
|
|
0 commit comments