broadcast should preserve like-columns

Author: jreback

pandas/core/apply.py

@@ -187,7 +187,7 @@ def apply_broadcast(self, target):
         # axis which we want to compare compliance
         result_compare = target.shape[0]
 
-        index = target.index
+        index = None
         for i, col in enumerate(target.columns):
             res = self.f(target[col])
             ares = np. asarray(res).ndim
@@ -208,6 +208,11 @@ def apply_broadcast(self, target):
 
             result_values[:, i] = res
 
+        # if we are returning a list-like
+        # then preserve the original index
+        if index is None:
+            index = target.index
+
         result = self.obj._constructor(result_values, index=index,
                                        columns=target.columns)
         return result

pandas/tests/frame/test_apply.py

@@ -163,6 +163,24 @@ def test_apply_broadcast(self):
                              index=self.frame.index)
         tm.assert_frame_equal(result, expected)
 
+        # preserve columns
+        df = DataFrame(np.tile(np.arange(3), 6).reshape(6, -1) + 1,
+                       columns=list('ABC'))
+        result = df.apply(lambda x: [1, 2, 3],
+                          axis=1,
+                          result_type='broadcast')
+        tm.assert_frame_equal(result, df)
+
+        # columms come from the returned Series
+        df = DataFrame(np.tile(np.arange(3), 6).reshape(6, -1) + 1,
+                       columns=list('ABC'))
+        result = df.apply(lambda x: Series([1, 2, 3], index=list('abc')),
+                          axis=1,
+                          result_type='broadcast')
+        expected = df.copy()
+        expected.columns = list('abc')
+        tm.assert_frame_equal(result, expected)
+
     def test_apply_broadcast_error(self):
         df = DataFrame(
             np.tile(np.arange(3, dtype='int64'), 6).reshape(6, -1) + 1,