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,