Add support for dask persist

sparsity/dask/core.py

@@ -89,6 +89,14 @@ def __dask_optimize__(dsk, keys, **kwargs):
         dsk2 = optimize(dsk, keys)
         return dsk2
 
+
+    def __dask_postpersist__(self):
+        def rebuild(dsk, *extra_args):
+            return SparseFrame(dsk, name=self._name,
+                               meta=self._meta,
+                               divisions=self.divisions)
+        return rebuild, ()
+
     def __dask_postcompute__(self):
         return finalize, ()
 
@@ -208,8 +216,8 @@ def join(self, other, on=None, how='left', lsuffix='',
 
     def to_npz(self, filename, blocksize=None,
                storage_options=None, compute=True):
-        from sparsity.dask.io_ import to_npz
-        to_npz(self, filename, blocksize, storage_options, compute)
+        import sparsity.dask.io_ as dsp_io
+        return dsp_io.to_npz(self, filename, blocksize, storage_options, compute)
 
     def groupby_sum(self, split_out=1, split_every=8):
         meta = self._meta

sparsity/test/test_dask_sparse_frame.py

@@ -1,5 +1,6 @@
 import datetime as dt
 import os
+from distributed import Client
 from uuid import uuid4
 
 import dask
@@ -460,3 +461,13 @@ def test_set_index(clickstream):
     res = dsf.set_index(level=1).compute().todense()
 
     pdt.assert_frame_equal(dense, res)
+
+
+def test_persist(dsf):
+    correct = dsf.compute().todense()
+    client = Client()
+    persisted = client.persist(dsf)
+
+    res = persisted.compute().todense()
+
+    pdt.assert_frame_equal(res, correct)