pydata / pydata/xarray

Faster unstacking of dask arrays

Open
#5,582 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement topic-performance
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

Recent dask version support assigning to a list of ints along one dimension. we can use this for unstacking (diff builds on #5577)

diff --git i/xarray/core/variable.py w/xarray/core/variable.py
index 222e8dab9..a50dfc574 100644
--- i/xarray/core/variable.py
+++ w/xarray/core/variable.py
@@ -1593,11 +1593,9 @@ class Variable(AbstractArray, NdimSizeLenMixin, VariableArithmetic):
         else:
             dtype = self.dtype
 
-        if sparse:
+        if sparse and not is_duck_dask_array(reordered):
             # 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
+            # Use the sparse.COO constructor since we cannot assign to sparse.COO
             from sparse import COO
 
             codes = zip(*index.codes)
@@ -1618,19 +1616,23 @@ class Variable(AbstractArray, NdimSizeLenMixin, VariableArithmetic):
             )
 
         else:
+            # dask supports assigning to a list of ints along one axis only.
+            # So we construct an array with the last dimension flattened,
+            # assign the values, then reshape to the final shape.
+            intermediate_shape = reordered.shape[:-1] + (np.prod(new_dim_sizes),)
+            indexer = np.ravel_multi_index(index.codes, new_dim_sizes)
             data = np.full_like(
                 self.data,
                 fill_value=fill_value,
-                shape=new_shape,
+                shape=intermediate_shape,
                 dtype=dtype,
             )
 
             # Indexer is a list of lists of locations. Each list is the locations
             # on the new dimension. This is robust to the data being sparse; in that
             # case the destinations will be NaN / zero.
-            # sparse doesn't support item assigment,
-            # https://github.com/pydata/sparse/issues/114
-            data[(..., *indexer)] = reordered
+            data[(..., indexer)] = reordered
+            data = data.reshape(new_shape)
 
         return self._replace(dims=new_dims, data=data)

This should be what alignment.reindex_variables is doing but I don't fully understand that function.

The annoying bit is figuring out when to use this version and what to do with things like dask wrapping sparse. I think we want to loop over each variable in Dataset.unstack calling Variable.unstack and dispatch based on the type of Variable.data to easily handle all the edge cases.

cc @Illviljan if you're interested in implementing this

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with Variable.unstack in xarray/core/variable.py and compare its current sparse and dask paths with alignment.reindex_variables. Trace Dataset.unstack to understand how variables are dispatched, then verify that unstacking handles dask arrays, sparse arrays, and dask-wrapped sparse data without regressing existing behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.