ENH: SparseDataFrame supports scipy.sparse.spmatrix in setitem

Author: kernc

pandas/core/internals.py

@@ -2571,6 +2571,16 @@ def _astype(self, dtype, copy=False, raise_on_error=True, values=None,
         return self.make_block_same_class(values=values,
                                           placement=self.mgr_locs)
 
+    def _can_hold_element(self, element):
+        element = np.asanyarray(element)
+        return np.issubdtype(element.dtype, self.sp_values.dtype)
+
+    def _try_cast(self, element):
+        try:
+            return np.asarray(element, dtype=self.sp_values.dtype)
+        except ValueError:
+            return element
+
     def __len__(self):
         try:
             return self.sp_index.length

pandas/core/sparse/frame.py

@@ -433,6 +433,15 @@ def __getitem__(self, key):
         else:
             return self._get_item_cache(key)
 
+    def __setitem__(self, key, value):
+        if is_scipy_sparse(value):
+            if any(ax == 1 for ax in value.shape):  # 1d spmatrix
+                value = SparseArray(value, fill_value=self._default_fill_value)
+            else:
+                # 2d; make it iterable
+                value = list(value.tocsc().T)
+        super().__setitem__(key, value)
+
     @Appender(DataFrame.get_value.__doc__, indents=0)
     def get_value(self, index, col, takeable=False):
         if takeable is True:

pandas/tests/sparse/test_frame.py

@@ -540,6 +540,33 @@ def test_setitem_array(self):
                                   self.frame['F'].reindex(index),
                                   check_names=False)
 
+    def test_setitem_spmatrix(self):
+        # GH-15634
+        tm.skip_if_no_package('scipy')
+        from scipy.sparse import csr_matrix
+
+        sdf = self.frame.copy(False)
+
+        # 1d -- column
+        spm = csr_matrix(np.arange(len(sdf))).T
+        sdf['X'] = spm
+        assert (sdf[['X']].to_coo() != spm.tocoo()).nnz == 0
+
+        # 1d -- existing column
+        sdf['A'] = spm.T
+        assert (sdf[['X']].to_coo() != spm.tocoo()).nnz == 0
+
+        # 1d row -- changing series contents not yet supported
+        spm = csr_matrix(np.arange(sdf.shape[1])).astype(float)
+        idx = np.r_[[False, True], np.full(sdf.shape[0] - 2, False)]
+        tm.assert_raises_regex(TypeError, 'assignment',
+                               lambda: sdf.__setitem__(idx, spm))
+
+        # 2d -- 2 columns
+        spm = csr_matrix(np.eye(len(sdf))[:, :2])
+        sdf[['X', 'A']] = spm
+        assert (sdf[['X', 'A']].to_coo() != spm.tocoo()).nnz == 0
+
     def test_delitem(self):
         A = self.frame['A']
         C = self.frame['C']