cleanups + add comments

pydata · Jul 5, 2021 · e4a6ec2 · e4a6ec2
1 parent e976ada
commit e4a6ec2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 15 deletions.
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -4133,23 +4133,21 @@ def 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)
-                # https://github.com/pydata/sparse/issues/422
+                # Sparse doesn't currently support advanced indexing
+                # https://github.com/pydata/sparse/issues/114
                 or any(
                     isinstance(v.data, sparse_array_type)
                     for v in self.variables.values()
                 )
-                # or sparse
                 # Until https://github.com/pydata/xarray/pull/4751 is resolved,
                 # we check explicitly whether it's a numpy array. Once that is
                 # resolved, explicitly exclude pint arrays.
-                # # pint doesn't implement `np.full_like` in a way that's
-                # # currently compatible.
-                # # https://github.com/pydata/xarray/pull/4746#issuecomment-753425173
-                # # or any(
-                # #     isinstance(v.data, pint_array_type) for v in self.variables.values()
-                # # )
+                # pint doesn't implement `np.full_like` in a way that's
+                # currently compatible.
+                # https://github.com/pydata/xarray/pull/4746#issuecomment-753425173
+                # or any(
+                #     isinstance(v.data, pint_array_type) for v in self.variables.values()
+                # )
                 or any(
                     not isinstance(v.data, np.ndarray) for v in self.variables.values()
                 )

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -1565,22 +1565,22 @@ def _unstack_once(
         index: pd.MultiIndex,
         dim: Hashable,
         fill_value=dtypes.NA,
-        sparse=False,
+        sparse: bool = False,
     ) -> "Variable":
         """
         Unstacks this variable given an index to unstack and the name of the
         dimension to which the index refers.
         """
 
         reordered = self.transpose(..., dim)
-        shape = reordered.shape
+
         new_dim_sizes = [lev.size for lev in index.levels]
         new_dim_names = index.names
         indexer = index.codes
 
         # Potentially we could replace `len(other_dims)` with just `-1`
         other_dims = [d for d in self.dims if d != dim]
-        new_shape = tuple(list(shape[: len(other_dims)]) + new_dim_sizes)
+        new_shape = tuple(list(reordered.shape[: len(other_dims)]) + new_dim_sizes)
         new_dims = reordered.dims[: len(other_dims)] + new_dim_names
 
         if fill_value is dtypes.NA:
@@ -1594,14 +1594,17 @@ def _unstack_once(
             dtype = self.dtype
 
         if sparse:
+            # unstacking a dense multitindexed array to a sparse array
+            # Use the sparse.COO constructor until sparse supports advanced indexing
+            # https://github.com/pydata/sparse/issues/114
             # TODO: how do we allow different sparse array types
             from sparse import COO
 
             codes = zip(*index.codes)
-            if not shape[:-1]:
+            if reordered.ndim == 1:
                 indexes = codes
             else:
-                sizes = itertools.product(range(*shape[:-1]))
+                sizes = itertools.product(range(*reordered.shape[:-1]))
                 tuple_indexes = itertools.product(sizes, codes)
                 indexes = map(lambda x: list(itertools.chain(*x)), tuple_indexes)  # type: ignore