chunk sparse arrays

xarray/core/duck_array_ops.py

@@ -186,7 +186,12 @@ def as_like_arrays(*data):
     elif any(isinstance(d, sparse_array_type) for d in data):
         from sparse import COO
 
-        return tuple(COO(d) for d in data)
+        out = []
+        for d in data:
+            if isinstance(d, dask_array_type):
+                d = d.compute()
+            out.append(COO(d))
+        return tuple(out)
     else:
         return tuple(np.asarray(d) for d in data)
 

xarray/core/variable.py

@@ -928,21 +928,26 @@ def chunk(self, chunks=None, name=None, lock=False):
         else:
             if utils.is_dict_like(chunks):
                 chunks = tuple(chunks.get(n, s) for n, s in enumerate(self.shape))
-            # da.from_array works by using lazily indexing with a tuple of
-            # slices. Using OuterIndexer is a pragmatic choice: dask does not
-            # yet handle different indexing types in an explicit way:
-            # https://github.com/dask/dask/issues/2883
-            data = indexing.ImplicitToExplicitIndexingAdapter(
-                data, indexing.OuterIndexer
-            )
 
-            # For now, assume that all arrays that we wrap with dask (including
-            # our lazily loaded backend array classes) should use NumPy array
-            # operations.
-            if LooseVersion(dask.__version__) > "1.2.2":
-                kwargs = dict(meta=np.ndarray)
+            if not isinstance(data, np.ndarray) and hasattr(data, "__array_function__"):
+                # This should always be true after Variable.as_compatible_data
+                assert IS_NEP18_ACTIVE
+                kwargs = {}
             else:
-                kwargs = dict()
+                # da.from_array works by using lazily indexing with a tuple of
+                # slices. Using OuterIndexer is a pragmatic choice: dask does not
+                # yet handle different indexing types in an explicit way:
+                # https://github.com/dask/dask/issues/2883
+                data = indexing.ImplicitToExplicitIndexingAdapter(
+                    data, indexing.OuterIndexer
+                )
+
+                if LooseVersion(dask.__version__) < "2.0.0":
+                    kwargs = {}
+                else:
+                    # All of our lazily loaded backend array classes) should use NumPy
+                    # array operations.
+                    kwargs = {"meta": np.ndarray}
 
             data = da.from_array(data, chunks, name=name, lock=lock, **kwargs)
 

xarray/tests/test_sparse.py

@@ -136,7 +136,6 @@ def test_variable_property(prop):
             True,
             marks=xfail(reason="'COO' object has no attribute 'argsort'"),
         ),
-        param(do("chunk", chunks=(5, 5)), True, marks=xfail),
         param(
             do(
                 "concat",
@@ -404,9 +403,6 @@ def test_dataarray_property(prop):
             False,
             marks=xfail(reason="Missing implementation for np.flip"),
         ),
-        param(
-            do("chunk", chunks=(5, 5)), False, marks=xfail(reason="Coercion to dense")
-        ),
         param(
             do("combine_first", make_xrarray({"x": 10, "y": 5})),
             True,
@@ -861,3 +857,17 @@ def test_sparse_coords(self):
             dims=["x"],
             coords={"x": COO.from_numpy([1, 2, 3, 4])},
         )
+
+
+def test_chunk():
+    s = sparse.COO.from_numpy(np.array([0, 0, 1, 2]))
+    a = DataArray(s)
+    ac = a.chunk(2)
+    assert ac.chunks == ((2, 2),)
+    assert isinstance(ac.data._meta, sparse.COO)
+    assert_identical(a, ac)
+
+    ds = a.to_dataset(name="a")
+    dsc = ds.chunk(2)
+    assert dsc.chunks == {"dim_0": (2, 2)}
+    assert_identical(ds, dsc)