add in split_and_operate

Author: jreback

pandas/core/internals.py

@@ -397,12 +397,75 @@ def fillna(self, value, limit=None, inplace=False, downcast=None,
             if not mask.any():
                 return self if inplace else self.copy()
 
-            # try again with a compatible block
-            block = self.coerce_to_target_dtype(original_value)
-            return block.fillna(original_value,
-                                limit=limit,
-                                inplace=inplace,
-                                downcast=False)
+            # operate column-by-column
+            def f(m, v):
+
+                # try again with a compatible block
+                block = self.coerce_to_target_dtype(original_value)
+                return block.fillna(original_value,
+                                    limit=limit,
+                                    inplace=inplace,
+                                    downcast=False)
+
+            return self.split_and_operate(mask, f, inplace)
+
+    def split_and_operate(self, mask, f, inplace):
+        """
+        split the block per-column. if masking
+        that block, promot if needed, then
+        apply the function f
+
+        Parameters
+        ----------
+        mask : 2-d boolean mask
+        f : function accepting (1-mask, 1-d values)
+        inplace : boolean
+
+        Returns
+        -------
+        list of blocks
+        """
+
+        new_values = self.values
+
+        def make_a_block(nv, ref_loc):
+            if isinstance(nv, Block):
+                block = nv
+            elif isinstance(nv, list):
+                block = nv[0]
+            else:
+                # Put back the dimension that was taken from it and make
+                # a block out of the result.
+                block = self.make_block(values=nv[np.newaxis],
+                                        placement=[ref_loc], fastpath=True)
+            return block
+
+        # ndim == 1
+        if self.ndim == 1:
+
+            if mask.any():
+                nv = f(mask, new_values)
+            else:
+                nv = new_values if inplace else new_values.copy()
+            block = make_a_block(nv, 0)
+            return block
+
+        # ndim > 1
+        new_blocks = []
+        for i, ref_loc in enumerate(self.mgr_locs):
+            m = mask[i]
+            v = new_values[i]
+
+            # need a new block
+            if m.any():
+                nv = f(m, v)
+            else:
+                nv = v if inplace else v.copy()
+
+            block = make_a_block(nv, ref_loc)
+            new_blocks.append(block)
+
+        return new_blocks
 
     def _maybe_downcast(self, blocks, downcast=None):
 
@@ -899,6 +962,9 @@ def putmask(self, mask, new, align=True, inplace=False, axis=0,
                     new_shape.insert(axis, 1)
                     new = new.reshape(tuple(new_shape))
 
+            # TODO(jreback)
+            # use split_and_operate
+
             # need to go column by column
             new_blocks = []
             if self.ndim > 1: