putmask to use split_and_operate

Author: jreback

pandas/core/internals.py

@@ -398,7 +398,7 @@ def fillna(self, value, limit=None, inplace=False, downcast=None,
                 return self if inplace else self.copy()
 
             # operate column-by-column
-            def f(m, v):
+            def f(m, v, i):
 
                 # try again with a compatible block
                 block = self.coerce_to_target_dtype(original_value)
@@ -411,14 +411,14 @@ def f(m, v):
 
     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
+        split the block per-column, and apply the callable f
+        per-column, return a new block for each. Handle
+        masking which will not change a block unless needed.
 
         Parameters
         ----------
         mask : 2-d boolean mask
-        f : function accepting (1-mask, 1-d values)
+        f : callable accepting (1d-mask, 1d values, indexer)
         inplace : boolean
 
         Returns
@@ -436,19 +436,19 @@ def make_a_block(nv, ref_loc):
             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)
+                nv = _block_shape(nv, ndim=self.ndim)
+                block = self.make_block(values=nv,
+                                        placement=ref_loc, fastpath=True)
             return block
 
         # ndim == 1
         if self.ndim == 1:
-
             if mask.any():
-                nv = f(mask, new_values)
+                nv = f(mask, new_values, None)
             else:
                 nv = new_values if inplace else new_values.copy()
-            block = make_a_block(nv, 0)
-            return block
+            block = make_a_block(nv, self.mgr_locs)
+            return [block]
 
         # ndim > 1
         new_blocks = []
@@ -458,11 +458,11 @@ def make_a_block(nv, ref_loc):
 
             # need a new block
             if m.any():
-                nv = f(m, v)
+                nv = f(m, v, i)
             else:
                 nv = v if inplace else v.copy()
 
-            block = make_a_block(nv, ref_loc)
+            block = make_a_block(nv, [ref_loc])
             new_blocks.append(block)
 
         return new_blocks
@@ -960,44 +960,31 @@ 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:
-                for i, ref_loc in enumerate(self.mgr_locs):
-                    m = mask[i]
-                    v = new_values[i]
-
-                    # need a new block
-                    if m.any():
-                        if isinstance(new, np.ndarray):
-                            n = np.squeeze(new[i % new.shape[0]])
-                        else:
-                            n = np.array(new)
-
-                        # type of the new block
-                        dtype, _ = maybe_promote(n.dtype)
+            # operate column-by-column
+            def f(m, v, i):
 
-                        # we need to explicitly astype here to make a copy
-                        n = n.astype(dtype)
+                # TODO(jreback)
+                # see if we can use coerce_to_target_dtype here instead
 
-                        nv = _putmask_smart(v, m, n)
+                if i is None:
+                    # ndim==1 case.
+                    n = new
+                else:
+                    if isinstance(new, np.ndarray):
+                        n = np.squeeze(new[i % new.shape[0]])
                     else:
-                        nv = v if inplace else v.copy()
+                        n = np.array(new)
 
-                    # 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)
+                    # type of the new block
+                    dtype, _ = maybe_promote(n.dtype)
 
-                    new_blocks.append(block)
+                    # we need to explicitly astype here to make a copy
+                    n = n.astype(dtype)
 
-            else:
-                nv = _putmask_smart(new_values, mask, new)
-                new_blocks.append(self.make_block(values=nv, fastpath=True))
+                nv = _putmask_smart(v, m, n)
+                return nv
 
+            new_blocks = self.split_and_operate(mask, f, inplace)
             return new_blocks
 
         if inplace: