Use _unstack_once for valid dask and sparse versions

xarray/core/dask_array_compat.py

@@ -1,13 +1,12 @@
 import warnings
-from distutils.version import LooseVersion
 
 import numpy as np
 
+from .pycompat import dask_version
+
 try:
     import dask.array as da
-    from dask import __version__ as dask_version
 except ImportError:
-    dask_version = "0.0.0"
     da = None
 
 
@@ -57,7 +56,7 @@ def pad(array, pad_width, mode="constant", **kwargs):
     return padded
 
 
-if LooseVersion(dask_version) > LooseVersion("2.30.0"):
+if dask_version > "2.30.0":
     ensure_minimum_chunksize = da.overlap.ensure_minimum_chunksize
 else:
 
@@ -114,7 +113,7 @@ def ensure_minimum_chunksize(size, chunks):
         return tuple(output)
 
 
-if LooseVersion(dask_version) > LooseVersion("2021.03.0"):
+if dask_version > "2021.03.0":
     sliding_window_view = da.lib.stride_tricks.sliding_window_view
 else:
 
@@ -180,7 +179,7 @@ def sliding_window_view(x, window_shape, axis=None):
             axis=axis,
         )
         # map_overlap's signature changed in https://github.com/dask/dask/pull/6165
-        if LooseVersion(dask_version) > "2.18.0":
+        if dask_version > "2.18.0":
             return map_overlap(_np_sliding_window_view, x, align_arrays=False, **kwargs)
         else:
             return map_overlap(x, _np_sliding_window_view, **kwargs)

xarray/core/dataset.py

@@ -78,7 +78,12 @@
 )
 from .missing import get_clean_interp_index
 from .options import OPTIONS, _get_keep_attrs
-from .pycompat import is_duck_dask_array, sparse_array_type
+from .pycompat import (
+    dask_version,
+    is_duck_dask_array,
+    sparse_array_type,
+    sparse_version,
+)
 from .utils import (
     Default,
     Frozen,
@@ -4030,17 +4035,22 @@ def unstack(
         for dim in dims:
 
             if (
-                # Dask arrays don't support assignment by index, which the fast unstack
-                # function requires.
-                # https://github.com/pydata/xarray/pull/4746#issuecomment-753282125
-                any(is_duck_dask_array(v.data) for v in self.variables.values())
-                # Sparse doesn't currently support (though we could special-case
-                # it)
+                # Dask arrays supports assignment by index,
+                # https://github.com/dask/dask/pull/7393
+                (
+                    dask_version < "2021.04.0"
+                    and any(is_duck_dask_array(v.data) for v in self.variables.values())
+                )
+                # Sparse now supports kwargs in full_like,
                 # https://github.com/pydata/sparse/issues/422
-                or any(
-                    isinstance(v.data, sparse_array_type)
-                    for v in self.variables.values()
+                or (
+                    sparse_version < "0.11.2"
+                    and any(
+                        isinstance(v.data, sparse_array_type)
+                        for v in self.variables.values()
+                    )
                 )
+                # Shifting the arrays to sparse requires _unstack_full_reindex:
                 or sparse
                 # Until https://github.com/pydata/xarray/pull/4751 is resolved,
                 # we check explicitly whether it's a numpy array. Once that is

xarray/core/indexing.py

@@ -2,22 +2,15 @@
 import functools
 import operator
 from collections import defaultdict
-from distutils.version import LooseVersion
 from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union
 
 import numpy as np
 import pandas as pd
 
-try:
-    import dask
-
-    DASK_VERSION = LooseVersion(dask.__version__)
-except ModuleNotFoundError:
-    DASK_VERSION = LooseVersion("0")
-
 from . import duck_array_ops, nputils, utils
 from .pycompat import (
     dask_array_type,
+    dask_version,
     integer_types,
     is_duck_dask_array,
     sparse_array_type,
@@ -1393,7 +1386,7 @@ def __getitem__(self, key):
                 return value
 
     def __setitem__(self, key, value):
-        if DASK_VERSION >= "2021.04.1":
+        if dask_version >= "2021.04.1":
             if isinstance(key, BasicIndexer):
                 self.array[key.tuple] = value
             elif isinstance(key, VectorizedIndexer):

xarray/core/pycompat.py

@@ -1,13 +1,18 @@
+from distutils.version import LooseVersion
+
 import numpy as np
 
 from .utils import is_duck_array
 
 integer_types = (int, np.integer)
 
 try:
+    import dask
     import dask.array
     from dask.base import is_dask_collection
 
+    dask_version = LooseVersion(dask.__version__)
+
     # solely for isinstance checks
     dask_array_type = (dask.array.Array,)
 
@@ -16,6 +21,7 @@ def is_duck_dask_array(x):
 
 
 except ImportError:  # pragma: no cover
+    dask_version = LooseVersion("0.0.0")
     dask_array_type = ()
     is_duck_dask_array = lambda _: False
     is_dask_collection = lambda _: False
@@ -24,14 +30,18 @@ def is_duck_dask_array(x):
     # solely for isinstance checks
     import sparse
 
+    sparse_version = LooseVersion(sparse.__version__)
     sparse_array_type = (sparse.SparseArray,)
 except ImportError:  # pragma: no cover
+    sparse_version = LooseVersion("0.0.0")
     sparse_array_type = ()
 
 try:
     # solely for isinstance checks
     import cupy
 
+    cupy_version = LooseVersion(sparse.__version__)
     cupy_array_type = (cupy.ndarray,)
 except ImportError:  # pragma: no cover
+    cupy_version = LooseVersion("0.0.0")
     cupy_array_type = ()

xarray/core/utils.py

@@ -31,8 +31,6 @@
 import numpy as np
 import pandas as pd
 
-from . import dtypes
-
 K = TypeVar("K")
 V = TypeVar("V")
 T = TypeVar("T")
@@ -83,9 +81,10 @@ def maybe_coerce_to_str(index, original_coords):
 
     pd.Index uses object-dtype to store str - try to avoid this for coords
     """
+    from .dtypes import result_type as dtypes_result_type
 
     try:
-        result_type = dtypes.result_type(*original_coords)
+        result_type = dtypes_result_type(*original_coords)
     except TypeError:
         pass
     else: