Handle DataFrames with multiple blocks correctly

Author: evanpw

pandas/core/generic.py

@@ -3401,7 +3401,8 @@ def where(self, cond, other=np.nan, inplace=False, axis=None, level=None,
             axis = 0
             align = True
         else:
-            align = False
+            axis = self._get_axis_number(axis)
+            align = (axis == 1)
 
         block_axis = self._get_block_manager_axis(axis)
 

pandas/core/internals.py

@@ -671,7 +671,7 @@ def putmask(self, mask, new, align=True, inplace=False,
             # direction, then explictly repeat and reshape new instead
             if getattr(new, 'ndim', 0) >= 1:
                 if self.ndim - 1 == new.ndim and axis == 1:
-                    new = np.repeat(new, self.shape[-1]).reshape(self.shape)
+                    new = np.repeat(new, new_values.shape[-1]).reshape(self.shape)
 
             np.putmask(new_values, mask, new)
 
@@ -2427,12 +2427,18 @@ def apply(self, f, axes=None, filter=None, do_integrity_check=False, **kwargs):
             else:
                 kwargs['filter'] = filter_locs
 
-        if f == 'where' and kwargs.get('align', True):
+        if f == 'where':
             align_copy = True
-            align_keys = ['other', 'cond']
-        elif f == 'putmask' and kwargs.get('align', True):
+            if kwargs.get('align', True):
+                align_keys = ['other', 'cond']
+            else:
+                align_keys = ['cond']
+        elif f == 'putmask':
             align_copy = False
-            align_keys = ['new', 'mask']
+            if kwargs.get('align', True):
+                align_keys = ['new', 'mask']
+            else:
+                align_keys = ['mask']
         elif f == 'eval':
             align_copy = False
             align_keys = ['other']

pandas/tests/test_frame.py

@@ -9883,6 +9883,36 @@ def test_where_axis(self):
         result.where(mask, s, axis='columns', inplace=True)
         assert_frame_equal(result, expected)
 
+        # Multiple dtypes (=> multiple Blocks)
+        df = pd.concat([DataFrame(np.random.randn(10, 2)),
+                     DataFrame(np.random.randint(0, 10, size=(10, 2)))],
+                     ignore_index=True, axis=1)
+        mask = DataFrame(False, columns=df.columns, index=df.index)
+        s1 = Series(1, index=df.columns)
+        s2 = Series(2, index=df.index)
+
+        result = df.where(mask, s1, axis='columns')
+        expected = DataFrame(1.0, columns=df.columns, index=df.index)
+        expected[2] = expected[2].astype(int)
+        expected[3] = expected[3].astype(int)
+        assert_frame_equal(result, expected)
+
+        result = df.copy()
+        result.where(mask, s1, axis='columns', inplace=True)
+        assert_frame_equal(result, expected)
+
+        result = df.where(mask, s2, axis='index')
+        expected = DataFrame(2.0, columns=df.columns, index=df.index)
+        expected[2] = expected[2].astype(int)
+        expected[3] = expected[3].astype(int)
+        assert_frame_equal(result, expected)
+
+        result = df.copy()
+        result.where(mask, s2, axis='index', inplace=True)
+        assert_frame_equal(result, expected)
+
+
+
 
     def test_mask(self):
         df = DataFrame(np.random.randn(5, 3))