DaskSparseFrame getitem, todense and bugfix

sparsity/dask/core.py

@@ -135,6 +135,22 @@ def index(self):
     def map_partitions(self, func, meta, *args, **kwargs):
         return map_partitions(func, self, meta, *args, **kwargs)
 
+    # noinspection PyTypeChecker
+    def todense(self):
+        """Convert into Dask DataFrame or Series
+
+        Returns
+        -------
+        res: dd.DataFrame | dd.Series
+        """
+        meta = dd.from_pandas(self._meta.todense(), npartitions=1)
+        res = self.map_partitions(
+            sp.SparseFrame.todense,
+            meta,
+            cls=dd.DataFrame if len(self.columns) > 1 else dd.Series,
+        )
+        return res
+
     def to_delayed(self):
         return [Delayed(k, self.dask) for k in self.__dask_keys__()]
 
@@ -644,7 +660,7 @@ def elemwise(op, *args, **kwargs):
     return SparseFrame(dsk, _name, meta, divisions)
 
 
-def map_partitions(func, ddf, meta, name=None, **kwargs):
+def map_partitions(func, ddf, meta, name=None, cls=SparseFrame, **kwargs):
     dsk = {}
     name = name or func.__name__
     token = tokenize(func, meta, **kwargs)
@@ -654,12 +670,11 @@ def map_partitions(func, ddf, meta, name=None, **kwargs):
         value = (ddf._name, i)
         dsk[(name, i)] = (apply_and_enforce, func, value, kwargs, meta)
 
-    return SparseFrame(merge(dsk, ddf.dask), name, meta, ddf.divisions)
+    return cls(merge(dsk, ddf.dask), name, meta, ddf.divisions)
 
 
 def apply_and_enforce(func, arg, kwargs, meta):
     sf = func(arg, **kwargs)
-    columns = meta.columns
     if isinstance(sf, sp.SparseFrame):
         if len(sf.data.data) == 0:
             assert meta.empty, \
@@ -668,6 +683,7 @@ def apply_and_enforce(func, arg, kwargs, meta):
                 "Computed a SparseFrame but meta is of type {}"\
                 .format(type(meta))
             return meta
+        columns = meta.columns
         if (len(columns) == len(sf.columns) and
                     type(columns) is type(sf.columns) and
                 columns.equals(sf.columns)):

sparsity/test/test_dask_sparse_frame.py

@@ -1,14 +1,14 @@
-import datetime as dt
 import os
-from uuid import uuid4
 
 import dask
 import dask.dataframe as dd
+import datetime as dt
 import numpy as np
 import pandas as pd
 import pandas.util.testing as pdt
 import pytest
 from distributed import Client
+from uuid import uuid4
 
 import sparsity as sp
 import sparsity.dask as dsp
@@ -45,6 +45,24 @@ def test_map_partitions():
     assert res.shape == (10, 2)
 
 
+def test_todense():
+    data = pd.DataFrame(np.random.rand(10, 2))
+    dsf = dsp.from_pandas(data, npartitions=3)
+    res = dsf.todense()
+    assert isinstance(res, dd.DataFrame)
+    computed = res.compute()
+    pdt.assert_frame_equal(computed, data, check_dtype=False)
+
+
+def test_todense_series():
+    data = pd.DataFrame(np.random.rand(10, 2))
+    dsf = dsp.from_pandas(data, npartitions=3)[0]
+    res = dsf.todense()
+    assert isinstance(res, dd.Series)
+    computed = res.compute()
+    pdt.assert_series_equal(computed, data[0], check_dtype=False)
+
+
 @pytest.mark.parametrize('item', [
     'X',
     ['X', 'Y'],